[CK_BUILDER] Add reflection for wmma and bwd weight instances to ck builder reflection

experimental/builder/include/ck_tile/builder/reflect/README.md

@@ -9,6 +9,7 @@ See the [main builder documentation](../README.md) for an overview.
 The reflection system works by extracting properties from a convolution kernel *type* and formatting them into a string. This is useful for debugging, performance tuning, and generating documentation.
 
 1. **Trait Extraction**: The `ConvTraits` template (in `conv_traits.hpp`) is specialized for each kernel instance. It extracts low-level details like tile sizes, data layouts, and pipeline versions from the kernel's type definition.
+This template is common for xld and wmma, fwd and backwards weight kernels. std::optional is used for parameters that are only used by some kernels
 
 2. **Description Generation**: The `describe<Instance>()` function (in `conv_description.hpp`) uses `ConvTraits` to populate a `ConvDescription` (`Description`) object.
 
@@ -48,6 +49,15 @@ The reflection system (`ckr::describe`) currently supports the following convolu
 - **Standard XDL Forward Convolution** (`DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle`)
 - **Large Tensor XDL Forward Convolution** (`DeviceGroupedConvFwdMultipleD_Xdl_CShuffle_Large_Tensor`)
 - **V3 XDL Forward Convolution** (`DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3`)
+- **V3 WMMA Forward Convolution** (`DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffleV3`)
+- **XDL Backward Weight Convolution** (`DeviceGroupedConvBwdWeight_Xdl_CShuffle`)
+- **V3 XDL Backward Weight Convolution** (`DeviceGroupedConvBwdWeight_Xdl_CShuffleV3`)
+- **XDL Multiple D Backward Weight Convolution** (`DeviceGroupedConvBwdWeightMultipleD_Xdl_CShuffle`)
+- **Two Stage XDL Backward Weight Convolution** (`DeviceGroupedConvBwdWeightTwoStage_Xdl_CShuffle`)
+- **V3 Two Stage XDL Backward Weight Convolution** (`DeviceGroupedConvBwdWeightTwoStage_Wmma_CShuffleV3`)
+- **Wmma Backward Weight Convolution** (`DeviceGroupedConvBwdWeight_Wmma_CShuffle`) 
+- **V3 Wmma Backward Weight Convolution** (`DeviceGroupedConvBwdWeight_Wmma_CShuffleV3`)
+- **V3 Wmma Multiple D Backward Weight Convolution** (`DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffleV3`)
 
 These variants all share similar template parameter structures and are compatible with the current `ConvTraits` implementation.
 
@@ -59,15 +69,6 @@ The following instance types are **not yet supported** by the reflection system:
   - Uses different internal structure with parameters like `K0PerBlock`, `K1`, `M1PerThread`, etc.
   - Missing standard members like `kKPerBlock`, `kMPerXDL`, `kAK1`
 
-- **WMMA Variants** (`DeviceGroupedConvFwdMultipleD_Wmma_CShuffle`)
-  - Uses WMMA-specific parameters like `MPerWmma`, `NPerWmma`, `MRepeat`, `NRepeat`
-  - Different tile transfer structure incompatible with current `ConvTraits`
-
-- **Backward Weight Convolution** (`DeviceGroupedConvBwdWeight_Xdl_CShuffle`)
-  - Uses different layout naming: `InLayout`, `WeiLayout`, `OutLayout` instead of `ALayout`, `BLayout`, `ELayout`
-  - Different specialization type: `ConvBackwardWeightSpecialization` vs `ConvForwardSpecialization`
-  - Missing several members expected by forward convolution traits
-
 ### Future Work
 
 To support these additional instance types, the reflection system would need:

experimental/builder/include/ck_tile/builder/reflect/conv_describe.hpp

