Skip to content

[CK_BUILDER] Tensor input/output reflection #3536

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Snektron merged 1 commit into from
Jan 12, 2026
Merged

[CK_BUILDER] Tensor input/output reflection #3536

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
96 changes: 9 additions & 87 deletions experimental/builder/include/ck_tile/builder/testing/conv_fwd.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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
11 changes: 11 additions & 0 deletions experimental/builder/include/ck_tile/builder/testing/tensor_buffer.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
41 changes: 26 additions & 15 deletions experimental/builder/include/ck_tile/builder/testing/testing.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
///
Expand All @@ -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.
///
Expand Down Expand Up @@ -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.
Expand All @@ -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>
Expand All @@ -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()
Expand All @@ -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,
Expand Down
Loading