Merge branch 'develop' into miopen_downstream_all

composable_kernel/include/problem_transform/transform_backward_data_convolution_into_gemm_v4r1r2_nhwc_kyxc_nhwk.hpp

@@ -21,8 +21,8 @@ template <typename... Wei,
           typename ConvDilations,
           typename InLeftPads,
           typename InRightPads,
-          index_t IYTildaValue,
-          index_t IXTildaValue,
+          typename IYTilda,
+          typename IXTilda,
           index_t GemmK1Value>
 __host__ __device__ constexpr auto
 transform_backward_data_convolution_into_gemm_v4r1r2_nhwc_kyxc_nhwk(
@@ -33,8 +33,8 @@ transform_backward_data_convolution_into_gemm_v4r1r2_nhwc_kyxc_nhwk(
     const ConvDilations& conv_dilations,
     const InLeftPads& in_left_pads,
     const InRightPads& in_right_pads,
-    Number<IYTildaValue>,
-    Number<IXTildaValue>,
+    IYTilda i_ytilda,
+    IXTilda i_xtilda,
     Number<GemmK1Value>)
 {
     constexpr auto I0 = Number<0>{};
@@ -43,8 +43,6 @@ transform_backward_data_convolution_into_gemm_v4r1r2_nhwc_kyxc_nhwk(
     constexpr auto I3 = Number<3>{};
 
     constexpr auto GemmK1 = Number<GemmK1Value>{};
-    constexpr auto IYTilda = Number<IYTildaValue>{};
-    constexpr auto IXTilda = Number<IXTildaValue>{};
 
     const auto N = in_n_hi_wi_c_grid_desc.GetLength(I0);
     const auto C = in_n_hi_wi_c_grid_desc.GetLength(I3);
@@ -98,8 +96,8 @@ transform_backward_data_convolution_into_gemm_v4r1r2_nhwc_kyxc_nhwk(
     const auto WTildaSlice = IWTildaSliceEnd - IWTildaSliceBegin;
 
     // GemmK is different for each GEMM
-    const auto YDotSlice = math::integer_divide_ceil(Y - IYTilda, YTilda);
-    const auto XDotSlice = math::integer_divide_ceil(X - IXTilda, XTilda);
+    const auto YDotSlice = math::integer_divide_ceil(Y - i_ytilda, YTilda);
+    const auto XDotSlice = math::integer_divide_ceil(X - i_xtilda, XTilda);
 
     const auto K1 = GemmK1;
     const auto K0 = K / K1;
@@ -183,8 +181,8 @@ transform_backward_data_convolution_into_gemm_v4r1r2_nhwc_kyxc_nhwk(
                                     make_tuple(make_unmerge_transform(make_tuple(K0, K1)),
                                                make_slice_transform(YDot, I0, YDotSlice),
                                                make_slice_transform(XDot, I0, XDotSlice),
-                                               make_freeze_transform(IYTilda),
-                                               make_freeze_transform(IXTilda),
+                                               make_freeze_transform(i_ytilda),
+                                               make_freeze_transform(i_xtilda),
                                                make_pass_through_transform(C)),
                                     make_tuple(Sequence<0>{},
                                                Sequence<1>{},
@@ -241,9 +239,9 @@ transform_backward_data_convolution_into_gemm_v4r1r2_nhwc_kyxc_nhwk(
     const auto in_n_htildaslice_wtildaslice_c_grid_desc = transform_tensor_descriptor(
         in_n_ytilda_htilda_xtilda_wtilda_c_grid_desc,
         make_tuple(make_pass_through_transform(N),
-                   make_freeze_transform(IYTilda),
+                   make_freeze_transform(i_ytilda),
                    make_slice_transform(HTilda, IHTildaSliceBegin, HTildaSlice),
-                   make_freeze_transform(IXTilda),
+                   make_freeze_transform(i_xtilda),
                    make_slice_transform(WTilda, IWTildaSliceBegin, WTildaSlice),
                    make_pass_through_transform(C)),
         make_tuple(Sequence<0>{},
@@ -271,5 +269,84 @@ transform_backward_data_convolution_into_gemm_v4r1r2_nhwc_kyxc_nhwk(
                       in_gemmm_gemmn_grid_desc);
 }
 
+// A: out
+// B: wei
+// C: in
+// Number of GEMMs = 1
+// GemmM = N * Ho * Wo
+// GemmN = C
+// GemmK = K
+template <typename... Wei,
+          typename... In,
+          typename... Out,
+          typename ConvStrides,
+          index_t GemmK1Value>
+__host__ __device__ constexpr auto
+transform_backward_data_convolution_into_gemm_v4r1r2_nhwc_kyxc_nhwk_1x1(
+    const TensorDescriptor<Out...>& out_n_ho_wo_k_grid_desc,
+    const TensorDescriptor<Wei...>& /* wei_k_y_x_c_grid_desc */,
+    const TensorDescriptor<In...>& in_n_hi_wi_c_grid_desc,
+    const ConvStrides& conv_strides,
+    Number<GemmK1Value>)
+{
+    constexpr auto I0 = Number<0>{};
+    constexpr auto I1 = Number<1>{};
+    constexpr auto I2 = Number<2>{};
+    constexpr auto I3 = Number<3>{};
+
+    constexpr auto GemmK1 = Number<GemmK1Value>{};
+
+    const auto N = in_n_hi_wi_c_grid_desc.GetLength(I0);
+    const auto C = in_n_hi_wi_c_grid_desc.GetLength(I3);
+    const auto K = out_n_ho_wo_k_grid_desc.GetLength(I3);
+
+    const auto Ho = out_n_ho_wo_k_grid_desc.GetLength(I1);
+    const auto Wo = out_n_ho_wo_k_grid_desc.GetLength(I2);
+
+    const auto ConvStrideH = conv_strides[I0];
+    const auto ConvStrideW = conv_strides[I1];
+
+    const auto K1 = GemmK1;
+    const auto K0 = K / K1;
+
+    // A: output tensor
+    const auto out_gemmk0_gemmm_gemmk1_grid_desc =
+        transform_tensor_descriptor(make_naive_tensor_descriptor_packed(make_tuple(N * Ho * Wo, K)),
+                                    make_tuple(make_pass_through_transform(N * Ho * Wo),
+                                               make_unmerge_transform(make_tuple(K0, K1))),
+                                    make_tuple(Sequence<0>{}, Sequence<1>{}),
+                                    make_tuple(Sequence<1>{}, Sequence<0, 2>{}));
+
+    // B: weight tensor
+    const auto wei_gemmk0_gemmn_gemmk1_grid_desc = transform_tensor_descriptor(
+        make_naive_tensor_descriptor_packed(make_tuple(K, C)),
+        make_tuple(make_unmerge_transform(make_tuple(K0, K1)), make_pass_through_transform(C)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
+
+    // C: input tensor
+    const auto in_n_y_ho_x_wo_c_grid_desc = transform_tensor_descriptor(
+        in_n_hi_wi_c_grid_desc,
+        make_tuple(make_pass_through_transform(N),
+                   make_embed_transform(make_tuple(I1, Ho), make_tuple(I1, ConvStrideH)),
+                   make_embed_transform(make_tuple(I1, Wo), make_tuple(I1, ConvStrideW)),
+                   make_pass_through_transform(C)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+        make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5>{}));
+
+    const auto in_gemmm_gemmn_grid_desc = transform_tensor_descriptor(
+        in_n_y_ho_x_wo_c_grid_desc,
+        make_tuple(make_freeze_transform(I0),
+                   make_freeze_transform(I0),
+                   make_merge_transform(make_tuple(N, Ho, Wo)),
+                   make_pass_through_transform(C)),
+        make_tuple(Sequence<1>{}, Sequence<3>{}, Sequence<0, 2, 4>{}, Sequence<5>{}),
+        make_tuple(Sequence<>{}, Sequence<>{}, Sequence<0>{}, Sequence<1>{}));
+
+    return make_tuple(out_gemmk0_gemmm_gemmk1_grid_desc,
+                      wei_gemmk0_gemmn_gemmk1_grid_desc,
+                      in_gemmm_gemmn_grid_desc);
+}
+
 } // namespace ck
 #endif

composable_kernel/include/problem_transform/transform_backward_weight_convolution_into_gemm_v4r4r2_atomic_nchw_kcyx_nkhw.hpp

+#ifndef CK_TRANSFORM_BACKWARD_WEIGHT_CONVOLUTION_INTO_GEMM_V4R4R2_ATOMIC_NCHW_KCYX_NKHW_HPP
+#define CK_TRANSFORM_BACKWARD_WEIGHT_CONVOLUTION_INTO_GEMM_V4R4R2_ATOMIC_NCHW_KCYX_NKHW_HPP
+
+#include "common_header.hpp"
+#include "tensor_descriptor.hpp"
+#include "tensor_descriptor_helper.hpp"
+
+namespace ck {
+
+// GemmM = K
+// GemmK = N * Ho * Wo
+// GemmN = C * Y * X
+template <typename... Wei,
+          typename... In,
+          typename... Out,
+          typename ConvStrides,
+          typename ConvDilations,
+          typename InLeftPads,
+          typename InRightPads,
+          index_t GemmK1Value,
+          typename GemmKBatchType,
+          typename GemmKPadType>
+__host__ __device__ constexpr auto
+transform_backward_weight_convolution_into_gemm_v4r4r2_atomic_nchw_kcyx_nkhw_pad(
+    const TensorDescriptor<Wei...>& wei_k_c_y_x_grid_desc,
+    const TensorDescriptor<In...>& in_n_c_hi_wi_grid_desc,
+    const TensorDescriptor<Out...>& out_n_k_ho_wo_grid_desc,
+    const ConvStrides& conv_strides,
+    const ConvDilations& conv_dilations,
+    const InLeftPads& in_left_pads,
+    const InRightPads& in_right_pads,
+    Number<GemmK1Value>,
+    GemmKBatchType GemmKBatch,
+    GemmKPadType GemmKPad)
+{
+    constexpr auto I0 = Number<0>{};
+    constexpr auto I1 = Number<1>{};
+    constexpr auto I2 = Number<2>{};
+    constexpr auto I3 = Number<3>{};
+
+    constexpr auto GemmK1 = Number<GemmK1Value>{};
+
+    const auto N = in_n_c_hi_wi_grid_desc.GetLength(I0);
+    const auto C = in_n_c_hi_wi_grid_desc.GetLength(I1);
+    const auto K = out_n_k_ho_wo_grid_desc.GetLength(I1);
+
+    const auto Hi = in_n_c_hi_wi_grid_desc.GetLength(I2);
+    const auto Wi = in_n_c_hi_wi_grid_desc.GetLength(I3);
+
+    const auto Ho = out_n_k_ho_wo_grid_desc.GetLength(I2);
+    const auto Wo = out_n_k_ho_wo_grid_desc.GetLength(I3);
+
+    const auto Y = wei_k_c_y_x_grid_desc.GetLength(I2);
+    const auto X = wei_k_c_y_x_grid_desc.GetLength(I3);
+
+    const auto ConvStrideH = conv_strides[I0];
+    const auto ConvStrideW = conv_strides[I1];
+
+    const auto ConvDilationH = conv_dilations[I0];
+    const auto ConvDilationW = conv_dilations[I1];
+
+    const auto InLeftPadH = in_left_pads[I0];
+    const auto InLeftPadW = in_left_pads[I1];
+
+    const auto InRightPadH = in_right_pads[I0];
+    const auto InRightPadW = in_right_pads[I1];
+
+    const auto GemmM      = K;
+    const auto GemmN      = C * Y * X;
+    const auto GemmKTotal = N * Ho * Wo;
+    const index_t GemmK0  = GemmKPad / (GemmKBatch * GemmK1);
+
+    // A: output tensor
+    const auto out_gemmktotal_gemmm_grid_desc = transform_tensor_descriptor(
+        make_naive_tensor_descriptor_packed(make_tuple(N, K, Ho * Wo)),
+        make_tuple(make_pass_through_transform(K), make_merge_transform(make_tuple(N, Ho * Wo))),
+        make_tuple(Sequence<1>{}, Sequence<0, 2>{}),
+        make_tuple(Sequence<1>{}, Sequence<0>{}));
+
+    const auto out_gemmkpad_gemmm_grid_desc = transform_tensor_descriptor(
+        out_gemmktotal_gemmm_grid_desc,
+        make_tuple(make_right_pad_transform(GemmKTotal, GemmKPad - GemmKTotal),
+                   make_pass_through_transform(GemmM)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+    const auto out_gemmkbatch_gemmk0_gemmm_gemmk1_grid_desc = transform_tensor_descriptor(
+        out_gemmkpad_gemmm_grid_desc,
+        make_tuple(make_unmerge_transform(make_tuple(GemmKBatch, GemmK0, GemmK1)),
+                   make_pass_through_transform(GemmM)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<0, 1, 3>{}, Sequence<2>{}));
+
+    // B: input tensor
+    const auto in_n_c_hip_wip_grid_desc = transform_tensor_descriptor(
+        in_n_c_hi_wi_grid_desc,
+        make_tuple(make_pass_through_transform(N),
+                   make_pass_through_transform(C),
+                   make_pad_transform(Hi, InLeftPadH, InRightPadH),
+                   make_pad_transform(Wi, InLeftPadW, InRightPadW)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
+
+    const auto in_n_c_y_ho_x_wo_grid_desc = transform_tensor_descriptor(
+        in_n_c_hip_wip_grid_desc,
+        make_tuple(make_pass_through_transform(N),
+                   make_pass_through_transform(C),
+                   make_embed_transform(make_tuple(Y, Ho), make_tuple(ConvDilationH, ConvStrideH)),
+                   make_embed_transform(make_tuple(X, Wo), make_tuple(ConvDilationW, ConvStrideW))),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2, 3>{}, Sequence<4, 5>{}));
+
+    const auto in_gemmktotal_gemmn_grid_desc =
+        transform_tensor_descriptor(in_n_c_y_ho_x_wo_grid_desc,
+                                    make_tuple(make_merge_transform(make_tuple(C, Y, X)),
+                                               make_merge_transform(make_tuple(N, Ho, Wo))),
+                                    make_tuple(Sequence<1, 2, 4>{}, Sequence<0, 3, 5>{}),
+                                    make_tuple(Sequence<1>{}, Sequence<0>{}));
+
+    const auto in_gemmkpad_gemmn_grid_desc = transform_tensor_descriptor(
+        in_gemmktotal_gemmn_grid_desc,
+        make_tuple(make_right_pad_transform(GemmKTotal, GemmKPad - GemmKTotal),
+                   make_pass_through_transform(GemmN)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+    const auto in_gemmkbatch_gemmk0_gemmn_gemmk1_grid_desc = transform_tensor_descriptor(
+        in_gemmkpad_gemmn_grid_desc,
+        make_tuple(make_unmerge_transform(make_tuple(GemmKBatch, GemmK0, GemmK1)),
+                   make_pass_through_transform(GemmN)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<0, 1, 3>{}, Sequence<2>{}));
+
+    // C: weight tensor
+    const auto wei_gemmm_gemmn_grid_desc = transform_tensor_descriptor(
+        make_naive_tensor_descriptor_packed(make_tuple(K, C * Y * X)),
+        make_tuple(make_pass_through_transform(K), make_pass_through_transform(C * Y * X)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+    return make_tuple(out_gemmkbatch_gemmk0_gemmm_gemmk1_grid_desc,
+                      in_gemmkbatch_gemmk0_gemmn_gemmk1_grid_desc,
+                      wei_gemmm_gemmn_grid_desc);
+}
+
+} // namespace ck
+#endif

composable_kernel/include/problem_transform/transform_backward_weight_convolution_into_gemm_v4r4r2_nchw_kcyx_nkhw.hpp

+#ifndef CK_TRANSFORM_BACKWARD_WEIGHT_CONVOLUTION_INTO_GEMM_V4R4R2_NCHW_KCYX_NKHW_HPP
+#define CK_TRANSFORM_BACKWARD_WEIGHT_CONVOLUTION_INTO_GEMM_V4R4R2_NCHW_KCYX_NKHW_HPP
+
+#include "common_header.hpp"
+#include "tensor_descriptor.hpp"
+#include "tensor_descriptor_helper.hpp"
+
+namespace ck {
+
+// GemmM = K
+// GemmK = N * Ho * Wo
+// GemmN = C * Y * X
+template <typename... Wei,
+          typename... In,
+          typename... Out,
+          typename ConvStrides,
+          typename ConvDilations,
+          typename InLeftPads,
+          typename InRightPads,
+          index_t GemmK1Value>
+__host__ __device__ constexpr auto
+transform_backward_weight_convolution_into_gemm_v4r4r2_nchw_kcyx_nkhw_pad(
+    const TensorDescriptor<Wei...>& wei_k_c_y_x_grid_desc,
+    const TensorDescriptor<In...>& in_n_c_hi_wi_grid_desc,
+    const TensorDescriptor<Out...>& out_n_k_ho_wo_grid_desc,
+    const ConvStrides& conv_strides,
+    const ConvDilations& conv_dilations,
+    const InLeftPads& in_left_pads,
+    const InRightPads& in_right_pads,
+    Number<GemmK1Value>)
+{
+    constexpr auto I0 = Number<0>{};
+    constexpr auto I1 = Number<1>{};
+    constexpr auto I2 = Number<2>{};
+    constexpr auto I3 = Number<3>{};
+
+    constexpr auto GemmK1 = Number<GemmK1Value>{};
+
+    const auto N = in_n_c_hi_wi_grid_desc.GetLength(I0);
+    const auto C = in_n_c_hi_wi_grid_desc.GetLength(I1);
+    const auto K = out_n_k_ho_wo_grid_desc.GetLength(I1);
+
+    const auto Hi = in_n_c_hi_wi_grid_desc.GetLength(I2);
+    const auto Wi = in_n_c_hi_wi_grid_desc.GetLength(I3);
+
+    const auto Ho = out_n_k_ho_wo_grid_desc.GetLength(I2);
+    const auto Wo = out_n_k_ho_wo_grid_desc.GetLength(I3);
+
+    const auto Y = wei_k_c_y_x_grid_desc.GetLength(I2);
+    const auto X = wei_k_c_y_x_grid_desc.GetLength(I3);
+
+    const auto ConvStrideH = conv_strides[I0];
+    const auto ConvStrideW = conv_strides[I1];
+
+    const auto ConvDilationH = conv_dilations[I0];
+    const auto ConvDilationW = conv_dilations[I1];
+
+    const auto InLeftPadH = in_left_pads[I0];
+    const auto InLeftPadW = in_left_pads[I1];
+
+    const auto InRightPadH = in_right_pads[I0];
+    const auto InRightPadW = in_right_pads[I1];
+
+    const auto GemmM  = K;
+    const auto GemmN  = C * Y * X;
+    const auto GemmK  = N * Ho * Wo;
+    const auto GemmK0 = GemmK / GemmK1;
+
+    // weight tensor
+    const auto wei_gemmm_gemmn_grid_desc = transform_tensor_descriptor(
+        make_naive_tensor_descriptor_packed(make_tuple(K, C * Y * X)),
+        make_tuple(make_pass_through_transform(K), make_pass_through_transform(C * Y * X)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+    // input tensor
+    const auto in_n_c_hip_wip_grid_desc = transform_tensor_descriptor(
+        in_n_c_hi_wi_grid_desc,
+        make_tuple(make_pass_through_transform(N),
+                   make_pass_through_transform(C),
+                   make_pad_transform(Hi, InLeftPadH, InRightPadH),
+                   make_pad_transform(Wi, InLeftPadW, InRightPadW)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
+
+    const auto in_n_c_y_ho_x_wo_grid_desc = transform_tensor_descriptor(
+        in_n_c_hip_wip_grid_desc,
+        make_tuple(make_pass_through_transform(N),
+                   make_pass_through_transform(C),
+                   make_embed_transform(make_tuple(Y, Ho), make_tuple(ConvDilationH, ConvStrideH)),
+                   make_embed_transform(make_tuple(X, Wo), make_tuple(ConvDilationW, ConvStrideW))),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2, 3>{}, Sequence<4, 5>{}));
+
+    const auto in_gemmk_gemmn_grid_desc =
+        transform_tensor_descriptor(in_n_c_y_ho_x_wo_grid_desc,
+                                    make_tuple(make_merge_transform(make_tuple(C, Y, X)),
+                                               make_merge_transform(make_tuple(N, Ho, Wo))),
+                                    make_tuple(Sequence<1, 2, 4>{}, Sequence<0, 3, 5>{}),
+                                    make_tuple(Sequence<1>{}, Sequence<0>{}));
+
+    const auto in_gemmk0_gemmn_gemmk1_grid_desc =
+        transform_tensor_descriptor(in_gemmk_gemmn_grid_desc,
+                                    make_tuple(make_unmerge_transform(make_tuple(GemmK0, GemmK1)),
+                                               make_pass_through_transform(GemmN)),
+                                    make_tuple(Sequence<0>{}, Sequence<1>{}),
+                                    make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
+
+    // output tensor
+    const auto out_gemmk_gemmm_grid_desc = transform_tensor_descriptor(
+        make_naive_tensor_descriptor_packed(make_tuple(N, K, Ho * Wo)),
+        make_tuple(make_pass_through_transform(K), make_merge_transform(make_tuple(N, Ho * Wo))),
+        make_tuple(Sequence<1>{}, Sequence<0, 2>{}),
+        make_tuple(Sequence<1>{}, Sequence<0>{}));
+
+    const auto out_gemmk0_gemmm_gemmk1_grid_desc =
+        transform_tensor_descriptor(out_gemmk_gemmm_grid_desc,
+                                    make_tuple(make_unmerge_transform(make_tuple(GemmK0, GemmK1)),
+                                               make_pass_through_transform(GemmM)),
+                                    make_tuple(Sequence<0>{}, Sequence<1>{}),
+                                    make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
+
+    return make_tuple(out_gemmk0_gemmm_gemmk1_grid_desc,
+                      in_gemmk0_gemmn_gemmk1_grid_desc,
+                      wei_gemmm_gemmn_grid_desc);
+}
+
+} // namespace ck
+#endif

composable_kernel/include/problem_transform/transform_backward_weight_convolution_into_gemm_v4r4r4_atomic_nhwc_kyxc_nhwk.hpp

+#ifndef CK_TRANSFORM_BACKWARD_WEIGHT_CONVOLUTION_INTO_GEMM_V4R4R4_ATOMIC_NHWC_KYXC_NHWK_HPP
+#define CK_TRANSFORM_BACKWARD_WEIGHT_CONVOLUTION_INTO_GEMM_V4R4R4_ATOMIC_NHWC_KYXC_NHWK_HPP
+
+#include "common_header.hpp"
+#include "tensor_descriptor.hpp"
+#include "tensor_descriptor_helper.hpp"
+
+namespace ck {
+
+// A: in
+// B: wei
+// C: out
+// GemmM = N * Ho * Wo
+// GemmN = K
+// GemmK = Y * X * C
+template <typename... In,
+          typename... Wei,
+          typename... Out,
+          typename ConvStrides,
+          typename ConvDilations,
+          typename InLeftPads,
+          typename InRightPads,
+          index_t GemmK1Value,
+          typename GemmKBatchType,
+          typename GemmKPadType>
+__host__ __device__ constexpr auto
+transform_backward_weight_convolution_into_gemm_v4r4r4_atomic_nhwc_kyxc_nhwk_pad(
+    const TensorDescriptor<In...>& in_n_hi_wi_c_grid_desc,
+    const TensorDescriptor<Wei...>& wei_k_y_x_c_grid_desc,
+    const TensorDescriptor<Out...>& out_n_ho_wo_k_grid_desc,
+    const ConvStrides& conv_strides,
+    const ConvDilations& conv_dilations,
+    const InLeftPads& in_left_pads,
+    const InRightPads& in_right_pads,
+    Number<GemmK1Value>,
+    GemmKBatchType GemmKBatch,
+    GemmKPadType GemmKPad)
+{
+    constexpr auto I0 = Number<0>{};
+    constexpr auto I1 = Number<1>{};
+    constexpr auto I2 = Number<2>{};
+    constexpr auto I3 = Number<3>{};
+
+    constexpr auto GemmK1 = Number<GemmK1Value>{};
+
+    const auto N = in_n_hi_wi_c_grid_desc.GetLength(I0);
+    const auto C = in_n_hi_wi_c_grid_desc.GetLength(I3);
+    const auto K = out_n_ho_wo_k_grid_desc.GetLength(I3);
+
+    const auto Hi = in_n_hi_wi_c_grid_desc.GetLength(I1);
+    const auto Wi = in_n_hi_wi_c_grid_desc.GetLength(I2);
+
+    const auto Ho = out_n_ho_wo_k_grid_desc.GetLength(I1);
+    const auto Wo = out_n_ho_wo_k_grid_desc.GetLength(I2);
+
+    const auto Y = wei_k_y_x_c_grid_desc.GetLength(I1);
+    const auto X = wei_k_y_x_c_grid_desc.GetLength(I2);
+
+    const auto ConvStrideH = conv_strides[I0];
+    const auto ConvStrideW = conv_strides[I1];
+
+    const auto ConvDilationH = conv_dilations[I0];
+    const auto ConvDilationW = conv_dilations[I1];
+
+    const auto InLeftPadH = in_left_pads[I0];
+    const auto InLeftPadW = in_left_pads[I1];
+
+    const auto InRightPadH = in_right_pads[I0];
+    const auto InRightPadW = in_right_pads[I1];
+
+    const auto GemmM      = Y * X * C;
+    const auto GemmN      = K;
+    const auto GemmKTotal = N * Ho * Wo;
+    const index_t GemmK0  = GemmKPad / (GemmKBatch * GemmK1);
+
+    // A: input tensor
+    const auto in_n_hip_wip_c_grid_desc = transform_tensor_descriptor(
+        in_n_hi_wi_c_grid_desc,
+        make_tuple(make_pass_through_transform(N),
+                   make_pad_transform(Hi, InLeftPadH, InRightPadH),
+                   make_pad_transform(Wi, InLeftPadW, InRightPadW),
+                   make_pass_through_transform(C)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
+
+    const auto in_n_y_ho_x_wo_c_grid_desc = transform_tensor_descriptor(
+        in_n_hip_wip_c_grid_desc,
+        make_tuple(make_pass_through_transform(N),
+                   make_embed_transform(make_tuple(Y, Ho), make_tuple(ConvDilationH, ConvStrideH)),
+                   make_embed_transform(make_tuple(X, Wo), make_tuple(ConvDilationW, ConvStrideW)),
+                   make_pass_through_transform(C)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+        make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5>{}));
+
+    const auto in_gemmktotal_gemmm_grid_desc =
+        transform_tensor_descriptor(in_n_y_ho_x_wo_c_grid_desc,
+                                    make_tuple(make_merge_transform(make_tuple(Y, X, C)),
+                                               make_merge_transform(make_tuple(N, Ho, Wo))),
+                                    make_tuple(Sequence<1, 3, 5>{}, Sequence<0, 2, 4>{}),
+                                    make_tuple(Sequence<1>{}, Sequence<0>{}));
+
+    const auto in_gemmkpad_gemmm_grid_desc = transform_tensor_descriptor(
+        in_gemmktotal_gemmm_grid_desc,
+        make_tuple(make_right_pad_transform(GemmKTotal, GemmKPad - GemmKTotal),
+                   make_pass_through_transform(GemmM)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+    const auto in_gemmkbatch_gemmk0_gemmm_gemmk1_grid_desc = transform_tensor_descriptor(
+        in_gemmkpad_gemmm_grid_desc,
+        make_tuple(make_unmerge_transform(make_tuple(GemmKBatch, GemmK0, GemmK1)),
+                   make_pass_through_transform(GemmM)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<0, 1, 3>{}, Sequence<2>{}));
+
+    // B: output tensor
+    const auto out_gemmktotal_gemmn_grid_desc =
+        make_naive_tensor_descriptor_packed(make_tuple(N * Ho * Wo, K));
+
+    const auto out_gemmkpad_gemmn_grid_desc = transform_tensor_descriptor(
+        out_gemmktotal_gemmn_grid_desc,
+        make_tuple(make_right_pad_transform(GemmKTotal, GemmKPad - GemmKTotal),
+                   make_pass_through_transform(GemmN)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+    const auto out_gemmkbatch_gemmk0_gemmn_gemmk1_grid_desc = transform_tensor_descriptor(
+        out_gemmkpad_gemmn_grid_desc,
+        make_tuple(make_unmerge_transform(make_tuple(GemmKBatch, GemmK0, GemmK1)),
+                   make_pass_through_transform(GemmN)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<0, 1, 3>{}, Sequence<2>{}));
+
+    // C: weight tensor
+    const auto wei_gemmm_gemmn_grid_desc = transform_tensor_descriptor(
+        make_naive_tensor_descriptor_packed(make_tuple(K, Y * X * C)),
+        make_tuple(make_pass_through_transform(K), make_pass_through_transform(Y * X * C)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<1>{}, Sequence<0>{}));
+
+    return make_tuple(in_gemmkbatch_gemmk0_gemmm_gemmk1_grid_desc,
+                      out_gemmkbatch_gemmk0_gemmn_gemmk1_grid_desc,
+                      wei_gemmm_gemmn_grid_desc);
+}
+
+} // namespace ck
+#endif

composable_kernel/include/problem_transform/transform_backward_weight_convolution_into_gemm_v4r4r4_nhwc_kyxc_nhwk.hpp

+#ifndef CK_TRANSFORM_BACKWARD_WEIGHT_CONVOLUTION_INTO_GEMM_V4R4R4_NHWC_KYXC_NHWK_HPP
+#define CK_TRANSFORM_BACKWARD_WEIGHT_CONVOLUTION_INTO_GEMM_V4R4R4_NHWC_KYXC_NHWK_HPP
+
+#include "common_header.hpp"
+#include "tensor_descriptor.hpp"
+#include "tensor_descriptor_helper.hpp"
+
+namespace ck {
+
+// A: in
+// B: wei
+// C: out
+// GemmM = N * Ho * Wo
+// GemmN = K
+// GemmK = Y * X * C
+template <typename... In,
+          typename... Wei,
+          typename... Out,
+          typename ConvStrides,
+          typename ConvDilations,
+          typename InLeftPads,
+          typename InRightPads,
+          index_t GemmK1Value>
+__host__ __device__ constexpr auto
+transform_backward_weight_convolution_into_gemm_v4r4r4_nhwc_kyxc_nhwk_pad(
+    const TensorDescriptor<In...>& in_n_hi_wi_c_grid_desc,
+    const TensorDescriptor<Wei...>& wei_k_y_x_c_grid_desc,
+    const TensorDescriptor<Out...>& out_n_ho_wo_k_grid_desc,
+    const ConvStrides& conv_strides,
+    const ConvDilations& conv_dilations,
+    const InLeftPads& in_left_pads,
+    const InRightPads& in_right_pads,
+    Number<GemmK1Value>)
+{
+    constexpr auto I0 = Number<0>{};
+    constexpr auto I1 = Number<1>{};
+    constexpr auto I2 = Number<2>{};
+    constexpr auto I3 = Number<3>{};
+
+    constexpr auto GemmK1 = Number<GemmK1Value>{};
+
+    const auto N = in_n_hi_wi_c_grid_desc.GetLength(I0);
+    const auto C = in_n_hi_wi_c_grid_desc.GetLength(I3);
+    const auto K = out_n_ho_wo_k_grid_desc.GetLength(I3);
+
+    const auto Hi = in_n_hi_wi_c_grid_desc.GetLength(I1);
+    const auto Wi = in_n_hi_wi_c_grid_desc.GetLength(I2);
+
+    const auto Ho = out_n_ho_wo_k_grid_desc.GetLength(I1);
+    const auto Wo = out_n_ho_wo_k_grid_desc.GetLength(I2);
+
+    const auto Y = wei_k_y_x_c_grid_desc.GetLength(I1);
+    const auto X = wei_k_y_x_c_grid_desc.GetLength(I2);
+
+    const auto ConvStrideH = conv_strides[I0];
+    const auto ConvStrideW = conv_strides[I1];
+
+    const auto ConvDilationH = conv_dilations[I0];
+    const auto ConvDilationW = conv_dilations[I1];
+
+    const auto InLeftPadH = in_left_pads[I0];
+    const auto InLeftPadW = in_left_pads[I1];
+
+    const auto InRightPadH = in_right_pads[I0];
+    const auto InRightPadW = in_right_pads[I1];
+
+    const auto GemmM  = Y * X * C;
+    const auto GemmN  = K;
+    const auto GemmK  = N * Ho * Wo;
+    const auto GemmK0 = GemmK / GemmK1;
+
+    // A: input tensor
+    const auto in_n_hip_wip_c_grid_desc = transform_tensor_descriptor(
+        in_n_hi_wi_c_grid_desc,
+        make_tuple(make_pass_through_transform(N),
+                   make_pad_transform(Hi, InLeftPadH, InRightPadH),
+                   make_pad_transform(Wi, InLeftPadW, InRightPadW),
+                   make_pass_through_transform(C)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
+
+    const auto in_n_y_ho_x_wo_c_grid_desc = transform_tensor_descriptor(
+        in_n_hip_wip_c_grid_desc,
+        make_tuple(make_pass_through_transform(N),
+                   make_embed_transform(make_tuple(Y, Ho), make_tuple(ConvDilationH, ConvStrideH)),
+                   make_embed_transform(make_tuple(X, Wo), make_tuple(ConvDilationW, ConvStrideW)),
+                   make_pass_through_transform(C)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+        make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5>{}));
+
+    const auto in_gemmk_gemmm_grid_desc =
+        transform_tensor_descriptor(in_n_y_ho_x_wo_c_grid_desc,
+                                    make_tuple(make_merge_transform(make_tuple(Y, X, C)),
+                                               make_merge_transform(make_tuple(N, Ho, Wo))),
+                                    make_tuple(Sequence<1, 3, 5>{}, Sequence<0, 2, 4>{}),
+                                    make_tuple(Sequence<1>{}, Sequence<0>{}));
+
+    const auto in_gemmk0_gemmm_gemmk1_grid_desc =
+        transform_tensor_descriptor(in_gemmk_gemmm_grid_desc,
+                                    make_tuple(make_unmerge_transform(make_tuple(GemmK0, GemmK1)),
+                                               make_pass_through_transform(GemmM)),
+                                    make_tuple(Sequence<0>{}, Sequence<1>{}),
+                                    make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
+
+    // B: output tensor
+    const auto out_gemmk_gemmn_grid_desc = transform_tensor_descriptor(
+        make_naive_tensor_descriptor_packed(make_tuple(N * Ho * Wo, K)),
+        make_tuple(make_pass_through_transform(N * Ho * Wo), make_pass_through_transform(K)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+    const auto out_gemmk0_gemmn_gemmk1_grid_desc =
+        transform_tensor_descriptor(out_gemmk_gemmn_grid_desc,
+                                    make_tuple(make_unmerge_transform(make_tuple(GemmK0, GemmK1)),
+                                               make_pass_through_transform(GemmN)),
+                                    make_tuple(Sequence<0>{}, Sequence<1>{}),
+                                    make_tuple(Sequence<0, 2>{}, Sequence<1>{}));
+
+    // C: weight tensor
+    const auto wei_gemmm_gemmn_grid_desc = transform_tensor_descriptor(
+        make_naive_tensor_descriptor_packed(make_tuple(K, Y * X * C)),
+        make_tuple(make_pass_through_transform(K), make_pass_through_transform(Y * X * C)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<1>{}, Sequence<0>{}));
+
+    return make_tuple(in_gemmk0_gemmm_gemmk1_grid_desc,
+                      out_gemmk0_gemmn_gemmk1_grid_desc,
+                      wei_gemmm_gemmn_grid_desc);
+}
+
+} // namespace ck
+#endif

composable_kernel/include/problem_transform/transform_backward_weight_convolution_into_gemm_v4r4r5_nhwc_kyxc_nhwk.hpp

+#ifndef CK_TRANSFORM_BACKWARD_WEIGHT_CONVOLUTION_INTO_GEMM_V4R4R5_NHWC_KYXC_NHWK_HPP
+#define CK_TRANSFORM_BACKWARD_WEIGHT_CONVOLUTION_INTO_GEMM_V4R4R5_NHWC_KYXC_NHWK_HPP
+
+#include "common_header.hpp"
+#include "tensor_descriptor.hpp"
+#include "tensor_descriptor_helper.hpp"
+
+namespace ck {
+
+// A: out
+// B: in
+// C: wei
+// GemmM = K
+// GemmN = Y * X * C
+// GemmKTotal = N * Ho * Wo
+template <typename... In,
+          typename... Wei,
+          typename... Out,
+          typename ConvStrides,
+          typename ConvDilations,
+          typename InLeftPads,
+          typename InRightPads,
+          index_t GemmK1Value,
+          typename GemmKBatchType,
+          typename GemmKPadType>
+__host__ __device__ constexpr auto
+transform_backward_weight_convolution_into_gemm_v4r4r5_nhwc_kyxc_nhwk_pad(
+    const TensorDescriptor<In...>& in_n_hi_wi_c_grid_desc,
+    const TensorDescriptor<Wei...>& wei_k_y_x_c_grid_desc,
+    const TensorDescriptor<Out...>& out_n_ho_wo_k_grid_desc,
+    const ConvStrides& conv_strides,
+    const ConvDilations& conv_dilations,
+    const InLeftPads& in_left_pads,
+    const InRightPads& in_right_pads,
+    Number<GemmK1Value>,
+    GemmKBatchType GemmKBatch,
+    GemmKPadType GemmKPad)
+{
+    constexpr auto I0 = Number<0>{};
+    constexpr auto I1 = Number<1>{};
+    constexpr auto I2 = Number<2>{};
+    constexpr auto I3 = Number<3>{};
+
+    constexpr auto GemmK1 = Number<GemmK1Value>{};
+
+    const auto N = in_n_hi_wi_c_grid_desc.GetLength(I0);
+    const auto C = in_n_hi_wi_c_grid_desc.GetLength(I3);
+    const auto K = out_n_ho_wo_k_grid_desc.GetLength(I3);
+
+    const auto Hi = in_n_hi_wi_c_grid_desc.GetLength(I1);
+    const auto Wi = in_n_hi_wi_c_grid_desc.GetLength(I2);
+
+    const auto Ho = out_n_ho_wo_k_grid_desc.GetLength(I1);
+    const auto Wo = out_n_ho_wo_k_grid_desc.GetLength(I2);
+
+    const auto Y = wei_k_y_x_c_grid_desc.GetLength(I1);
+    const auto X = wei_k_y_x_c_grid_desc.GetLength(I2);
+
+    const auto ConvStrideH = conv_strides[I0];
+    const auto ConvStrideW = conv_strides[I1];
+
+    const auto ConvDilationH = conv_dilations[I0];
+    const auto ConvDilationW = conv_dilations[I1];
+
+    const auto InLeftPadH = in_left_pads[I0];
+    const auto InLeftPadW = in_left_pads[I1];
+
+    const auto InRightPadH = in_right_pads[I0];
+    const auto InRightPadW = in_right_pads[I1];
+
+    const auto GemmM      = K;
+    const auto GemmN      = Y * X * C;
+    const auto GemmKTotal = N * Ho * Wo;
+    const index_t GemmK0  = GemmKPad / (GemmKBatch * GemmK1);
+
+    // A: output tensor
+    const auto out_gemmktotal_gemmm_grid_desc =
+        make_naive_tensor_descriptor_packed(make_tuple(N * Ho * Wo, K));
+
+    const auto out_gemmkpad_gemmm_grid_desc = transform_tensor_descriptor(
+        out_gemmktotal_gemmm_grid_desc,
+        make_tuple(make_right_pad_transform(GemmKTotal, GemmKPad - GemmKTotal),
+                   make_pass_through_transform(GemmM)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+    const auto out_gemmkbatch_gemmk0_gemmm_gemmk1_grid_desc = transform_tensor_descriptor(
+        out_gemmkpad_gemmm_grid_desc,
+        make_tuple(make_unmerge_transform(make_tuple(GemmKBatch, GemmK0, GemmK1)),
+                   make_pass_through_transform(GemmM)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<0, 1, 3>{}, Sequence<2>{}));
+
+    // B: input tensor
+    const auto in_n_hip_wip_c_grid_desc = transform_tensor_descriptor(
+        in_n_hi_wi_c_grid_desc,
+        make_tuple(make_pass_through_transform(N),
+                   make_pad_transform(Hi, InLeftPadH, InRightPadH),
+                   make_pad_transform(Wi, InLeftPadW, InRightPadW),
+                   make_pass_through_transform(C)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
+
+    const auto in_n_y_ho_x_wo_c_grid_desc = transform_tensor_descriptor(
+        in_n_hip_wip_c_grid_desc,
+        make_tuple(make_pass_through_transform(N),
+                   make_embed_transform(make_tuple(Y, Ho), make_tuple(ConvDilationH, ConvStrideH)),
+                   make_embed_transform(make_tuple(X, Wo), make_tuple(ConvDilationW, ConvStrideW)),
+                   make_pass_through_transform(C)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+        make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5>{}));
+
+    const auto in_gemmktotal_gemmn_grid_desc =
+        transform_tensor_descriptor(in_n_y_ho_x_wo_c_grid_desc,
+                                    make_tuple(make_merge_transform(make_tuple(Y, X, C)),
+                                               make_merge_transform(make_tuple(N, Ho, Wo))),
+                                    make_tuple(Sequence<1, 3, 5>{}, Sequence<0, 2, 4>{}),
+                                    make_tuple(Sequence<1>{}, Sequence<0>{}));
+
+    const auto in_gemmkpad_gemmn_grid_desc = transform_tensor_descriptor(
+        in_gemmktotal_gemmn_grid_desc,
+        make_tuple(make_right_pad_transform(GemmKTotal, GemmKPad - GemmKTotal),
+                   make_pass_through_transform(GemmN)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+    const auto in_gemmkbatch_gemmk0_gemmn_gemmk1_grid_desc = transform_tensor_descriptor(
+        in_gemmkpad_gemmn_grid_desc,
+        make_tuple(make_unmerge_transform(make_tuple(GemmKBatch, GemmK0, GemmK1)),
+                   make_pass_through_transform(GemmN)),
+        make_tuple(Sequence<0>{}, Sequence<1>{}),
+        make_tuple(Sequence<0, 1, 3>{}, Sequence<2>{}));
+
+    // C: weight tensor
+    const auto wei_gemmm_gemmn_grid_desc =
+        make_naive_tensor_descriptor_packed(make_tuple(K, Y * X * C));
+
+    return make_tuple(out_gemmkbatch_gemmk0_gemmm_gemmk1_grid_desc,
+                      in_gemmkbatch_gemmk0_gemmn_gemmk1_grid_desc,
+                      wei_gemmm_gemmn_grid_desc);
+}
+
+} // namespace ck
+#endif

composable_kernel/include/tensor_description/multi_index_transform.hpp

@@ -1327,6 +1327,129 @@ struct Merge_v2r2_magic_division
     }
 };
 
+// Implementation of "Merge" transformation primitive that uses division and mod. It is supposed to
+// be used for low_lengths that are known at compile time and are power of 2, otherwise performance
+// will be very bad
+template <typename LowLengths>
+struct Merge_v3_division_mod
+{
+    static constexpr index_t NDimLow = LowLengths::Size();
+
+    using LowerIndex = MultiIndex<NDimLow>;
+    using UpperIndex = MultiIndex<1>;
+
+    using LowLengthsScan =
+        decltype(container_reverse_exclusive_scan(LowLengths{}, math::multiplies{}, Number<1>{}));
+
+    using UpLengths =
+        decltype(make_tuple(container_reduce(LowLengths{}, math::multiplies{}, Number<1>{})));
+
+    LowLengths low_lengths_;
+    LowLengthsScan low_lengths_scan_;
+    UpLengths up_lengths_;
+
+    __host__ __device__ constexpr Merge_v3_division_mod() = default;
+
+    __host__ __device__ constexpr Merge_v3_division_mod(const LowLengths& low_lengths)
+        : low_lengths_{low_lengths},
+          low_lengths_scan_{
+              container_reverse_exclusive_scan(low_lengths, math::multiplies{}, Number<1>{})},
+          up_lengths_{make_tuple(container_reduce(low_lengths, math::multiplies{}, Number<1>{}))}
+    {
+        static_assert(LowerIndex::Size() == NDimLow, "wrong!");
+    }
+
+    __host__ __device__ static constexpr index_t GetNumOfLowerDimension() { return NDimLow; }
+
+    __host__ __device__ static constexpr index_t GetNumOfUpperDimension() { return 1; }
+
+    __host__ __device__ constexpr const auto& GetUpperLengths() const { return up_lengths_; }
+
+    template <typename LowIdx, typename UpIdx>
+    __host__ __device__ constexpr void CalculateLowerIndex(LowIdx& idx_low,
+                                                           const UpIdx& idx_up) const
+    {
+        static_assert(LowIdx::Size() == NDimLow && UpIdx::Size() == 1,
+                      "wrong! inconsistent # of dimension");
+
+        index_t tmp = idx_up[Number<0>{}];
+
+        // division and mod
+        static_for<0, NDimLow - 1, 1>{}([&](auto i) {
+            idx_low(i) = tmp / this->low_lengths_scan_[i];
+            tmp %= this->low_lengths_scan_[i];
+        });
+
+        idx_low(Number<NDimLow - 1>{}) = tmp;
+    }
+
+    template <typename LowIdxDiff,
+              typename UpIdxDiff,
+              typename LowIdx,
+              typename UpIdx,
+              index_t Hack>
+    __host__ __device__ void UpdateLowerIndex(LowIdxDiff& idx_diff_low,
+                                              const UpIdxDiff&,
+                                              LowIdx& idx_low,
+                                              const UpIdx& idx_up_new,
+                                              Number<Hack>) const
+    {
+        static_assert(LowIdxDiff::Size() == NDimLow && UpIdxDiff::Size() == 1 &&
+                          LowIdx::Size() == NDimLow && UpIdx::Size() == 1,
+                      "wrong! inconsistent # of dimension");
+
+        constexpr auto I0   = Number<0>{};
+        constexpr auto INm1 = Number<NDimLow - 1>{};
+
+        index_t tmp = idx_up_new[I0];
+
+        static_for<0, NDimLow - 1, 1>{}([&](auto i) {
+            const index_t tmp2 = idx_low[i];
+            idx_low(i)         = tmp / this->low_lengths_scan_[i];
+            idx_diff_low(i)    = idx_low[i] - tmp2;
+            tmp %= this->low_lengths_scan_[i];
+        });
+
+        const index_t tmp2 = idx_low[INm1];
+        idx_low(INm1)      = tmp;
+        idx_diff_low(INm1) = idx_low[INm1] - tmp2;
+    }
+
+    __host__ __device__ static constexpr bool IsLinearTransform() { return false; }
+
+    __host__ __device__ static constexpr bool IsValidUpperIndexAlwaysMappedToValidLowerIndex()
+    {
+        return true;
+    }
+
+    __host__ __device__ static constexpr bool IsKnownAtCompileTime()
+    {
+        return is_known_at_compile_time<LowLengths>::value &&
+               is_known_at_compile_time<LowLengthsScan>::value &&
+               is_known_at_compile_time<UpLengths>::value;
+    }
+
+    template <typename UpIdx>
+    __host__ __device__ static constexpr bool
+    IsValidUpperIndexMappedToValidLowerIndex(const UpIdx& /* idx_up */)
+    {
+        return true;
+    }
+
+    __host__ __device__ void Print() const
+    {
+        printf("{");
+        printf("Merge_v3_direct_division_mod, ");
+        printf("low_lengths_ ");
+        print_multi_index(low_lengths_);
+        printf("low_lengths_scan_ ");
+        print_multi_index(low_lengths_scan_);
+        printf("up_lengths_ ");
+        print_multi_index(up_lengths_);
+        printf("}");
+    }
+};
+
 template <typename UpLengths, bool Use24BitIntegerCalculation>
 struct UnMerge
 {

composable_kernel/include/tensor_description/multi_index_transform_helper.hpp

@@ -31,7 +31,7 @@ __host__ __device__ constexpr auto make_left_pad_transform(
     return LeftPad<LowLength, LeftPadLength, SkipIsValidCheck>{low_length, left_pad};
 }
 
-template <typename LowLength, typename RightPadLength, bool SkipIsValidCheck>
+template <typename LowLength, typename RightPadLength, bool SkipIsValidCheck = false>
 __host__ __device__ constexpr auto make_right_pad_transform(
     const LowLength& low_length,
     const RightPadLength& right_pad,
@@ -52,22 +52,36 @@ __host__ __device__ constexpr auto make_embed_transform(const UpLengths& up_leng
 template <typename LowLengths>
 __host__ __device__ constexpr auto make_merge_transform(const LowLengths& low_lengths)
 {
-#if !CK_EXPERIMENTAL_MERGE_USE_MAGIC_DIVISION
-    return Merge_v1_carry_check<LowLengths>{low_lengths};
+#if CK_EXPERIMENTAL_MERGE_USE_MAGIC_DIVISION
+    return make_merge_transform_v2_magic_division(low_lengths);
 #else
+    return make_merge_transform_v1_carry_check(low_lengths);
+#endif
+}
+
+template <typename LowLengths>
+__host__ __device__ constexpr auto
+make_merge_transform_v1_carry_check(const LowLengths& low_lengths)
+{
+    return Merge_v1_carry_check<LowLengths>{low_lengths};
+}
+
+template <typename LowLengths>
+__host__ __device__ constexpr auto
+make_merge_transform_v2_magic_division(const LowLengths& low_lengths)
+{
 #if 1
     return Merge_v2_magic_division<LowLengths>{low_lengths};
 #else
     return Merge_v2r2_magic_division<LowLengths>{low_lengths};
 #endif
-#endif
 }
 
 template <typename LowLengths>
 __host__ __device__ constexpr auto
-make_merge_transform_v2_magic_division(const LowLengths& low_lengths)
+make_merge_transform_v3_division_mod(const LowLengths& low_lengths)
 {
-    return Merge_v2_magic_division<LowLengths>{low_lengths};
+    return Merge_v3_division_mod<LowLengths>{low_lengths};
 }
 
 template <typename UpLengths, bool Use24BitIntegerCalculation = false>

composable_kernel/include/tensor_description/tensor_adaptor.hpp

@@ -189,8 +189,7 @@ struct TensorAdaptor
         bool is_known = true;
 
         static_for<0, Transforms::Size(), 1>{}([&](auto i) {
-            is_known &=
-                remove_cv_t<remove_reference_t<decltype(Transforms{}[i])>>::IsKnownAtCompileTime();
+            is_known &= remove_cvref_t<decltype(Transforms{}[i])>::IsKnownAtCompileTime();
         });
 
         return is_known && is_known_at_compile_time<ElementSize>::value;

composable_kernel/include/tensor_description/tensor_descriptor.hpp

@@ -185,8 +185,7 @@ struct TensorDescriptor
         bool is_known = true;
 
         static_for<0, Transforms::Size(), 1>{}([&](auto i) {
-            is_known &=
-                remove_cv_t<remove_reference_t<decltype(Transforms{}[i])>>::IsKnownAtCompileTime();
+            is_known &= remove_cvref_t<decltype(Transforms{}[i])>::IsKnownAtCompileTime();
         });
 
         return is_known && is_known_at_compile_time<ElementSize>::value &&
@@ -587,11 +586,11 @@ __host__ __device__ constexpr bool coordinate_has_valid_offset(const TensorDesc&
 
 template <typename TensorDesc>
 using TensorCoordinate_t = decltype(make_tensor_coordinate(
-    TensorDesc{}, MultiIndex<remove_cv_t<remove_reference_t<TensorDesc>>::GetNumOfDimension()>{}));
+    TensorDesc{}, MultiIndex<remove_cvref_t<TensorDesc>::GetNumOfDimension()>{}));
 
 template <typename TensorDesc>
 using TensorCoordinateStep_t = decltype(make_tensor_coordinate_step(
-    TensorDesc{}, MultiIndex<remove_cv_t<remove_reference_t<TensorDesc>>::GetNumOfDimension()>{}));
+    TensorDesc{}, MultiIndex<remove_cvref_t<TensorDesc>::GetNumOfDimension()>{}));
 
 } // namespace ck
 #endif

composable_kernel/include/tensor_operation/blockwise_gemm_dlops_v3.hpp

@@ -110,12 +110,10 @@ struct BlockwiseGemmDlops_km_kn_m0m1n0n1_v3
                         const BThreadBuffer& b_thread_buf,
                         CThreadBuffer& c_thread_buf) const
     {
-        static_assert(is_same<remove_cv_t<remove_reference_t<typename ABlockBuffer::type>>,
-                              remove_cv_t<remove_reference_t<FloatA>>>::value &&
-                      is_same<remove_cv_t<remove_reference_t<typename BThreadBuffer::type>>,
-                              remove_cv_t<remove_reference_t<FloatB>>>::value &&
-                      is_same<remove_cv_t<remove_reference_t<typename CThreadBuffer::type>>,
-                              remove_cv_t<remove_reference_t<FloatC>>>::value &&
+        static_assert(
+            is_same<remove_cvref_t<typename ABlockBuffer::type>, remove_cvref_t<FloatA>>::value &&
+            is_same<remove_cvref_t<typename BThreadBuffer::type>, remove_cvref_t<FloatB>>::value &&
+            is_same<remove_cvref_t<typename CThreadBuffer::type>, remove_cvref_t<FloatC>>::value &&
             "wrong! inconsistent type");
 
         constexpr auto I0 = Number<0>{};

composable_kernel/include/tensor_operation/blockwise_gemm_xdlops.hpp

@@ -4,21 +4,22 @@
 #include "common_header.hpp"
 #include "threadwise_tensor_slice_transfer.hpp"
 #include "xdlops_gemm.hpp"
+#include "tensor_adaptor.hpp"
 
 namespace ck {
 
 template <index_t BlockSize,
           typename FloatAB,
-          class ABlockDesc,
-          class BBlockDesc,
-          index_t MPerWave,
-          index_t NPerWave,
+          typename FloatAcc,
+          typename AK0MK1BlockDesc,
+          typename BK0NK1BlockDesc,
+          index_t MPerXDL,
+          index_t NPerXDL,
+          index_t MRepeat,
+          index_t NRepeat,
           index_t K1>
-struct BlockwiseGemmXdlops_km_kn_m0m1m2n_v1
+struct BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1
 {
-
-    using CIndex = MultiIndex<2>;
-
     static constexpr auto I0 = Number<0>{};
     static constexpr auto I1 = Number<1>{};
     static constexpr auto I2 = Number<2>{};
@@ -26,329 +27,169 @@ struct BlockwiseGemmXdlops_km_kn_m0m1m2n_v1
 
     static constexpr index_t WaveSize = 64;
 
-    static constexpr index_t M0 = ABlockDesc{}.GetLength(I1);
-    static constexpr index_t M1 = ABlockDesc{}.GetLength(I2);
-
-    static constexpr index_t N0 = BBlockDesc{}.GetLength(I1);
-    static constexpr index_t N1 = BBlockDesc{}.GetLength(I2);
+    static constexpr index_t MPerBlock = AK0MK1BlockDesc{}.GetLength(I1);
+    static constexpr index_t NPerBlock = BK0NK1BlockDesc{}.GetLength(I1);
 
-    static constexpr auto xdlops_gemm = XdlopsGemm<FloatAB, MPerWave, NPerWave, K1>{};
+    static constexpr index_t K0 = BK0NK1BlockDesc{}.GetLength(I0);
 
-    static constexpr index_t MWaves = M1 / MPerWave;
-    static constexpr index_t NWaves = N1 / NPerWave;
+    static constexpr auto xdlops_gemm = XdlopsGemm<FloatAB, MPerXDL, NPerXDL, K1>{};
 
-    static constexpr index_t MRepeat = M0;
-    static constexpr index_t NRepeat = N0;
+    static constexpr index_t MWaves = MPerBlock / (MRepeat * MPerXDL);
+    static constexpr index_t NWaves = NPerBlock / (NRepeat * NPerXDL);
 
-    __device__ constexpr auto GetCLayout() const { return xdlops_gemm.GetCLayout(); }
+    StaticBufferV2<AddressSpaceEnum_t::Vgpr, vector_type<FloatAcc, 16>, MRepeat * NRepeat, true>
+        c_thread_buf_;
 
-    __device__ constexpr auto GetNumBlks() const { return xdlops_gemm.GetCLayout().GetNumBlks(); }
+    __host__ __device__ constexpr auto& GetCThreadBuffer() { return c_thread_buf_; }
 
-    __device__ constexpr auto GetBlkSize() const { return xdlops_gemm.GetCLayout().GetBlkSize(); }
-
-    __device__ static auto CalculateAThreadOriginDataIndex()
+    __device__ static auto GetWaveIdx()
     {
         const index_t thread_id = get_thread_local_1d_id();
-        const index_t waveId    = thread_id / WaveSize;
-        const index_t laneId    = thread_id % WaveSize;
-        const index_t waveId_m  = waveId / NWaves;
 
-        if constexpr(xdlops_gemm.IsKReduction)
-        {
-            const index_t m_offset = waveId_m * MPerWave + xdlops_gemm.GetBlkTd(laneId);
-            const index_t k_offset = xdlops_gemm.GetBlkId(laneId);
-            return make_tuple(k_offset, 0, m_offset, 0);
+        constexpr auto threadid_to_wave_idx_adaptor = make_single_stage_tensor_adaptor(
+            make_tuple(make_merge_transform(make_tuple(MWaves, NWaves, WaveSize))),
+            make_tuple(Sequence<0, 1, 2>{}),
+            make_tuple(Sequence<0>{}));
+
+        return threadid_to_wave_idx_adaptor.CalculateBottomIndex(make_multi_index(thread_id));
     }
-        else
+
+    __device__ static auto CalculateAThreadOriginDataIndex()
     {
-            const index_t m_offset = waveId_m * MPerWave + laneId;
-            const index_t k_offset = 0;
-            return make_tuple(k_offset, 0, m_offset, 0);
-        }
+        const auto wave_idx = GetWaveIdx();
+
+        const auto waveId_m = wave_idx[I0];
+
+        const auto xdlops_a_idx = xdlops_gemm.CalculateAThreadOriginDataIndex();
+
+        return make_tuple(xdlops_a_idx[I0], 0, waveId_m, xdlops_a_idx[I1], 0);
     }
 
     __device__ static auto CalculateBThreadOriginDataIndex()
     {
-        const index_t thread_id = get_thread_local_1d_id();
-        const index_t waveId    = thread_id / WaveSize;
-        const index_t laneId    = thread_id % WaveSize;
-        const index_t waveId_n  = waveId % NWaves;
+        const auto wave_idx = GetWaveIdx();
 
-        if constexpr(xdlops_gemm.IsKReduction)
-        {
-            const index_t n_offset = waveId_n * NPerWave + xdlops_gemm.GetBlkTd(laneId);
-            const index_t k_offset = xdlops_gemm.GetBlkId(laneId);
-            return make_tuple(k_offset, 0, n_offset, 0);
-        }
-        else
-        {
-            const index_t n_offset = waveId_n * NPerWave + laneId;
-            const index_t k_offset = 0;
-            return make_tuple(k_offset, 0, n_offset, 0);
-        }
+        const auto waveId_n = wave_idx[I1];
+
+        const auto xdlops_b_idx = xdlops_gemm.CalculateBThreadOriginDataIndex();
+
+        return make_tuple(xdlops_b_idx[I0], 0, waveId_n, xdlops_b_idx[I1], 0);
     }
 
     template <index_t m0, index_t n0, index_t xdlops_i, index_t blk_i>
-    __device__ static CIndex
+    __device__ static auto
         CalculateCThreadOriginDataIndex(Number<m0>, Number<n0>, Number<xdlops_i>, Number<blk_i>)
     {
+        const auto wave_idx = GetWaveIdx();
+
+        const auto waveId_m = wave_idx[I0];
+        const auto waveId_n = wave_idx[I1];
 
-        const index_t waveId = get_thread_local_1d_id() / WaveSize;
+        const auto blk_idx = xdlops_gemm.GetBeginOfThreadBlk(xdlops_i, blk_i);
 
-        const auto thread_mtx_on_blk = xdlops_gemm.GetBeginOfThreadBlk(xdlops_i, blk_i);
+        constexpr auto mrepeat_mwave_mperxdl_to_m_adaptor = make_single_stage_tensor_adaptor(
+            make_tuple(make_unmerge_transform(make_tuple(MRepeat, MWaves, MPerXDL))),
+            make_tuple(Sequence<0>{}),
+            make_tuple(Sequence<0, 1, 2>{}));
 
-        const index_t waveId_m = waveId / NWaves;
-        const index_t waveId_n = waveId % NWaves;
+        constexpr auto nrepeat_nwave_nperxdl_to_n_adaptor = make_single_stage_tensor_adaptor(
+            make_tuple(make_unmerge_transform(make_tuple(NRepeat, NWaves, NPerXDL))),
+            make_tuple(Sequence<0>{}),
+            make_tuple(Sequence<0, 1, 2>{}));
 
-        const index_t m_offset = m0 * M1 + waveId_m * MPerWave + thread_mtx_on_blk[I0];
-        const index_t n_offset = n0 * N1 + waveId_n * NPerWave + thread_mtx_on_blk[I1];
+        const index_t c_thread_m = mrepeat_mwave_mperxdl_to_m_adaptor.CalculateBottomIndex(
+            make_tuple(m0, waveId_m, blk_idx[I0]))[I0];
+        const index_t c_thread_n = nrepeat_nwave_nperxdl_to_n_adaptor.CalculateBottomIndex(
+            make_tuple(n0, waveId_n, blk_idx[I1]))[I0];
 
-        return CIndex{m_offset, n_offset};
+        return make_tuple(c_thread_m, c_thread_n);
     }
 
-    __device__ BlockwiseGemmXdlops_km_kn_m0m1m2n_v1()
-        : a_thread_copy_{CalculateAThreadOriginDataIndex()},
-          b_thread_copy_{CalculateBThreadOriginDataIndex()}
+    __host__ __device__ BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1()
     {
-        static_assert(ABlockDesc::IsKnownAtCompileTime() && BBlockDesc::IsKnownAtCompileTime(),
+        static_assert(AK0MK1BlockDesc::IsKnownAtCompileTime() &&
+                          BK0NK1BlockDesc::IsKnownAtCompileTime(),
                       "wrong! Desc should be known at compile-time");
 
-        static_assert(ABlockDesc{}.GetLength(I0) == BBlockDesc{}.GetLength(I0),
-                      "wrong! K dimension not consistent");
+        static_assert(AK0MK1BlockDesc{}.GetLength(I0) == BK0NK1BlockDesc{}.GetLength(I0),
+                      "wrong! K0 dimension not consistent");
 
-        static_assert(ABlockDesc{}.GetLength(I3) == BBlockDesc{}.GetLength(I3),
+        static_assert(AK0MK1BlockDesc{}.GetLength(I2) == BK0NK1BlockDesc{}.GetLength(I2),
                       "wrong! K1 dimension not consistent");
 
         static_assert(BlockSize == MWaves * NWaves * WaveSize,
                       "BlockSize != MWaves * NWaves * WaveSize\n");
 
-        static_assert(K1 == BBlockDesc{}.GetLength(I3), "K1 is wrong!");
-
-        constexpr index_t KPerBlock = ABlockDesc{}.GetLength(I0);
-
-        static_assert(KPerBlock % xdlops_gemm.KPerXdlops == 0, "KPerBlock is wrong!");
-
-        static_assert(K1 % xdlops_gemm.mfma_type.k_base == 0, "K1 is wrong!");
+        static_assert(MPerBlock % (MPerXDL * MRepeat) == 0 && NPerBlock % (NPerXDL * NRepeat) == 0,
+                      "wrong!");
     }
 
-    template <typename ABlockBuffer, typename BBlockBuffer, typename CThreadBuffer>
-    __device__ void Run(const ABlockBuffer& a_block_buf,
-                        const BBlockBuffer& b_block_buf,
-                        CThreadBuffer& c_thread_buf) const
+    __host__ __device__ static constexpr auto GetCM0N0M1N1M2M3M4N2ThreadDescriptor()
     {
-        auto a_thread_buf = make_static_buffer<AddressSpaceEnum_t::Vgpr, FloatAB>(
-            a_thread_desc_.GetElementSpaceSize());
-        auto b_thread_buf = make_static_buffer<AddressSpaceEnum_t::Vgpr, FloatAB>(
-            b_thread_desc_.GetElementSpaceSize());
-
-        constexpr index_t KPerBlock = ABlockDesc{}.GetLength(I0);
-
-        vector_type<FloatAB, a_thread_desc_.GetElementSpaceSize()> a_thread_vec;
-
-        vector_type<FloatAB, b_thread_desc_.GetElementSpaceSize()> b_thread_vec;
-
-        static_for<0, KPerBlock, xdlops_gemm.KPerXdlops>{}([&](auto k) {
-            // read A
-            a_thread_copy_.Run(ABlockDesc{},
-                               make_tuple(k, I0, I0, I0),
-                               a_block_buf,
-                               a_thread_desc_,
-                               make_tuple(I0, I0, I0, I0),
-                               a_thread_buf);
-
-            // read B
-            b_thread_copy_.Run(BBlockDesc{},
-                               make_tuple(k, I0, I0, I0),
-                               b_block_buf,
-                               b_thread_desc_,
-                               make_tuple(I0, I0, I0, I0),
-                               b_thread_buf);
+        constexpr auto c_m0_m1_m2_n_tblk_lens = xdlops_gemm.GetCM0M1M2NThreadBlkLengths();
 
-            using mfma_input_type =
-                typename vector_type<FloatAB, xdlops_gemm.mfma_type.k_base>::type;
-
-            static_for<0, a_thread_desc_.GetElementSpaceSize(), 1>{}([&](auto i) {
-                a_thread_vec.template AsType<FloatAB>()(Number<i>{}) = a_thread_buf[Number<i>{}];
-            });
+        constexpr auto M0 = c_m0_m1_m2_n_tblk_lens[I0];
+        constexpr auto M1 = c_m0_m1_m2_n_tblk_lens[I1];
+        constexpr auto M2 = c_m0_m1_m2_n_tblk_lens[I2];
+        constexpr auto N  = c_m0_m1_m2_n_tblk_lens[I3];
 
-            static_for<0, b_thread_desc_.GetElementSpaceSize(), 1>{}([&](auto i) {
-                b_thread_vec.template AsType<FloatAB>()(Number<i>{}) = b_thread_buf[Number<i>{}];
-            });
-
-            static_for<0, MRepeat, 1>{}([&](auto m0) {
-                static_for<0, NRepeat, 1>{}([&](auto n0) {
-                    xdlops_gemm.template Run<decltype(a_thread_desc_),
-                                             decltype(b_thread_desc_),
-                                             decltype(c_thread_desc_),
-                                             m0,
-                                             n0>(a_thread_vec.template AsType<mfma_input_type>(),
-                                                 b_thread_vec.template AsType<mfma_input_type>(),
-                                                 c_thread_buf);
-                });
-            });
-        });
+        return make_naive_tensor_descriptor_packed(make_tuple(I1, I1, I1, I1, M0, M1, M2, N));
     }
 
-    private:
-    // A[K, M]
-    static constexpr auto a_thread_desc_ =
-        make_naive_tensor_descriptor_packed(make_tuple(I1, Number<MRepeat>{}, I1, Number<K1>{}));
-
-    // B[K, N]
-    static constexpr auto b_thread_desc_ =
-        make_naive_tensor_descriptor_packed(make_tuple(I1, Number<NRepeat>{}, I1, Number<K1>{}));
-
-    static constexpr auto c_thread_desc_ =
-        make_naive_tensor_descriptor_packed(make_tuple(Number<MRepeat>{}, Number<NRepeat>{}));
-
-    using AThreadCopy = ThreadwiseTensorSliceTransfer_v4<FloatAB,
-                                                         FloatAB,
-                                                         ABlockDesc,
-                                                         decltype(a_thread_desc_),
-                                                         Sequence<1, MRepeat, 1, K1>,
-                                                         Sequence<0, 1, 2, 3>,
-                                                         3,
-                                                         K1,
-                                                         1>;
-
-    using BThreadCopy = ThreadwiseTensorSliceTransfer_v4<FloatAB,
-                                                         FloatAB,
-                                                         BBlockDesc,
-                                                         decltype(b_thread_desc_),
-                                                         Sequence<1, NRepeat, 1, K1>,
-                                                         Sequence<0, 1, 2, 3>,
-                                                         3,
-                                                         K1,
-                                                         1>;
-
-    AThreadCopy a_thread_copy_;
-    BThreadCopy b_thread_copy_;
-};
-
-template <index_t BlockSize,
-          typename FloatAB,
-          class ABlockDesc,
-          class BBlockDesc,
-          index_t MPerWave,
-          index_t NPerWave,
-          index_t K1>
-struct BlockwiseGemmXdlops_km_kn_m0m1m2n_v1_2x2pipeline
-{
-
-    using CIndex = MultiIndex<2>;
-
-    static constexpr auto I0 = Number<0>{};
-    static constexpr auto I1 = Number<1>{};
-    static constexpr auto I2 = Number<2>{};
-    static constexpr auto I3 = Number<3>{};
-
-    static constexpr auto xdlops_gemm = XdlopsGemm<float, MPerWave, NPerWave, K1>{};
-
-    static constexpr index_t WaveSize = 64;
-
-    static constexpr index_t M0 = ABlockDesc{}.GetLength(I1);
-    static constexpr index_t M1 = ABlockDesc{}.GetLength(I2);
-
-    static constexpr index_t N0 = BBlockDesc{}.GetLength(I1);
-    static constexpr index_t N1 = BBlockDesc{}.GetLength(I2);
-
-    static constexpr index_t MWaves = M1 / MPerWave;
-    static constexpr index_t NWaves = N1 / NPerWave;
-
-    static constexpr index_t MRepeat = M0;
-    static constexpr index_t NRepeat = N0;
-
-    __device__ constexpr auto GetCLayout() const { return xdlops_gemm.GetCLayout(); }
-
-    __device__ constexpr auto GetNumBlks() const { return xdlops_gemm.GetCLayout().GetNumBlks(); }
-
-    __device__ constexpr auto GetBlkSize() const { return xdlops_gemm.GetCLayout().GetBlkSize(); }
-
-    __device__ static auto CalculateAThreadOriginDataIndex()
+    __host__ __device__ static constexpr auto GetCM0N0M1N1M2M3M4N2BlockDescriptor()
     {
-        const index_t thread_id = get_thread_local_1d_id();
-        const index_t waveId    = thread_id / WaveSize;
-        const index_t laneId    = thread_id % WaveSize;
-        const index_t waveId_m  = waveId / NWaves;
-
-        if constexpr(xdlops_gemm.IsKReduction)
-        {
-            const index_t m_offset = waveId_m * MPerWave + xdlops_gemm.GetBlkTd(laneId);
-            const index_t k_offset = xdlops_gemm.GetBlkId(laneId);
-            return make_tuple(k_offset, 0, m_offset, 0);
-        }
-        else
-        {
-            const index_t m_offset = waveId_m * MPerWave + laneId;
-            const index_t k_offset = 0;
-            return make_tuple(k_offset, 0, m_offset, 0);
-        }
+        constexpr auto c_m0_n0_m1_n1_m2_n2_block_desc =
+            make_naive_tensor_descriptor_packed(make_tuple(Number<MRepeat>{},
+                                                           Number<NRepeat>{},
+                                                           Number<MWaves>{},
+                                                           Number<NWaves>{},
+                                                           Number<MPerXDL>{},
+                                                           Number<NPerXDL>{}));
+
+        return xdlops_gemm.MakeCM0N0M1N1M2M3M4N2Descriptor(c_m0_n0_m1_n1_m2_n2_block_desc);
     }
 
-    __device__ static auto CalculateBThreadOriginDataIndex()
-    {
-        const index_t thread_id = get_thread_local_1d_id();
-        const index_t waveId    = thread_id / WaveSize;
-        const index_t laneId    = thread_id % WaveSize;
-        const index_t waveId_n  = waveId % NWaves;
-
-        if constexpr(xdlops_gemm.IsKReduction)
-        {
-            const index_t n_offset = waveId_n * NPerWave + xdlops_gemm.GetBlkTd(laneId);
-            const index_t k_offset = xdlops_gemm.GetBlkId(laneId);
-            return make_tuple(k_offset, 0, n_offset, 0);
-        }
-        else
+    template <typename CMNGridDesc>
+    __host__ __device__ static constexpr auto
+    MakeCM0N0M1N1M2M3M4N2GridDescriptor(const CMNGridDesc& c_m_n_grid_desc)
     {
-            const index_t n_offset = waveId_n * NPerWave + laneId;
-            const index_t k_offset = 0;
-            return make_tuple(k_offset, 0, n_offset, 0);
-        }
+        const auto c_m0_n0_m1_n1_m2_n2_grid_desc = transform_tensor_descriptor(
+            c_m_n_grid_desc,
+            make_tuple(make_unmerge_transform(make_tuple(MRepeat, MWaves, MPerXDL)),
+                       make_unmerge_transform(make_tuple(NRepeat, NWaves, NPerXDL))),
+            make_tuple(Sequence<0>{}, Sequence<1>{}),
+            make_tuple(Sequence<0, 2, 4>{}, Sequence<1, 3, 5>{}));
+
+        return xdlops_gemm.MakeCM0N0M1N1M2M3M4N2Descriptor(c_m0_n0_m1_n1_m2_n2_grid_desc);
     }
 
-    template <index_t m0, index_t n0, index_t xdlops_i, index_t blk_i>
-    __device__ static CIndex
-        CalculateCThreadOriginDataIndex(Number<m0>, Number<n0>, Number<xdlops_i>, Number<blk_i>)
+    __host__ __device__ static constexpr auto MakeAK0M0M1M2K1BlockDescriptor()
     {
-
-        const index_t waveId = get_thread_local_1d_id() / WaveSize;
-
-        const auto thread_mtx_on_blk = xdlops_gemm.GetBeginOfThreadBlk(xdlops_i, blk_i);
-
-        const index_t waveId_m = waveId / NWaves;
-        const index_t waveId_n = waveId % NWaves;
-
-        const index_t m_offset = m0 * M1 + waveId_m * MPerWave + thread_mtx_on_blk[I0];
-        const index_t n_offset = n0 * N1 + waveId_n * NPerWave + thread_mtx_on_blk[I1];
-
-        return CIndex{m_offset, n_offset};
+        return transform_tensor_descriptor(
+            AK0MK1BlockDesc{},
+            make_tuple(make_pass_through_transform(Number<K0>{}),
+                       make_unmerge_transform(
+                           make_tuple(Number<MRepeat>{}, Number<MWaves>{}, Number<MPerXDL>{})),
+                       make_pass_through_transform(Number<K1>{})),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
+            make_tuple(Sequence<0>{}, Sequence<1, 2, 3>{}, Sequence<4>{}));
     }
 
-    __device__ BlockwiseGemmXdlops_km_kn_m0m1m2n_v1_2x2pipeline()
-        : a_thread_copy_{CalculateAThreadOriginDataIndex()},
-          b_thread_copy_{CalculateBThreadOriginDataIndex()}
+    __host__ __device__ static constexpr auto MakeBK0N0N1N2K1BlockDescriptor()
     {
-        static_assert(ABlockDesc::IsKnownAtCompileTime() && BBlockDesc::IsKnownAtCompileTime(),
-                      "wrong! Desc should be known at compile-time");
-
-        static_assert(ABlockDesc{}.GetLength(I0) == BBlockDesc{}.GetLength(I0),
-                      "wrong! K dimension not consistent");
-
-        static_assert(ABlockDesc{}.GetLength(I3) == BBlockDesc{}.GetLength(I3),
-                      "wrong! K1 dimension not consistent");
-
-        static_assert(BlockSize == MWaves * NWaves * WaveSize,
-                      "BlockSize != MWaves * NWaves * WaveSize\n");
-
-        static_assert(K1 == BBlockDesc{}.GetLength(I3), "K1 is wrong!");
-
-        constexpr index_t KPerBlock = ABlockDesc{}.GetLength(I0);
-
-        static_assert(KPerBlock % xdlops_gemm.KPerXdlops == 0, "KPerBlock is wrong!");
-
-        static_assert(K1 % xdlops_gemm.mfma_type.k_base == 0, "K1 is wrong!");
+        return transform_tensor_descriptor(
+            BK0NK1BlockDesc{},
+            make_tuple(make_pass_through_transform(Number<K0>{}),
+                       make_unmerge_transform(
+                           make_tuple(Number<NRepeat>{}, Number<NWaves>{}, Number<NPerXDL>{})),
+                       make_pass_through_transform(Number<K1>{})),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
+            make_tuple(Sequence<0>{}, Sequence<1, 2, 3>{}, Sequence<4>{}));
     }
 
+    static constexpr auto a_k0_m0_m1_m2_k1_block_desc = MakeAK0M0M1M2K1BlockDescriptor();
+    static constexpr auto b_k0_n0_n1_n2_k1_block_desc = MakeBK0N0N1N2K1BlockDescriptor();
+
     template <typename ABlockBuffer, typename BBlockBuffer, typename CThreadBuffer>
     __device__ void Run(const ABlockBuffer& a_block_buf,
                         const BBlockBuffer& b_block_buf,
@@ -359,165 +200,82 @@ struct BlockwiseGemmXdlops_km_kn_m0m1m2n_v1_2x2pipeline
         auto b_thread_buf = make_static_buffer<AddressSpaceEnum_t::Vgpr, FloatAB>(
             b_thread_desc_.GetElementSpaceSize());
 
-        constexpr index_t KPerBlock = ABlockDesc{}.GetLength(I0);
-
-        // read A_sub_0
-        a_thread_copy_.Run(ABlockDesc{},
-                           make_tuple(I0, I0, I0, I0),
-                           a_block_buf,
-                           a_thread_desc_,
-                           make_tuple(I0, I0, I0, I0),
-                           a_thread_buf);
-
-        // read B_sub_0
-        b_thread_copy_.Run(BBlockDesc{},
-                           make_tuple(I0, I0, I0, I0),
-                           b_block_buf,
-                           b_thread_desc_,
-                           make_tuple(I0, I0, I0, I0),
-                           b_thread_buf);
-
-        // read B_sub_1
-        b_thread_copy_.Run(BBlockDesc{},
-                           make_tuple(I0, I1, I0, I0),
-                           b_block_buf,
-                           b_thread_desc_,
-                           make_tuple(I0, I1, I0, I0),
-                           b_thread_buf);
-
-        // read A_sub_1
-        a_thread_copy_.Run(ABlockDesc{},
-                           make_tuple(I0, I1, I0, I0),
-                           a_block_buf,
-                           a_thread_desc_,
-                           make_tuple(I0, I1, I0, I0),
-                           a_thread_buf);
-
-        // C_sub_00 += transpose(A_sub_0) * B_sub_0
-        xdlops_gemm.template Run<decltype(a_thread_desc_),
-                                 decltype(b_thread_desc_),
-                                 decltype(c_thread_desc_),
-                                 0,
-                                 0>(a_thread_buf, b_thread_buf, c_thread_buf);
-
-        // C_sub_01 += transpose(A_sub_0) * B_sub_1
-        xdlops_gemm.template Run<decltype(a_thread_desc_),
-                                 decltype(b_thread_desc_),
-                                 decltype(c_thread_desc_),
-                                 0,
-                                 1>(a_thread_buf, b_thread_buf, c_thread_buf);
-
-        static_for<xdlops_gemm.KPerXdlops, KPerBlock, xdlops_gemm.KPerXdlops>{}([&](auto k) {
-            // read A_sub_0
-            a_thread_copy_.Run(ABlockDesc{},
-                               make_tuple(k, I0, I0, I0),
+        static_for<0, MRepeat, 1>{}([&](auto m0) {
+            // read A
+            a_thread_copy_.Run(a_k0_m0_m1_m2_k1_block_desc,
+                               make_tuple(I0, m0, I0, I0, I0),
                                a_block_buf,
                                a_thread_desc_,
-                               make_tuple(I0, I0, I0, I0),
+                               make_tuple(I0, I0, I0, I0, I0),
                                a_thread_buf);
 
-            // C_sub_10 += transpose(A_sub_1) * B_sub_0
-            xdlops_gemm.template Run<decltype(a_thread_desc_),
-                                     decltype(b_thread_desc_),
-                                     decltype(c_thread_desc_),
-                                     1,
-                                     0>(a_thread_buf, b_thread_buf, c_thread_buf);
-
-            // read B_sub_0
-            b_thread_copy_.Run(BBlockDesc{},
-                               make_tuple(k, I0, I0, I0),
+            static_for<0, NRepeat, 1>{}([&](auto n0) {
+                // read B
+                b_thread_copy_.Run(b_k0_n0_n1_n2_k1_block_desc,
+                                   make_tuple(I0, n0, I0, I0, I0),
                                    b_block_buf,
                                    b_thread_desc_,
-                               make_tuple(I0, I0, I0, I0),
+                                   make_tuple(I0, I0, I0, I0, I0),
                                    b_thread_buf);
 
-            // C_sub_11 += transpose(A_sub_1) * B_sub_1
-            xdlops_gemm.template Run<decltype(a_thread_desc_),
-                                     decltype(b_thread_desc_),
-                                     decltype(c_thread_desc_),
-                                     1,
-                                     1>(a_thread_buf, b_thread_buf, c_thread_buf);
+                static_for<0, K0, xdlops_gemm.K0PerXdlops>{}([&](auto k0) {
+                    vector_type<FloatAB, K1> a_thread_vec;
+                    vector_type<FloatAB, K1> b_thread_vec;
 
-            // read B_sub_1
-            b_thread_copy_.Run(BBlockDesc{},
-                               make_tuple(k, I1, I0, I0),
-                               b_block_buf,
-                               b_thread_desc_,
-                               make_tuple(I0, I1, I0, I0),
-                               b_thread_buf);
+                    static_for<0, K1, 1>{}([&](auto i) {
+                        a_thread_vec.template AsType<FloatAB>()(i) = a_thread_buf
+                            [Number<a_thread_desc_.CalculateOffset(make_tuple(k0, 0, 0, 0, i))>{}];
+                        b_thread_vec.template AsType<FloatAB>()(i) = b_thread_buf
+                            [Number<b_thread_desc_.CalculateOffset(make_tuple(k0, 0, 0, 0, i))>{}];
+                    });
 
-            // read A_sub_1
-            a_thread_copy_.Run(ABlockDesc{},
-                               make_tuple(k, I1, I0, I0),
-                               a_block_buf,
-                               a_thread_desc_,
-                               make_tuple(I0, I1, I0, I0),
-                               a_thread_buf);
+                    using mfma_input_type =
+                        typename vector_type<FloatAB, xdlops_gemm.K1PerXdlops>::type;
 
-            // C_sub_00 += transpose(A_sub_0) * B_sub_0
-            xdlops_gemm.template Run<decltype(a_thread_desc_),
-                                     decltype(b_thread_desc_),
-                                     decltype(c_thread_desc_),
-                                     0,
-                                     0>(a_thread_buf, b_thread_buf, c_thread_buf);
+                    constexpr index_t c_offset = c_thread_desc_.CalculateOffset(make_tuple(m0, n0));
 
-            // C_sub_01 += transpose(A_sub_0) * B_sub_1
-            xdlops_gemm.template Run<decltype(a_thread_desc_),
-                                     decltype(b_thread_desc_),
-                                     decltype(c_thread_desc_),
-                                     0,
-                                     1>(a_thread_buf, b_thread_buf, c_thread_buf);
+                    xdlops_gemm.template Run(a_thread_vec.template AsType<mfma_input_type>(),
+                                             b_thread_vec.template AsType<mfma_input_type>(),
+                                             c_thread_buf.GetVector(Number<c_offset>{}));
+                });
+            });
         });
-
-        // C_sub_10 += transpose(A_sub_1) * B_sub_0
-        xdlops_gemm.template Run<decltype(a_thread_desc_),
-                                 decltype(b_thread_desc_),
-                                 decltype(c_thread_desc_),
-                                 1,
-                                 0>(a_thread_buf, b_thread_buf, c_thread_buf);
-
-        // C_sub_11 += transpose(A_sub_1) * B_sub_1
-        xdlops_gemm.template Run<decltype(a_thread_desc_),
-                                 decltype(b_thread_desc_),
-                                 decltype(c_thread_desc_),
-                                 1,
-                                 1>(a_thread_buf, b_thread_buf, c_thread_buf);
     }
 
     private:
     // A[K, M]
     static constexpr auto a_thread_desc_ =
-        make_naive_tensor_descriptor_packed(make_tuple(I1, Number<MRepeat>{}, I1, Number<K1>{}));
+        make_naive_tensor_descriptor_packed(make_tuple(Number<K0>{}, I1, I1, I1, Number<K1>{}));
 
     // B[K, N]
     static constexpr auto b_thread_desc_ =
-        make_naive_tensor_descriptor_packed(make_tuple(I1, Number<NRepeat>{}, I1, Number<K1>{}));
+        make_naive_tensor_descriptor_packed(make_tuple(Number<K0>{}, I1, I1, I1, Number<K1>{}));
 
     static constexpr auto c_thread_desc_ =
         make_naive_tensor_descriptor_packed(make_tuple(Number<MRepeat>{}, Number<NRepeat>{}));
 
     using AThreadCopy = ThreadwiseTensorSliceTransfer_v4<FloatAB,
                                                          FloatAB,
-                                                         ABlockDesc,
+                                                         decltype(a_k0_m0_m1_m2_k1_block_desc),
                                                          decltype(a_thread_desc_),
-                                                         Sequence<1, 1, 1, K1>,
-                                                         Sequence<0, 1, 2, 3>,
-                                                         3,
-                                                         1, // K1,
-                                                         1>;
+                                                         Sequence<K0, 1, 1, 1, K1>,
+                                                         Sequence<0, 1, 2, 3, 4>,
+                                                         4,
+                                                         K1,
+                                                         K1>;
 
     using BThreadCopy = ThreadwiseTensorSliceTransfer_v4<FloatAB,
                                                          FloatAB,
-                                                         BBlockDesc,
+                                                         decltype(b_k0_n0_n1_n2_k1_block_desc),
                                                          decltype(b_thread_desc_),
-                                                         Sequence<1, 1, 1, K1>,
-                                                         Sequence<0, 1, 2, 3>,
-                                                         3,
-                                                         1, // K1,
-                                                         1>;
-
-    AThreadCopy a_thread_copy_;
-    BThreadCopy b_thread_copy_;
+                                                         Sequence<K0, 1, 1, 1, K1>,
+                                                         Sequence<0, 1, 2, 3, 4>,
+                                                         4,
+                                                         K1,
+                                                         K1>;
+
+    AThreadCopy a_thread_copy_{CalculateAThreadOriginDataIndex()};
+    BThreadCopy b_thread_copy_{CalculateBThreadOriginDataIndex()};
 };
 
 } // namespace ck

composable_kernel/include/tensor_operation/gridwise_gemm_xdlops_v2r3.hpp

@@ -18,8 +18,9 @@ template <typename GridwiseGemm,
           typename FloatC,
           typename AK0MK1GridDesc,
           typename BK0NK1GridDesc,
-          typename CM0M1M2NGridDesc,
-          typename CBlockClusterAdaptor>
+          typename CM0N0M1N1M2M3M4N2GridDesc,
+          typename CBlockClusterAdaptor,
+          bool HasMainKBlockLoop>
 __global__ void
 #if CK_USE_LAUNCH_BOUNDS
     __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, CK_MIN_BLOCK_PER_CU)
@@ -29,7 +30,7 @@ __global__ void
                                 FloatC* __restrict__ p_c_grid,
                                 const AK0MK1GridDesc a_k0_m_k1_grid_desc,
                                 const BK0NK1GridDesc b_k0_n_k1_grid_desc,
-                                const CM0M1M2NGridDesc c_m0_m1_m2_n_grid_desc,
+                                const CM0N0M1N1M2M3M4N2GridDesc c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
                                 const CBlockClusterAdaptor c_block_cluster_adaptor)
 {
     constexpr index_t shared_block_size =
@@ -37,13 +38,13 @@ __global__ void
 
     __shared__ FloatAB p_shared_block[shared_block_size];
 
-    GridwiseGemm::Run(p_a_grid,
+    GridwiseGemm::template Run<HasMainKBlockLoop>(p_a_grid,
                                                   p_b_grid,
                                                   p_c_grid,
                                                   p_shared_block,
                                                   a_k0_m_k1_grid_desc,
                                                   b_k0_n_k1_grid_desc,
-                      c_m0_m1_m2_n_grid_desc,
+                                                  c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
                                                   c_block_cluster_adaptor);
 }
 #elif CK_EXPERIMENTAL_PASS_TENSOR_DESCRIPTOR_BY_VOID_POINTER
@@ -52,7 +53,7 @@ template <typename GridwiseGemm,
           typename FloatC,
           typename AK0MK1GridDesc,
           typename BK0NK1GridDesc,
-          typename CM0M1M2NGridDesc,
+          typename CM0N0M1N1M2M3M4N2GridDesc,
           typename CBlockClusterAdaptor>
 __global__ void
 #if CK_USE_LAUNCH_BOUNDS
@@ -63,7 +64,7 @@ __global__ void
                                 FloatC* __restrict__ p_c_grid,
                                 const void CONSTANT* p_a_k0_m_k1_grid_desc,
                                 const void CONSTANT* p_b_k0_n_k1_grid_desc,
-                                const void CONSTANT* p_c_m0_m1_m2_n_grid_desc,
+                                const void CONSTANT* p_c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
                                 const void CONSTANT* p_c_block_cluster_adaptor)
 {
     constexpr index_t shared_block_size =
@@ -73,20 +74,21 @@ __global__ void
         cast_pointer_to_generic_address_space(p_a_k0_m_k1_grid_desc));
     const auto b_k0_n_k1_grid_desc = *reinterpret_cast<const BK0NK1GridDesc*>(
         cast_pointer_to_generic_address_space(p_b_k0_n_k1_grid_desc));
-    const auto c_m0_m1_m2_n_grid_desc = *reinterpret_cast<const CM0M1M2NGridDesc*>(
-        cast_pointer_to_generic_address_space(p_c_m0_m1_m2_n_grid_desc));
+    const auto c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc =
+        *reinterpret_cast<const CM0N0M1N1M2M3M4N2GridDesc*>(
+            cast_pointer_to_generic_address_space(p_c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc));
     const auto c_block_cluster_adaptor = *reinterpret_cast<const CBlockClusterAdaptor*>(
         cast_pointer_to_generic_address_space(p_c_block_cluster_adaptor));
 
     __shared__ FloatAB p_shared_block[shared_block_size];
 
-    GridwiseGemm::Run(p_a_grid,
+    GridwiseGemm::template Run<HasMainKBlockLoop>(p_a_grid,
                                                   p_b_grid,
                                                   p_c_grid,
                                                   p_shared_block,
                                                   a_k0_m_k1_grid_desc,
                                                   b_k0_n_k1_grid_desc,
-                      c_m0_m1_m2_n_grid_desc,
+                                                  c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
                                                   c_block_cluster_adaptor);
 }
 #endif
@@ -101,9 +103,9 @@ template <index_t BlockSize,
           typename CMNGridDesc,
           index_t MPerBlock,
           index_t NPerBlock,
-          index_t KPerBlock,
-          index_t MPerWave,
-          index_t NPerWave,
+          index_t K0PerBlock,
+          index_t MPerXDL,
+          index_t NPerXDL,
           index_t K1Value,
           index_t MRepeat,
           index_t NRepeat,
@@ -131,13 +133,19 @@ template <index_t BlockSize,
           typename CGridStepHacks,
           typename AGridMoveSliceWindowStepHacks,
           typename BGridMoveSliceWindowStepHacks,
-          bool CAccessOrderMRepeatNRepeat>
+          bool CAccessOrderMRepeatNRepeat,
+          bool ABlockLdsExtraM,
+          bool BBlockLdsExtraN>
 struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3
 {
     static constexpr auto I0 = Number<0>{};
     static constexpr auto I1 = Number<1>{};
     static constexpr auto I2 = Number<2>{};
     static constexpr auto I3 = Number<3>{};
+    static constexpr auto I4 = Number<4>{};
+    static constexpr auto I5 = Number<5>{};
+    static constexpr auto I6 = Number<6>{};
+    static constexpr auto I7 = Number<7>{};
 
     // K1 should be Number<...>
     static constexpr auto K1 = Number<K1Value>{};
@@ -147,14 +155,34 @@ struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3
         constexpr auto max_lds_align = K1;
 
         // A matrix in LDS memory, dst of blockwise copy
-        //   be careful of LDS alignment
-        constexpr auto a_k0_m_k1_block_desc = make_naive_tensor_descriptor_aligned(
-            make_tuple(Number<KPerBlock>{}, Number<MPerBlock>{}, K1), max_lds_align);
+        constexpr auto a_k0_m_k1_block_desc = [&]() {
+            if constexpr(ABlockLdsExtraM)
+            {
+                return make_naive_tensor_descriptor(
+                    make_tuple(Number<K0PerBlock>{}, Number<MPerBlock>{}, K1),
+                    make_tuple(Number<MPerBlock + 1>{} * K1, K1, I1));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor_aligned(
+                    make_tuple(Number<K0PerBlock>{}, Number<MPerBlock>{}, K1), max_lds_align);
+            }
+        }();
 
         // B matrix in LDS memory, dst of blockwise copy
-        //   be careful of LDS alignment
-        constexpr auto b_k0_n_k1_block_desc = make_naive_tensor_descriptor_aligned(
-            make_tuple(Number<KPerBlock>{}, Number<NPerBlock>{}, K1), max_lds_align);
+        constexpr auto b_k0_n_k1_block_desc = [&]() {
+            if constexpr(BBlockLdsExtraN)
+            {
+                return make_naive_tensor_descriptor(
+                    make_tuple(Number<K0PerBlock>{}, Number<NPerBlock>{}, K1),
+                    make_tuple(Number<NPerBlock + 1>{} * K1, K1, I1));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor_aligned(
+                    make_tuple(Number<K0PerBlock>{}, Number<NPerBlock>{}, K1), max_lds_align);
+            }
+        }();
 
         // LDS allocation for A and B: be careful of alignment
         constexpr auto a_block_space_size =
@@ -166,27 +194,45 @@ struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3
         return (a_block_space_size + b_block_space_size) * sizeof(FloatAB);
     }
 
+    // block_id to matrix tile idx (m0, n0) mapping are controlled by {M01, N01}
     __host__ __device__ static constexpr bool
     CheckValidity(const AK0MK1GridDesc& a_k0_m_k1_grid_desc,
                   const BK0NK1GridDesc& b_k0_n_k1_grid_desc,
-                  const CMNGridDesc& c_m_n_grid_desc)
+                  const CMNGridDesc& c_m_n_grid_desc,
+                  index_t M01,
+                  index_t N01)
     {
-        // TODO: turn on this
         static_assert(is_known_at_compile_time<remove_cv_t<decltype(K1)>>::value,
                       "wrong! K1 need to be known at compile-time");
 
+        static_assert((MPerBlock % (MPerXDL * MRepeat) == 0) &&
+                          (NPerBlock % (NRepeat * NPerXDL)) == 0,
+                      "Invalid tuning param!");
+
         const auto M  = a_k0_m_k1_grid_desc.GetLength(I1);
         const auto N  = b_k0_n_k1_grid_desc.GetLength(I1);
         const auto K0 = a_k0_m_k1_grid_desc.GetLength(I0);
 
-        // TODO: also check validity of all components (blockwise-copy, threadwise-copy, etc)
+        if(!(M == c_m_n_grid_desc.GetLength(I0) && N == c_m_n_grid_desc.GetLength(I1) &&
+             K0 == b_k0_n_k1_grid_desc.GetLength(I0) && K1 == a_k0_m_k1_grid_desc.GetLength(I2) &&
+             K1 == b_k0_n_k1_grid_desc.GetLength(I2)))
+            return false;
+
+        if(!(M % MPerBlock == 0 && N % NPerBlock == 0 && K0 % K0PerBlock == 0))
+            return false;
+
+        // check M01, N01
+        constexpr auto M1 = Number<MPerBlock>{};
+        constexpr auto N1 = Number<NPerBlock>{};
+
+        const auto M0 = M / M1;
+        const auto N0 = N / N1;
 
-        return (M == c_m_n_grid_desc.GetLength(I0) && N == c_m_n_grid_desc.GetLength(I1) &&
-                K0 == b_k0_n_k1_grid_desc.GetLength(I0) &&
-                K1 == a_k0_m_k1_grid_desc.GetLength(I2) &&
-                K1 == b_k0_n_k1_grid_desc.GetLength(I2)) &&
-               (M % MPerBlock == 0 && N % NPerBlock == 0 && K0 % KPerBlock == 0) &&
-               (MPerBlock % MPerWave == 0 && NPerBlock % NPerWave == 0);
+        if(!(M0 % M01 == 0 && N0 % N01 == 0))
+            return false;
+
+        // TODO: also check validity of all components (blockwise-copy, threadwise-copy, etc)
+        return true;
     }
 
     __host__ __device__ static constexpr index_t
@@ -200,34 +246,66 @@ struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3
         return grid_size;
     }
 
-    __host__ __device__ static constexpr auto
-    MakeCM0M1M2NGridDescriptor(const CMNGridDesc& c_m_n_grid_desc)
+    __host__ __device__ static constexpr bool CalculateHasMainK0BlockLoop(index_t K0)
     {
-        constexpr auto xdlops_gemm = XdlopsGemm<FloatAB, MPerWave, NPerWave, K1>{};
+        const bool has_main_k0_block_loop = (K0 / K0PerBlock) > 1;
 
-        constexpr auto CLayout = xdlops_gemm.GetCLayout();
+        return has_main_k0_block_loop;
+    }
 
-        constexpr auto M0 = Number<CLayout.M1()>{};
-        constexpr auto M1 = Number<CLayout.N1()>{};
-        constexpr auto M2 = Number<CLayout.M0()>{};
+    __host__ __device__ static constexpr auto
+    MakeCM0N0M1N1M2M3M4N2GridDescriptor(const CMNGridDesc& c_m_n_grid_desc)
+    {
+        constexpr auto max_lds_align = K1;
 
-        constexpr index_t MWaves = MPerBlock / (MPerWave * MRepeat);
-        constexpr index_t NWaves = NPerBlock / (NPerWave * NRepeat);
+        // A matrix in LDS memory, dst of blockwise copy
+        constexpr auto a_k0_m_k1_block_desc = [&]() {
+            if constexpr(ABlockLdsExtraM)
+            {
+                return make_naive_tensor_descriptor(
+                    make_tuple(Number<K0PerBlock>{}, Number<MPerBlock>{}, K1),
+                    make_tuple(Number<MPerBlock + 1>{} * K1, K1, I1));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor_aligned(
+                    make_tuple(Number<K0PerBlock>{}, Number<MPerBlock>{}, K1), max_lds_align);
+            }
+        }();
 
-        constexpr auto N1 = Number<CLayout.N0()>{};
+        // B matrix in LDS memory, dst of blockwise copy
+        constexpr auto b_k0_n_k1_block_desc = [&]() {
+            if constexpr(BBlockLdsExtraN)
+            {
+                return make_naive_tensor_descriptor(
+                    make_tuple(Number<K0PerBlock>{}, Number<NPerBlock>{}, K1),
+                    make_tuple(Number<NPerBlock + 1>{} * K1, K1, I1));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor_aligned(
+                    make_tuple(Number<K0PerBlock>{}, Number<NPerBlock>{}, K1), max_lds_align);
+            }
+        }();
 
-        const auto c_m0_m1_m2_n_grid_desc = transform_tensor_descriptor(
-            c_m_n_grid_desc,
-            make_tuple(make_unmerge_transform(make_tuple(MRepeat, MWaves, M0, M1, M2)),
-                       make_unmerge_transform(make_tuple(NRepeat, NWaves, N1))),
-            make_tuple(Sequence<0>{}, Sequence<1>{}),
-            make_tuple(Sequence<0, 2, 4, 5, 6>{}, Sequence<1, 3, 7>{}));
+        using BlockwiseGemm =
+            BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1<BlockSize,
+                                                                FloatAB,
+                                                                FloatAcc,
+                                                                decltype(a_k0_m_k1_block_desc),
+                                                                decltype(b_k0_n_k1_block_desc),
+                                                                MPerXDL,
+                                                                NPerXDL,
+                                                                MRepeat,
+                                                                NRepeat,
+                                                                K1>;
 
-        return c_m0_m1_m2_n_grid_desc;
+        return BlockwiseGemm::MakeCM0N0M1N1M2M3M4N2GridDescriptor(c_m_n_grid_desc);
     }
 
+    // return block_id to C matrix tile idx (m0, n0) mapping
     __host__ __device__ static constexpr auto
-    MakeCBlockClusterAdaptor(const CMNGridDesc& c_m_n_grid_desc)
+    MakeCBlockClusterAdaptor(const CMNGridDesc& c_m_n_grid_desc, index_t M01, index_t N01)
     {
         const auto M = c_m_n_grid_desc.GetLength(I0);
         const auto N = c_m_n_grid_desc.GetLength(I1);
@@ -238,31 +316,40 @@ struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3
         const auto M0 = M / M1;
         const auto N0 = N / N1;
 
-#if 1
-        const auto c_blockid_to_m0_n0_block_cluster_adaptor =
-            make_single_stage_tensor_adaptor(make_tuple(make_merge_transform(make_tuple(M0, N0))),
-                                             make_tuple(Sequence<0, 1>{}),
+        const auto M00 = M0 / M01;
+        const auto N00 = N0 / N01;
+
+        const auto m00_m01_n00_n01_to_m0_n0_block_cluster_adaptor =
+            make_single_stage_tensor_adaptor(
+                make_tuple(make_unmerge_transform(make_tuple(M00, M01)),
+                           make_unmerge_transform(make_tuple(N00, N01))),
+                make_tuple(Sequence<0>{}, Sequence<1>{}),
+                make_tuple(Sequence<0, 2>{}, Sequence<1, 3>{}));
+
+        const auto c_blockid_to_m00_m01_n00_n01_block_cluster_adaptor =
+            make_single_stage_tensor_adaptor(
+                make_tuple(make_merge_transform(make_tuple(M00, N00, M01, N01))),
+                make_tuple(Sequence<0, 1, 2, 3>{}),
                 make_tuple(Sequence<0>{}));
-#elif 1
+
         const auto c_blockid_to_m0_n0_block_cluster_adaptor =
-            make_single_stage_tensor_adaptor(make_tuple(make_merge_transform(make_tuple(N0, M0))),
-                                             make_tuple(Sequence<1, 0>{}),
-                                             make_tuple(Sequence<0>{}));
-#endif
+            chain_tensor_adaptors(m00_m01_n00_n01_to_m0_n0_block_cluster_adaptor,
+                                  c_blockid_to_m00_m01_n00_n01_block_cluster_adaptor);
 
         return c_blockid_to_m0_n0_block_cluster_adaptor;
     }
 
-    using CM0M1M2NGridDesc     = decltype(MakeCM0M1M2NGridDescriptor(CMNGridDesc{}));
-    using CBlockClusterAdaptor = decltype(MakeCBlockClusterAdaptor(CMNGridDesc{}));
+    using CM0N0M1N1M2M3M4N2GridDesc = decltype(MakeCM0N0M1N1M2M3M4N2GridDescriptor(CMNGridDesc{}));
+    using CBlockClusterAdaptor      = decltype(MakeCBlockClusterAdaptor(CMNGridDesc{}, 1, 1));
 
+    template <bool HasMainKBlockLoop>
     __device__ static void Run(const FloatAB* __restrict__ p_a_grid,
                                const FloatAB* __restrict__ p_b_grid,
                                FloatC* __restrict__ p_c_grid,
                                FloatAB* __restrict__ p_shared_block,
                                const AK0MK1GridDesc& a_k0_m_k1_grid_desc,
                                const BK0NK1GridDesc& b_k0_n_k1_grid_desc,
-                               const CM0M1M2NGridDesc& c_m0_m1_m2_n_grid_desc,
+                               const CM0N0M1N1M2M3M4N2GridDesc& c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
                                const CBlockClusterAdaptor& c_block_cluster_adaptor)
     {
         const auto a_grid_buf = make_dynamic_buffer<AddressSpaceEnum_t::Global>(
@@ -270,7 +357,7 @@ struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3
         const auto b_grid_buf = make_dynamic_buffer<AddressSpaceEnum_t::Global>(
             p_b_grid, b_k0_n_k1_grid_desc.GetElementSpaceSize());
         auto c_grid_buf = make_dynamic_buffer<AddressSpaceEnum_t::Global>(
-            p_c_grid, c_m0_m1_m2_n_grid_desc.GetElementSpaceSize());
+            p_c_grid, c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc.GetElementSpaceSize());
 
         const auto K0 = a_k0_m_k1_grid_desc.GetLength(I0);
 
@@ -289,20 +376,40 @@ struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3
         constexpr auto max_lds_align = K1;
 
         // A matrix in LDS memory, dst of blockwise copy
-        //   be careful of LDS alignment
-        constexpr auto a_k0_m_k1_block_desc = make_naive_tensor_descriptor_aligned(
-            make_tuple(Number<KPerBlock>{}, Number<MPerBlock>{}, K1), max_lds_align);
+        constexpr auto a_k0_m_k1_block_desc = [&]() {
+            if constexpr(ABlockLdsExtraM)
+            {
+                return make_naive_tensor_descriptor(
+                    make_tuple(Number<K0PerBlock>{}, Number<MPerBlock>{}, K1),
+                    make_tuple(Number<MPerBlock + 1>{} * K1, K1, I1));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor_aligned(
+                    make_tuple(Number<K0PerBlock>{}, Number<MPerBlock>{}, K1), max_lds_align);
+            }
+        }();
 
         // B matrix in LDS memory, dst of blockwise copy
-        //   be careful of LDS alignment
-        constexpr auto b_k0_n_k1_block_desc = make_naive_tensor_descriptor_aligned(
-            make_tuple(Number<KPerBlock>{}, Number<NPerBlock>{}, K1), max_lds_align);
+        constexpr auto b_k0_n_k1_block_desc = [&]() {
+            if constexpr(BBlockLdsExtraN)
+            {
+                return make_naive_tensor_descriptor(
+                    make_tuple(Number<K0PerBlock>{}, Number<NPerBlock>{}, K1),
+                    make_tuple(Number<NPerBlock + 1>{} * K1, K1, I1));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor_aligned(
+                    make_tuple(Number<K0PerBlock>{}, Number<NPerBlock>{}, K1), max_lds_align);
+            }
+        }();
 
         // A matrix blockwise copy
         auto a_blockwise_copy =
             BlockwiseTensorSliceTransfer_v4<BlockSize,
                                             InMemoryDataOperationEnum_t::Set,
-                                            Sequence<KPerBlock, MPerBlock, K1>,
+                                            Sequence<K0PerBlock, MPerBlock, K1>,
                                             ABlockTransferThreadSliceLengths_K0_M_K1,
                                             ABlockTransferThreadClusterLengths_K0_M_K1,
                                             ABlockTransferThreadClusterArrangeOrder,
@@ -328,7 +435,7 @@ struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3
         auto b_blockwise_copy =
             BlockwiseTensorSliceTransfer_v4<BlockSize,
                                             InMemoryDataOperationEnum_t::Set,
-                                            Sequence<KPerBlock, NPerBlock, K1>,
+                                            Sequence<K0PerBlock, NPerBlock, K1>,
                                             BBlockTransferThreadSliceLengths_K0_N_K1,
                                             BBlockTransferThreadClusterLengths_K0_N_K1,
                                             BBlockTransferThreadClusterArrangeOrder,
@@ -352,59 +459,25 @@ struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3
 
         // GEMM definition
         //   c_mtx += transpose(a_mtx) * b_mtx
-        //     a_mtx[KPerBlock, MPerBlock] is in LDS
-        //     b_mtx[KPerBlock, NPerBlock] is in LDS
+        //     a_mtx[K0PerBlock, MPerBlock] is in LDS
+        //     b_mtx[K0PerBlock, NPerBlock] is in LDS
         //     c_mtx[MPerBlock, NPerBlock] is distributed among threads, and saved in
         //       register
         // sanity check
 
-        static_assert(MPerBlock % (MPerWave * MRepeat) == 0 &&
-                          NPerBlock % (NPerWave * NRepeat) == 0,
-                      "wrong!");
-
-        constexpr auto a_k0_m0_m1_k1_block_desc = transform_tensor_descriptor(
-            a_k0_m_k1_block_desc,
-            make_tuple(make_pass_through_transform(Number<KPerBlock>{}),
-                       make_unmerge_transform(
-                           make_tuple(Number<MRepeat>{}, Number<MPerBlock / MRepeat>{})),
-                       make_pass_through_transform(K1)),
-            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
-            make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3>{}));
-
-        constexpr auto b_k0_n0_n1_k1_block_desc = transform_tensor_descriptor(
-            b_k0_n_k1_block_desc,
-            make_tuple(make_pass_through_transform(Number<KPerBlock>{}),
-                       make_unmerge_transform(
-                           make_tuple(Number<NRepeat>{}, Number<NPerBlock / NRepeat>{})),
-                       make_pass_through_transform(K1)),
-            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
-            make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3>{}));
-
-        const auto blockwise_gemm =
-            BlockwiseGemmXdlops_km_kn_m0m1m2n_v1<BlockSize,
+        auto blockwise_gemm =
+            BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1<BlockSize,
                                                                 FloatAB,
-                                                 decltype(a_k0_m0_m1_k1_block_desc),
-                                                 decltype(b_k0_n0_n1_k1_block_desc),
-                                                 MPerWave,
-                                                 NPerWave,
+                                                                FloatAcc,
+                                                                decltype(a_k0_m_k1_block_desc),
+                                                                decltype(b_k0_n_k1_block_desc),
+                                                                MPerXDL,
+                                                                NPerXDL,
+                                                                MRepeat,
+                                                                NRepeat,
                                                                 K1>{};
 
-        constexpr auto CLayout = blockwise_gemm.GetCLayout();
-
-        constexpr index_t BlkSize   = CLayout.GetBlkSize();
-        constexpr index_t NumBlks   = CLayout.GetNumBlks();
-        constexpr index_t NumXdlops = CLayout.GetNumXdlops();
-
-        static_assert(NumBlks == 1 && NumXdlops == 1, "K Reduction Mfma only");
-
-        constexpr auto c_mr_nr_blk_desc =
-            make_naive_tensor_descriptor_packed(make_tuple(Number<MRepeat>{}, Number<NRepeat>{}));
-
-        StaticBuffer<AddressSpaceEnum_t::Vgpr,
-                     vector_type<FloatAcc, BlkSize>,
-                     c_mr_nr_blk_desc.GetElementSpaceSize(),
-                     true>
-            c_thread_buf;
+        auto c_thread_buf = blockwise_gemm.GetCThreadBuffer();
 
         // LDS allocation for A and B: be careful of alignment
         constexpr auto a_block_space_size =
@@ -413,8 +486,8 @@ struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3
         FloatAB* p_a_block = p_shared_block;
         FloatAB* p_b_block = p_shared_block + a_block_space_size;
 
-        constexpr auto a_block_slice_copy_step = make_multi_index(KPerBlock, 0, 0);
-        constexpr auto b_block_slice_copy_step = make_multi_index(KPerBlock, 0, 0);
+        constexpr auto a_block_slice_copy_step = make_multi_index(K0PerBlock, 0, 0);
+        constexpr auto b_block_slice_copy_step = make_multi_index(K0PerBlock, 0, 0);
 
         // hack to control index calculation when iterating over A and B matrix for threadwise copy
         constexpr auto a_k0_m_k1_grid_step_hacks = AGridStepHacks{};
@@ -440,8 +513,10 @@ struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3
         }
 
         // main body
-        index_t k_block_data_begin = 0;
+        index_t k0_block_data_begin = 0;
 
+        if constexpr(HasMainKBlockLoop)
+        {
             do
             {
                 a_blockwise_copy.MoveSrcSliceWindow(a_k0_m_k1_grid_desc,
@@ -451,11 +526,13 @@ struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3
                                                     b_block_slice_copy_step,
                                                     b_k0_n_k1_grid_move_slice_window_step_hack);
 
-            a_blockwise_copy.RunRead(a_k0_m_k1_grid_desc, a_grid_buf, a_k0_m_k1_grid_step_hacks);
+                a_blockwise_copy.RunRead(
+                    a_k0_m_k1_grid_desc, a_grid_buf, a_k0_m_k1_grid_step_hacks);
 
                 block_sync_lds();
 
-            b_blockwise_copy.RunRead(b_k0_n_k1_grid_desc, b_grid_buf, b_k0_n_k1_grid_step_hacks);
+                b_blockwise_copy.RunRead(
+                    b_k0_n_k1_grid_desc, b_grid_buf, b_k0_n_k1_grid_step_hacks);
 
                 blockwise_gemm.Run(a_block_buf, b_block_buf, c_thread_buf);
 
@@ -464,8 +541,9 @@ struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3
                 a_blockwise_copy.RunWrite(a_k0_m_k1_block_desc, a_block_buf);
                 b_blockwise_copy.RunWrite(b_k0_n_k1_block_desc, b_block_buf);
 
-            k_block_data_begin += KPerBlock;
-        } while(k_block_data_begin < (K0 - KPerBlock));
+                k0_block_data_begin += K0PerBlock;
+            } while(k0_block_data_begin < (K0 - K0PerBlock));
+        }
 
         // tail
         {
@@ -474,41 +552,23 @@ struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3
             blockwise_gemm.Run(a_block_buf, b_block_buf, c_thread_buf);
         }
 
-#if 0
         // output: register to global memory
         {
-            constexpr index_t M0 = CLayout.M1();
-            constexpr index_t M1 = CLayout.N1();
-            constexpr index_t M2 = CLayout.M0();
-
-            constexpr index_t N0 = CLayout.N1();
-            constexpr index_t N1 = CLayout.N0();
-
-            constexpr auto c_m0_m1_m2_n_thread_desc =
-                make_naive_tensor_descriptor_packed(make_tuple(Number<MRepeat>{},
-                                                                          Number<NRepeat>{},
-                                                                          Number<1>{},
-                                                                          Number<1>{},
-                                                                          Number<M0>{},
-                                                                          Number<1>{},
-                                                                          Number<M2>{},
-                                                                          Number<1>{}));
-
-            StaticBuffer<AddressSpaceEnum_t::Vgpr, FloatC, c_m0_m1_m2_n_thread_desc.GetElementSpaceSize(), true>
-                c_blk_buf_;
-
-            static_for<0, MRepeat, 1>{}([&](auto mr_i) {
-                static_for<0, NRepeat, 1>{}([&](auto nr_i) {
-                    constexpr auto blk_off =
-                        c_mr_nr_blk_desc.CalculateOffset(make_tuple(mr_i, nr_i));
-
-                    static_for<0, BlkSize, 1>{}([&](auto j) {
-                        c_blk_buf_(Number<blk_off * BlkSize + j>{}) =
-                            c_thread_buf[Number<blk_off>{}]
-                                .template AsType<FloatAcc>()[Number<j>{}];
-                    });
-                });
-            });
+            constexpr auto c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc =
+                blockwise_gemm.GetCM0N0M1N1M2M3M4N2BlockDescriptor();
+
+            constexpr auto M0 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I0);
+            constexpr auto N0 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I1);
+            constexpr auto M1 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I2);
+            constexpr auto N1 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I3);
+            constexpr auto M2 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I4);
+            constexpr auto M3 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I5);
+            constexpr auto M4 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I6);
+            constexpr auto N2 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I7);
+
+            constexpr auto c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc =
+                make_naive_tensor_descriptor_packed(make_tuple(
+                    Number<M0>{}, Number<N0>{}, I1, I1, Number<M2>{}, I1, Number<M4>{}, I1));
 
             // calculate origin of thread output tensor on global memory
             //     blockwise GEMM c matrix starting index
@@ -521,278 +581,58 @@ struct GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3
             const index_t n_thread_data_on_grid =
                 n_block_data_idx_on_grid + c_thread_mtx_on_block[I1];
 
-            constexpr auto c_m0_m1_m2_n_grid_tensor_step_hacks = CGridStepHacks{};
-
-            constexpr index_t MWaves = MPerBlock / (MPerWave * MRepeat);
-            constexpr index_t NWaves = NPerBlock / (NPerWave * NRepeat);
+            constexpr auto c_m0_n0_m1_n1_m2_m3_m4_n2_grid_tensor_step_hacks = CGridStepHacks{};
 
-            ThreadwiseTensorSliceTransfer_v1r3<
-                FloatC,
-                FloatC,
-                decltype(c_m0_m1_m2_n_thread_desc),
-                decltype(c_m0_m1_m2_n_grid_desc),
-                Sequence<MRepeat, NRepeat, 1, 1, M0, 1, M2, 1>,
-                CThreadTransferSrcDstAccessOrder,
-                CThreadTransferSrcDstVectorDim,
-                CThreadTransferDstScalarPerVector,
-                CGlobalMemoryDataOperation,
-                1,
-                true>{
-                c_m0_m1_m2_n_grid_desc,
-                make_multi_index(m_thread_data_on_grid / (M2 * M1 * M0 * MWaves),
-                                 n_thread_data_on_grid / (N1 * NWaves),
-                                 m_thread_data_on_grid % (M2 * M1 * M0 * MWaves) / (M2 * M1 * M0),
-                                 n_thread_data_on_grid % (N1 * NWaves) / N1,
-                                 m_thread_data_on_grid % (M2 * M1 * M0) / (M2 * M1),
-                                 m_thread_data_on_grid % (M2 * M1) / M2,
-                                 m_thread_data_on_grid % M2,
-                                 n_thread_data_on_grid % N1)}
-                .Run(c_m0_m1_m2_n_thread_desc,
-                     make_tuple(I0, I0, I0, I0, I0, I0, I0, I0),
-                     c_blk_buf_,
-                     c_m0_m1_m2_n_grid_desc,
-                     c_grid_buf,
-                     c_m0_m1_m2_n_grid_tensor_step_hacks);
-        }
-#else
-        {
-            constexpr index_t M0 = CLayout.M1();
-            constexpr index_t M1 = CLayout.N1();
-            constexpr index_t M2 = CLayout.M0();
-
-            constexpr auto c_m0_m1_m2_n_thread_desc = make_naive_tensor_descriptor_packed(
-                make_tuple(I1, I1, I1, I1, Number<M0>{}, Number<1>{}, Number<M2>{}, Number<1>{}));
-
-            // calculate origin of thread output tensor on global memory
-            //     blockwise GEMM c matrix starting index
-            const auto c_thread_mtx_on_block =
-                blockwise_gemm.CalculateCThreadOriginDataIndex(I0, I0, I0, I0);
+            const auto m_thread_data_on_grid_to_m0_m1_m2_m3_m4_adaptor =
+                make_single_stage_tensor_adaptor(
+                    make_tuple(make_merge_transform(make_tuple(M0, M1, M2, M3, M4))),
+                    make_tuple(Sequence<0, 1, 2, 3, 4>{}),
+                    make_tuple(Sequence<0>{}));
 
-            const index_t m_thread_data_on_grid =
-                m_block_data_idx_on_grid + c_thread_mtx_on_block[I0];
+            const auto m_thread_data_on_grid_idx =
+                m_thread_data_on_grid_to_m0_m1_m2_m3_m4_adaptor.CalculateBottomIndex(
+                    make_multi_index(m_thread_data_on_grid));
 
-            const index_t n_thread_data_on_grid =
-                n_block_data_idx_on_grid + c_thread_mtx_on_block[I1];
+            const auto n_thread_data_on_grid_to_n0_n1_n2_adaptor = make_single_stage_tensor_adaptor(
+                make_tuple(make_merge_transform(make_tuple(N0, N1, N2))),
+                make_tuple(Sequence<0, 1, 2>{}),
+                make_tuple(Sequence<0>{}));
 
-            constexpr auto c_m0_m1_m2_n_grid_tensor_step_hacks = CGridStepHacks{};
+            const auto n_thread_data_on_grid_idx =
+                n_thread_data_on_grid_to_n0_n1_n2_adaptor.CalculateBottomIndex(
+                    make_multi_index(n_thread_data_on_grid));
 
             auto c_thread_copy =
-                ThreadwiseTensorSliceTransfer_v1r3<FloatC,
+                ThreadwiseTensorSliceTransfer_v1r3<FloatAcc,
                                                    FloatC,
-                                                   decltype(c_m0_m1_m2_n_thread_desc),
-                                                   decltype(c_m0_m1_m2_n_grid_desc),
-                                                   Sequence<1, 1, 1, 1, M0, 1, M2, 1>,
+                                                   decltype(c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc),
+                                                   decltype(c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc),
+                                                   Sequence<M0, N0, I1, I1, M2, I1, M4, I1>,
                                                    CThreadTransferSrcDstAccessOrder,
                                                    CThreadTransferSrcDstVectorDim,
                                                    CThreadTransferDstScalarPerVector,
                                                    CGlobalMemoryDataOperation,
                                                    1,
                                                    true>{
-                    c_m0_m1_m2_n_grid_desc,
-                    make_multi_index(0,
-                                     0,
-                                     0,
-                                     0,
-                                     m_thread_data_on_grid / (M2 * M1),
-                                     m_thread_data_on_grid % (M2 * M1) / M2,
-                                     m_thread_data_on_grid % M2,
-                                     n_thread_data_on_grid)};
-
-            auto init_copy = [&](auto c_thread_idx_) {
-                constexpr auto blk_off = c_mr_nr_blk_desc.CalculateOffset(c_thread_idx_);
-                c_thread_copy.Run(c_m0_m1_m2_n_thread_desc,
-                                  make_tuple(I0, I0, I0, I0, I0, I0, I0, I0),
-                                  c_thread_buf[Number<blk_off>{}].template AsType<FloatAcc>(),
-                                  c_m0_m1_m2_n_grid_desc,
-                                  c_grid_buf,
-                                  c_m0_m1_m2_n_grid_tensor_step_hacks);
-
-                return c_thread_idx_;
-            };
-
-            auto mrepeat_plus_copy = [&](auto c_thread_idx_) {
-                constexpr auto mrepeat_step_plus = make_multi_index(1, 0, 0, 0, 0, 0, 0, 0);
-                c_thread_copy.MoveDstSliceWindow(c_m0_m1_m2_n_grid_desc, mrepeat_step_plus);
-
-                constexpr auto blk_off = c_mr_nr_blk_desc.CalculateOffset(c_thread_idx_);
-                c_thread_copy.Run(c_m0_m1_m2_n_thread_desc,
-                                  make_tuple(I0, I0, I0, I0, I0, I0, I0, I0),
-                                  c_thread_buf[Number<blk_off>{}].template AsType<FloatAcc>(),
-                                  c_m0_m1_m2_n_grid_desc,
-                                  c_grid_buf,
-                                  c_m0_m1_m2_n_grid_tensor_step_hacks);
-            };
-
-            auto nrepeat_plus_copy = [&](auto c_thread_idx_) {
-                constexpr auto nrepeat_step_plus = make_multi_index(0, 1, 0, 0, 0, 0, 0, 0);
-                c_thread_copy.MoveDstSliceWindow(c_m0_m1_m2_n_grid_desc, nrepeat_step_plus);
 
-                constexpr auto blk_off = c_mr_nr_blk_desc.CalculateOffset(c_thread_idx_);
-                c_thread_copy.Run(c_m0_m1_m2_n_thread_desc,
+                    c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                    make_multi_index(m_thread_data_on_grid_idx[I0],
+                                     n_thread_data_on_grid_idx[I0],
+                                     m_thread_data_on_grid_idx[I1],
+                                     n_thread_data_on_grid_idx[I1],
+                                     m_thread_data_on_grid_idx[I2],
+                                     m_thread_data_on_grid_idx[I3],
+                                     m_thread_data_on_grid_idx[I4],
+                                     n_thread_data_on_grid_idx[I2])};
+
+            c_thread_copy.Run(c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc,
                               make_tuple(I0, I0, I0, I0, I0, I0, I0, I0),
-                                  c_thread_buf[Number<blk_off>{}].template AsType<FloatAcc>(),
-                                  c_m0_m1_m2_n_grid_desc,
+                              c_thread_buf,
+                              c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
                               c_grid_buf,
-                                  c_m0_m1_m2_n_grid_tensor_step_hacks);
-            };
-
-            auto mrepeat_minus_copy = [&](auto c_thread_idx_) {
-                constexpr auto mrepeat_step_plus = make_multi_index(-1, 0, 0, 0, 0, 0, 0, 0);
-                c_thread_copy.MoveDstSliceWindow(c_m0_m1_m2_n_grid_desc, mrepeat_step_plus);
-
-                constexpr auto blk_off = c_mr_nr_blk_desc.CalculateOffset(c_thread_idx_);
-                c_thread_copy.Run(c_m0_m1_m2_n_thread_desc,
-                                  make_tuple(I0, I0, I0, I0, I0, I0, I0, I0),
-                                  c_thread_buf[Number<blk_off>{}].template AsType<FloatAcc>(),
-                                  c_m0_m1_m2_n_grid_desc,
-                                  c_grid_buf,
-                                  c_m0_m1_m2_n_grid_tensor_step_hacks);
-            };
-
-            auto nrepeat_minus_copy = [&](auto c_thread_idx_) {
-                constexpr auto nrepeat_step_minus = make_multi_index(0, -1, 0, 0, 0, 0, 0, 0);
-                c_thread_copy.MoveDstSliceWindow(c_m0_m1_m2_n_grid_desc, nrepeat_step_minus);
-
-                constexpr auto blk_off = c_mr_nr_blk_desc.CalculateOffset(c_thread_idx_);
-                c_thread_copy.Run(c_m0_m1_m2_n_thread_desc,
-                                  make_tuple(I0, I0, I0, I0, I0, I0, I0, I0),
-                                  c_thread_buf[Number<blk_off>{}].template AsType<FloatAcc>(),
-                                  c_m0_m1_m2_n_grid_desc,
-                                  c_grid_buf,
-                                  c_m0_m1_m2_n_grid_tensor_step_hacks);
-            };
-
-            static_assert((MRepeat == 4 && NRepeat == 4) or (MRepeat == 4 && NRepeat == 2) or
-                              (MRepeat == 2 && NRepeat == 4) or (MRepeat == 2 && NRepeat == 2) or
-                              (MRepeat == 2 && NRepeat == 1) or (MRepeat == 1 && NRepeat == 2) or
-                              (MRepeat == 1 && NRepeat == 1),
-                          "wrong");
-
-            if constexpr(MRepeat == 4 && NRepeat == 4)
-            {
-                init_copy(make_tuple(I0, I0));
-
-                if constexpr(CAccessOrderMRepeatNRepeat)
-                {
-                    nrepeat_plus_copy(make_tuple(I0, I1));
-                    nrepeat_plus_copy(make_tuple(I0, I2));
-                    nrepeat_plus_copy(make_tuple(I0, I3));
-                    mrepeat_plus_copy(make_tuple(I1, I3));
-                    nrepeat_minus_copy(make_tuple(I1, I2));
-                    nrepeat_minus_copy(make_tuple(I1, I1));
-                    nrepeat_minus_copy(make_tuple(I1, I0));
-                    mrepeat_plus_copy(make_tuple(I2, I0));
-                    nrepeat_plus_copy(make_tuple(I2, I1));
-                    nrepeat_plus_copy(make_tuple(I2, I2));
-                    nrepeat_plus_copy(make_tuple(I2, I3));
-                    mrepeat_plus_copy(make_tuple(I3, I3));
-                    nrepeat_minus_copy(make_tuple(I3, I2));
-                    nrepeat_minus_copy(make_tuple(I3, I1));
-                    nrepeat_minus_copy(make_tuple(I3, I0));
-                }
-                else
-                {
-                    mrepeat_plus_copy(make_tuple(I1, I0));
-                    mrepeat_plus_copy(make_tuple(I2, I0));
-                    mrepeat_plus_copy(make_tuple(I3, I0));
-                    nrepeat_plus_copy(make_tuple(I3, I1));
-                    mrepeat_minus_copy(make_tuple(I2, I1));
-                    mrepeat_minus_copy(make_tuple(I1, I1));
-                    mrepeat_minus_copy(make_tuple(I0, I1));
-                    nrepeat_plus_copy(make_tuple(I0, I2));
-                    mrepeat_plus_copy(make_tuple(I1, I2));
-                    mrepeat_plus_copy(make_tuple(I2, I2));
-                    mrepeat_plus_copy(make_tuple(I3, I2));
-                    nrepeat_plus_copy(make_tuple(I3, I3));
-                    mrepeat_minus_copy(make_tuple(I2, I3));
-                    mrepeat_minus_copy(make_tuple(I1, I3));
-                    mrepeat_minus_copy(make_tuple(I0, I3));
-                }
-            }
-            else if constexpr(MRepeat == 4 && NRepeat == 2)
-            {
-                init_copy(make_tuple(I0, I0));
-
-                if constexpr(CAccessOrderMRepeatNRepeat)
-                {
-                    nrepeat_plus_copy(make_tuple(I0, I1));
-                    mrepeat_plus_copy(make_tuple(I1, I1));
-                    nrepeat_minus_copy(make_tuple(I1, I0));
-                    mrepeat_plus_copy(make_tuple(I2, I0));
-                    nrepeat_plus_copy(make_tuple(I2, I1));
-                    mrepeat_plus_copy(make_tuple(I3, I1));
-                    nrepeat_minus_copy(make_tuple(I3, I0));
-                }
-                else
-                {
-                    mrepeat_plus_copy(make_tuple(I1, I0));
-                    mrepeat_plus_copy(make_tuple(I2, I0));
-                    mrepeat_plus_copy(make_tuple(I3, I0));
-                    nrepeat_plus_copy(make_tuple(I3, I1));
-                    mrepeat_minus_copy(make_tuple(I2, I1));
-                    mrepeat_minus_copy(make_tuple(I1, I1));
-                    mrepeat_minus_copy(make_tuple(I0, I1));
-                }
-            }
-            else if constexpr(MRepeat == 2 && NRepeat == 4)
-            {
-                init_copy(make_tuple(I0, I0));
-
-                if constexpr(CAccessOrderMRepeatNRepeat)
-                {
-                    nrepeat_plus_copy(make_tuple(I0, I1));
-                    nrepeat_plus_copy(make_tuple(I0, I2));
-                    nrepeat_plus_copy(make_tuple(I0, I3));
-                    mrepeat_plus_copy(make_tuple(I1, I3));
-                    nrepeat_minus_copy(make_tuple(I1, I2));
-                    nrepeat_minus_copy(make_tuple(I1, I1));
-                    nrepeat_minus_copy(make_tuple(I1, I0));
-                }
-                else
-                {
-                    mrepeat_plus_copy(make_tuple(I1, I0));
-                    nrepeat_plus_copy(make_tuple(I1, I1));
-                    mrepeat_minus_copy(make_tuple(I0, I1));
-                    nrepeat_plus_copy(make_tuple(I0, I2));
-                    mrepeat_plus_copy(make_tuple(I1, I2));
-                    nrepeat_plus_copy(make_tuple(I1, I3));
-                    mrepeat_minus_copy(make_tuple(I0, I3));
-                }
-            }
-            else if constexpr(MRepeat == 2 && NRepeat == 2)
-            {
-                init_copy(make_tuple(I0, I0));
-
-                if constexpr(CAccessOrderMRepeatNRepeat)
-                {
-                    nrepeat_plus_copy(make_tuple(I0, I1));
-                    mrepeat_plus_copy(make_tuple(I1, I1));
-                    nrepeat_minus_copy(make_tuple(I1, I0));
-                }
-                else
-                {
-                    mrepeat_plus_copy(make_tuple(I1, I0));
-                    nrepeat_plus_copy(make_tuple(I1, I1));
-                    mrepeat_minus_copy(make_tuple(I0, I1));
+                              c_m0_n0_m1_n1_m2_m3_m4_n2_grid_tensor_step_hacks);
         }
     }
-            else if constexpr(MRepeat == 2 && NRepeat == 1)
-            {
-                init_copy(make_tuple(I0, I0));
-                mrepeat_plus_copy(make_tuple(I1, I0));
-            }
-            else if constexpr(MRepeat == 1 && NRepeat == 2)
-            {
-                init_copy(make_tuple(I0, I0));
-                nrepeat_plus_copy(make_tuple(I0, I1));
-            }
-            else if constexpr(MRepeat == 1 && NRepeat == 1)
-            {
-                init_copy(make_tuple(I0, I0));
-            }
-        }
-#endif
-    }
 }; // namespace ck
 
 } // namespace ck

composable_kernel/include/tensor_operation/gridwise_gemm_xdlops_v2r4.hpp

+#ifndef CK_GRIDWISE_GEMM_XDLOPS_V2R4_HPP
+#define CK_GRIDWISE_GEMM_XDLOPS_V2R4_HPP
+
+#include "common_header.hpp"
+#include "multi_index_transform_helper.hpp"
+#include "tensor_descriptor.hpp"
+#include "tensor_descriptor_helper.hpp"
+#include "blockwise_gemm_xdlops.hpp"
+#include "blockwise_tensor_slice_transfer.hpp"
+#include "threadwise_tensor_slice_transfer.hpp"
+#include "threadwise_tensor_slice_set.hpp"
+
+namespace ck {
+
+#if CK_EXPERIMENTAL_PASS_TENSOR_DESCRIPTOR_BY_VALUE
+template <typename GridwiseGemm,
+          typename FloatAB,
+          typename FloatC,
+          typename ABK0MK1GridDesc,
+          typename BBK0NK1GridDesc,
+          typename CM0N0M1N1M2M3M4N2GridDesc,
+          typename CBlockClusterAdaptor>
+__global__ void
+#if CK_USE_LAUNCH_BOUNDS
+    __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, CK_MIN_BLOCK_PER_CU)
+#endif
+        kernel_gemm_xdlops_v2r4(const FloatAB* __restrict__ p_a_grid,
+                                const FloatAB* __restrict__ p_b_grid,
+                                FloatC* __restrict__ p_c_grid,
+                                const ABK0MK1GridDesc a_b_k0_m_k1_grid_desc,
+                                const BBK0NK1GridDesc b_b_k0_n_k1_grid_desc,
+                                const CM0N0M1N1M2M3M4N2GridDesc c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                                const CBlockClusterAdaptor c_block_cluster_adaptor)
+{
+    constexpr index_t shared_block_size =
+        GridwiseGemm::GetSharedMemoryNumberOfByte() / sizeof(FloatAB);
+
+    __shared__ FloatAB p_shared_block[shared_block_size];
+
+    GridwiseGemm::Run(p_a_grid,
+                      p_b_grid,
+                      p_c_grid,
+                      p_shared_block,
+                      a_b_k0_m_k1_grid_desc,
+                      b_b_k0_n_k1_grid_desc,
+                      c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                      c_block_cluster_adaptor);
+}
+#elif CK_EXPERIMENTAL_PASS_TENSOR_DESCRIPTOR_BY_VOID_POINTER
+template <typename GridwiseGemm,
+          typename FloatAB,
+          typename FloatC,
+          typename ABK0MK1GridDesc,
+          typename BBK0NK1GridDesc,
+          typename CM0N0M1N1M2M3M4N2GridDesc,
+          typename CBlockClusterAdaptor>
+__global__ void
+#if CK_USE_LAUNCH_BOUNDS
+    __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, CK_MIN_BLOCK_PER_CU)
+#endif
+        kernel_gemm_xdlops_v2r4(const FloatAB* __restrict__ p_a_grid,
+                                const FloatAB* __restrict__ p_b_grid,
+                                FloatC* __restrict__ p_c_grid,
+                                const void CONSTANT* p_a_b_k0_m_k1_grid_desc,
+                                const void CONSTANT* p_b_b_k0_n_k1_grid_desc,
+                                const void CONSTANT* p_c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                                const void CONSTANT* p_c_block_cluster_adaptor)
+{
+    constexpr index_t shared_block_size =
+        GridwiseGemm::GetSharedMemoryNumberOfByte() / sizeof(FloatAB);
+
+    const auto a_b_k0_m_k1_grid_desc = *reinterpret_cast<const ABK0MK1GridDesc*>(
+        cast_pointer_to_generic_address_space(p_a_b_k0_m_k1_grid_desc));
+    const auto b_b_k0_n_k1_grid_desc = *reinterpret_cast<const BBK0NK1GridDesc*>(
+        cast_pointer_to_generic_address_space(p_b_b_k0_n_k1_grid_desc));
+    const auto c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc =
+        *reinterpret_cast<const CM0N0M1N1M2M3M4N2GridDesc*>(
+            cast_pointer_to_generic_address_space(p_c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc));
+    const auto c_block_cluster_adaptor = *reinterpret_cast<const CBlockClusterAdaptor*>(
+        cast_pointer_to_generic_address_space(p_c_block_cluster_adaptor));
+
+    __shared__ FloatAB p_shared_block[shared_block_size];
+
+    GridwiseGemm::Run(p_a_grid,
+                      p_b_grid,
+                      p_c_grid,
+                      p_shared_block,
+                      a_b_k0_m_k1_grid_desc,
+                      b_b_k0_n_k1_grid_desc,
+                      c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                      c_block_cluster_adaptor);
+}
+#endif
+
+template <index_t BlockSize,
+          typename FloatAB,
+          typename FloatAcc,
+          typename FloatC,
+          InMemoryDataOperationEnum_t CGlobalMemoryDataOperation,
+          typename ABK0MK1GridDesc,
+          typename BBK0NK1GridDesc,
+          typename CMNGridDesc,
+          index_t MPerBlock,
+          index_t NPerBlock,
+          index_t KPerBlock,
+          index_t MPerXDL,
+          index_t NPerXDL,
+          index_t K1Value,
+          index_t MRepeat,
+          index_t NRepeat,
+          typename ABlockTransferThreadSliceLengths_K0_M_K1,
+          typename ABlockTransferThreadClusterLengths_K0_M_K1,
+          typename ABlockTransferThreadClusterArrangeOrder,
+          typename ABlockTransferSrcAccessOrder,
+          index_t ABlockTransferSrcVectorDim,
+          index_t ABlockTransferSrcScalarPerVector,
+          index_t ABlockTransferDstScalarPerVector_K1,
+          bool AThreadTransferSrcResetCoordinateAfterRun,
+          typename BBlockTransferThreadSliceLengths_K0_N_K1,
+          typename BBlockTransferThreadClusterLengths_K0_N_K1,
+          typename BBlockTransferThreadClusterArrangeOrder,
+          typename BBlockTransferSrcAccessOrder,
+          index_t BBlockTransferSrcVectorDim,
+          index_t BBlockTransferSrcScalarPerVector,
+          index_t BBlockTransferDstScalarPerVector_K1,
+          bool BThreadTransferSrcResetCoordinateAfterRun,
+          typename CThreadTransferSrcDstAccessOrder,
+          index_t CThreadTransferSrcDstVectorDim,
+          index_t CThreadTransferDstScalarPerVector,
+          typename AGridStepHacks,
+          typename BGridStepHacks,
+          typename CGridStepHacks,
+          typename AGridMoveSliceWindowStepHacks,
+          typename BGridMoveSliceWindowStepHacks,
+          bool CAccessOrderMRepeatNRepeat,
+          bool ABlockLdsExtraM,
+          bool BBlockLdsExtraN>
+struct GridwiseGemm_bk0mk1_bk0nk1_mn_xdlops_v2r4
+{
+    static constexpr auto I0 = Number<0>{};
+    static constexpr auto I1 = Number<1>{};
+    static constexpr auto I2 = Number<2>{};
+    static constexpr auto I3 = Number<3>{};
+    static constexpr auto I4 = Number<4>{};
+    static constexpr auto I5 = Number<5>{};
+    static constexpr auto I6 = Number<6>{};
+    static constexpr auto I7 = Number<7>{};
+
+    // K1 should be Number<...>
+    static constexpr auto K1 = Number<K1Value>{};
+
+    __host__ __device__ static constexpr index_t GetSharedMemoryNumberOfByte()
+    {
+        constexpr auto max_lds_align = K1;
+
+        // A matrix in LDS memory, dst of blockwise copy
+        constexpr auto a_k0_m_k1_block_desc = [&]() {
+            if constexpr(ABlockLdsExtraM)
+            {
+                return make_naive_tensor_descriptor(
+                    make_tuple(Number<KPerBlock>{}, Number<MPerBlock>{}, K1),
+                    make_tuple(Number<MPerBlock + 1>{} * K1, K1, I1));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor_aligned(
+                    make_tuple(Number<KPerBlock>{}, Number<MPerBlock>{}, K1), max_lds_align);
+            }
+        }();
+
+        // B matrix in LDS memory, dst of blockwise copy
+        constexpr auto b_k0_n_k1_block_desc = [&]() {
+            if constexpr(BBlockLdsExtraN)
+            {
+                return make_naive_tensor_descriptor(
+                    make_tuple(Number<KPerBlock>{}, Number<NPerBlock>{}, K1),
+                    make_tuple(Number<NPerBlock + 1>{} * K1, K1, I1));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor_aligned(
+                    make_tuple(Number<KPerBlock>{}, Number<NPerBlock>{}, K1), max_lds_align);
+            }
+        }();
+
+        // LDS allocation for A and B: be careful of alignment
+        constexpr auto a_block_space_size =
+            math::integer_least_multiple(a_k0_m_k1_block_desc.GetElementSpaceSize(), max_lds_align);
+
+        constexpr auto b_block_space_size =
+            math::integer_least_multiple(b_k0_n_k1_block_desc.GetElementSpaceSize(), max_lds_align);
+
+        return (a_block_space_size + b_block_space_size) * sizeof(FloatAB);
+    }
+
+    // block_id to matrix tile idx (m0, n0) mapping are controlled by {M01, N01}
+    __host__ __device__ static constexpr bool
+    CheckValidity(const ABK0MK1GridDesc& a_b_k0_m_k1_grid_desc,
+                  const BBK0NK1GridDesc& b_b_k0_n_k1_grid_desc,
+                  const CMNGridDesc& c_m_n_grid_desc,
+                  index_t M01,
+                  index_t N01)
+    {
+        static_assert(is_known_at_compile_time<remove_cv_t<decltype(K1)>>::value,
+                      "wrong! K1 need to be known at compile-time");
+
+        static_assert((MPerBlock % (MPerXDL * MRepeat) == 0) &&
+                          (NPerBlock % (NRepeat * NPerXDL)) == 0,
+                      "Invalid tuning param!");
+
+        const auto M      = a_b_k0_m_k1_grid_desc.GetLength(I2);
+        const auto N      = b_b_k0_n_k1_grid_desc.GetLength(I2);
+        const auto K0     = a_b_k0_m_k1_grid_desc.GetLength(I1);
+        const auto KBatch = a_b_k0_m_k1_grid_desc.GetLength(I0);
+
+        if(!(M == c_m_n_grid_desc.GetLength(I0) && N == c_m_n_grid_desc.GetLength(I1) &&
+             K0 == b_b_k0_n_k1_grid_desc.GetLength(I1) &&
+             K1 == a_b_k0_m_k1_grid_desc.GetLength(I3) &&
+             K1 == b_b_k0_n_k1_grid_desc.GetLength(I3) &&
+             KBatch == b_b_k0_n_k1_grid_desc.GetLength(I0)))
+            return false;
+
+        if(!(M % MPerBlock == 0 && N % NPerBlock == 0 && K0 % KPerBlock == 0))
+            return false;
+
+        // check M01, N01
+        constexpr auto M1 = Number<MPerBlock>{};
+        constexpr auto N1 = Number<NPerBlock>{};
+
+        const auto M0 = M / M1;
+        const auto N0 = N / N1;
+
+        if(!(M0 % M01 == 0 && N0 % N01 == 0))
+            return false;
+
+        // TODO: also check validity of all components (blockwise-copy, threadwise-copy, etc)
+        return true;
+    }
+
+    __host__ __device__ static constexpr index_t
+    CalculateGridSize(const CMNGridDesc& c_m_n_grid_desc, index_t KBatch)
+    {
+        const auto M = c_m_n_grid_desc.GetLength(I0);
+        const auto N = c_m_n_grid_desc.GetLength(I1);
+
+        const index_t grid_size = (M / MPerBlock) * (N / NPerBlock) * KBatch;
+
+        return grid_size;
+    }
+
+    __host__ __device__ static constexpr auto
+    MakeCM0N0M1N1M2M3M4N2GridDescriptor(const CMNGridDesc& c_m_n_grid_desc)
+    {
+        constexpr auto max_lds_align = K1;
+
+        // A matrix in LDS memory, dst of blockwise copy
+        constexpr auto a_k0_m_k1_block_desc = [&]() {
+            if constexpr(ABlockLdsExtraM)
+            {
+                return make_naive_tensor_descriptor(
+                    make_tuple(Number<KPerBlock>{}, Number<MPerBlock>{}, K1),
+                    make_tuple(Number<MPerBlock + 1>{} * K1, K1, I1));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor_aligned(
+                    make_tuple(Number<KPerBlock>{}, Number<MPerBlock>{}, K1), max_lds_align);
+            }
+        }();
+
+        // B matrix in LDS memory, dst of blockwise copy
+        constexpr auto b_k0_n_k1_block_desc = [&]() {
+            if constexpr(BBlockLdsExtraN)
+            {
+                return make_naive_tensor_descriptor(
+                    make_tuple(Number<KPerBlock>{}, Number<NPerBlock>{}, K1),
+                    make_tuple(Number<NPerBlock + 1>{} * K1, K1, I1));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor_aligned(
+                    make_tuple(Number<KPerBlock>{}, Number<NPerBlock>{}, K1), max_lds_align);
+            }
+        }();
+
+        using BlockwiseGemm =
+            BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1<BlockSize,
+                                                                FloatAB,
+                                                                FloatAcc,
+                                                                decltype(a_k0_m_k1_block_desc),
+                                                                decltype(b_k0_n_k1_block_desc),
+                                                                MPerXDL,
+                                                                NPerXDL,
+                                                                MRepeat,
+                                                                NRepeat,
+                                                                K1>;
+
+        return BlockwiseGemm::MakeCM0N0M1N1M2M3M4N2GridDescriptor(c_m_n_grid_desc);
+    }
+
+    // return block_id to C matrix tile idx (m0, n0) mapping
+    __host__ __device__ static constexpr auto MakeCBlockClusterAdaptor(
+        const CMNGridDesc& c_m_n_grid_desc, index_t M01, index_t N01, index_t KBatch)
+    {
+        const auto M = c_m_n_grid_desc.GetLength(I0);
+        const auto N = c_m_n_grid_desc.GetLength(I1);
+
+        constexpr auto M1 = Number<MPerBlock>{};
+        constexpr auto N1 = Number<NPerBlock>{};
+
+        const auto M0 = M / M1;
+        const auto N0 = N / N1;
+
+        const auto M00 = M0 / M01;
+        const auto N00 = N0 / N01;
+
+        const auto kbatch_m00_m01_n00_n01_to_m0_n0_block_cluster_adaptor =
+            make_single_stage_tensor_adaptor(
+                make_tuple(make_pass_through_transform(KBatch),
+                           make_unmerge_transform(make_tuple(M00, M01)),
+                           make_unmerge_transform(make_tuple(N00, N01))),
+                make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
+                make_tuple(Sequence<0>{}, Sequence<1, 3>{}, Sequence<2, 4>{}));
+
+        const auto c_blockid_to_kbatch_m00_m01_n00_n01_block_cluster_adaptor =
+            make_single_stage_tensor_adaptor(
+                make_tuple(make_merge_transform(make_tuple(KBatch, M00, N00, M01, N01))),
+                make_tuple(Sequence<0, 1, 2, 3, 4>{}),
+                make_tuple(Sequence<0>{}));
+
+        const auto c_blockid_to_kbatch_m0_n0_block_cluster_adaptor =
+            chain_tensor_adaptors(kbatch_m00_m01_n00_n01_to_m0_n0_block_cluster_adaptor,
+                                  c_blockid_to_kbatch_m00_m01_n00_n01_block_cluster_adaptor);
+
+        return c_blockid_to_kbatch_m0_n0_block_cluster_adaptor;
+    }
+
+    using CM0N0M1N1M2M3M4N2GridDesc = decltype(MakeCM0N0M1N1M2M3M4N2GridDescriptor(CMNGridDesc{}));
+    using CBlockClusterAdaptor      = decltype(MakeCBlockClusterAdaptor(CMNGridDesc{}, 1, 1, 1));
+
+    __device__ static void Run(const FloatAB* __restrict__ p_a_grid,
+                               const FloatAB* __restrict__ p_b_grid,
+                               FloatC* __restrict__ p_c_grid,
+                               FloatAB* __restrict__ p_shared_block,
+                               const ABK0MK1GridDesc& a_b_k0_m_k1_grid_desc,
+                               const BBK0NK1GridDesc& b_b_k0_n_k1_grid_desc,
+                               const CM0N0M1N1M2M3M4N2GridDesc& c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                               const CBlockClusterAdaptor& c_block_cluster_adaptor)
+    {
+        const auto a_grid_buf = make_dynamic_buffer<AddressSpaceEnum_t::Global>(
+            p_a_grid, a_b_k0_m_k1_grid_desc.GetElementSpaceSize());
+        const auto b_grid_buf = make_dynamic_buffer<AddressSpaceEnum_t::Global>(
+            p_b_grid, b_b_k0_n_k1_grid_desc.GetElementSpaceSize());
+        auto c_grid_buf = make_dynamic_buffer<AddressSpaceEnum_t::Global>(
+            p_c_grid, c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc.GetElementSpaceSize());
+
+        const auto K0 = a_b_k0_m_k1_grid_desc.GetLength(I1);
+
+        // divide block work by [M, N]
+        const auto block_work_idx =
+            c_block_cluster_adaptor.CalculateBottomIndex(make_multi_index(get_block_1d_id()));
+
+        const index_t k_batch_id = block_work_idx[I0];
+        // HACK: this force m/n_block_data_idx_on_grid into SGPR
+        const index_t m_block_data_idx_on_grid =
+            __builtin_amdgcn_readfirstlane(block_work_idx[I1] * MPerBlock);
+
+        const index_t n_block_data_idx_on_grid =
+            __builtin_amdgcn_readfirstlane(block_work_idx[I2] * NPerBlock);
+
+        // lds max alignment
+        constexpr auto max_lds_align = K1;
+
+        // A matrix in LDS memory, dst of blockwise copy
+        constexpr auto a_k0_m_k1_block_desc = [&]() {
+            if constexpr(ABlockLdsExtraM)
+            {
+                return make_naive_tensor_descriptor(
+                    make_tuple(Number<KPerBlock>{}, Number<MPerBlock>{}, K1),
+                    make_tuple(Number<MPerBlock + 1>{} * K1, K1, I1));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor_aligned(
+                    make_tuple(Number<KPerBlock>{}, Number<MPerBlock>{}, K1), max_lds_align);
+            }
+        }();
+
+        constexpr auto a_b_k0_m_k1_block_desc = [&]() {
+            if constexpr(ABlockLdsExtraM)
+            {
+                return make_naive_tensor_descriptor(
+                    make_tuple(Number<1>{}, Number<KPerBlock>{}, Number<MPerBlock>{}, K1),
+                    make_tuple(Number<KPerBlock>{} * Number<MPerBlock + 1>{} * K1,
+                               Number<MPerBlock + 1>{} * K1,
+                               K1,
+                               I1));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor_aligned(
+                    make_tuple(Number<1>{}, Number<KPerBlock>{}, Number<MPerBlock>{}, K1),
+                    max_lds_align);
+            }
+        }();
+        // B matrix in LDS memory, dst of blockwise copy
+        constexpr auto b_k0_n_k1_block_desc = [&]() {
+            if constexpr(BBlockLdsExtraN)
+            {
+                return make_naive_tensor_descriptor(
+                    make_tuple(Number<KPerBlock>{}, Number<NPerBlock>{}, K1),
+                    make_tuple(Number<NPerBlock + 1>{} * K1, K1, I1));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor_aligned(
+                    make_tuple(Number<KPerBlock>{}, Number<NPerBlock>{}, K1), max_lds_align);
+            }
+        }();
+
+        constexpr auto b_b_k0_n_k1_block_desc = [&]() {
+            if constexpr(BBlockLdsExtraN)
+            {
+                return make_naive_tensor_descriptor(
+                    make_tuple(Number<1>{}, Number<KPerBlock>{}, Number<NPerBlock>{}, K1),
+                    make_tuple(Number<KPerBlock>{} * Number<NPerBlock + 1>{} * K1,
+                               Number<NPerBlock + 1>{} * K1,
+                               K1,
+                               I1));
+            }
+            else
+            {
+                return make_naive_tensor_descriptor_aligned(
+                    make_tuple(Number<1>{}, Number<KPerBlock>{}, Number<NPerBlock>{}, K1),
+                    max_lds_align);
+            }
+        }();
+        // A matrix blockwise copy
+        auto a_blockwise_copy =
+            BlockwiseTensorSliceTransfer_v4<BlockSize,
+                                            InMemoryDataOperationEnum_t::Set,
+                                            Sequence<1, KPerBlock, MPerBlock, K1>,
+                                            ABlockTransferThreadSliceLengths_K0_M_K1,
+                                            ABlockTransferThreadClusterLengths_K0_M_K1,
+                                            ABlockTransferThreadClusterArrangeOrder,
+                                            FloatAB,
+                                            FloatAB,
+                                            decltype(a_b_k0_m_k1_grid_desc),
+                                            decltype(a_b_k0_m_k1_block_desc),
+                                            ABlockTransferSrcAccessOrder,
+                                            Sequence<0, 2, 1, 3>,
+                                            ABlockTransferSrcVectorDim,
+                                            3,
+                                            ABlockTransferSrcScalarPerVector,
+                                            ABlockTransferDstScalarPerVector_K1,
+                                            1,
+                                            1,
+                                            AThreadTransferSrcResetCoordinateAfterRun,
+                                            true>(
+                a_b_k0_m_k1_grid_desc,
+                make_multi_index(k_batch_id, 0, m_block_data_idx_on_grid, 0),
+                a_b_k0_m_k1_block_desc,
+                make_multi_index(0, 0, 0, 0));
+
+        // B matrix blockwise copy
+        auto b_blockwise_copy =
+            BlockwiseTensorSliceTransfer_v4<BlockSize,
+                                            InMemoryDataOperationEnum_t::Set,
+                                            Sequence<1, KPerBlock, NPerBlock, K1>,
+                                            BBlockTransferThreadSliceLengths_K0_N_K1,
+                                            BBlockTransferThreadClusterLengths_K0_N_K1,
+                                            BBlockTransferThreadClusterArrangeOrder,
+                                            FloatAB,
+                                            FloatAB,
+                                            decltype(b_b_k0_n_k1_grid_desc),
+                                            decltype(b_b_k0_n_k1_block_desc),
+                                            BBlockTransferSrcAccessOrder,
+                                            Sequence<0, 2, 1, 3>,
+                                            BBlockTransferSrcVectorDim,
+                                            3,
+                                            BBlockTransferSrcScalarPerVector,
+                                            BBlockTransferDstScalarPerVector_K1,
+                                            1,
+                                            1,
+                                            BThreadTransferSrcResetCoordinateAfterRun,
+                                            true>(
+                b_b_k0_n_k1_grid_desc,
+                make_multi_index(k_batch_id, 0, n_block_data_idx_on_grid, 0),
+                b_b_k0_n_k1_block_desc,
+                make_multi_index(0, 0, 0, 0));
+
+        // GEMM definition
+        //   c_mtx += transpose(a_mtx) * b_mtx
+        //     a_mtx[KPerBlock, MPerBlock] is in LDS
+        //     b_mtx[KPerBlock, NPerBlock] is in LDS
+        //     c_mtx[MPerBlock, NPerBlock] is distributed among threads, and saved in
+        //       register
+        // sanity check
+
+        auto blockwise_gemm =
+            BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1<BlockSize,
+                                                                FloatAB,
+                                                                FloatAcc,
+                                                                decltype(a_k0_m_k1_block_desc),
+                                                                decltype(b_k0_n_k1_block_desc),
+                                                                MPerXDL,
+                                                                NPerXDL,
+                                                                MRepeat,
+                                                                NRepeat,
+                                                                K1>{};
+
+        auto c_thread_buf = blockwise_gemm.GetCThreadBuffer();
+
+        // LDS allocation for A and B: be careful of alignment
+        constexpr auto a_block_space_size =
+            math::integer_least_multiple(a_k0_m_k1_block_desc.GetElementSpaceSize(), max_lds_align);
+
+        FloatAB* p_a_block = p_shared_block;
+        FloatAB* p_b_block = p_shared_block + a_block_space_size;
+
+        constexpr auto a_block_slice_copy_step = make_multi_index(0, KPerBlock, 0, 0);
+        constexpr auto b_block_slice_copy_step = make_multi_index(0, KPerBlock, 0, 0);
+
+        // hack to control index calculation when iterating over A and B matrix for threadwise copy
+        constexpr auto a_k0_m_k1_grid_step_hacks = AGridStepHacks{};
+        constexpr auto b_k0_n_k1_grid_step_hacks = BGridStepHacks{};
+
+        // hack to control index calculation when move slice window for A and B matrix for
+        // threadwise copy
+        constexpr auto a_k0_m_k1_grid_move_slice_window_step_hack = AGridMoveSliceWindowStepHacks{};
+        constexpr auto b_k0_n_k1_grid_move_slice_window_step_hack = BGridMoveSliceWindowStepHacks{};
+
+        auto a_block_buf = make_dynamic_buffer<AddressSpaceEnum_t::Lds>(
+            p_a_block, a_k0_m_k1_block_desc.GetElementSpaceSize());
+        auto b_block_buf = make_dynamic_buffer<AddressSpaceEnum_t::Lds>(
+            p_b_block, b_k0_n_k1_block_desc.GetElementSpaceSize());
+
+        // preload data into LDS
+        {
+            a_blockwise_copy.RunRead(a_b_k0_m_k1_grid_desc, a_grid_buf, a_k0_m_k1_grid_step_hacks);
+            b_blockwise_copy.RunRead(b_b_k0_n_k1_grid_desc, b_grid_buf, b_k0_n_k1_grid_step_hacks);
+
+            a_blockwise_copy.RunWrite(a_b_k0_m_k1_block_desc, a_block_buf);
+            b_blockwise_copy.RunWrite(b_b_k0_n_k1_block_desc, b_block_buf);
+        }
+
+        // main body
+        index_t k_block_data_begin = 0;
+
+        do
+        {
+            a_blockwise_copy.MoveSrcSliceWindow(a_b_k0_m_k1_grid_desc,
+                                                a_block_slice_copy_step,
+                                                a_k0_m_k1_grid_move_slice_window_step_hack);
+            b_blockwise_copy.MoveSrcSliceWindow(b_b_k0_n_k1_grid_desc,
+                                                b_block_slice_copy_step,
+                                                b_k0_n_k1_grid_move_slice_window_step_hack);
+
+            a_blockwise_copy.RunRead(a_b_k0_m_k1_grid_desc, a_grid_buf, a_k0_m_k1_grid_step_hacks);
+
+            block_sync_lds();
+
+            b_blockwise_copy.RunRead(b_b_k0_n_k1_grid_desc, b_grid_buf, b_k0_n_k1_grid_step_hacks);
+
+            blockwise_gemm.Run(a_block_buf, b_block_buf, c_thread_buf);
+
+            block_sync_lds();
+
+            a_blockwise_copy.RunWrite(a_b_k0_m_k1_block_desc, a_block_buf);
+            b_blockwise_copy.RunWrite(b_b_k0_n_k1_block_desc, b_block_buf);
+
+            k_block_data_begin += KPerBlock;
+        } while(k_block_data_begin < (K0 - KPerBlock));
+
+        // tail
+        {
+            block_sync_lds();
+
+            blockwise_gemm.Run(a_block_buf, b_block_buf, c_thread_buf);
+        }
+
+        // output: register to global memory
+        {
+            constexpr auto c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc =
+                blockwise_gemm.GetCM0N0M1N1M2M3M4N2BlockDescriptor();
+
+            constexpr auto M0 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I0);
+            constexpr auto N0 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I1);
+            constexpr auto M1 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I2);
+            constexpr auto N1 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I3);
+            constexpr auto M2 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I4);
+            constexpr auto M3 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I5);
+            constexpr auto M4 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I6);
+            constexpr auto N2 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I7);
+
+            constexpr auto c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc =
+                make_naive_tensor_descriptor_packed(make_tuple(
+                    Number<M0>{}, Number<N0>{}, I1, I1, Number<M2>{}, I1, Number<M4>{}, I1));
+
+            // calculate origin of thread output tensor on global memory
+            //     blockwise GEMM c matrix starting index
+            const auto c_thread_mtx_on_block =
+                blockwise_gemm.CalculateCThreadOriginDataIndex(I0, I0, I0, I0);
+
+            const index_t m_thread_data_on_grid =
+                m_block_data_idx_on_grid + c_thread_mtx_on_block[I0];
+
+            const index_t n_thread_data_on_grid =
+                n_block_data_idx_on_grid + c_thread_mtx_on_block[I1];
+
+            constexpr auto c_m0_n0_m1_n1_m2_m3_m4_n2_grid_tensor_step_hacks = CGridStepHacks{};
+
+            const auto m_thread_data_on_grid_to_m0_m1_m2_m3_m4_adaptor =
+                make_single_stage_tensor_adaptor(
+                    make_tuple(make_merge_transform(make_tuple(M0, M1, M2, M3, M4))),
+                    make_tuple(Sequence<0, 1, 2, 3, 4>{}),
+                    make_tuple(Sequence<0>{}));
+
+            const auto m_thread_data_on_grid_idx =
+                m_thread_data_on_grid_to_m0_m1_m2_m3_m4_adaptor.CalculateBottomIndex(
+                    make_multi_index(m_thread_data_on_grid));
+
+            const auto n_thread_data_on_grid_to_n0_n1_n2_adaptor = make_single_stage_tensor_adaptor(
+                make_tuple(make_merge_transform(make_tuple(N0, N1, N2))),
+                make_tuple(Sequence<0, 1, 2>{}),
+                make_tuple(Sequence<0>{}));
+
+            const auto n_thread_data_on_grid_idx =
+                n_thread_data_on_grid_to_n0_n1_n2_adaptor.CalculateBottomIndex(
+                    make_multi_index(n_thread_data_on_grid));
+
+            auto c_thread_copy =
+                ThreadwiseTensorSliceTransfer_v1r3<FloatAcc,
+                                                   FloatC,
+                                                   decltype(c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc),
+                                                   decltype(c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc),
+                                                   Sequence<M0, N0, I1, I1, M2, I1, M4, I1>,
+                                                   CThreadTransferSrcDstAccessOrder,
+                                                   CThreadTransferSrcDstVectorDim,
+                                                   CThreadTransferDstScalarPerVector,
+                                                   CGlobalMemoryDataOperation,
+                                                   1,
+                                                   true>{
+
+                    c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                    make_multi_index(m_thread_data_on_grid_idx[I0],
+                                     n_thread_data_on_grid_idx[I0],
+                                     m_thread_data_on_grid_idx[I1],
+                                     n_thread_data_on_grid_idx[I1],
+                                     m_thread_data_on_grid_idx[I2],
+                                     m_thread_data_on_grid_idx[I3],
+                                     m_thread_data_on_grid_idx[I4],
+                                     n_thread_data_on_grid_idx[I2])};
+
+            c_thread_copy.Run(c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc,
+                              make_tuple(I0, I0, I0, I0, I0, I0, I0, I0),
+                              c_thread_buf,
+                              c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                              c_grid_buf,
+                              c_m0_n0_m1_n1_m2_m3_m4_n2_grid_tensor_step_hacks);
+        }
+    }
+}; // namespace ck
+
+} // namespace ck
+#endif

composable_kernel/include/tensor_operation/threadwise_contraction_dlops.hpp

@@ -55,18 +55,15 @@ struct ThreadwiseGemmDlops_km0m1_kn0n1_m0m1n0n1
                                CBuffer& c_buf,
                                COriginIdx)
     {
-        static_assert(
-            is_known_at_compile_time<remove_cv_t<remove_reference_t<AOriginIdx>>>::value &&
-                is_known_at_compile_time<remove_cv_t<remove_reference_t<BOriginIdx>>>::value &&
-                is_known_at_compile_time<remove_cv_t<remove_reference_t<COriginIdx>>>::value,
+        static_assert(is_known_at_compile_time<remove_cvref_t<AOriginIdx>>::value &&
+                          is_known_at_compile_time<remove_cvref_t<BOriginIdx>>::value &&
+                          is_known_at_compile_time<remove_cvref_t<COriginIdx>>::value,
                       "wrong! AOriginIdx, BOriginIdx, COringinIdx should be known at compile-time");
 
-        static_assert(is_same<remove_cv_t<remove_reference_t<typename ABuffer::type>>,
-                              remove_cv_t<remove_reference_t<FloatA>>>::value &&
-                      is_same<remove_cv_t<remove_reference_t<typename BBuffer::type>>,
-                              remove_cv_t<remove_reference_t<FloatB>>>::value &&
-                      is_same<remove_cv_t<remove_reference_t<typename CBuffer::type>>,
-                              remove_cv_t<remove_reference_t<FloatC>>>::value &&
+        static_assert(
+            is_same<remove_cvref_t<typename ABuffer::type>, remove_cvref_t<FloatA>>::value &&
+            is_same<remove_cvref_t<typename BBuffer::type>, remove_cvref_t<FloatB>>::value &&
+            is_same<remove_cvref_t<typename CBuffer::type>, remove_cvref_t<FloatC>>::value &&
             "wrong! inconsistent type");
 
         constexpr auto I0 = Number<0>{};
@@ -157,18 +154,15 @@ struct ThreadwiseContractionDlops_A_TK0_TM0_TM1_TK1_B_TK0_TN0_TN1_TK1_C_TM0_TM1_
                                CBuffer& c_buf,
                                COriginIdx)
     {
-        static_assert(
-            is_known_at_compile_time<remove_cv_t<remove_reference_t<AOriginIdx>>>::value &&
-                is_known_at_compile_time<remove_cv_t<remove_reference_t<BOriginIdx>>>::value &&
-                is_known_at_compile_time<remove_cv_t<remove_reference_t<COriginIdx>>>::value,
+        static_assert(is_known_at_compile_time<remove_cvref_t<AOriginIdx>>::value &&
+                          is_known_at_compile_time<remove_cvref_t<BOriginIdx>>::value &&
+                          is_known_at_compile_time<remove_cvref_t<COriginIdx>>::value,
                       "wrong! AOriginIdx, BOriginIdx, COringinIdx should be known at compile-time");
 
-        static_assert(is_same<remove_cv_t<remove_reference_t<typename ABuffer::type>>,
-                              remove_cv_t<remove_reference_t<FloatA>>>::value &&
-                      is_same<remove_cv_t<remove_reference_t<typename BBuffer::type>>,
-                              remove_cv_t<remove_reference_t<FloatB>>>::value &&
-                      is_same<remove_cv_t<remove_reference_t<typename CBuffer::type>>,
-                              remove_cv_t<remove_reference_t<FloatC>>>::value &&
+        static_assert(
+            is_same<remove_cvref_t<typename ABuffer::type>, remove_cvref_t<FloatA>>::value &&
+            is_same<remove_cvref_t<typename BBuffer::type>, remove_cvref_t<FloatB>>::value &&
+            is_same<remove_cvref_t<typename CBuffer::type>, remove_cvref_t<FloatC>>::value &&
             "wrong! inconsistent type");
 
         constexpr auto I0 = Number<0>{};

composable_kernel/include/tensor_operation/threadwise_gemm_dlops_v3.hpp

@@ -41,18 +41,15 @@ struct ThreadwiseGemmDlops_km_kn_mn_v3
                           CDesc::IsKnownAtCompileTime(),
                       "wrong! Desc should be known at compile-time");
 
-        static_assert(
-            is_known_at_compile_time<remove_cv_t<remove_reference_t<AOriginIdx>>>::value &&
-                is_known_at_compile_time<remove_cv_t<remove_reference_t<BOriginIdx>>>::value &&
-                is_known_at_compile_time<remove_cv_t<remove_reference_t<COriginIdx>>>::value,
+        static_assert(is_known_at_compile_time<remove_cvref_t<AOriginIdx>>::value &&
+                          is_known_at_compile_time<remove_cvref_t<BOriginIdx>>::value &&
+                          is_known_at_compile_time<remove_cvref_t<COriginIdx>>::value,
                       "wrong! AOriginIdx, BOriginIdx, COringinIdx should be known at compile-time");
 
-        static_assert(is_same<remove_cv_t<remove_reference_t<typename ABuffer::type>>,
-                              remove_cv_t<remove_reference_t<FloatA>>>::value &&
-                      is_same<remove_cv_t<remove_reference_t<typename BBuffer::type>>,
-                              remove_cv_t<remove_reference_t<FloatB>>>::value &&
-                      is_same<remove_cv_t<remove_reference_t<typename CBuffer::type>>,
-                              remove_cv_t<remove_reference_t<FloatC>>>::value &&
+        static_assert(
+            is_same<remove_cvref_t<typename ABuffer::type>, remove_cvref_t<FloatA>>::value &&
+            is_same<remove_cvref_t<typename BBuffer::type>, remove_cvref_t<FloatB>>::value &&
+            is_same<remove_cvref_t<typename CBuffer::type>, remove_cvref_t<FloatC>>::value &&
             "wrong! inconsistent type");
 
         constexpr auto I0 = Number<0>{};

composable_kernel/include/tensor_operation/threadwise_tensor_slice_set.hpp

@@ -30,11 +30,11 @@ struct ThreadwiseTensorSliceSet_v1
 
         static_assert(Buffer::IsStaticBuffer(), "wrong! DstBuffer need to be StaticBuffer");
 
-        static_assert(is_known_at_compile_time<remove_cv_t<remove_reference_t<OriginIdx>>>::value,
+        static_assert(is_known_at_compile_time<remove_cvref_t<OriginIdx>>::value,
                       "wrong! OriginIdx need to be known at compile-time");
 
         // Desc is known at compile-time
-        constexpr auto desc = remove_cv_t<remove_reference_t<Desc>>{};
+        constexpr auto desc = remove_cvref_t<Desc>{};
 
         // OriginIdx is known at compile-time
         constexpr auto origin_idx = to_multi_index(OriginIdx{});

composable_kernel/include/tensor_operation/threadwise_tensor_slice_transfer.hpp

@@ -95,18 +95,13 @@ struct ThreadwiseTensorSliceTransfer_v1r3
         static_assert(SrcDesc::IsKnownAtCompileTime(),
                       "wrong! SrcDesc need to known at compile-time");
 
-        static_assert(
-            is_known_at_compile_time<remove_cv_t<remove_reference_t<SrcSliceOriginIdx>>>::value,
+        static_assert(is_known_at_compile_time<remove_cvref_t<SrcSliceOriginIdx>>::value,
                       "wrong! SrcSliceOrigin need to known at compile-time");
 
         static_assert(SrcBuffer::IsStaticBuffer(), "wrong! SrcBuffer need to be StaticBuffer");
 
-        // static_assert(is_same<remove_cv_t<remove_reference_t<typename SrcBuffer::type>>,
-        // remove_cv_t<remove_reference_t<SrcData>>>::value,
-        //"wrong! SrcBuffer data type is wrong");
-
         // SrcDesc and src_slice_origin_idx are known at compile-time
-        constexpr auto src_desc             = remove_cv_t<remove_reference_t<SrcDesc>>{};
+        constexpr auto src_desc             = remove_cvref_t<SrcDesc>{};
         constexpr auto src_slice_origin_idx = to_multi_index(SrcSliceOriginIdx{});
 
         constexpr auto I0 = Number<0>{};
@@ -208,10 +203,20 @@ struct ThreadwiseTensorSliceTransfer_v1r3
                 coordinate_has_valid_offset_assuming_visible_index_is_valid(dst_desc, dst_coord_);
 
             // copy data from dst_vector into dst_buf
+            if constexpr(DstInMemOp == InMemoryDataOperationEnum_t::Set)
+            {
                 dst_buf.template Set<dst_vector_t>(
                     dst_coord_.GetOffset(),
                     is_dst_valid,
                     dst_vector.template AsType<dst_vector_t>()[Number<0>{}]);
+            }
+            else if constexpr(DstInMemOp == InMemoryDataOperationEnum_t::AtomicAdd)
+            {
+                dst_buf.template AtomicAdd<dst_vector_t>(
+                    dst_coord_.GetOffset(),
+                    is_dst_valid,
+                    dst_vector.template AsType<dst_vector_t>()[Number<0>{}]);
+            }
 
             constexpr auto move_on_dim = [&]() constexpr
             {
@@ -392,7 +397,7 @@ struct ThreadwiseTensorSliceTransfer_v2
                       "wrong! SrcDesc need to known at compile-time");
     }
 
-    __device__ void SetDstSliceOrigin(const SrcDesc& src_desc, const Index& src_slice_origin_idx)
+    __device__ void SetSrcSliceOrigin(const SrcDesc& src_desc, const Index& src_slice_origin_idx)
     {
         src_coord_ = make_tensor_coordinate(src_desc, src_slice_origin_idx);
     }
@@ -411,16 +416,15 @@ struct ThreadwiseTensorSliceTransfer_v2
         static_assert(DstDesc::IsKnownAtCompileTime(),
                       "wrong! DstDesc need to known at compile-time");
 
-        static_assert(
-            is_known_at_compile_time<remove_cv_t<remove_reference_t<DstSliceOriginIdx>>>::value,
+        static_assert(is_known_at_compile_time<remove_cvref_t<DstSliceOriginIdx>>::value,
                       "wrong! DstSliceOrigin need to known at compile-time");
 
-        static_assert(is_same<remove_cv_t<remove_reference_t<typename DstBuffer::type>>,
-                              remove_cv_t<remove_reference_t<DstData>>>::value &&
+        static_assert(
+            is_same<remove_cvref_t<typename DstBuffer::type>, remove_cvref_t<DstData>>::value &&
             "wrong! inconsistent type");
 
         // DstDesc and dst_slice_origin_idx are known at compile-time
-        constexpr auto dst_desc             = remove_cv_t<remove_reference_t<DstDesc>>{};
+        constexpr auto dst_desc             = remove_cvref_t<DstDesc>{};
         constexpr auto dst_slice_origin_idx = DstSliceOriginIdx{};
 
         constexpr auto I0 = Number<0>{};
@@ -729,8 +733,8 @@ struct ThreadwiseTensorSliceTransfer_v3
                           SrcBuffer::GetAddressSpace() == AddressSpaceEnum_t::Lds,
                       "wrong!");
 
-        static_assert(is_same<remove_cv_t<remove_reference_t<typename SrcBuffer::type>>,
-                              remove_cv_t<remove_reference_t<SrcData>>>::value,
+        static_assert(
+            is_same<remove_cvref_t<typename SrcBuffer::type>, remove_cvref_t<SrcData>>::value,
             "wrong! SrcBuffer and SrcData data type are inconsistent");
 
         constexpr auto I0 = Number<0>{};
@@ -886,8 +890,8 @@ struct ThreadwiseTensorSliceTransfer_v3
                           DstBuffer::GetAddressSpace() == AddressSpaceEnum_t::Lds,
                       "wrong!");
 
-        static_assert(is_same<remove_cv_t<remove_reference_t<typename DstBuffer::type>>,
-                              remove_cv_t<remove_reference_t<DstData>>>::value,
+        static_assert(
+            is_same<remove_cvref_t<typename DstBuffer::type>, remove_cvref_t<DstData>>::value,
             "wrong! SrcBuffer or DstBuffer data type is wrong");
 
         constexpr auto I0 = Number<0>{};
@@ -1303,24 +1307,21 @@ struct ThreadwiseTensorSliceTransfer_v4
         static_assert(SrcDesc::IsKnownAtCompileTime() && DstDesc::IsKnownAtCompileTime(),
                       "wrong! SrcDesc and DstDesc need to known at compile-time");
 
-        static_assert(is_same<remove_cv_t<remove_reference_t<typename SrcBuffer::type>>,
-                              remove_cv_t<remove_reference_t<SrcData>>>::value &&
-                          is_same<remove_cv_t<remove_reference_t<typename DstBuffer::type>>,
-                                  remove_cv_t<remove_reference_t<DstData>>>::value,
+        static_assert(
+            is_same<remove_cvref_t<typename SrcBuffer::type>, remove_cvref_t<SrcData>>::value &&
+                is_same<remove_cvref_t<typename DstBuffer::type>, remove_cvref_t<DstData>>::value,
             "wrong! SrcBuffer or DstBuffer data type is wrong");
 
         static_assert(DstBuffer::IsStaticBuffer(), "wrong! DstBuffer need to be StaticBuffer");
 
-        static_assert(
-            is_known_at_compile_time<
-                remove_cv_t<remove_reference_t<SrcRefToOriginDisplacement>>>::value &&
-                is_known_at_compile_time<remove_cv_t<remove_reference_t<DstOriginIdx>>>::value,
+        static_assert(is_known_at_compile_time<remove_cvref_t<SrcRefToOriginDisplacement>>::value &&
+                          is_known_at_compile_time<remove_cvref_t<DstOriginIdx>>::value,
                       "wrong! SrcOriginToRefDistance and DstOriginToRefDistance need to be known "
                       "at compile-time");
 
         // SrcDesc and DstDesc are known at compile-time
-        constexpr auto src_desc = remove_cv_t<remove_reference_t<SrcDesc>>{};
-        constexpr auto dst_desc = remove_cv_t<remove_reference_t<DstDesc>>{};
+        constexpr auto src_desc = remove_cvref_t<SrcDesc>{};
+        constexpr auto dst_desc = remove_cvref_t<DstDesc>{};
 
         // SrcOriginToRefDisttance and DstOriginToRefDistance are known at compile-time
         constexpr auto src_ref_to_origin_disp_idx = to_multi_index(SrcRefToOriginDisplacement{});

composable_kernel/include/tensor_operation/threadwise_tensor_slice_transfer_v2.hpp

@@ -80,8 +80,8 @@ struct ThreadwiseTensorSliceTransfer_v3r1
                           SrcBuffer::GetAddressSpace() == AddressSpaceEnum_t::Lds,
                       "wrong!");
 
-        static_assert(is_same<remove_cv_t<remove_reference_t<typename SrcBuffer::type>>,
-                              remove_cv_t<remove_reference_t<SrcData>>>::value,
+        static_assert(
+            is_same<remove_cvref_t<typename SrcBuffer::type>, remove_cvref_t<SrcData>>::value,
             "wrong! SrcBuffer and SrcData data type are inconsistent");
 
         // tensor descriptor for src_vector
@@ -248,8 +248,8 @@ struct ThreadwiseTensorSliceTransfer_v3r1
                           DstBuffer::GetAddressSpace() == AddressSpaceEnum_t::Lds,
                       "wrong!");
 
-        static_assert(is_same<remove_cv_t<remove_reference_t<typename DstBuffer::type>>,
-                              remove_cv_t<remove_reference_t<DstData>>>::value,
+        static_assert(
+            is_same<remove_cvref_t<typename DstBuffer::type>, remove_cvref_t<DstData>>::value,
             "wrong! SrcBuffer or DstBuffer data type is wrong");
 
         // tensor descriptor for dst_vector
@@ -669,24 +669,21 @@ struct ThreadwiseTensorSliceTransfer_v4r1
         static_assert(SrcDesc::IsKnownAtCompileTime() && DstDesc::IsKnownAtCompileTime(),
                       "wrong! SrcDesc and DstDesc need to known at compile-time");
 
-        static_assert(is_same<remove_cv_t<remove_reference_t<typename SrcBuffer::type>>,
-                              remove_cv_t<remove_reference_t<SrcData>>>::value &&
-                          is_same<remove_cv_t<remove_reference_t<typename DstBuffer::type>>,
-                                  remove_cv_t<remove_reference_t<DstData>>>::value,
+        static_assert(
+            is_same<remove_cvref_t<typename SrcBuffer::type>, remove_cvref_t<SrcData>>::value &&
+                is_same<remove_cvref_t<typename DstBuffer::type>, remove_cvref_t<DstData>>::value,
             "wrong! SrcBuffer or DstBuffer data type is wrong");
 
         static_assert(DstBuffer::IsStaticBuffer(), "wrong! DstBuffer need to be StaticBuffer");
 
-        static_assert(
-            is_known_at_compile_time<
-                remove_cv_t<remove_reference_t<SrcRefToOriginDisplacement>>>::value &&
-                is_known_at_compile_time<remove_cv_t<remove_reference_t<DstOriginIdx>>>::value,
+        static_assert(is_known_at_compile_time<remove_cvref_t<SrcRefToOriginDisplacement>>::value &&
+                          is_known_at_compile_time<remove_cvref_t<DstOriginIdx>>::value,
                       "wrong! SrcOriginToRefDistance and DstOriginToRefDistance need to be known "
                       "at compile-time");
 
         // SrcDesc and DstDesc are known at compile-time
-        constexpr auto src_desc = remove_cv_t<remove_reference_t<SrcDesc>>{};
-        constexpr auto dst_desc = remove_cv_t<remove_reference_t<DstDesc>>{};
+        constexpr auto src_desc = remove_cvref_t<SrcDesc>{};
+        constexpr auto dst_desc = remove_cvref_t<DstDesc>{};
 
         // SrcOriginToRefDisttance and DstOriginToRefDistance are known at compile-time
         constexpr auto src_ref_to_origin_disp_idx = to_multi_index(SrcRefToOriginDisplacement{});

composable_kernel/include/tensor_operation/xdlops_gemm.hpp

@@ -7,21 +7,18 @@
 
 namespace ck {
 
-enum struct mfma_instr
+enum struct MfmaInstr
 {
-    /// fp32
     mfma_f32_32x32x1xf32 = 0,
     mfma_f32_16x16x1xf32,
     mfma_f32_4x4x1xf32,
     mfma_f32_32x32x2xf32, // k reduction
     mfma_f32_16x16x4xf32, // k reduction
-                          /// fp16
     mfma_f32_32x32x4f16,
     mfma_f32_16x16x4f16,
     mfma_f32_4x4x4f16,
     mfma_f32_32x32x8f16,  // k reduction
     mfma_f32_16x16x16f16, // k reduction
-                          /// bfp16
     mfma_f32_32x32x2bf16,
     mfma_f32_16x16x2bf16,
     mfma_f32_4x4x2bf16,
@@ -29,317 +26,245 @@ enum struct mfma_instr
     mfma_f32_16x16x8bf16, // k reduction
 };
 
-template <mfma_instr instr>
-struct mfma_info;
+template <MfmaInstr instr>
+struct mfma_type;
 
 template <>
-struct mfma_info<mfma_instr::mfma_f32_32x32x1xf32>
+struct mfma_type<MfmaInstr::mfma_f32_32x32x1xf32>
 {
     static constexpr index_t group_size          = 4;
-    static constexpr index_t num_groups_blk  = 4;
-    static constexpr index_t num_regs_blk    = group_size * num_groups_blk;
-    static constexpr index_t num_threads_blk = 32;
+    static constexpr index_t num_groups_per_blk  = 4;
+    static constexpr index_t num_regs_per_blk    = 16;
+    static constexpr index_t num_threads_per_blk = 32;
     static constexpr index_t wave_size           = 64;
-    static constexpr index_t num_input_blks  = wave_size / num_threads_blk;
+    static constexpr index_t num_input_blks      = 2;
     static constexpr index_t num_output_blks     = 2;
-    static constexpr index_t num_regs_xdlops = num_regs_blk * num_output_blks;
-    static constexpr index_t m               = 32;
-    static constexpr index_t n               = 32;
-    static constexpr index_t k               = 1;
-    static constexpr index_t cycles          = 64;
-    static constexpr index_t k_base          = 1;
+    static constexpr index_t m_per_blk           = 32;
+    static constexpr index_t n_per_blk           = 32;
+    static constexpr index_t k_per_blk           = 1;
+    static constexpr bool is_k_reduction         = false;
 
-    template <index_t MPerXdlops,
-              index_t NPerXdlops,
-              index_t COffset,
-              class FloatA,
-              class FloatB,
-              class FloatC>
+    template <index_t MPerXdlops, index_t NPerXdlops, class FloatA, class FloatB, class FloatC>
     __device__ void run(const FloatA& a, const FloatB& b, FloatC& reg_c) const
     {
-        intrin_mfma_f32_32x32x1f32<MPerXdlops, NPerXdlops, COffset>::Run(a, b, reg_c);
+        intrin_mfma_f32_32x32x1f32<MPerXdlops, NPerXdlops>::Run(a, b, reg_c);
     }
 };
 
 template <>
-struct mfma_info<mfma_instr::mfma_f32_32x32x2xf32>
+struct mfma_type<MfmaInstr::mfma_f32_32x32x2xf32>
 {
     static constexpr index_t group_size          = 4;
-    static constexpr index_t num_groups_blk  = 4;
-    static constexpr index_t num_regs_blk    = group_size * num_groups_blk;
-    static constexpr index_t num_threads_blk = 32;
+    static constexpr index_t num_groups_per_blk  = 4;
+    static constexpr index_t num_regs_per_blk    = 16;
+    static constexpr index_t num_threads_per_blk = 32;
     static constexpr index_t wave_size           = 64;
-    static constexpr index_t num_input_blks  = wave_size / num_threads_blk;
+    static constexpr index_t num_input_blks      = 2;
     static constexpr index_t num_output_blks     = 1;
-    static constexpr index_t num_regs_xdlops = num_regs_blk * num_output_blks;
-    static constexpr index_t m               = 32;
-    static constexpr index_t n               = 32;
-    static constexpr index_t k               = 2;
-    static constexpr index_t cycles          = 64;
-    static constexpr index_t k_base          = 1;
+    static constexpr index_t m_per_blk           = 32;
+    static constexpr index_t n_per_blk           = 32;
+    static constexpr index_t k_per_blk           = 1;
+    static constexpr bool is_k_reduction         = true;
 
-    template <index_t MPerXdlops,
-              index_t NPerXdlops,
-              index_t COffset,
-              class FloatA,
-              class FloatB,
-              class FloatC>
+    template <index_t MPerXdlops, index_t NPerXdlops, class FloatA, class FloatB, class FloatC>
     __device__ void run(const FloatA& a, const FloatB& b, FloatC& reg_c) const
     {
-        intrin_mfma_f32_32x32x2f32<MPerXdlops, NPerXdlops, COffset>::Run(a, b, reg_c);
+        intrin_mfma_f32_32x32x2f32<MPerXdlops, NPerXdlops>::Run(a, b, reg_c);
     }
 };
 
 template <>
-struct mfma_info<mfma_instr::mfma_f32_16x16x4xf32>
+struct mfma_type<MfmaInstr::mfma_f32_16x16x4xf32>
 {
     static constexpr index_t group_size          = 4;
-    static constexpr index_t num_groups_blk  = 1;
-    static constexpr index_t num_regs_blk    = group_size * num_groups_blk;
-    static constexpr index_t num_threads_blk = 16;
+    static constexpr index_t num_groups_per_blk  = 1;
+    static constexpr index_t num_regs_per_blk    = 4;
+    static constexpr index_t num_threads_per_blk = 16;
     static constexpr index_t wave_size           = 64;
-    static constexpr index_t num_input_blks  = wave_size / num_threads_blk;
+    static constexpr index_t num_input_blks      = 4;
     static constexpr index_t num_output_blks     = 1;
-    static constexpr index_t num_regs_xdlops = num_regs_blk * num_output_blks;
-    static constexpr index_t m               = 16;
-    static constexpr index_t n               = 16;
-    static constexpr index_t k               = 4;
-    static constexpr index_t cycles          = 32;
-    static constexpr index_t k_base          = 1;
+    static constexpr index_t m_per_blk           = 16;
+    static constexpr index_t n_per_blk           = 16;
+    static constexpr index_t k_per_blk           = 1;
+    static constexpr bool is_k_reduction         = true;
 
-    template <index_t MPerXdlops,
-              index_t NPerXdlops,
-              index_t COffset,
-              class FloatA,
-              class FloatB,
-              class FloatC>
+    template <index_t MPerXdlops, index_t NPerXdlops, class FloatA, class FloatB, class FloatC>
     __device__ void run(const FloatA& a, const FloatB& b, FloatC& reg_c) const
     {
-        intrin_mfma_f32_16x16x4f32<MPerXdlops, NPerXdlops, COffset>::Run(a, b, reg_c);
+        intrin_mfma_f32_16x16x4f32<MPerXdlops, NPerXdlops>::Run(a, b, reg_c);
     }
 };
 
 template <>
-struct mfma_info<mfma_instr::mfma_f32_16x16x1xf32>
+struct mfma_type<MfmaInstr::mfma_f32_16x16x1xf32>
 {
     static constexpr index_t group_size          = 4;
-    static constexpr index_t num_groups_blk  = 1;
-    static constexpr index_t num_regs_blk    = group_size * num_groups_blk;
-    static constexpr index_t num_threads_blk = 16;
+    static constexpr index_t num_groups_per_blk  = 1;
+    static constexpr index_t num_regs_per_blk    = 4;
+    static constexpr index_t num_threads_per_blk = 16;
     static constexpr index_t wave_size           = 64;
-    static constexpr index_t num_input_blks  = wave_size / num_threads_blk;
+    static constexpr index_t num_input_blks      = 4;
     static constexpr index_t num_output_blks     = 4;
-    static constexpr index_t num_regs_xdlops = num_regs_blk * num_output_blks;
-    static constexpr index_t m               = 16;
-    static constexpr index_t n               = 16;
-    static constexpr index_t k               = 1;
-    static constexpr index_t cycles          = 32;
-    static constexpr index_t k_base          = 1;
+    static constexpr index_t m_per_blk           = 16;
+    static constexpr index_t n_per_blk           = 16;
+    static constexpr index_t k_per_blk           = 1;
+    static constexpr bool is_k_reduction         = false;
 
-    template <index_t MPerXdlops,
-              index_t NPerXdlops,
-              index_t COffset,
-              class FloatA,
-              class FloatB,
-              class FloatC>
+    template <index_t MPerXdlops, index_t NPerXdlops, class FloatA, class FloatB, class FloatC>
     __device__ void run(const FloatA& a, const FloatB& b, FloatC& reg_c) const
     {
-        intrin_mfma_f32_16x16x1f32<MPerXdlops, NPerXdlops, COffset>::Run(a, b, reg_c);
+        intrin_mfma_f32_16x16x1f32<MPerXdlops, NPerXdlops>::Run(a, b, reg_c);
     }
 };
 
 // treat 4x4x1 as a single-blk 4x64 mfma
 template <>
-struct mfma_info<mfma_instr::mfma_f32_4x4x1xf32>
+struct mfma_type<MfmaInstr::mfma_f32_4x4x1xf32>
 {
     static constexpr index_t group_size          = 4;
-    static constexpr index_t num_groups_blk  = 1;
-    static constexpr index_t num_regs_blk    = group_size * num_groups_blk;
-    static constexpr index_t num_threads_blk = 64;
+    static constexpr index_t num_groups_per_blk  = 1;
+    static constexpr index_t num_regs_per_blk    = 4;
+    static constexpr index_t num_threads_per_blk = 64;
     static constexpr index_t wave_size           = 64;
     static constexpr index_t num_input_blks      = 1;
     static constexpr index_t num_output_blks     = 1;
-    static constexpr index_t num_regs_xdlops = 4;
-    static constexpr index_t m               = 4;
-    static constexpr index_t n               = 64;
-    static constexpr index_t k               = 1;
-    static constexpr index_t cycles          = 8;
-    static constexpr index_t k_base          = 1;
+    static constexpr index_t m_per_blk           = 4;
+    static constexpr index_t n_per_blk           = 64;
+    static constexpr index_t k_per_blk           = 1;
+    static constexpr bool is_k_reduction         = false;
 
-    template <index_t MPerXdlops,
-              index_t NPerXdlops,
-              index_t COffset,
-              class FloatA,
-              class FloatB,
-              class FloatC>
+    template <index_t MPerXdlops, index_t NPerXdlops, class FloatA, class FloatB, class FloatC>
     __device__ void run(const FloatA& a, const FloatB& b, FloatC& reg_c) const
     {
-        intrin_mfma_f32_4x4x1f32<MPerXdlops, NPerXdlops, COffset>::Run(a, b, reg_c);
+        intrin_mfma_f32_4x4x1f32<MPerXdlops, NPerXdlops>::Run(a, b, reg_c);
     }
 };
 
 template <>
-struct mfma_info<mfma_instr::mfma_f32_32x32x4f16>
+struct mfma_type<MfmaInstr::mfma_f32_32x32x4f16>
 {
     static constexpr index_t group_size          = 4;
-    static constexpr index_t num_groups_blk  = 4;
-    static constexpr index_t num_regs_blk    = group_size * num_groups_blk;
-    static constexpr index_t num_threads_blk = 32;
+    static constexpr index_t num_groups_per_blk  = 4;
+    static constexpr index_t num_regs_per_blk    = 16;
+    static constexpr index_t num_threads_per_blk = 32;
     static constexpr index_t wave_size           = 64;
-    static constexpr index_t num_input_blks  = wave_size / num_threads_blk;
+    static constexpr index_t num_input_blks      = 2;
     static constexpr index_t num_output_blks     = 2;
-    static constexpr index_t num_regs_xdlops = num_regs_blk * num_output_blks;
-    static constexpr index_t m               = 32;
-    static constexpr index_t n               = 32;
-    static constexpr index_t k               = 4;
-    static constexpr index_t cycles          = 64;
-    static constexpr index_t k_base          = 4;
+    static constexpr index_t m_per_blk           = 32;
+    static constexpr index_t n_per_blk           = 32;
+    static constexpr index_t k_per_blk           = 4;
+    static constexpr bool is_k_reduction         = false;
 
-    template <index_t MPerXdlops,
-              index_t NPerXdlops,
-              index_t COffset,
-              class FloatA,
-              class FloatB,
-              class FloatC>
+    template <index_t MPerXdlops, index_t NPerXdlops, class FloatA, class FloatB, class FloatC>
     __device__ void run(const FloatA& a, const FloatB& b, FloatC& reg_c) const
     {
-        intrin_mfma_f32_32x32x4f16<MPerXdlops, NPerXdlops, COffset>::Run(a, b, reg_c);
+        intrin_mfma_f32_32x32x4f16<MPerXdlops, NPerXdlops>::Run(a, b, reg_c);
     }
 };
 
 template <>
-struct mfma_info<mfma_instr::mfma_f32_32x32x8f16>
+struct mfma_type<MfmaInstr::mfma_f32_32x32x8f16>
 {
     static constexpr index_t group_size          = 4;
-    static constexpr index_t num_groups_blk  = 4;
-    static constexpr index_t num_regs_blk    = group_size * num_groups_blk;
-    static constexpr index_t num_threads_blk = 32;
+    static constexpr index_t num_groups_per_blk  = 4;
+    static constexpr index_t num_regs_per_blk    = 16;
+    static constexpr index_t num_threads_per_blk = 32;
     static constexpr index_t wave_size           = 64;
-    static constexpr index_t num_input_blks  = wave_size / num_threads_blk;
+    static constexpr index_t num_input_blks      = 2;
     static constexpr index_t num_output_blks     = 1;
-    static constexpr index_t num_regs_xdlops = num_regs_blk * num_output_blks;
-    static constexpr index_t m               = 32;
-    static constexpr index_t n               = 32;
-    static constexpr index_t k               = 8;
-    static constexpr index_t cycles          = 64;
-    static constexpr index_t k_base          = 4;
+    static constexpr index_t m_per_blk           = 32;
+    static constexpr index_t n_per_blk           = 32;
+    static constexpr index_t k_per_blk           = 4;
+    static constexpr bool is_k_reduction         = true;
 
-    template <index_t MPerXdlops,
-              index_t NPerXdlops,
-              index_t COffset,
-              class FloatA,
-              class FloatB,
-              class FloatC>
+    template <index_t MPerXdlops, index_t NPerXdlops, class FloatA, class FloatB, class FloatC>
     __device__ void run(const FloatA& a, const FloatB& b, FloatC& reg_c) const
     {
-        intrin_mfma_f32_32x32x8f16<MPerXdlops, NPerXdlops, COffset>::Run(a, b, reg_c);
+        intrin_mfma_f32_32x32x8f16<MPerXdlops, NPerXdlops>::Run(a, b, reg_c);
     }
 };
 
 template <>
-struct mfma_info<mfma_instr::mfma_f32_16x16x16f16>
+struct mfma_type<MfmaInstr::mfma_f32_16x16x16f16>
 {
     static constexpr index_t group_size          = 4;
-    static constexpr index_t num_groups_blk  = 1;
-    static constexpr index_t num_regs_blk    = group_size * num_groups_blk;
-    static constexpr index_t num_threads_blk = 16;
+    static constexpr index_t num_groups_per_blk  = 1;
+    static constexpr index_t num_regs_per_blk    = 4;
+    static constexpr index_t num_threads_per_blk = 16;
     static constexpr index_t wave_size           = 64;
-    static constexpr index_t num_input_blks  = wave_size / num_threads_blk;
+    static constexpr index_t num_input_blks      = 4;
     static constexpr index_t num_output_blks     = 1;
-    static constexpr index_t num_regs_xdlops = num_regs_blk * num_output_blks;
-    static constexpr index_t m               = 16;
-    static constexpr index_t n               = 16;
-    static constexpr index_t k               = 16;
-    static constexpr index_t cycles          = 32;
-    static constexpr index_t k_base          = 4;
+    static constexpr index_t m_per_blk           = 16;
+    static constexpr index_t n_per_blk           = 16;
+    static constexpr index_t k_per_blk           = 4;
+    static constexpr bool is_k_reduction         = true;
 
-    template <index_t MPerXdlops,
-              index_t NPerXdlops,
-              index_t COffset,
-              class FloatA,
-              class FloatB,
-              class FloatC>
+    template <index_t MPerXdlops, index_t NPerXdlops, class FloatA, class FloatB, class FloatC>
     __device__ void run(const FloatA& a, const FloatB& b, FloatC& reg_c) const
     {
-        intrin_mfma_f32_16x16x16f16<MPerXdlops, NPerXdlops, COffset>::Run(a, b, reg_c);
+        intrin_mfma_f32_16x16x16f16<MPerXdlops, NPerXdlops>::Run(a, b, reg_c);
     }
 };
 
 template <>
-struct mfma_info<mfma_instr::mfma_f32_16x16x4f16>
+struct mfma_type<MfmaInstr::mfma_f32_16x16x4f16>
 {
     static constexpr index_t group_size          = 4;
-    static constexpr index_t num_groups_blk  = 1;
-    static constexpr index_t num_regs_blk    = group_size * num_groups_blk;
-    static constexpr index_t num_threads_blk = 16;
+    static constexpr index_t num_groups_per_blk  = 1;
+    static constexpr index_t num_regs_per_blk    = 4;
+    static constexpr index_t num_threads_per_blk = 16;
     static constexpr index_t wave_size           = 64;
-    static constexpr index_t num_input_blks  = wave_size / num_threads_blk;
+    static constexpr index_t num_input_blks      = 4;
     static constexpr index_t num_output_blks     = 4;
-    static constexpr index_t num_regs_xdlops = num_regs_blk * num_output_blks;
-    static constexpr index_t m               = 16;
-    static constexpr index_t n               = 16;
-    static constexpr index_t k               = 4;
-    static constexpr index_t cycles          = 32;
-    static constexpr index_t k_base          = 4;
+    static constexpr index_t m_per_blk           = 16;
+    static constexpr index_t n_per_blk           = 16;
+    static constexpr index_t k_per_blk           = 4;
+    static constexpr bool is_k_reduction         = false;
 
-    template <index_t MPerXdlops,
-              index_t NPerXdlops,
-              index_t COffset,
-              class FloatA,
-              class FloatB,
-              class FloatC>
+    template <index_t MPerXdlops, index_t NPerXdlops, class FloatA, class FloatB, class FloatC>
     __device__ void run(const FloatA& a, const FloatB& b, FloatC& reg_c) const
     {
-        intrin_mfma_f32_16x16x4f16<MPerXdlops, NPerXdlops, COffset>::Run(a, b, reg_c);
+        intrin_mfma_f32_16x16x4f16<MPerXdlops, NPerXdlops>::Run(a, b, reg_c);
     }
 };
 
 template <>
-struct mfma_info<mfma_instr::mfma_f32_4x4x4f16>
+struct mfma_type<MfmaInstr::mfma_f32_4x4x4f16>
 {
     static constexpr index_t group_size          = 4;
-    static constexpr index_t num_groups_blk  = 1;
-    static constexpr index_t num_regs_blk    = group_size * num_groups_blk;
-    static constexpr index_t num_threads_blk = 64;
+    static constexpr index_t num_groups_per_blk  = 1;
+    static constexpr index_t num_regs_per_blk    = 4;
+    static constexpr index_t num_threads_per_blk = 64;
     static constexpr index_t wave_size           = 64;
     static constexpr index_t num_input_blks      = 1;
     static constexpr index_t num_output_blks     = 1;
-    static constexpr index_t num_regs_xdlops = 4;
-    static constexpr index_t m               = 4;
-    static constexpr index_t n               = 64;
-    static constexpr index_t k               = 4;
-    static constexpr index_t cycles          = 8;
-    static constexpr index_t k_base          = 4;
+    static constexpr index_t m_per_blk           = 4;
+    static constexpr index_t n_per_blk           = 64;
+    static constexpr index_t k_per_blk           = 4;
+    static constexpr bool is_k_reduction         = false;
 
-    template <index_t MPerXdlops,
-              index_t NPerXdlops,
-              index_t COffset,
-              class FloatA,
-              class FloatB,
-              class FloatC>
+    template <index_t MPerXdlops, index_t NPerXdlops, class FloatA, class FloatB, class FloatC>
     __device__ void run(const FloatA& a, const FloatB& b, FloatC& reg_c) const
     {
-        intrin_mfma_f32_4x4x4f16<MPerXdlops, NPerXdlops, COffset>::Run(a, b, reg_c);
+        intrin_mfma_f32_4x4x4f16<MPerXdlops, NPerXdlops>::Run(a, b, reg_c);
     }
 };
 
 #if 0
 template <>
-struct mfma_info<mfma_instr::mfma_f32_32x32x2bf16>
+struct mfma_type<MfmaInstr::mfma_f32_32x32x2bf16>
 {
     static constexpr index_t group_size          = 4;
-    static constexpr index_t num_groups_blk  = 4;
-    static constexpr index_t num_regs_blk    = group_size * num_groups_blk;
-    static constexpr index_t num_threads_blk = 32;
+    static constexpr index_t num_groups_per_blk  = 4;
+    static constexpr index_t num_regs_per_blk    = 16;
+    static constexpr index_t num_threads_per_blk = 32;
     static constexpr index_t wave_size           = 64;
-    static constexpr index_t num_input_blks  = wave_size / num_threads_blk;
+    static constexpr index_t num_input_blks      = 2;
     static constexpr index_t num_output_blks     = 2;
-    static constexpr index_t num_regs_xdlops = num_regs_blk * num_output_blks;
-    static constexpr index_t m               = 32;
-    static constexpr index_t n               = 32;
-    static constexpr index_t k               = 2;
-    static constexpr index_t cycles          = 64;
-    static constexpr index_t k_base          = 2;
+    static constexpr index_t m_per_blk           = 32;
+    static constexpr index_t n_per_blk           = 32;
+    static constexpr index_t k_per_blk           = 2;
+    static constexpr bool is_k_reduction         = false;
 
     template <index_t MPerXdlops,
               index_t NPerXdlops,
@@ -359,21 +284,19 @@ struct mfma_info<mfma_instr::mfma_f32_32x32x2bf16>
 };
 
 template <>
-struct mfma_info<mfma_instr::mfma_f32_32x32x4bf16>
+struct mfma_type<MfmaInstr::mfma_f32_32x32x4bf16>
 {
     static constexpr index_t group_size          = 4;
-    static constexpr index_t num_groups_blk  = 4;
-    static constexpr index_t num_regs_blk    = group_size * num_groups_blk;
-    static constexpr index_t num_threads_blk = 32;
+    static constexpr index_t num_groups_per_blk  = 4;
+    static constexpr index_t num_regs_per_blk    = 16;
+    static constexpr index_t num_threads_per_blk = 32;
     static constexpr index_t wave_size           = 64;
-    static constexpr index_t num_input_blks  = wave_size / num_threads_blk;
+    static constexpr index_t num_input_blks      = 2;
     static constexpr index_t num_output_blks     = 1;
-    static constexpr index_t num_regs_xdlops = num_regs_blk * num_output_blks;
-    static constexpr index_t m               = 32;
-    static constexpr index_t n               = 32;
-    static constexpr index_t k               = 4;
-    static constexpr index_t cycles          = 64;
-    static constexpr index_t k_base          = 2;
+    static constexpr index_t m_per_blk           = 32;
+    static constexpr index_t n_per_blk           = 32;
+    static constexpr index_t k_per_blk           = 2;
+    static constexpr bool is_k_reduction         = true;
 
     template <index_t MPerXdlops,
               index_t NPerXdlops,
@@ -392,21 +315,19 @@ struct mfma_info<mfma_instr::mfma_f32_32x32x4bf16>
 };
 
 template <>
-struct mfma_info<mfma_instr::mfma_f32_16x16x8bf16>
+struct mfma_type<MfmaInstr::mfma_f32_16x16x8bf16>
 {
     static constexpr index_t group_size          = 4;
-    static constexpr index_t num_groups_blk  = 1;
-    static constexpr index_t num_regs_blk    = group_size * num_groups_blk;
-    static constexpr index_t num_threads_blk = 16;
+    static constexpr index_t num_groups_per_blk  = 1;
+    static constexpr index_t num_regs_per_blk    = 4;
+    static constexpr index_t num_threads_per_blk = 16;
     static constexpr index_t wave_size           = 64;
-    static constexpr index_t num_input_blks  = wave_size / num_threads_blk;
+    static constexpr index_t num_input_blks      = 4;
     static constexpr index_t num_output_blks     = 1;
-    static constexpr index_t num_regs_xdlops = num_regs_blk * num_output_blks;
-    static constexpr index_t m               = 16;
-    static constexpr index_t n               = 16;
-    static constexpr index_t k               = 8;
-    static constexpr index_t cycles          = 32;
-    static constexpr index_t k_base          = 2;
+    static constexpr index_t m_per_blk           = 16;
+    static constexpr index_t n_per_blk           = 16;
+    static constexpr index_t k_per_blk           = 2;
+    static constexpr bool is_k_reduction         = true;
 
     template <index_t MPerXdlops,
               index_t NPerXdlops,
@@ -425,21 +346,19 @@ struct mfma_info<mfma_instr::mfma_f32_16x16x8bf16>
 };
 
 template <>
-struct mfma_info<mfma_instr::mfma_f32_16x16x2bf16>
+struct mfma_type<MfmaInstr::mfma_f32_16x16x2bf16>
 {
     static constexpr index_t group_size          = 4;
-    static constexpr index_t num_groups_blk  = 1;
-    static constexpr index_t num_regs_blk    = group_size * num_groups_blk;
-    static constexpr index_t num_threads_blk = 16;
+    static constexpr index_t num_groups_per_blk  = 1;
+    static constexpr index_t num_regs_per_blk    = 4;
+    static constexpr index_t num_threads_per_blk = 16;
     static constexpr index_t wave_size           = 64;
-    static constexpr index_t num_input_blks  = wave_size / num_threads_blk;
+    static constexpr index_t num_input_blks      = 4;
     static constexpr index_t num_output_blks     = 4;
-    static constexpr index_t num_regs_xdlops = num_regs_blk * num_output_blks;
-    static constexpr index_t m               = 16;
-    static constexpr index_t n               = 16;
-    static constexpr index_t k               = 2;
-    static constexpr index_t cycles          = 32;
-    static constexpr index_t k_base          = 2;
+    static constexpr index_t m_per_blk           = 16;
+    static constexpr index_t n_per_blk           = 16;
+    static constexpr index_t k_per_blk           = 2;
+    static constexpr bool is_k_reduction         = false;
 
     template <index_t MPerXdlops,
               index_t NPerXdlops,
@@ -458,21 +377,19 @@ struct mfma_info<mfma_instr::mfma_f32_16x16x2bf16>
 };
 
 template <>
-struct mfma_info<mfma_instr::mfma_f32_4x4x2bf16>
+struct mfma_type<MfmaInstr::mfma_f32_4x4x2bf16>
 {
     static constexpr index_t group_size          = 4;
-    static constexpr index_t num_groups_blk  = 1;
-    static constexpr index_t num_regs_blk    = group_size * num_groups_blk;
-    static constexpr index_t num_threads_blk = 64;
+    static constexpr index_t num_groups_per_blk  = 1;
+    static constexpr index_t num_regs_per_blk    = 4;
+    static constexpr index_t num_threads_per_blk = 64;
     static constexpr index_t wave_size           = 64;
     static constexpr index_t num_input_blks      = 1;
     static constexpr index_t num_output_blks     = 1;
-    static constexpr index_t num_regs_xdlops = 4;
-    static constexpr index_t m               = 4;
-    static constexpr index_t n               = 64;
-    static constexpr index_t k               = 2;
-    static constexpr index_t cycles          = 8;
-    static constexpr index_t k_base          = 2;
+    static constexpr index_t m_per_blk           = 4;
+    static constexpr index_t n_per_blk           = 64;
+    static constexpr index_t k_per_blk           = 2;
+    static constexpr bool is_k_reduction         = false;
 
     template <index_t MPerXdlops,
               index_t NPerXdlops,
@@ -491,200 +408,227 @@ struct mfma_info<mfma_instr::mfma_f32_4x4x2bf16>
 };
 #endif
 
-template <mfma_instr instr, index_t MPerXdlops_, index_t NPerXdlops_>
-struct xdlops_info
+template <typename base_type, index_t MPerXdlops, index_t NPerXdlops>
+struct MfmaSelector
 {
-    static constexpr auto mfma_type = mfma_info<instr>{};
-
-    static constexpr index_t MPerXdlops = MPerXdlops_;
-    static constexpr index_t NPerXdlops = NPerXdlops_;
+    template <typename base_type_, index_t MPerXdlops_, index_t NPerXdlops_>
+    static constexpr auto GetMfma();
 
-    static constexpr bool IsABroadcast()
+    template <>
+    static constexpr auto GetMfma<float, 64, 64>()
     {
-        static_assert(NPerXdlops >= MPerXdlops, "only support ABroadcast");
-        return true;
+        return MfmaInstr::mfma_f32_32x32x1xf32;
     }
 
-    static constexpr bool IsKReduction()
+    template <>
+    static constexpr auto GetMfma<float, 32, 64>()
     {
-        return (mfma_type.num_output_blks == 1) && (mfma_type.num_input_blks > 1);
+        return MfmaInstr::mfma_f32_32x32x1xf32;
     }
 
-    static constexpr index_t GetKPerXdlops()
+    template <>
+    static constexpr auto GetMfma<float, 16, 64>()
     {
-        return IsKReduction() ? mfma_type.num_input_blks : 1;
+        return MfmaInstr::mfma_f32_16x16x1xf32;
     }
 
-    static constexpr index_t GetNumCRegs() { return MPerXdlops * NPerXdlops / mfma_type.wave_size; }
-};
-
-template <class base_type, index_t MPerWave, index_t NPerWave, index_t KPack>
-struct XdlopsGemm
-{
-    template <class base_type_  = base_type,
-              index_t MPerWave_ = MPerWave,
-              index_t NPerWave_ = NPerWave>
-    static constexpr auto GetXdlopsInfo();
-
     template <>
-    static constexpr auto GetXdlopsInfo<float, 64, 64>()
+    static constexpr auto GetMfma<float, 8, 64>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_32x32x1xf32, 64, 64>{};
+        return MfmaInstr::mfma_f32_4x4x1xf32;
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<float, 32, 64>()
+    static constexpr auto GetMfma<float, 4, 64>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_32x32x1xf32, 32, 64>{};
+        return MfmaInstr::mfma_f32_4x4x1xf32;
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<float, 16, 64>()
+    static constexpr auto GetMfma<float, 32, 32>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_16x16x1xf32, 16, 64>{};
+        return MfmaInstr::mfma_f32_32x32x2xf32;
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<float, 8, 64>()
+    static constexpr auto GetMfma<float, 16, 16>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_4x4x1xf32, 8, 64>{};
+        return MfmaInstr::mfma_f32_16x16x4xf32;
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<float, 4, 64>()
+    static constexpr auto GetMfma<half_t, 64, 64>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_4x4x1xf32, 4, 64>{};
+        return MfmaInstr::mfma_f32_32x32x4f16;
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<float, 32, 32>()
+    static constexpr auto GetMfma<half_t, 32, 64>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_32x32x2xf32, 32, 32>{};
+        return MfmaInstr::mfma_f32_32x32x4f16;
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<float, 16, 16>()
+    static constexpr auto GetMfma<half_t, 32, 32>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_16x16x4xf32, 16, 16>{};
+        return MfmaInstr::mfma_f32_32x32x8f16;
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<half_t, 64, 64>()
+    static constexpr auto GetMfma<half_t, 16, 16>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_32x32x4f16, 64, 64>{};
+        return MfmaInstr::mfma_f32_16x16x16f16;
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<half_t, 32, 64>()
+    static constexpr auto GetMfma<half_t, 16, 64>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_32x32x4f16, 32, 64>{};
+        return MfmaInstr::mfma_f32_16x16x4f16;
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<half_t, 32, 32>()
+    static constexpr auto GetMfma<half_t, 8, 64>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_32x32x8f16, 32, 32>{};
+        return MfmaInstr::mfma_f32_4x4x4f16;
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<half_t, 16, 16>()
+    static constexpr auto GetMfma<half_t, 4, 64>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_16x16x16f16, 16, 16>{};
+        return MfmaInstr::mfma_f32_4x4x4f16;
     }
 
+#if 0
     template <>
-    static constexpr auto GetXdlopsInfo<half_t, 16, 64>()
+    static constexpr auto GetMfma<ushort, 128, 64>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_16x16x4f16, 16, 64>{};
+        return xdlops_info<MfmaInstr::mfma_f32_32x32x2bf16, 64, 64, 2, 1, c_vec32_4_t>{};
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<half_t, 8, 64>()
+    static constexpr auto GetMfma<ushort, 64, 128>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_4x4x4f16, 8, 64>{};
+        return xdlops_info<MfmaInstr::mfma_f32_32x32x2bf16, 64, 64, 1, 2, c_vec32_4_t>{};
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<half_t, 4, 64>()
+    static constexpr auto GetMfma<ushort, 64, 64>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_4x4x4f16, 4, 64>{};
+        return xdlops_info<MfmaInstr::mfma_f32_32x32x2bf16, 64, 64, 1, 1, c_vec32_2_t>{};
     }
 
-#if 0
     template <>
-    static constexpr auto GetXdlopsInfo<ushort, 128, 64>()
+    static constexpr auto GetMfma<ushort, 64, 32>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_32x32x2bf16, 64, 64, 2, 1, c_vec32_4_t>{};
+        return xdlops_info<MfmaInstr::mfma_f32_32x32x2bf16, 64, 32, 1, 1, c_vec32_1_t>{};
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<ushort, 64, 128>()
+    static constexpr auto GetMfma<ushort, 32, 64>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_32x32x2bf16, 64, 64, 1, 2, c_vec32_4_t>{};
+        return xdlops_info<MfmaInstr::mfma_f32_32x32x2bf16, 32, 64, 1, 1, c_vec32_1_t>{};
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<ushort, 64, 64>()
+    static constexpr auto GetMfma<ushort, 64, 16>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_32x32x2bf16, 64, 64, 1, 1, c_vec32_2_t>{};
+        return xdlops_info<MfmaInstr::mfma_f32_16x16x2bf16, 64, 16, 1, 1, c_vec16_1_t>{};
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<ushort, 64, 32>()
+    static constexpr auto GetMfma<ushort, 16, 64>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_32x32x2bf16, 64, 32, 1, 1, c_vec32_1_t>{};
+        return xdlops_info<MfmaInstr::mfma_f32_16x16x2bf16, 16, 64, 1, 1, c_vec16_1_t>{};
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<ushort, 32, 64>()
+    static constexpr auto GetMfma<ushort, 8, 64>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_32x32x2bf16, 32, 64, 1, 1, c_vec32_1_t>{};
+        return xdlops_info<MfmaInstr::mfma_f32_4x4x2bf16, 8, 64, 1, 1, c_vec4_2_t>{};
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<ushort, 64, 16>()
+    static constexpr auto GetMfma<ushort, 4, 64>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_16x16x2bf16, 64, 16, 1, 1, c_vec16_1_t>{};
+        return xdlops_info<MfmaInstr::mfma_f32_4x4x2bf16, 4, 64, 1, 1, c_vec4_1_t>{};
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<ushort, 16, 64>()
+    static constexpr auto GetMfma<ushort, 32, 32>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_16x16x2bf16, 16, 64, 1, 1, c_vec16_1_t>{};
+        return xdlops_info<MfmaInstr::mfma_f32_32x32x4bf16, 32, 32, 1, 1, c_vec16_1_t>{};
     }
 
     template <>
-    static constexpr auto GetXdlopsInfo<ushort, 8, 64>()
+    static constexpr auto GetMfma<ushort, 16, 16>()
     {
-        return xdlops_info<mfma_instr::mfma_f32_4x4x2bf16, 8, 64, 1, 1, c_vec4_2_t>{};
+        return xdlops_info<MfmaInstr::mfma_f32_16x16x8bf16, 16, 16, 1, 1, c_vec4_1_t>{};
     }
+#endif
 
-    template <>
-    static constexpr auto GetXdlopsInfo<ushort, 4, 64>()
+    static constexpr auto selected_mfma = mfma_type<GetMfma<base_type, MPerXdlops, NPerXdlops>()>{};
+
+    __host__ __device__ static constexpr void mfma_check()
     {
-        return xdlops_info<mfma_instr::mfma_f32_4x4x2bf16, 4, 64, 1, 1, c_vec4_1_t>{};
+        static_assert(selected_mfma.group_size * selected_mfma.num_groups_per_blk ==
+                          selected_mfma.num_regs_per_blk,
+                      "wrong! num_regs_per_blk");
+
+        static_assert(selected_mfma.num_threads_per_blk == selected_mfma.n_per_blk,
+                      "n_per_blk != num_threads_per_blk");
+
+        static_assert(selected_mfma.num_regs_per_blk * selected_mfma.num_input_blks ==
+                          selected_mfma.m_per_blk,
+                      "m_per_blk != num_input_blks * num_regs_per_blk");
+
+        static_assert(selected_mfma.num_output_blks == selected_mfma.num_input_blks ||
+                          selected_mfma.num_output_blks == 1,
+                      "incorrect num_output_blks");
+
+        static_assert(selected_mfma.num_regs_per_blk * selected_mfma.wave_size ==
+                          selected_mfma.m_per_blk * selected_mfma.n_per_blk,
+                      "num_regs_per_blk incorrect");
+
+        static_assert(selected_mfma.is_k_reduction ||
+                          (selected_mfma.num_input_blks == selected_mfma.num_output_blks),
+                      "is_k_reduction wrong!");
     }
 
-    template <>
-    static constexpr auto GetXdlopsInfo<ushort, 32, 32>()
+    __host__ __device__ constexpr MfmaSelector() { mfma_check(); }
+
+    static constexpr bool IsABroadcast()
     {
-        return xdlops_info<mfma_instr::mfma_f32_32x32x4bf16, 32, 32, 1, 1, c_vec16_1_t>{};
+        static_assert(NPerXdlops >= MPerXdlops, "only support ABroadcast");
+        return true;
     }
 
-    template <>
-    static constexpr auto GetXdlopsInfo<ushort, 16, 16>()
+    static constexpr index_t GetKPerXdlops()
     {
-        return xdlops_info<mfma_instr::mfma_f32_16x16x8bf16, 16, 16, 1, 1, c_vec4_1_t>{};
+        return (selected_mfma.is_k_reduction ? selected_mfma.num_input_blks : 1) *
+               selected_mfma.k_per_blk;
     }
-#endif
+
+    static constexpr index_t GetKPerThread() { return selected_mfma.k_per_blk; }
+};
+
+template <typename base_type, index_t MPerXdlops, index_t NPerXdlops, index_t KPack>
+struct XdlopsGemm
+{
+    static constexpr auto I0 = Number<0>{};
+    static constexpr auto I1 = Number<1>{};
+    static constexpr auto I2 = Number<2>{};
+    static constexpr auto I3 = Number<3>{};
+    static constexpr auto I4 = Number<4>{};
+    static constexpr auto I5 = Number<5>{};
 
     using CIndex = MultiIndex<2>;
 
-    __device__ static constexpr index_t GetNumBlks() { return mfma_type.num_output_blks; }
+    __device__ static constexpr index_t GetNumBlks() { return mfma_instr.num_output_blks; }
 
     __device__ static constexpr index_t GetNumXdlops()
     {
-        return MPerXdlops * NPerXdlops / (mfma_type.m * mfma_type.n * mfma_type.num_output_blks);
+        return MPerXdlops * NPerXdlops /
+               (mfma_instr.m_per_blk * mfma_instr.n_per_blk * mfma_instr.num_output_blks);
     }
 
     __host__ __device__ constexpr XdlopsGemm()
@@ -697,104 +641,142 @@ struct XdlopsGemm
                           MPerXdlops == 64,
                       "Only support GemmMPerXdlops == 4, 8, 16, 32 or 64 for xdlops");
 
-        static_assert(mfma_type.num_threads_blk == mfma_type.n, "n != num_threads_blk");
-        static_assert(mfma_type.num_regs_blk * mfma_type.num_input_blks == mfma_type.m,
-                      "m != num_input_blks * num_regs_blk");
-        static_assert(mfma_type.num_output_blks == mfma_type.num_input_blks ||
-                          mfma_type.num_output_blks == 1,
-                      "incorrect num_output_blks");
-        static_assert(mfma_type.num_regs_blk * mfma_type.wave_size == mfma_type.m * mfma_type.n,
-                      "num_regs_blk incorrect");
-
-        static_assert(mfma_type.k % mfma_type.k_base == 0, "k % kbase != 0!");
+        static_assert(KPack % mfma_instr.k_per_blk == 0, "KPack cannot be divided by k_per_blk");
+    }
+
+    template <typename CM0N0M1N1M2N2Desc>
+    __host__ __device__ static constexpr auto
+    MakeCM0N0M1N1M2M3M4N2Descriptor(const CM0N0M1N1M2N2Desc& c_m0_n0_m1_n1_m2_n2_desc)
+    {
+        const auto M0 = c_m0_n0_m1_n1_m2_n2_desc.GetLength(I0);
+        const auto N0 = c_m0_n0_m1_n1_m2_n2_desc.GetLength(I1);
+        const auto M1 = c_m0_n0_m1_n1_m2_n2_desc.GetLength(I2);
+        const auto N1 = c_m0_n0_m1_n1_m2_n2_desc.GetLength(I3);
+
+        return transform_tensor_descriptor(
+            c_m0_n0_m1_n1_m2_n2_desc,
+            make_tuple(make_pass_through_transform(M0),
+                       make_pass_through_transform(N0),
+                       make_pass_through_transform(M1),
+                       make_pass_through_transform(N1),
+                       make_unmerge_transform(make_tuple(mfma_instr.num_groups_per_blk,
+                                                         mfma_instr.num_input_blks,
+                                                         mfma_instr.group_size)),
+                       make_pass_through_transform(mfma_instr.num_threads_per_blk)),
+            make_tuple(Sequence<0>{},
+                       Sequence<1>{},
+                       Sequence<2>{},
+                       Sequence<3>{},
+                       Sequence<4>{},
+                       Sequence<5>{}),
+            make_tuple(Sequence<0>{},
+                       Sequence<1>{},
+                       Sequence<2>{},
+                       Sequence<3>{},
+                       Sequence<4, 5, 6>{},
+                       Sequence<7>{}));
     }
 
     __device__ static constexpr index_t GetRegSizePerXdlops()
     {
-        return MPerXdlops * NPerXdlops / mfma_type.wave_size;
+        return MPerXdlops * NPerXdlops / mfma_instr.wave_size;
     }
 
-    template <class ADesc,
-              class BDesc,
-              class CDesc,
-              index_t m0,
-              index_t n0,
-              class FloatA,
-              class FloatB,
-              class FloatC>
+    template <class FloatA, class FloatB, class FloatC>
     __device__ void Run(const FloatA& p_a_wave, const FloatB& p_b_wave, FloatC& p_c_thread) const
     {
         static_assert(is_same<base_type, float>::value || is_same<base_type, half_t>::value ||
                           is_same<base_type, ushort>::value,
                       "base base_type must be float, half, ushort!");
 
-        static_assert(KPack % mfma_type.k_base == 0, "KPack cannot be divided by k_base");
+        static_for<0, KPack / mfma_instr.k_per_blk, 1>{}([&](auto k) {
+            mfma_instr.template run<MPerXdlops, NPerXdlops>(p_a_wave[k], p_b_wave[k], p_c_thread);
+        });
+    }
+
+    __device__ static auto GetLaneId() { return get_thread_local_1d_id() % mfma_instr.wave_size; }
 
-        constexpr index_t c_offset = CDesc{}.CalculateOffset(make_tuple(m0, n0)) * GetNumXdlops();
+    __device__ static auto GetBlkIdx()
+    {
+        const auto laneId = GetLaneId();
 
-        static_for<0, KPack, mfma_type.k_base>{}([&](auto k) {
-            constexpr index_t a_offset = ADesc{}.CalculateOffset(make_tuple(0, m0, 0, k));
-            constexpr index_t b_offset = BDesc{}.CalculateOffset(make_tuple(0, n0, 0, k));
+        constexpr auto threadidx_to_blk_idx_adaptor = make_single_stage_tensor_adaptor(
+            make_tuple(make_merge_transform(
+                make_tuple(1, mfma_instr.num_input_blks, mfma_instr.num_threads_per_blk))),
+            make_tuple(Sequence<0, 1, 2>{}),
+            make_tuple(Sequence<0>{}));
 
-            mfma_type.template run<MPerXdlops, NPerXdlops, c_offset>(
-                p_a_wave[Number<a_offset / mfma_type.k_base>{}],
-                p_b_wave[Number<b_offset / mfma_type.k_base>{}],
-                p_c_thread);
-        });
+        const auto blk_idx =
+            threadidx_to_blk_idx_adaptor.CalculateBottomIndex(make_multi_index(laneId));
+
+        const auto blk_id = blk_idx[I1];
+        const auto blk_td = blk_idx[I2];
+
+        return make_tuple(blk_id, blk_td);
     }
 
-    __device__ static CIndex GetBeginOfThreadBlk(index_t xdlops_i, index_t blk_i)
+    __host__ __device__ static auto CalculateAThreadOriginDataIndex()
     {
-        const index_t laneId = get_thread_local_1d_id() % mfma_type.wave_size;
-        const index_t blk_id = laneId / mfma_type.num_threads_blk;
-        const index_t blk_td = laneId % mfma_type.num_threads_blk;
+        const auto laneId  = GetLaneId();
+        const auto blk_idx = GetBlkIdx();
 
-        index_t n_offset = blk_i * mfma_type.n + blk_td;
-        index_t m_offset = xdlops_i * mfma_type.m + blk_id * mfma_type.group_size;
+        const auto blk_id = blk_idx[I0];
+        const auto blk_td = blk_idx[I1];
 
-        return CIndex{m_offset, n_offset};
+        if constexpr(mfma_instr.is_k_reduction)
+        {
+            return make_tuple(blk_id, blk_td);
+        }
+        else
+        {
+            return make_tuple(0, laneId);
+        }
     }
 
-    static constexpr index_t MRepeats   = GetXdlopsInfo().MRepeats;
-    static constexpr index_t NRepeats   = GetXdlopsInfo().NRepeats;
-    static constexpr index_t MPerXdlops = GetXdlopsInfo().MPerXdlops;
-    static constexpr index_t NPerXdlops = GetXdlopsInfo().NPerXdlops;
+    __host__ __device__ static auto CalculateBThreadOriginDataIndex()
+    {
+        const auto laneId  = GetLaneId();
+        const auto blk_idx = GetBlkIdx();
 
-    static constexpr bool IsKReduction  = GetXdlopsInfo().IsKReduction();
-    static constexpr bool IsABroadcast  = GetXdlopsInfo().IsABroadcast();
-    static constexpr index_t KPerXdlops = GetXdlopsInfo().GetKPerXdlops();
+        const auto blk_id = blk_idx[I0];
+        const auto blk_td = blk_idx[I1];
 
-    static constexpr auto GetBlkId(const index_t lane_id)
+        if constexpr(mfma_instr.is_k_reduction)
         {
-        return lane_id / mfma_type.num_threads_blk;
+            return make_tuple(blk_id, blk_td);
         }
-
-    static constexpr auto GetBlkTd(const index_t lane_id)
+        else
         {
-        return lane_id % mfma_type.num_threads_blk;
+            return make_tuple(0, laneId);
+        }
     }
 
-    static constexpr auto mfma_type = GetXdlopsInfo().mfma_type;
-
-    struct CLayout
+    __device__ static CIndex GetBeginOfThreadBlk(index_t xdlops_i, index_t blk_i)
     {
-        __host__ __device__ static constexpr index_t M1() { return mfma_type.num_groups_blk; }
-        __host__ __device__ static constexpr index_t M0() { return mfma_type.group_size; }
-        __host__ __device__ static constexpr index_t N1() { return mfma_type.num_input_blks; }
-        __host__ __device__ static constexpr index_t N0() { return mfma_type.num_threads_blk; }
+        const auto blk_idx = GetBlkIdx();
 
-        __device__ static constexpr index_t GetBlkSize() { return mfma_type.num_regs_blk; }
+        const auto blk_id = blk_idx[I0];
+        const auto blk_td = blk_idx[I1];
 
-        __device__ static constexpr index_t GetNumBlks() { return mfma_type.num_output_blks; }
+        index_t n_offset = blk_i * mfma_instr.n_per_blk + blk_td;
+        index_t m_offset = xdlops_i * mfma_instr.m_per_blk + blk_id * mfma_instr.group_size;
 
-        __device__ static constexpr index_t GetNumXdlops()
-        {
-            return MPerXdlops * NPerXdlops /
-                   (mfma_type.m * mfma_type.n * mfma_type.num_output_blks);
+        return CIndex{m_offset, n_offset};
     }
-    };
 
-    __host__ __device__ static constexpr auto GetCLayout() { return CLayout{}; }
+    static constexpr auto mfma = MfmaSelector<base_type, MPerXdlops, NPerXdlops>{};
+
+    static constexpr auto mfma_instr = mfma.selected_mfma;
+
+    static constexpr auto KPerXdlops  = mfma.GetKPerXdlops();
+    static constexpr auto K1PerXdlops = mfma.GetKPerThread();
+    static constexpr auto K0PerXdlops = KPerXdlops / K1PerXdlops;
+
+    __host__ __device__ static constexpr auto GetCM0M1M2NThreadBlkLengths()
+    {
+        return make_tuple(
+            Number<mfma_instr.num_groups_per_blk>{}, I1, Number<mfma_instr.group_size>{}, I1);
+    }
 };
 
 } // namespace ck