@@ -41,16 +41,21 @@ conv::ConvDescription describe()
             .output_element_op = traits.output_element_op,
         },
         conv::GemmAlgorithmInfo{
-            .thread_block_size   = traits.thread_block_size,
-            .tile_dims           = traits.tile_dims,
-            .warp_gemm           = traits.warp_gemm,
-            .a_tile_transfer     = traits.a_tile_transfer,
-            .b_tile_transfer     = traits.b_tile_transfer,
-            .c_tile_transfer     = traits.c_tile_transfer,
-            .pipeline_version    = traits.pipeline_version,
-            .pipeline_scheduler  = traits.pipeline_scheduler,
-            .conv_specialization = traits.conv_specialization,
-            .padding             = traits.gemm_padding,
+            .thread_block_size         = traits.thread_block_size,
+            .tile_dims                 = traits.tile_dims,
+            .warp_gemm                 = traits.warp_gemm,
+            .a_tile_transfer           = traits.a_tile_transfer,
+            .b_tile_transfer           = traits.b_tile_transfer,
+            .c_tile_transfer           = traits.c_tile_transfer,
+            .pipeline_version          = traits.pipeline_version,
+            .pipeline_scheduler        = traits.pipeline_scheduler,
+            .conv_specialization       = traits.conv_specialization,
+            .padding                   = traits.gemm_padding,
+            .num_gemm_k_prefetch_stage = traits.num_gemm_k_prefetch_stage,
+            .max_transpose_transfer_src_scalar_per_vector =
+                traits.max_transpose_transfer_src_scalar_per_vector,
+            .max_transpose_dst_scalar_per_vector = traits.max_transpose_dst_scalar_per_vector,
+            .num_groups_to_merge                 = traits.num_groups_to_merge,
         },
         []<typename T = Instance>() { return reflect::instance_string<T>(); });
 }

experimental/builder/include/ck_tile/builder/reflect/conv_description.hpp

@@ -64,7 +64,11 @@ struct GemmAlgorithmInfo
     builder::PipelineVersion pipeline_version;
     builder::PipelineScheduler pipeline_scheduler;
     builder::ConvSpecialization conv_specialization;
-    builder::GemmPadding padding;
+    std::optional<builder::GemmPadding> padding;
+    std::optional<int> num_gemm_k_prefetch_stage;
+    std::optional<int> max_transpose_transfer_src_scalar_per_vector;
+    std::optional<int> max_transpose_dst_scalar_per_vector;
+    std::optional<int> num_groups_to_merge;
 };
 
 /// @brief Provides human-readable descriptions of convolution kernel instances
@@ -121,7 +125,11 @@ class ConvDescription : public Description
                     algorithm_.tile_dims.n,
                     "×",
                     algorithm_.tile_dims.k);
-        f.writeLine(2, "Gemm padding: ", algorithm_.padding);
+        if(algorithm_.padding)
+            f.writeLine(
+                2, "Gemm padding: ", algorithm_.padding.value_or(builder::GemmPadding::DEFAULT));
+        else
+            f.writeLine(2, "Struct does not contain optional padding argument");
         f.writeLine(2, "Convolution specialization: ", algorithm_.conv_specialization);
         // Pipeline section
         f.writeLine(2, "Pipeline version: ", algorithm_.pipeline_version);
@@ -231,9 +239,39 @@ class ConvDescription : public Description
                     algorithm_.c_tile_transfer.thread_cluster_dims[2],
                     "×",
                     algorithm_.c_tile_transfer.thread_cluster_dims[3]);
-        f.writeLast(4,
+        f.writeLine(4,
                     "Vector access (GMEM write) instruction size: ",
                     algorithm_.c_tile_transfer.scalar_per_vector);
+        if(algorithm_.num_gemm_k_prefetch_stage)
+            f.writeLine(2,
+                        "Max Transpose transfer scr scalar per vector: ",
+                        algorithm_.num_gemm_k_prefetch_stage.value_or(0));
+        else
+            f.writeLine(2,
+                        "Struct does not contain optional "
+                        "num_gemm_k_prefetch_stage parameter");
+
+        if(algorithm_.max_transpose_transfer_src_scalar_per_vector)
+            f.writeLine(2,
+                        "Max Transpose transfer scr scalar per vector: ",
+                        algorithm_.max_transpose_transfer_src_scalar_per_vector.value_or(0));
+        else
+            f.writeLine(2,
+                        "Struct does not contain optional "
+                        "max_transpose_transfer_src_scalar_per_vector parameter");
+        if(algorithm_.max_transpose_dst_scalar_per_vector)
+            f.writeLine(2,
+                        "Max Transpose dst scalar per vector: ",
+                        algorithm_.max_transpose_dst_scalar_per_vector.value_or(0));
+        else
+            f.writeLine(
+                2,
+                "Struct does not contain optional max_transpose_dst_scalar_per_vector parameter");
+        if(algorithm_.num_groups_to_merge)
+            f.writeLast(2, "Num groups to merge: ", algorithm_.num_groups_to_merge.value_or(0));
+        else
+            f.writeLast(2, "Struct does not contain optional num_groups_to_merge parameter");
+
         return f.getString();
     }
 

