ck-builder: tensor input/output reflection (#3536)

This adds some utilities to automatically generate UniqueInputs,
UniqueOutputs, alloc_inputs, alloc_outputs, and validate, based
on a Inputs::reflect() and Outputs::reflect().

Author: Snektron

experimental/builder/include/ck_tile/builder/testing/conv_fwd.hpp

@@ -7,6 +7,7 @@
 #include "ck_tile/builder/factory/helpers/ck/conv_tensor_layout.hpp"
 #include "ck_tile/builder/factory/helpers/ck/conv_elementwise_op.hpp"
 #include "ck_tile/builder/testing/testing.hpp"
+#include "ck_tile/builder/testing/testing_reflect.hpp"
 #include "ck_tile/builder/testing/filter_extent.hpp"
 #include "ck_tile/builder/testing/tensor_buffer.hpp"
 #include "ck_tile/builder/testing/tensor_initialization.hpp"
@@ -182,6 +183,12 @@ struct Inputs<SIGNATURE>
 {
     void* input;
     void* weight;
+
+    static void reflect(const Args<SIGNATURE>& args, const auto& inspect)
+    {
+        inspect("input", args.make_input_descriptor(), &Inputs<SIGNATURE>::input);
+        inspect("weight", args.make_weight_descriptor(), &Inputs<SIGNATURE>::weight);
+    }
 };
 
 /// @brief `Outputs` specialization for forward convolution.
@@ -194,68 +201,13 @@ template <auto SIGNATURE>
 struct Outputs<SIGNATURE>
 {
     void* output;
-};
-
-/// @brief `UniqueInputs` specialization for forward convolution.
-///
-/// @tparam SIGNATURE Forward convolution signature.
-///
-/// @see UniqueInputs
-/// @see ValidUniqueInputs
-template <auto SIGNATURE>
-    requires ValidConvSignature<SIGNATURE> && ConvDirectionIsForward<SIGNATURE>
-struct UniqueInputs<SIGNATURE>
-{
-    DeviceBuffer input_buf;
-    DeviceBuffer weight_buf;
-
-    /// @see ValidUniqueInputs
-    Inputs<SIGNATURE> get()
-    {
-        return {
-            .input  = input_buf.get(),
-            .weight = weight_buf.get(),
-        };
-    }
-};
-
-/// @brief `UniqueOutputs` specialization for forward convolution.
-///
-/// @tparam SIGNATURE Forward convolution signature.
-///
-/// @see UniqueOutputs
-/// @see ValidUniqueOutputs
-template <auto SIGNATURE>
-    requires ValidConvSignature<SIGNATURE> && ConvDirectionIsForward<SIGNATURE>
-struct UniqueOutputs<SIGNATURE>
-{
-    DeviceBuffer output_buf;
 
-    /// @see ValidUniqueOutputs
-    Outputs<SIGNATURE> get()
+    static void reflect(const Args<SIGNATURE>& args, const auto& inspect)
     {
-        return {
-            .output = output_buf.get(),
-        };
+        inspect("output", args.make_output_descriptor(), &Outputs<SIGNATURE>::output);
     }
 };
 
-/// @brief `alloc_inputs()` specialization for forward convolution.
-///
-/// @tparam SIGNATURE Forward convolution signature.
-///
-/// @see alloc_inputs()
-template <auto SIGNATURE>
-    requires ValidConvSignature<SIGNATURE> && ConvDirectionIsForward<SIGNATURE> &&
-             ValidUniqueInputs<SIGNATURE>
-UniqueInputs<SIGNATURE> alloc_inputs(const Args<SIGNATURE>& args)
-{
-    return {
-        .input_buf  = alloc_tensor_buffer(args.make_input_descriptor()),
-        .weight_buf = alloc_tensor_buffer(args.make_weight_descriptor()),
-    };
-}
-
 /// @brief `init_inputs()` specialization for forward convolution.
 ///
 /// @tparam SIGNATURE Forward convolution signature.
@@ -269,34 +221,4 @@ void init_inputs(const Args<SIGNATURE>& args, Inputs<SIGNATURE> inputs)
     init_tensor_buffer_uniform_fp(inputs.weight, args.make_weight_descriptor(), -2.0f, 2.0f);
 }
 
-/// @brief `alloc_outputs()` specialization for forward convolution.
-///
-/// @tparam SIGNATURE Forward convolution signature.
-///
-/// @see alloc_outputs()
-template <auto SIGNATURE>
-    requires ValidConvSignature<SIGNATURE> && ConvDirectionIsForward<SIGNATURE> &&
-             ValidUniqueOutputs<SIGNATURE>
-UniqueOutputs<SIGNATURE> alloc_outputs(const Args<SIGNATURE>& args)
-{
-    return {
-        .output_buf = alloc_tensor_buffer(args.make_output_descriptor()),
-    };
-}
-
-/// @brief `validate()` specialization for forward convolution.
-///
-/// @tparam SIGNATURE Forward convolution signature.
-///
-/// @see validate()
-template <auto SIGNATURE>
-    requires ValidConvSignature<SIGNATURE> && ConvDirectionIsForward<SIGNATURE>
-ValidationReport
-validate(const Args<SIGNATURE>& args, Outputs<SIGNATURE> actual, Outputs<SIGNATURE> expected)
-{
-    ValidationReport report;
-    report.check("output", args.make_output_descriptor(), actual.output, expected.output);
-    return report;
-}
-
 } // namespace ck_tile::builder::test

experimental/builder/include/ck_tile/builder/testing/tensor_buffer.hpp

