ck-builder: test system initial prototype

Author: Snektron

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

@@ -4,6 +4,7 @@
 #pragma once
 
 #include "ck_tile/builder/conv_signature_concepts.hpp"
+#include "ck/library/utility/convolution_parameter.hpp"
 
 namespace ck_tile::builder::test {
 
@@ -12,8 +13,53 @@ struct FilterExtent
     ck::index_t width  = 1;
     ck::index_t height = 1;
     ck::index_t depth  = 1;
+
+    template <int SPATIAL_DIM>
+    std::vector<ck::index_t> to_vector() const
+    {
+        if constexpr(SPATIAL_DIM == 1)
+        {
+            return {std::initializer_list<ck::index_t>{this->width}};
+        }
+        else if constexpr(SPATIAL_DIM == 2)
+        {
+            return {{this->height, this->width}};
+        }
+        else if constexpr(SPATIAL_DIM == 3)
+        {
+            return {{this->depth, this->height, this->width}};
+        }
+    }
 };
 
+template <int SPATIAL_DIM>
+std::array<ck::index_t, SPATIAL_DIM + 3> to_ck_lengths(const std::array<ck::index_t, 3>& gnc,
+                                                       const FilterExtent& whd)
+{
+    std::array<ck::index_t, SPATIAL_DIM + 3> result = {0};
+    result[0]                                       = gnc[0];
+    result[1]                                       = gnc[1];
+    result[2]                                       = gnc[2];
+
+    if constexpr(SPATIAL_DIM == 1)
+    {
+        result[3] = whd.width;
+    }
+    else if constexpr(SPATIAL_DIM == 2)
+    {
+        result[3] = whd.height;
+        result[4] = whd.width;
+    }
+    else if constexpr(SPATIAL_DIM == 3)
+    {
+        result[3] = whd.depth;
+        result[4] = whd.height;
+        result[5] = whd.width;
+    }
+
+    return result;
+}
+
 struct TensorExtent
 {
     ck::index_t batch_size      = 1;  // N
@@ -39,6 +85,21 @@ struct ConvArgs
     FilterExtent filter_dilation;
     FilterExtent input_left_pad;
     FilterExtent input_right_pad;
+
+    ck::utils::conv::ConvParam to_conv_param() const
+    {
+        return ck::utils::conv::ConvParam(SPATIAL_DIM,
+                                          this->lengths.groups,
+                                          this->lengths.batch_size,
+                                          this->lengths.output_channels,
+                                          this->lengths.input_channels,
+                                          this->lengths.filter.to_vector<SPATIAL_DIM>(),
+                                          this->lengths.image.to_vector<SPATIAL_DIM>(),
+                                          this->filter_strides.to_vector<SPATIAL_DIM>(),
+                                          this->filter_dilation.to_vector<SPATIAL_DIM>(),
+                                          this->input_left_pad.to_vector<SPATIAL_DIM>(),
+                                          this->input_right_pad.to_vector<SPATIAL_DIM>(), );
+    }
 };
 
 } // namespace ck_tile::builder::test

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

@@ -4,8 +4,13 @@
 #pragma once
 
 #include <memory>
+#include <numeric>
 #include <hip/hip_runtime.h>
 #include "ck_tile/builder/conv_signature_concepts.hpp"