experimental/builder/include/ck_tile/builder/reflect/conv_traits.hpp

@@ -88,7 +88,7 @@ struct ConvTraits
     builder::ElementwiseOperation weight_element_op;
     builder::ElementwiseOperation output_element_op;
 
-    builder::GemmPadding gemm_padding;
+    std::optional<builder::GemmPadding> gemm_padding = std::nullopt;
     builder::ConvSpecialization conv_specialization;
 
     // --- Algorithm Information ---
@@ -102,8 +102,14 @@ struct ConvTraits
 
     OutputTileTransferInfo c_tile_transfer;
 
+    std::optional<int> num_gemm_k_prefetch_stage = std::nullopt;
+
     builder::PipelineVersion pipeline_version;
     builder::PipelineScheduler pipeline_scheduler;
+
+    std::optional<int> max_transpose_transfer_src_scalar_per_vector = std::nullopt;
+    std::optional<int> max_transpose_dst_scalar_per_vector          = std::nullopt;
+    std::optional<int> num_groups_to_merge                          = std::nullopt;
 };
 
 } // namespace ck_tile::reflect::conv

experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_multiple_d_wmma_cshuffle_v3.hpp

@@ -0,0 +1,46 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <concepts>
+
+#include "ck_tile/builder/reflect/conv_traits.hpp"
+#include "ck_tile/builder/reflect/conv_traits_helpers.hpp"
+#include "ck_tile/builder/reflect/instance_traits.hpp"
+#include "ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_multiple_d_wmma_cshuffle_v3.hpp"
+
+namespace ck_tile::reflect::conv {
+
+/// @brief Tag dispatch implementation for DeviceGroupedConvBwdWeightMultipleD_Wmma_CShuffle_Tag
+template <typename Instance>
+    requires HasInstanceTraits<Instance> &&
+             std::same_as<typename InstanceTraits<Instance>::device_kernel_tag,
+                          DeviceGroupedConvBwdWeight_multiple_d_Wmma_CShuffle_V3_Tag>
+constexpr ConvTraits instance_to_conv_traits()
+{
+    using InstTraits = InstanceTraits<Instance>;
+
+    return ConvTraits{
+        .spatial_dim         = InstTraits::kSpatialDim,
+        .direction           = conv_direction<Instance>(),
+        .layout              = bwd_wei_conv_layout<Instance>(),
+        .data_type           = conv_data_type<typename InstTraits::InDataType>(),
+        .input_element_op    = elementwise_op<typename InstTraits::InElementwiseOperation>(),
+        .weight_element_op   = elementwise_op<typename InstTraits::WeiElementwiseOperation>(),
+        .output_element_op   = elementwise_op<typename InstTraits::OutElementwiseOperation>(),
+        .conv_specialization = conv_spec<Instance>(),
+        .thread_block_size   = InstTraits::kBlockSize,
+        .tile_dims           = conv_traits_data_tile<InstTraits>(InstTraits::kKPerBlock),
+        .a_tile_transfer =
+            conv_traits_a_transfer_params<InstTraits>(InstTraits::kK1, InstTraits::kKPerBlock),
+        .b_tile_transfer =
+            conv_traits_b_transfer_params<InstTraits>(InstTraits::kK1, InstTraits::kKPerBlock),
+        .warp_gemm          = conv_traits_wmma_warp_gemm_params<InstTraits>(),
+        .c_tile_transfer    = conv_traits_wmma_c_tile_transfer<InstTraits>(),
+        .pipeline_version   = get_pipeline_version<InstTraits>(),
+        .pipeline_scheduler = get_pipeline_scheduler<InstTraits>(),
+    };
+}
+
+} // namespace ck_tile::reflect::conv

experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp

@@ -0,0 +1,53 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <concepts>
+
+#include "ck_tile/builder/reflect/conv_traits.hpp"
+#include "ck_tile/builder/reflect/conv_traits_helpers.hpp"
+#include "ck_tile/builder/reflect/instance_traits.hpp"
+#include "ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_multiple_d_xdl_cshuffle.hpp"
+
+namespace ck_tile::reflect::conv {
+
+/// @brief Tag dispatch implementation for DeviceGroupedConvBwdWeight_Xdl_CShuffle_Tag
+template <typename Instance>
+    requires HasInstanceTraits<Instance> &&
+             std::same_as<typename InstanceTraits<Instance>::device_kernel_tag,
+                          DeviceGroupedConvBwdWeight_multiple_d_Xdl_CShuffle_Tag>
+constexpr ConvTraits instance_to_conv_traits()
+{
+    using InstTraits = InstanceTraits<Instance>;
+
+    return ConvTraits{
+        .spatial_dim         = InstTraits::kSpatialDim,
+        .direction           = conv_direction<Instance>(),
+        .layout              = bwd_wei_conv_layout<Instance>(),
+        .data_type           = conv_data_type<typename InstTraits::InDataType>(),
+        .input_element_op    = elementwise_op<typename InstTraits::InElementwiseOperation>(),
+        .weight_element_op   = elementwise_op<typename InstTraits::WeiElementwiseOperation>(),
+        .output_element_op   = elementwise_op<typename InstTraits::OutElementwiseOperation>(),
+        .conv_specialization = conv_spec<Instance>(),
+        .thread_block_size   = InstTraits::kBlockSize,
+        .tile_dims           = conv_traits_data_tile<InstTraits>(InstTraits::kK0PerBlock),
+        .a_tile_transfer =
+            conv_traits_a_transfer_params<InstTraits>(InstTraits::kK1, InstTraits::kK0PerBlock),
+        .b_tile_transfer =
+            conv_traits_b_transfer_params<InstTraits>(InstTraits::kK1, InstTraits::kK0PerBlock),
+        .warp_gemm = conv_traits_xdl_warp_gemm_params<InstTraits>(),
+        .c_tile_transfer =
+            {.shuffle_params = {.m_gemms_per_shuffle = InstTraits::kCShuffleMXdlPerWavePerShuffle,
+                                .n_gemms_per_shuffle = InstTraits::kCShuffleNXdlPerWavePerShuffle},
+             .thread_cluster_dims = {InstTraits::kCThreadClusterLengths[0],
+                                     InstTraits::kCThreadClusterLengths[1],
+                                     InstTraits::kCThreadClusterLengths[2],
+                                     InstTraits::kCThreadClusterLengths[3]},
+             .scalar_per_vector   = InstTraits::kCBlockTransferScalarPerVector_NWaveNPerXdl},
+        .pipeline_version   = get_pipeline_version<InstTraits>(),
+        .pipeline_scheduler = get_pipeline_scheduler<InstTraits>(),
+    };
+}
+
+} // namespace ck_tile::reflect::conv

experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_two_stage_wmma_cshuffle_v3.hpp

@@ -0,0 +1,50 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <concepts>
+
+#include "ck_tile/builder/reflect/conv_traits.hpp"
+#include "ck_tile/builder/reflect/conv_traits_helpers.hpp"
+#include "ck_tile/builder/reflect/instance_traits.hpp"
+#include "ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_two_stage_wmma_cshuffle_v3.hpp"
+
+namespace ck_tile::reflect::conv {
+
+/// @brief Tag dispatch implementation for DeviceGroupedConvBwdWeight_wmma_CShuffle_Tag
+template <typename Instance>
+    requires HasInstanceTraits<Instance> &&
+             std::same_as<typename InstanceTraits<Instance>::device_kernel_tag,
+                          DeviceGroupedConvBwdWeight_two_stage_Wmma_CShuffle_Tag>
+constexpr ConvTraits instance_to_conv_traits()
+{
+    using InstTraits = InstanceTraits<Instance>;
+
+    return ConvTraits{
+        .spatial_dim         = InstTraits::kSpatialDim,
+        .direction           = conv_direction<Instance>(),
+        .layout              = bwd_wei_conv_layout<Instance>(),
+        .data_type           = conv_data_type<typename InstTraits::InDataType>(),
+        .input_element_op    = elementwise_op<typename InstTraits::InElementwiseOperation>(),
+        .weight_element_op   = elementwise_op<typename InstTraits::WeiElementwiseOperation>(),
+        .output_element_op   = elementwise_op<typename InstTraits::OutElementwiseOperation>(),
+        .conv_specialization = conv_spec<Instance>(),
+        .thread_block_size   = InstTraits::kBlockSize,
+        .tile_dims           = conv_traits_data_tile<InstTraits>(InstTraits::kKPerBlock),
+        .a_tile_transfer =
+            conv_traits_a_transfer_params<InstTraits>(InstTraits::kABK1, InstTraits::kKPerBlock),
+        .b_tile_transfer =
+            conv_traits_b_transfer_params<InstTraits>(InstTraits::kABK1, InstTraits::kKPerBlock),
+        .warp_gemm          = conv_traits_wmma_warp_gemm_params<InstTraits>(),
+        .c_tile_transfer    = conv_traits_wmma_c_tile_transfer<InstTraits>(),
+        .pipeline_version   = get_pipeline_version<InstTraits>(),
+        .pipeline_scheduler = get_pipeline_scheduler<InstTraits>(),
+        .max_transpose_transfer_src_scalar_per_vector =
+            InstTraits::kTransposeTransferSrcScalarPerVector,
+        .max_transpose_dst_scalar_per_vector = InstTraits::kTransposeTransferDstScalarPerVector,
+        .num_groups_to_merge                 = InstTraits::kNumGroupsToMerge,
+    };
+}
+
+} // namespace ck_tile::reflect::conv