@@ -81,4 +81,15 @@ inline DeviceBuffer alloc_buffer(size_t size)
     return DeviceBuffer(d_buf);
 }
 
+/// @brief "Align" an offset to a multiple of a particular alignment.
+///
+/// Returns `addr` aligned to the next multiple of `alignment`.
+///
+/// @param addr The address to align.
+/// @param alignment The alignment.
+inline size_t align_fwd(size_t addr, size_t alignment)
+{
+    return addr % alignment == 0 ? addr : addr - addr % alignment + alignment;
+}
+
 } // namespace ck_tile::builder::test

experimental/builder/include/ck_tile/builder/testing/testing.hpp

@@ -5,6 +5,8 @@
 
 #include <concepts>
 
+#include "ck_tile/builder/testing/tensor_descriptor.hpp"
+#include "ck_tile/builder/testing/tensor_buffer.hpp"
 #include "ck_tile/builder/testing/validation.hpp"
 
 /// This file is the main header for the CK-Builder testing system. A high-level
@@ -132,8 +134,8 @@ struct Outputs;
 /// be created using `alloc_inputs()` and that an instance of the corresponding
 /// `Inputs` structure can be obtained using `.get()`.
 ///
-/// @note The easiest way to implement this type is to use the `DeviceBuffer`
-/// type to allocate individual device buffers for each input tensor.
+/// @note A default implementation is provided for this type if `Inputs`
+/// supports `TensorReflectable`.
 ///
 /// @tparam SIGNATURE The signature to specialize the structure for.
 ///
@@ -151,8 +153,8 @@ struct UniqueInputs;
 /// be created using `alloc_outputs()` and that an instance of the corresponding
 /// `Outputs` structure can be obtained using `.get()`.
 ///
-/// @note The easiest way to implement this type is to use the `DeviceBuffer`
-/// type to allocate individual device buffers for each output tensor.
+/// @note A default implementation is provided for this type if `Outputs`
+/// supports `TensorReflectable`.
 ///
 /// @tparam SIGNATURE The signature to specialize the structure for.
 ///
@@ -197,6 +199,12 @@ concept ValidUniqueOutputs = requires(UniqueOutputs<SIGNATURE>& inputs) {
 /// amount of memory required and then allocate it on the device, for example
 /// using `alloc_buffer` or `alloc_tensor_buffer`.
 ///
+/// @note This function is explicitly deleted to generate compile errors
+/// for missing implementations.
+///
+/// @note A default implementation is provided for this function if `Inputs`
+/// supports `TensorReflectable`.
+///
 /// @tparam SIGNATURE The signature to specialize the structure for.
 ///
 /// @param args The run-time arguments of the operation.
@@ -207,22 +215,22 @@ concept ValidUniqueOutputs = requires(UniqueOutputs<SIGNATURE>& inputs) {
 /// @see alloc_tensor_buffer()
 template <auto SIGNATURE>
     requires ValidUniqueInputs<SIGNATURE>
-UniqueInputs<SIGNATURE> alloc_inputs(const Args<SIGNATURE>& args);
+UniqueInputs<SIGNATURE> alloc_inputs(const Args<SIGNATURE>& args) = delete;
 
-/// @brief Allocate inputs corresponding to a signature.
+/// @brief Initialize inputs corresponding to a signature.
 ///
 /// The `init_inputs()` function is used to initialize pseudo-random data
 /// to the tensors specified in the Inputs structure. Implementors should
 /// fill each of the tensors in `inputs` with appropriate random data.
 ///
+/// @note This function is explicitly deleted to generate compile errors
+/// for missing implementations.
+///
 /// @tparam SIGNATURE the signature to specialize the structure for.
 ///
 /// @param args The run-time arguments of the operation.
 /// @param inputs The operation inputs to initialize with random data.
 ///
-/// @note This function is explicitly deleted to generate compile errors
-/// for missing implementations.
-///
 /// @see Inputs
 /// @see tensor_initialization
 template <auto SIGNATURE>
@@ -235,13 +243,16 @@ void init_inputs(const Args<SIGNATURE>& args, Inputs<SIGNATURE> inputs) = delete
 /// amount of memory required and then allocate it on the device, for example
 /// using `alloc_buffer` or `alloc_tensor_buffer`.
 ///
+/// @note This function is explicitly deleted to generate compile errors
+/// for missing implementations.
+///
+/// @note A default implementation is provided for this function if `Outputs`
+/// supports `TensorReflectable`.
+///
 /// @tparam SIGNATURE The signature to specialize the structure for.
 ///
 /// @param args The run-time arguments of the operation.
 ///
-/// @note This function is explicitly deleted to generate compile errors
-/// for missing implementations.
-///
 /// @see Outputs
 /// @see UniqueOutputs
 /// @see alloc_buffer()
@@ -262,15 +273,15 @@ UniqueInputs<SIGNATURE> alloc_outputs(const Args<SIGNATURE>& args) = delete;
 /// were incorrect, and where (a subset of) those elements are located within
 /// the tensor. See `ValidationReport` for more information about the report.
 ///
+/// @note This function is explicitly deleted to generate compile errors
+/// for missing implementations.
+///
 /// @tparam SIGNATURE The signature to specialize the structure for.
 ///
 /// @param args The run-time arguments of the operation.
 /// @param actual The actual results, the results of the operation to-be-tested.
 /// @param expected The expected results, the results of the reference implementation.
 ///
-/// @note This function is explicitly deleted to generate compile errors
-/// for missing implementations.
-///
 /// @see ValidationReport
 template <auto SIGNATURE>
 ValidationReport validate(const Args<SIGNATURE>& args,