+#include "ck_tile/builder/testing/conv_args.hpp"
+#include "ck_tile/builder/testing/type_traits.hpp"
+#include "ck_tile/builder/conv_factory.hpp"
+#include "ck/library/utility/convolution_host_tensor_descriptor_helper.hpp"
 
 namespace ck_tile::builder::test {
 
@@ -20,17 +25,58 @@ struct DeviceMemoryDeleter
 
 using DeviceBuffer = std::unique_ptr<std::byte[], DeviceMemoryDeleter>;
 
+template <DataType DT>
+DeviceBuffer alloc_tensor(ck::HostTensorDescriptor descriptor)
+{
+    const auto total_elements = descriptor.GetElementSpaceSize();
+    const auto total_size     = total_elements * sizeof(typename DataTypeTraits<DT>::Type);
+
+    std::byte* d_buf  = nullptr;
+    const auto status = hipMalloc(&d_buf, total_size);
+    // TODO(Robin): How to check error without relying on google test?
+    // Ideally we get some sort of trace here, but thats not possible until c++23.
+    // For now just throw a runtime error.
+    if(status != hipSuccess)
+    {
+        throw std::runtime_error("failed to load hip memory");
+    }
+    return DeviceBuffer(d_buf);
+}
+
 template <auto SIGNATURE>
     requires ValidConvSignature<SIGNATURE>
 struct TensorMemoryManager
 {
     // Type aliases for tensor data types
     // For now, all tensors use the same data type from the signature
-    using InputDataType  = decltype(SIGNATURE.data_type);
-    using WeightDataType = decltype(SIGNATURE.data_type);
-    using OutputDataType = decltype(SIGNATURE.data_type);
+    using InputDataType  = DataTypeTraits<SIGNATURE.data_type>::Type;
+    using WeightDataType = DataTypeTraits<SIGNATURE.data_type>::Type;
+    using OutputDataType = DataTypeTraits<SIGNATURE.data_type>::Type;
+
+    using Layouts =
+        decltype(ck_tile::builder::factory_internal::GetTensorLayout<SIGNATURE.layout,
+                                                                     SIGNATURE.spatial_dim,
+                                                                     ConvDirection::FORWARD>());
+
+    TensorMemoryManager(const ConvArgs<SIGNATURE>& args)
+        : param(args.to_conv_param()),
+          input_descriptor(ck::utils::conv::make_input_host_tensor_descriptor_g_n_c_wis_packed<
+                           typename Layouts::ALayout>(this->param)),
+          weight_descriptor(ck::utils::conv::make_weight_host_tensor_descriptor_g_k_c_xs_packed<
+                            typename Layouts::BLayout>(this->param)),
+          output_descriptor(ck::utils::conv::make_output_host_tensor_descriptor_g_n_k_wos_packed<
+                            typename Layouts::ELayout>(this->param)),
+          input_buf(alloc_tensor<SIGNATURE.data_type>(this->input_descriptor)),
+          weight_buf(alloc_tensor<SIGNATURE.data_type>(this->weight_descriptor)),
+          output_buf(alloc_tensor<SIGNATURE.data_type>(this->output_descriptor))
+    {
+    }
+
+    ck::utils::conv::ConvParam param;
 
-    TensorMemoryManager() = default;
+    ck::HostTensorDescriptor input_descriptor;
+    ck::HostTensorDescriptor weight_descriptor;
+    ck::HostTensorDescriptor output_descriptor;
 
     // Device memory buffers
     DeviceBuffer input_buf  = nullptr;

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

@@ -0,0 +1,57 @@
+// Copyright (C) Advanced Micro Devices, Inc. All rights reserved.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <cstddef>
+#include "ck_tile/builder/types.hpp"
+#include "ck_tile/ops/common/tensor_layout.hpp"
+
+// TODO(Robin): Test to check that all DataType variants are covered?
+// TODO(Robin): Put this file somewhere else?
+
+namespace ck_tile::builder::test {
+
+/// This structure contains some useful traits for CK-Builder's DataType
+/// type. Its main usecase is to convert a CK-Builder DataType into an
+/// equivalent C++ type.
+template <DataType DT>
+struct DataTypeTraits;
+
+template <>
+struct DataTypeTraits<DataType::FP32>
+{
+    using Type = float;
+};
+
+template <>
+struct DataTypeTraits<DataType::FP16>
+{
+    using Type = ck::half_t;
+};
+
+template <>
+struct DataTypeTraits<DataType::BF16>
+{
+    using Type = ck::bhalf_t;
+};
+
+template <>
+struct DataTypeTraits<DataType::FP8>
+{
+    using Type = ck::f8_t;
+};
+
+template <>
+struct DataTypeTraits<DataType::I8>
+{
+    using Type = int8_t;
+};
+
+template <>
+struct DataTypeTraits<DataType::U8>
+{
+    using Type = uint8_t;
+};
+
+} // namespace ck_tile::builder::test

experimental/builder/test/CMakeLists.txt

@@ -23,7 +23,6 @@ add_ck_builder_test(test_ckb_conv_builder
     test_bwd_weight_instance_traits.cpp
     test_instance_traits_util.cpp
     test_memory_manager.cpp)
-
 add_ck_builder_test(test_ckb_inline_diff test_inline_diff.cpp)
 
 # Testing the virtual GetInstanceString methods requires kernel compilation.
@@ -51,6 +50,7 @@ add_ck_builder_test(test_ckb_build_fwd_instances
     conv/test_ckb_conv_fwd_3d_fp16.cpp
     conv/test_ckb_conv_fwd_3d_fp32.cpp
     )
+target_link_libraries(test_ckb_build_fwd_instances PRIVATE utility)
 
 function(add_ck_factory_test test_name)
     add_ck_builder_test(${test_name} ${ARGN})

experimental/builder/test/conv/test_ckb_conv_fwd_2d_dl_fp16.cpp

@@ -3,11 +3,16 @@
 
 #include "utils/ckb_conv_test_configs.hpp"
 #include "utils/ckb_conv_test_utils.hpp"
+#include "ck_tile/builder/testing/tensor_memory_manager.hpp"
+#include "ck_tile/builder/testing/conv_args.hpp"
 
 namespace {
 
 using namespace ck_tile::builder::test_utils;
 
+namespace ckb = ck_tile::builder;
+namespace ckt = ck_tile::builder::test;
+
 TEST(FwdConvInstances, Create_DeviceGroupedConvFwdDlMultipleD_NHWC_KYXC_NHWK_Instance_2D_FP16_GNHWC)
 {
     constexpr ConvSignature FwdConvSignature{.spatial_dim = 2,
@@ -54,4 +59,104 @@ TEST(FwdConvInstances,
         {"DeviceGroupedConvFwdDlMultipleD_NHWC_KYXC_NHWK", "256, 128, 128, 16", "Filter1x1Pad0"});
 }
 
+TEST(FwdConvInstances, Fp16_2D_DL_EndToEndBasic)
+{
+    constexpr ConvSignature signature{.spatial_dim           = 2,
+                                      .direction             = ConvDirection::FORWARD,
+                                      .layout                = GroupConvLayout2D::GNHWC_GKYXC_GNHWK,
+                                      .data_type             = DataType::FP16,
+                                      .elementwise_operation = ElementwiseOperation::PASS_THROUGH};
+
+    constexpr auto algorithm =
+        ConvAlgorithm_DeviceGroupedConvFwdDlMultipleD_NHWC_KYXC_NHWK{}
+            .with_thread_block(FwdThreadBlock_256_128x128x16)
+            .with_specializations(ConvFwdSpecialization::DEFAULT, GemmSpecialization::MNKPadding)
+            .with_dl_thread_config(DlThreadConfig_16x2x4x4x1)
+            .with_dl_thread_cluster(DlThreadCluster_8x2)
+            .with_dl_transfer(DlFwdTransfer);
+
+    auto args = ckt::ConvArgs<signature>{
+        .lengths =
+            {
+                .batch_size      = 16,
+                .groups          = 1,
+                .input_channels  = 32,
+                .output_channels = 16,
+                .image =
+                    {
+                        .width  = 56,
+                        .height = 56,
+                    },
+                .filter =
+                    {
+                        .width  = 3,
+                        .height = 3,
+                    },
+            },
+        .filter_strides  = {.width = 1, .height = 1},
+        .filter_dilation = {.width = 1, .height = 1},
+        .input_left_pad  = {.width = 0, .height = 0},
+        .input_right_pad = {.width = 0, .height = 0},
+    };
+
+    auto tmm = ckt::TensorMemoryManager<signature>(args);
+
+    auto conv = ConvBuilder<signature, algorithm>::Instance{};
+
+    auto invoker = conv.MakeInvoker();
+
+    const auto input_desc  = tmm.input_descriptor;
+    const auto weight_desc = tmm.weight_descriptor;
+    const auto output_desc = tmm.output_descriptor;
+
+    std::array<ck::index_t, 2 + 3> input_lengths;
+    std::array<ck::index_t, 2 + 3> input_strides;
+    std::array<ck::index_t, 2 + 3> weight_lengths;
+    std::array<ck::index_t, 2 + 3> weight_strides;
+    std::array<ck::index_t, 2 + 3> output_lengths;
+    std::array<ck::index_t, 2 + 3> output_strides;
+    std::array<ck::index_t, 2> conv_filter_strides;
+    std::array<ck::index_t, 2> conv_filter_dilations;
+    std::array<ck::index_t, 2> input_left_pads;
+    std::array<ck::index_t, 2> input_right_pads;
+
+    auto copy = [](auto& src, auto& dst) { std::copy(src.begin(), src.end(), dst.begin()); };
+
+    copy(input_desc.GetLengths(), input_lengths);
+    copy(input_desc.GetStrides(), input_strides);
+    copy(weight_desc.GetLengths(), weight_lengths);
+    copy(weight_desc.GetStrides(), weight_strides);
+    copy(output_desc.GetLengths(), output_lengths);
+    copy(output_desc.GetStrides(), output_strides);
+
+    copy(tmm.param.conv_filter_strides_, conv_filter_strides);
+    copy(tmm.param.conv_filter_dilations_, conv_filter_dilations);
+    copy(tmm.param.input_left_pads_, input_left_pads);
+    copy(tmm.param.input_right_pads_, input_right_pads);
+
+    auto argument = conv.MakeArgument(tmm.input_buf.get(),
+                                      tmm.weight_buf.get(),
+                                      {},
+                                      tmm.output_buf.get(),
+                                      input_lengths,
+                                      input_strides,
+                                      weight_lengths,
+                                      weight_strides,
+                                      {},
+                                      {},
+                                      output_lengths,
+                                      output_strides,
+                                      conv_filter_strides,
+                                      conv_filter_dilations,
+                                      input_left_pads,
+                                      input_right_pads,
+                                      ck::tensor_operation::element_wise::PassThrough{},
+                                      ck::tensor_operation::element_wise::PassThrough{},
+                                      ck::tensor_operation::element_wise::PassThrough{});
+
+    ASSERT_THAT(conv.IsSupportedArgument(argument), testing::IsTrue());
+
+    invoker.Run(argument, {});
+}
+
 } // namespace