experimental/builder/include/ck_tile/builder/reflect/conv_traits_device_grouped_conv_bwd_weight_two_stage_xdl_cshuffle.hpp

@@ -0,0 +1,57 @@
+// Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <concepts>
+
+#include "ck_tile/builder/reflect/conv_traits.hpp"
+#include "ck_tile/builder/reflect/conv_traits_helpers.hpp"
+#include "ck_tile/builder/reflect/instance_traits.hpp"
+#include "ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_two_stage_xdl_cshuffle.hpp"
+
+namespace ck_tile::reflect::conv {
+
+/// @brief Tag dispatch implementation for DeviceGroupedConvBwdTwoStage_Xdl_CShuffle_Tag
+template <typename Instance>
+    requires HasInstanceTraits<Instance> &&
+             std::same_as<typename InstanceTraits<Instance>::device_kernel_tag,
+                          DeviceGroupedConvBwdWeight_two_stage_Xdl_CShuffle_Tag>
+constexpr ConvTraits instance_to_conv_traits()
+{
+    using InstTraits = InstanceTraits<Instance>;
+
+    return ConvTraits{
+        .spatial_dim         = InstTraits::kSpatialDim,
+        .direction           = conv_direction<Instance>(),
+        .layout              = bwd_wei_conv_layout<Instance>(),
+        .data_type           = conv_data_type<typename InstTraits::InDataType>(),
+        .input_element_op    = elementwise_op<typename InstTraits::InElementwiseOperation>(),
+        .weight_element_op   = elementwise_op<typename InstTraits::WeiElementwiseOperation>(),
+        .output_element_op   = elementwise_op<typename InstTraits::OutElementwiseOperation>(),
+        .conv_specialization = conv_spec<Instance>(),
+        .thread_block_size   = InstTraits::kBlockSize,
+        .tile_dims           = conv_traits_data_tile<InstTraits>(InstTraits::kKPerBlock),
+        .a_tile_transfer =
+            conv_traits_a_transfer_params<InstTraits>(InstTraits::kK1, InstTraits::kKPerBlock),
+        .b_tile_transfer =
+            conv_traits_b_transfer_params<InstTraits>(InstTraits::kK1, InstTraits::kKPerBlock),
+        .warp_gemm = conv_traits_xdl_warp_gemm_params<InstTraits>(),
+        .c_tile_transfer =
+            {.shuffle_params = {.m_gemms_per_shuffle = InstTraits::kCShuffleMXdlPerWavePerShuffle,
+                                .n_gemms_per_shuffle = InstTraits::kCShuffleNXdlPerWavePerShuffle},
+             .thread_cluster_dims = {InstTraits::kCThreadClusterLengths[0],
+                                     InstTraits::kCThreadClusterLengths[1],
+                                     InstTraits::kCThreadClusterLengths[2],
+                                     InstTraits::kCThreadClusterLengths[3]},
+             .scalar_per_vector   = InstTraits::kCBlockTransferScalarPerVector_NWaveNPerXdl},
+        .pipeline_version   = get_pipeline_version<InstTraits>(),
+        .pipeline_scheduler = get_pipeline_scheduler<InstTraits>(),
+        .max_transpose_transfer_src_scalar_per_vector =
+            InstTraits::kTransposeTransferSrcScalarPerVector,
+        .max_transpose_dst_scalar_per_vector = InstTraits::kTransposeTransferDstScalarPerVector,
+        .num_groups_to_merge                 = InstTraits::kNumGroupsToMerge,
+    };
+}
+
+} // namespace ck_tile::reflect::conv