Dynamic tensor descriptor (#24)

* support dynamic tensor descriptor

* use buffer load OOB feature for padding case

* add navi support

* add int8x4 inference kernel

Co-authored-by: Chao Liu <chao@ixt-rack-81.local.lan>
Co-authored-by: Jing Zhang <jizhan@amd.com>

composable_kernel/include/driver/driver_dynamic_convolution_forward_implicit_gemm_v5r1_nchw_kcyx_nkhw.hpp

@@ -0,0 +1,353 @@
+#ifndef CK_DRIVER_DYNAMIC_CONVOLUTION_FORWARD_IMPLICIT_GEMM_V5R1_NCHW_KCYX_NKHW_HPP
+#define CK_DRIVER_DYNAMIC_CONVOLUTION_FORWARD_IMPLICIT_GEMM_V5R1_NCHW_KCYX_NKHW_HPP
+
+#include "common_header.hpp"
+#include "dynamic_tensor_descriptor.hpp"
+#include "dynamic_tensor_descriptor_helper.hpp"
+#include "gridwise_dynamic_gemm_v2.hpp"
+#include "gridwise_operation_wrapper.hpp"
+
+namespace ck {
+
+template <index_t BlockSize,
+          typename FloatAB,
+          typename FloatAcc,
+          typename FloatC,
+          index_t KPerBlock,
+          index_t HoPerBlock,
+          index_t WoPerBlock,
+          index_t EPerBlock,
+          index_t KPerThread,
+          index_t HoPerThread,
+          index_t WoPerThread,
+          index_t EPerThread,
+          typename ABlockTransferThreadSliceLengths_E_K,
+          typename ABlockTransferThreadClusterLengths_E_K,
+          index_t ABlockTransferSrcScalarPerVector_E,
+          index_t ABlockTransferDstScalarPerVector_K,
+          index_t BThreadTransferSrcScalarPerVector_W,
+          index_t CThreadTransferDstScalarPerVector_W>
+struct DriverDynamicConvolutionForwardImplicitGemm_v5r1_nchw_kcyx_nkhw_pad
+{
+    template <typename... Wei,
+              typename... In,
+              typename... Out,
+              typename ConvStrides,
+              typename ConvDilations,
+              typename InLeftPads,
+              typename InRightPads>
+    __host__ void Run(const DynamicTensorDescriptor<Wei...>& wei_k_c_y_x_global_desc,
+                      const DynamicTensorDescriptor<In...>& in_n_c_hi_wi_global_desc,
+                      const DynamicTensorDescriptor<Out...>& out_n_k_ho_wo_global_desc,
+                      const ConvStrides& conv_strides,
+                      const ConvDilations& conv_dilations,
+                      const InLeftPads& in_left_pads,
+                      const InRightPads& in_right_pads,
+                      const FloatAB* __restrict__ p_wei_global,
+                      const FloatAB* __restrict__ p_in_global,
+                      FloatC* __restrict__ p_out_global) const
+    {
+        constexpr auto I0 = Number<0>{};
+        constexpr auto I1 = Number<1>{};
+        constexpr auto I2 = Number<2>{};
+        constexpr auto I3 = Number<3>{};
+
+        const auto N = in_n_c_hi_wi_global_desc.GetLength(I0);
+        const auto C = in_n_c_hi_wi_global_desc.GetLength(I1);
+        const auto K = out_n_k_ho_wo_global_desc.GetLength(I1);
+
+        const auto Hi = in_n_c_hi_wi_global_desc.GetLength(I2);
+        const auto Wi = in_n_c_hi_wi_global_desc.GetLength(I3);
+
+        const auto Ho = out_n_k_ho_wo_global_desc.GetLength(I2);
+        const auto Wo = out_n_k_ho_wo_global_desc.GetLength(I3);
+
+        const auto Y = wei_k_c_y_x_global_desc.GetLength(I2);
+        const auto X = wei_k_c_y_x_global_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];
+
+        // weight tensor
+        const auto wei_gemmk_gemmm_global_desc = transform_dynamic_tensor_descriptor(
+            make_dynamic_naive_tensor_descriptor_packed_v2(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<1>{}, Sequence<0>{}));
+
+        // input tensor
+        const auto in_n_c_hip_wip_global_desc = transform_dynamic_tensor_descriptor(
+            in_n_c_hi_wi_global_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_global_desc = transform_dynamic_tensor_descriptor(
+            in_n_c_hip_wip_global_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_n_ho_wo_global_desc = transform_dynamic_tensor_descriptor(
+            in_n_c_y_ho_x_wo_global_desc,
+            make_tuple(make_merge_transform(make_tuple(C, Y, X)),
+                       make_pass_through_transform(N),
+                       make_pass_through_transform(Ho),
+                       make_pass_through_transform(Wo)),
+            make_tuple(Sequence<1, 2, 4>{}, Sequence<0>{}, Sequence<3>{}, Sequence<5>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
+
+        // output tensor
+        const auto out_gemmm_n_ho_wo_global_desc = transform_dynamic_tensor_descriptor(
+            make_dynamic_naive_tensor_descriptor_packed_v2(make_tuple(N, K, Ho, Wo)),
+            make_tuple(make_pass_through_transform(K),
+                       make_pass_through_transform(N),
+                       make_pass_through_transform(Ho),
+                       make_pass_through_transform(Wo)),
+            make_tuple(Sequence<1>{}, Sequence<0>{}, Sequence<2>{}, Sequence<3>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
+
+        const auto E = C * Y * X;
+
+        if(!(K % KPerBlock == 0 && Ho % HoPerBlock == 0 && Wo % WoPerBlock == 0 &&
+             E % EPerBlock == 0))
+        {
+            throw std::runtime_error("wrong! GEMM size no divisible");
+        }
+
+        // hack to control index calculation when iterating over a_k_m_global tensor
+        constexpr auto a_k_m_global_iterator_hacks =
+            make_tuple(make_tuple(Sequence<0, 0, 0>{}, Sequence<0, 0, 0>{}),
+                       make_tuple(Sequence<0, 0, 0>{}, Sequence<0, 0, 0>{}));
+
+        constexpr auto a_k_m_global_move_slice_window_iterator_hack = Sequence<0, 0, 0>{};
+
+        constexpr auto b_k_n_global_iterator_hacks =
+            make_tuple(make_tuple(Sequence<0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0>{},
+                                  Sequence<0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0>{},
+                                  Sequence<0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0>{},
+                                  Sequence<0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0>{}),
+                       make_tuple(Sequence<0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0>{},
+                                  Sequence<0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0>{},
+                                  Sequence<0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0>{},
+                                  Sequence<0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0>{}));
+
+        constexpr auto b_k_n_global_move_slice_window_iterator_hack =
+            Sequence<0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0>{};
+
+        // hack to control index calculation when iterating over c_m0_m1_n0_n1_global tensor
+        // hack for NKHW format
+        constexpr auto c_k_n_h_w_global_tensor_iterator_hacks =
+            make_tuple(make_tuple(Sequence<0, 0, 0, 0, 0>{},
+                                  Sequence<0, 0, 0, 0, 0>{},
+                                  Sequence<0, 0, 0, 0, 0>{},
+                                  Sequence<0, 0, 0, 0, 0>{}),
+                       make_tuple(Sequence<0, 0, 0, 0, 0>{},
+                                  Sequence<0, 0, 0, 0, 0>{},
+                                  Sequence<0, 0, 0, 0, 0>{},
+                                  Sequence<0, 0, 0, 0, 0>{}));
+
+#if 1
+        // GEMM
+        using gridwise_gemm = GridwiseDynamicGemm_km_kn_mn_v3<
+            BlockSize,
+            FloatAB,
+            FloatAcc,
+            FloatC,
+            InMemoryDataOperation::Set,
+            decltype(wei_gemmk_gemmm_global_desc),
+            decltype(in_gemmk_n_ho_wo_global_desc),
+            decltype(out_gemmm_n_ho_wo_global_desc),
+            KPerBlock,
+            HoPerBlock,
+            WoPerBlock,
+            EPerBlock,
+            KPerThread,
+            HoPerThread,
+            WoPerThread,
+            EPerThread,
+            ABlockTransferThreadSliceLengths_E_K,
+            ABlockTransferThreadClusterLengths_E_K,
+            Sequence<1, 0>,
+            Sequence<1, 0>,
+            0,
+            ABlockTransferSrcScalarPerVector_E,
+            ABlockTransferDstScalarPerVector_K,
+            false, // don't move back src coordinate after threadwise copy
+            Sequence<0, 2, 3, 1>,
+            3,
+            BThreadTransferSrcScalarPerVector_W,
+            false, // don't move back src coordinate after threadwise copy, which will be fused with
+                   // MoveSrcSliceWindow() to save addr computation
+            Sequence<0, 2, 3, 1>,
+            3,
+            CThreadTransferDstScalarPerVector_W,
+            decltype(a_k_m_global_iterator_hacks),
+            decltype(b_k_n_global_iterator_hacks),
+            decltype(c_k_n_h_w_global_tensor_iterator_hacks),
+            decltype(a_k_m_global_move_slice_window_iterator_hack),
+            decltype(b_k_n_global_move_slice_window_iterator_hack)>;
+
+        const auto GridSize = (K / KPerBlock) * (Ho / HoPerBlock) * (Wo / WoPerBlock) * N;
+
+        const bool has_main_k_block_loop = (E + EPerBlock) / (2 * EPerBlock) > 1;
+
+        const bool has_double_tail_k_block_loop = (E / EPerBlock) % 2 == 0;
+
+        index_t nrepeat = 100;
+
+        for(index_t i = 0; i < 5; ++i)
+        {
+            std::cout << "Start running " << nrepeat << " times..." << std::endl;
+
+            KernelTimer timer;
+            timer.Start();
+            std::cout << "has_main_k_block_loop: " << has_main_k_block_loop
+                      << " has_double_tail_k_block_loop: " << has_double_tail_k_block_loop
+                      << std::endl;
+
+            for(index_t j = 0; j < nrepeat; ++j)
+            {
+                if(has_main_k_block_loop && has_double_tail_k_block_loop)
+                {
+                    const auto kernel =
+                        run_gridwise_operation<gridwise_gemm,
+                                               decltype(wei_gemmk_gemmm_global_desc),
+                                               const FloatAB*,
+                                               decltype(in_gemmk_n_ho_wo_global_desc),
+                                               const FloatAB*,
+                                               decltype(out_gemmm_n_ho_wo_global_desc),
+                                               FloatC*,
+                                               integral_constant<bool, true>,
+                                               integral_constant<bool, true>>;
+
+                    launch_kernel(kernel,
+                                  dim3(GridSize),
+                                  dim3(BlockSize),
+                                  0,
+                                  0,
+                                  wei_gemmk_gemmm_global_desc,
+                                  p_wei_global,
+                                  in_gemmk_n_ho_wo_global_desc,
+                                  p_in_global,
+                                  out_gemmm_n_ho_wo_global_desc,
+                                  p_out_global,
+                                  integral_constant<bool, true>{},
+                                  integral_constant<bool, true>{});
+                }
+                else if(has_main_k_block_loop && !has_double_tail_k_block_loop)
+                {
+                    const auto kernel =
+                        run_gridwise_operation<gridwise_gemm,
+                                               decltype(wei_gemmk_gemmm_global_desc),
+                                               const FloatAB*,
+                                               decltype(in_gemmk_n_ho_wo_global_desc),
+                                               const FloatAB*,
+                                               decltype(out_gemmm_n_ho_wo_global_desc),
+                                               FloatC*,
+                                               integral_constant<bool, true>,
+                                               integral_constant<bool, false>>;
+
+                    launch_kernel(kernel,
+                                  dim3(GridSize),
+                                  dim3(BlockSize),
+                                  0,
+                                  0,
+                                  wei_gemmk_gemmm_global_desc,
+                                  p_wei_global,
+                                  in_gemmk_n_ho_wo_global_desc,
+                                  p_in_global,
+                                  out_gemmm_n_ho_wo_global_desc,
+                                  p_out_global,
+                                  integral_constant<bool, true>{},
+                                  integral_constant<bool, false>{});
+                }
+                else if(!has_main_k_block_loop && has_double_tail_k_block_loop)
+                {
+                    const auto kernel =
+                        run_gridwise_operation<gridwise_gemm,
+                                               decltype(wei_gemmk_gemmm_global_desc),
+                                               const FloatAB*,
+                                               decltype(in_gemmk_n_ho_wo_global_desc),
+                                               const FloatAB*,
+                                               decltype(out_gemmm_n_ho_wo_global_desc),
+                                               FloatC*,
+                                               integral_constant<bool, false>,
+                                               integral_constant<bool, true>>;
+
+                    launch_kernel(kernel,
+                                  dim3(GridSize),
+                                  dim3(BlockSize),
+                                  0,
+                                  0,
+                                  wei_gemmk_gemmm_global_desc,
+                                  p_wei_global,
+                                  in_gemmk_n_ho_wo_global_desc,
+                                  p_in_global,
+                                  out_gemmm_n_ho_wo_global_desc,
+                                  p_out_global,
+                                  integral_constant<bool, false>{},
+                                  integral_constant<bool, true>{});
+                }
+                else
+                {
+                    const auto kernel =
+                        run_gridwise_operation<gridwise_gemm,
+                                               decltype(wei_gemmk_gemmm_global_desc),
+                                               const FloatAB*,
+                                               decltype(in_gemmk_n_ho_wo_global_desc),
+                                               const FloatAB*,
+                                               decltype(out_gemmm_n_ho_wo_global_desc),
+                                               FloatC*,
+                                               integral_constant<bool, false>,
+                                               integral_constant<bool, false>>;
+
+                    launch_kernel(kernel,
+                                  dim3(GridSize),
+                                  dim3(BlockSize),
+                                  0,
+                                  0,
+                                  wei_gemmk_gemmm_global_desc,
+                                  p_wei_global,
+                                  in_gemmk_n_ho_wo_global_desc,
+                                  p_in_global,
+                                  out_gemmm_n_ho_wo_global_desc,
+                                  p_out_global,
+                                  integral_constant<bool, false>{},
+                                  integral_constant<bool, false>{});
+                }
+            }
+
+            timer.End();
+
+            float ave_time = timer.GetElapsedTime() / nrepeat;
+
+            float perf = (float)calculate_convolution_flops(in_n_c_hi_wi_global_desc,
+                                                            wei_k_c_y_x_global_desc,
+                                                            out_n_k_ho_wo_global_desc) /
+                         (std::size_t(1000) * 1000 * 1000) / ave_time;
+
+            std::cout << "Average time : " << ave_time << " ms, " << perf << " TFlop/s"
+                      << std::endl;
+        }
+#endif
+    }
+};
+} // namespace ck
+#endif

composable_kernel/include/gridwise_operation_wrapper.hpp

@@ -2,7 +2,11 @@
 #define CK_GRIDWISE_OPERATION_KERNEL_WRAPPER
 
 template <typename GridwiseOp, typename... Xs>
-__global__ void run_gridwise_operation(Xs... xs)
+__global__ void
+#if CK_USE_LAUNCH_BOUNDS
+    __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, CK_MIN_BLOCK_PER_CU)
+#endif
+        run_gridwise_operation(Xs... xs)
 {
     GridwiseOp{}.Run(xs...);
 }

composable_kernel/include/kernel_algorithm/gridwise_convolution_backward_data_implicit_gemm_v1r2_nchw_kcyx_nkhw_lds_double_buffer.hpp

@@ -107,8 +107,8 @@ struct GridwiseConvolutionBackwardDataImplicitGemm_v1r2_nchw_kcyx_nkhw_lds_doubl
 
         const auto block_work_id = block_work_desc.CalculateClusterIndex(get_block_1d_id());
 
-        const index_t e_block_data_on_global = block_work_id[0] * EPerBlock;
-        const index_t b_block_data_on_global = block_work_id[1] * BPerBlock;
+        const index_t e_block_data_on_global = block_work_id[Number<0>{}] * EPerBlock;
+        const index_t b_block_data_on_global = block_work_id[Number<1>{}] * BPerBlock;
 
         // output tensor
         //     global tensor in global memory, src of blockwise copy
@@ -151,7 +151,7 @@ struct GridwiseConvolutionBackwardDataImplicitGemm_v1r2_nchw_kcyx_nkhw_lds_doubl
                                                AddressSpace::Vgpr,
                                                AddressSpace::Lds,
                                                InMemoryDataOperation::Set>(
-                {0, b_block_data_on_global, 0}, {0, 0, 0});
+                make_multi_index(0, b_block_data_on_global, 0), make_multi_index(0, 0, 0));
 
         // weight tensor
         //     global tensor in global memory, src of blockwise copy
@@ -191,7 +191,7 @@ struct GridwiseConvolutionBackwardDataImplicitGemm_v1r2_nchw_kcyx_nkhw_lds_doubl
                                                AddressSpace::Vgpr,
                                                AddressSpace::Lds,
                                                InMemoryDataOperation::Set>(
-                {0, e_block_data_on_global, 0}, {0, 0, 0});
+                make_multi_index(0, e_block_data_on_global, 0), make_multi_index(0, 0, 0));
 
         // GEMM definition
         //   c_mtx += transpose(a_mtx) * b_mtx
@@ -354,7 +354,7 @@ struct GridwiseConvolutionBackwardDataImplicitGemm_v1r2_nchw_kcyx_nkhw_lds_doubl
 
         {
 #if 1 // debug
-            // input: register to global memory, atomic add
+      // input: register to global memory, atomic add
             constexpr auto in_memory_op = (Y <= ConvStrideH && X <= ConvStrideW)
                                               ? InMemoryDataOperation::Set
                                               : InMemoryDataOperation::AtomicAdd;
@@ -434,13 +434,13 @@ struct GridwiseConvolutionBackwardDataImplicitGemm_v1r2_nchw_kcyx_nkhw_lds_doubl
                 InThreadCopyDstDataPerWrite_B,
                 AddressSpace::Vgpr,
                 AddressSpace::Global,
-                in_memory_op>({0, 0, 0, 0, 0, 0},
-                              {e_thread_data_on_global / E1,
-                               e_thread_data_on_global % E1,
-                               0,
-                               b_thread_data_on_global / B1,
-                               b_thread_data_on_global % B1,
-                               0})
+                in_memory_op>(make_multi_index(0, 0, 0, 0, 0, 0),
+                              make_multi_index(e_thread_data_on_global / E1,
+                                               e_thread_data_on_global % E1,
+                                               0,
+                                               b_thread_data_on_global / B1,
+                                               b_thread_data_on_global % B1,
+                                               0))
                 .Run(p_in_thread, p_in_global);
         }
     }

composable_kernel/include/kernel_algorithm/gridwise_convolution_backward_data_implicit_gemm_v5r1_nhwc_kyxc_nhwk.hpp

@@ -125,7 +125,7 @@ struct GridwiseConvolutionBackwardDataImplicitGemm_v5r1_nhwc_kyxc_nhwk
         index_t GemmK1 = XDotSlice;
         index_t GemmK2 = K;
 
-        return Array<index_t, 5>{GemmM, GemmN, GemmK0, GemmK1, GemmK2};
+        return make_multi_index(GemmM, GemmN, GemmK0, GemmK1, GemmK2);
     }
 
     __host__ __device__ static constexpr auto GetGemmSize(index_t gemm_id)

composable_kernel/include/kernel_algorithm/gridwise_convolution_forward_implicit_gemm_v4r1_nchw_kcyx_nkhw_lds_double_buffer.hpp

@@ -1,5 +1,5 @@
-#ifndef CK_GRIDWISE_CONVOLUTION_IMPLICIT_GEMM_V4R1_NCHW_KCYX_NKHW_LDS_DOUBLE_BUFFER_HPP
-#define CK_GRIDWISE_CONVOLUTION_IMPLICIT_GEMM_V4R1_NCHW_KCYX_NKHW_LDS_DOUBLE_BUFFER_HPP
+#ifndef CK_GRIDWISE_CONVOLUTION_FORWARD_IMPLICIT_GEMM_V4R1_NCHW_KCYX_NKHW_LDS_DOUBLE_BUFFER_HPP
+#define CK_GRIDWISE_CONVOLUTION_FORWARD_IMPLICIT_GEMM_V4R1_NCHW_KCYX_NKHW_LDS_DOUBLE_BUFFER_HPP
 
 #include "common_header.hpp"
 #include "tensor_descriptor.hpp"
@@ -49,7 +49,7 @@ template <index_t GridSize,
           typename WeiBlockCopyDstAccessOrder,
           index_t WeiBlockCopySrcDataPerRead_E,
           index_t WeiBlockCopyDstDataPerWrite_K>
-struct GridwiseConvolutionImplicitGemm_v4r1_nchw_kcyx_nkhw_lds_double_buffer
+struct GridwiseConvolutionForwardImplicitGemm_v4r1_nchw_kcyx_nkhw_lds_double_buffer
 {
     __device__ void Run(const Float* const __restrict__ p_in_global,
                         const Float* const __restrict__ p_wei_global,
@@ -119,8 +119,8 @@ struct GridwiseConvolutionImplicitGemm_v4r1_nchw_kcyx_nkhw_lds_double_buffer
 
         const auto block_work_id = block_work_desc.CalculateClusterIndex(get_block_1d_id());
 
-        const index_t k_block_data_on_global = block_work_id[0] * KPerBlock;
-        const index_t b_block_data_on_global = block_work_id[1] * BPerBlock;
+        const index_t k_block_data_on_global = block_work_id[I0] * KPerBlock;
+        const index_t b_block_data_on_global = block_work_id[I1] * BPerBlock;
 
         // input tensor
         //     global tensor in global memory
@@ -183,7 +183,7 @@ struct GridwiseConvolutionImplicitGemm_v4r1_nchw_kcyx_nkhw_lds_double_buffer
                                                AddressSpace::Vgpr,
                                                AddressSpace::Lds,
                                                InMemoryDataOperation::Set>(
-                {0, 0, b_block_data_on_global, 0}, {0, 0, 0, 0});
+                make_multi_index(0, 0, b_block_data_on_global, 0), make_multi_index(0, 0, 0, 0));
 
         // weight tensor
         //     global tensor in global memory, src of blockwise copy
@@ -226,7 +226,7 @@ struct GridwiseConvolutionImplicitGemm_v4r1_nchw_kcyx_nkhw_lds_double_buffer
                                                AddressSpace::Vgpr,
                                                AddressSpace::Lds,
                                                InMemoryDataOperation::Set>(
-                {0, k_block_data_on_global}, {0, 0});
+                make_multi_index(0, k_block_data_on_global), make_multi_index(0, 0));
 
         // GEMM definition
         //   c_mtx += transpose(a_mtx) * b_mtx
@@ -439,12 +439,12 @@ struct GridwiseConvolutionImplicitGemm_v4r1_nchw_kcyx_nkhw_lds_double_buffer
                 1,
                 AddressSpace::Vgpr,
                 AddressSpace::Global,
-                InMemoryDataOperation::Set>({0, 0, 0, 0, 0},
-                                            {k_thread_data_on_global / K1,
-                                             k_thread_data_on_global % K1,
-                                             0,
-                                             b_thread_data_on_global,
-                                             0})
+                InMemoryDataOperation::Set>(make_multi_index(0, 0, 0, 0, 0),
+                                            make_multi_index(k_thread_data_on_global / K1,
+                                                             k_thread_data_on_global % K1,
+                                                             0,
+                                                             b_thread_data_on_global,
+                                                             0))
                 .Run(p_out_thread, p_out_global);
         }
     }

composable_kernel/include/kernel_algorithm/gridwise_convolution_forward_implicit_gemm_v4r4_nchw_kcyx_nkhw.hpp

@@ -1,5 +1,5 @@
-#ifndef CK_GRIDWISE_CONVOLUTION_IMPLICIT_GEMM_V4R4_NCHW_KCYX_NKHW_HPP
-#define CK_GRIDWISE_CONVOLUTION_IMPLICIT_GEMM_V4R4_NCHW_KCYX_NKHW_HPP
+#ifndef CK_GRIDWISE_CONVOLUTION_FORWARD_IMPLICIT_GEMM_V4R4_NCHW_KCYX_NKHW_HPP
+#define CK_GRIDWISE_CONVOLUTION_FORWARD_IMPLICIT_GEMM_V4R4_NCHW_KCYX_NKHW_HPP
 
 #include "common_header.hpp"
 #include "tensor_descriptor.hpp"
@@ -43,7 +43,7 @@ template <index_t GridSize,
           index_t GemmBBlockCopySrcDataPerRead_GemmN,
           index_t GemmBBlockCopyDstDataPerWrite_GemmN,
           index_t GemmCThreadCopyDstDataPerWrite_GemmN1>
-struct GridwiseConvolutionImplicitGemm_v4r4_nchw_kcyx_nkhw
+struct GridwiseConvolutionForwardImplicitGemm_v4r4_nchw_kcyx_nkhw
 {
     __device__ void Run(const Float* const __restrict__ p_in_global,
                         const Float* const __restrict__ p_wei_global,

composable_kernel/include/kernel_algorithm/gridwise_convolution_forward_implicit_gemm_v4r4_nhwc_kyxc_nhwk.hpp

@@ -0,0 +1,162 @@
+#ifndef CK_GRIDWISE_CONVOLUTION_FORWARD_IMPLICIT_GEMM_V4R4_NHWC_KYXC_NHWK_HPP
+#define CK_GRIDWISE_CONVOLUTION_FORWARD_IMPLICIT_GEMM_V4R4_NHWC_KYXC_NHWK_HPP
+
+#include "common_header.hpp"
+#include "tensor_descriptor.hpp"
+#include "tensor_descriptor_helper.hpp"
+#include "gridwise_gemm.hpp"
+
+namespace ck {
+
+// GemmM = K
+// GemmN = N * Ho * Wo
+// GemmK = C * Y * X
+template <index_t GridSize,
+          index_t BlockSize,
+          typename Float,
+          typename AccFloat,
+          typename InGlobalDesc,
+          typename WeiGlobalDesc,
+          typename OutGlobalDesc,
+          typename ConvStrides,
+          typename ConvDilations,
+          typename InLeftPads,
+          typename InRightPads,
+          index_t GemmMPerBlock,
+          index_t GemmNPerBlock,
+          index_t GemmKPerBlock,
+          index_t GemmMPerThread,
+          index_t GemmNPerThread,
+          index_t GemmKPerThread,
+          index_t GemmMLevel0Cluster,
+          index_t GemmNLevel0Cluster,
+          index_t GemmMLevel1Cluster,
+          index_t GemmNLevel1Cluster,
+          index_t ThreadGemmDataPerRead_GemmM,
+          index_t ThreadGemmDataPerRead_GemmN,
+          typename GemmABlockCopyThreadSliceLengths_GemmK_GemmM,
+          typename GemmABlockCopyThreadClusterLengths_GemmK_GemmM,
+          index_t GemmABlockCopySrcDataPerRead_GemmK,
+          index_t GemmABlockCopyDstDataPerWrite_GemmM,
+          typename GemmBBlockCopyThreadSliceLengths_GemmK_GemmN,
+          typename GemmBBlockCopyThreadClusterLengths_GemmK_GemmN,
+          index_t GemmBBlockCopySrcDataPerRead_GemmK,
+          index_t GemmBBlockCopyDstDataPerWrite_GemmN,
+          index_t GemmCThreadCopyDstDataPerWrite_GemmM1>
+struct GridwiseConvolutionForwardImplicitGemm_v4r4_nhwc_kyxc_nhwk
+{
+    __device__ void Run(const Float* const __restrict__ p_in_global,
+                        const Float* const __restrict__ p_wei_global,
+                        Float* const __restrict__ p_out_global) const
+    {
+        constexpr auto I0 = Number<0>{};
+        constexpr auto I1 = Number<1>{};
+        constexpr auto I2 = Number<2>{};
+        constexpr auto I3 = Number<3>{};
+
+        constexpr auto in_n_hi_wi_c_global_desc  = InGlobalDesc{};
+        constexpr auto wei_k_y_x_c_global_desc   = WeiGlobalDesc{};
+        constexpr auto out_n_ho_wo_k_global_desc = OutGlobalDesc{};
+
+        constexpr index_t N  = in_n_hi_wi_c_global_desc.GetLengths()[I0];
+        constexpr index_t Hi = in_n_hi_wi_c_global_desc.GetLengths()[I1];
+        constexpr index_t Wi = in_n_hi_wi_c_global_desc.GetLengths()[I2];
+        constexpr index_t C  = in_n_hi_wi_c_global_desc.GetLengths()[I3];
+
+        constexpr index_t K  = out_n_ho_wo_k_global_desc.GetLengths()[I3];
+        constexpr index_t Ho = out_n_ho_wo_k_global_desc.GetLengths()[I1];
+        constexpr index_t Wo = out_n_ho_wo_k_global_desc.GetLengths()[I2];
+
+        constexpr index_t Y = wei_k_y_x_c_global_desc.GetLengths()[I1];
+        constexpr index_t X = wei_k_y_x_c_global_desc.GetLengths()[I2];
+
+        constexpr index_t ConvStrideH = ConvStrides{}[I0];
+        constexpr index_t ConvStrideW = ConvStrides{}[I1];
+
+        constexpr index_t ConvDilationH = ConvDilations{}[I0];
+        constexpr index_t ConvDilationW = ConvDilations{}[I1];
+
+        // weight tensor
+        constexpr auto wei_gemmk_gemmm_global_desc = reorder_tensor_descriptor_given_upper2lower(
+            unfold_tensor_descriptor(wei_k_y_x_c_global_desc, I1, I3), Sequence<1, 0>{});
+
+        // input tensor
+        constexpr auto in_n_hip_wip_c_global_desc =
+            transform_tensor_descriptor(in_n_hi_wi_c_global_desc,
+                                        make_tuple(PassThrough<N>{},
+                                                   Pad<Sequence<Hi, Wi>, InLeftPads, InRightPads>{},
+                                                   PassThrough<C>{}),
+                                        make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3>{}),
+                                        make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3>{}));
+
+        constexpr index_t Hip = in_n_hip_wip_c_global_desc.GetLengths()[I1];
+        constexpr index_t Wip = in_n_hip_wip_c_global_desc.GetLengths()[I2];
+
+        constexpr auto in_n_y_ho_x_wo_c_global_desc = transform_tensor_descriptor(
+            in_n_hip_wip_c_global_desc,
+            make_tuple(PassThrough<N>{},
+                       Embed<Hip, Sequence<Y, Ho>, Sequence<ConvDilationH, ConvStrideH, 0>>{},
+                       Embed<Wip, Sequence<X, Wo>, Sequence<ConvDilationW, ConvStrideW, 0>>{},
+                       PassThrough<C>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+            make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5>{}));
+
+        constexpr auto in_gemmk_gemmn_global_desc = transform_tensor_descriptor(
+            in_n_y_ho_x_wo_c_global_desc,
+            make_tuple(Merge<Sequence<Y, X, C>>{}, Merge<Sequence<N, Ho, Wo>>{}),
+            make_tuple(Sequence<1, 3, 5>{}, Sequence<0, 2, 4>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+        // output tensor
+        constexpr auto out_gemmm_gemmn_global_desc = transform_tensor_descriptor(
+            unfold_tensor_descriptor(out_n_ho_wo_k_global_desc, I0, I2),
+            make_tuple(PassThrough<K>{}, Merge<Sequence<N * Ho * Wo>>{}),
+            make_tuple(Sequence<1>{}, Sequence<0>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+        // GEMM
+        constexpr auto gridwise_gemm =
+            GridwiseGemmTransposedANormalBNormalC_v1<GridSize,
+                                                     BlockSize,
+                                                     Float,
+                                                     AccFloat,
+                                                     decltype(wei_gemmk_gemmm_global_desc),
+                                                     decltype(in_gemmk_gemmn_global_desc),
+                                                     decltype(out_gemmm_gemmn_global_desc),
+                                                     InMemoryDataOperation::Set,
+                                                     GemmMPerBlock,
+                                                     GemmNPerBlock,
+                                                     GemmKPerBlock,
+                                                     GemmMPerThread,
+                                                     GemmNPerThread,
+                                                     GemmKPerThread,
+                                                     GemmMLevel0Cluster,
+                                                     GemmNLevel0Cluster,
+                                                     GemmMLevel1Cluster,
+                                                     GemmNLevel1Cluster,
+                                                     ThreadGemmDataPerRead_GemmM,
+                                                     ThreadGemmDataPerRead_GemmN,
+                                                     GemmABlockCopyThreadSliceLengths_GemmK_GemmM,
+                                                     GemmABlockCopyThreadClusterLengths_GemmK_GemmM,
+                                                     Sequence<1, 0>,
+                                                     Sequence<1, 0>,
+                                                     0,
+                                                     GemmABlockCopySrcDataPerRead_GemmK,
+                                                     GemmABlockCopyDstDataPerWrite_GemmM,
+                                                     GemmBBlockCopyThreadSliceLengths_GemmK_GemmN,
+                                                     GemmBBlockCopyThreadClusterLengths_GemmK_GemmN,
+                                                     Sequence<1, 0>,
+                                                     Sequence<1, 0>,
+                                                     0,
+                                                     GemmBBlockCopySrcDataPerRead_GemmK,
+                                                     GemmBBlockCopyDstDataPerWrite_GemmN,
+                                                     Sequence<2, 3, 0, 1>,
+                                                     1,
+                                                     GemmCThreadCopyDstDataPerWrite_GemmM1>{};
+
+        gridwise_gemm.Run(p_wei_global, p_in_global, p_out_global);
+    }
+};
+
+} // namespace ck
+#endif

composable_kernel/include/tensor_description/array_multi_index.hpp

@@ -0,0 +1,77 @@
+#ifndef CK_ARRAY_MULTI_INDEX_HPP
+#define CK_ARRAY_MULTI_INDEX_HPP
+
+#include "common_header.hpp"
+
+namespace ck {
+
+template <index_t N>
+using MultiIndex = Array<index_t, N>;
+
+template <typename... Xs>
+__host__ __device__ constexpr auto make_multi_index(Xs&&... xs)
+{
+    return make_array<index_t>(index_t{xs}...);
+}
+
+template <index_t NSize>
+__host__ __device__ constexpr auto make_zero_multi_index()
+{
+    return unpack([](auto... xs) { return make_multi_index(xs...); },
+                  typename uniform_sequence_gen<NSize, 0>::type{});
+}
+
+template <typename T>
+__host__ __device__ constexpr auto to_multi_index(const T& x)
+{
+    return unpack([](auto... ys) { return make_multi_index(ys...); }, x);
+}
+
+template <index_t NSize, typename X>
+__host__ __device__ constexpr auto operator+=(MultiIndex<NSize>& y, const X& x)
+{
+    static_assert(X::Size() == NSize, "wrong! size not the same");
+    static_for<0, NSize, 1>{}([&](auto i) { y(i) += x[i]; });
+    return y;
+}
+
+template <index_t NSize, typename X>
+__host__ __device__ constexpr auto operator-=(MultiIndex<NSize>& y, const X& x)
+{
+    static_assert(X::Size() == NSize, "wrong! size not the same");
+    static_for<0, NSize, 1>{}([&](auto i) { y(i) -= x[i]; });
+    return y;
+}
+
+template <index_t NSize, typename T>
+__host__ __device__ constexpr auto operator+(const MultiIndex<NSize>& a, const T& b)
+{
+    using type = MultiIndex<NSize>;
+    static_assert(T::Size() == NSize, "wrong! size not the same");
+    type r;
+    static_for<0, NSize, 1>{}([&](auto i) { r(i) = a[i] + b[i]; });
+    return r;
+}
+
+template <index_t NSize, typename T>
+__host__ __device__ constexpr auto operator-(const MultiIndex<NSize>& a, const T& b)
+{
+    using type = MultiIndex<NSize>;
+    static_assert(T::Size() == NSize, "wrong! size not the same");
+    type r;
+    static_for<0, NSize, 1>{}([&](auto i) { r(i) = a[i] - b[i]; });
+    return r;
+}
+
+template <index_t NSize, typename T>
+__host__ __device__ constexpr auto operator*(const MultiIndex<NSize>& a, const T& b)
+{
+    using type = MultiIndex<NSize>;
+    static_assert(T::Size() == NSize, "wrong! size not the same");
+    type r;
+    static_for<0, NSize, 1>{}([&](auto i) { r(i) = a[i] * b[i]; });
+    return r;
+}
+
+} // namespace ck
+#endif

composable_kernel/include/tensor_description/cluster_descriptor.hpp

@@ -0,0 +1,48 @@
+#ifndef CK_CLUSTER_DESCRIPTOR_HPP
+#define CK_CLUSTER_DESCRIPTOR_HPP
+
+#include "common_header.hpp"
+
+// TODO remove dependency on deprecated tensor descriptor
+#include "tensor_descriptor.hpp"
+
+namespace ck {
+
+// a cluster map 1d index to N-d index
+template <typename Lengths, typename ArrangeOrder>
+struct ClusterDescriptor
+{
+    static constexpr index_t nDim = Lengths::Size();
+
+    static constexpr auto mDesc = transform_tensor_descriptor(
+        make_native_tensor_descriptor_packed(Lengths{}),
+        make_tuple(Merge<decltype(Lengths::ReorderGivenNew2Old(ArrangeOrder{}))>{}),
+        make_tuple(ArrangeOrder{}),
+        make_tuple(Sequence<0>{}));
+
+    __host__ __device__ constexpr ClusterDescriptor()
+    {
+        static_assert(Lengths::Size() == nDim && ArrangeOrder::Size() == nDim,
+                      "wrong! size not the same");
+
+        static_assert(is_valid_sequence_map<ArrangeOrder>{}, "wrong! ArrangeOrder is wrong");
+    }
+
+    __host__ __device__ static constexpr index_t GetElementSize() { return mDesc.GetElementSize(); }
+
+    __host__ __device__ static constexpr auto CalculateClusterIndex(index_t idx_1d)
+    {
+        return mDesc.CalculateLowerIndex(MultiIndex<1>{idx_1d});
+    }
+};
+
+template <typename Lengths,
+          typename ArrangeOrder = typename arithmetic_sequence_gen<0, Lengths::Size(), 1>::type>
+__host__ __device__ constexpr auto make_cluster_descriptor(
+    Lengths, ArrangeOrder order = typename arithmetic_sequence_gen<0, Lengths::Size(), 1>::type{})
+{
+    return ClusterDescriptor<Lengths, decltype(order)>{};
+}
+
+} // namespace ck
+#endif

composable_kernel/include/tensor_description/dynamic_multi_index_transform_helper.hpp

@@ -0,0 +1,74 @@
+#ifndef CK_DYNAMIC_MULTI_INDEX_TRANSFORM_HELPER_HPP
+#define CK_DYNAMIC_MULTI_INDEX_TRANSFORM_HELPER_HPP
+
+#include "common_header.hpp"
+#include "dynamic_multi_index_transform.hpp"
+
+namespace ck {
+
+template <typename LowLength>
+__host__ __device__ constexpr auto make_pass_through_transform(const LowLength& low_length)
+{
+    return DynamicPassThrough<LowLength>{low_length};
+}
+
+template <typename LowLength, typename LeftPad, typename RightPad, bool SkipIsValidCheck = false>
+__host__ __device__ constexpr auto
+make_pad_transform(const LowLength& low_length,
+                   const LeftPad& left_pad,
+                   const RightPad& right_pad,
+                   integral_constant<bool, SkipIsValidCheck> = integral_constant<bool, false>{})
+{
+    return DynamicPad<LowLength, LeftPad, RightPad, SkipIsValidCheck>{
+        low_length, left_pad, right_pad};
+}
+
+template <typename LowLength, typename LeftPad, bool SkipIsValidCheck = false>
+__host__ __device__ constexpr auto make_left_pad_transform(
+    const LowLength& low_length,
+    const LeftPad& left_pad,
+    integral_constant<bool, SkipIsValidCheck> = integral_constant<bool, false>{})
+{
+    return DynamicLeftPad<LowLength, LeftPad, SkipIsValidCheck>{low_length, left_pad};
+}
+
+template <typename LowLength, typename RightPad, bool SkipIsValidCheck>
+__host__ __device__ constexpr auto make_right_pad_transform(
+    const LowLength& low_length,
+    const RightPad& right_pad,
+    integral_constant<bool, SkipIsValidCheck> = integral_constant<bool, false>{})
+{
+    return DynamicRightPad<LowLength, RightPad, SkipIsValidCheck>{low_length, right_pad};
+}
+
+template <typename UpLengths,
+          typename Coefficients,
+          typename std::enable_if<UpLengths::Size() == Coefficients::Size(), bool>::type = false>
+__host__ __device__ constexpr auto make_embed_transform(const UpLengths& up_lengths,
+                                                        const Coefficients& coefficients)
+{
+    return DynamicEmbed<UpLengths, Coefficients>{up_lengths, coefficients};
+}
+
+template <typename LowLengths>
+__host__ __device__ constexpr auto make_merge_transform(const LowLengths& low_lengths)
+{
+    return DynamicMerge<LowLengths>{low_lengths};
+}
+
+template <typename UpLengths, bool Use24BitIntegerCalculation = false>
+__host__ __device__ constexpr auto make_unmerge_transform(
+    const UpLengths& up_lengths,
+    integral_constant<bool, Use24BitIntegerCalculation> = integral_constant<bool, false>{})
+{
+    return DynamicUnMerge<UpLengths, Use24BitIntegerCalculation>{up_lengths};
+}
+
+template <typename LowerIndex>
+__host__ __device__ constexpr auto make_freeze_transform(const LowerIndex& low_idx)
+{
+    return DynamicFreeze<LowerIndex>{low_idx};
+}
+
+} // namespace ck
+#endif

composable_kernel/include/tensor_description/dynamic_tensor_descriptor.hpp

@@ -0,0 +1,608 @@
+#ifndef CK_DYNAMIC_TENSOR_DESCRIPTOR_HPP
+#define CK_DYNAMIC_TENSOR_DESCRIPTOR_HPP
+
+#include "common_header.hpp"
+#include "dynamic_multi_index_transform.hpp"
+
+namespace ck {
+
+template <index_t NDimHidden, typename VisibleDimensionIds>
+struct DynamicTensorCoordinate;
+
+template <index_t NTransform, index_t NDimVisible, typename UpdateLowerIndexHack>
+struct DynamicTensorCoordinateIterator;
+
+template <typename LowerDimensionIdss, typename UpperDimensionIdss>
+__host__ __device__ constexpr index_t GetNumOfHiddenDimension(LowerDimensionIdss,
+                                                              UpperDimensionIdss)
+{
+    constexpr auto all_low_dim_ids =
+        unpack([](auto&&... xs) constexpr { return merge_sequences(xs...); }, LowerDimensionIdss{});
+
+    constexpr auto all_up_dim_ids =
+        unpack([](auto&&... xs) constexpr { return merge_sequences(xs...); }, UpperDimensionIdss{});
+
+    constexpr auto all_dim_ids = merge_sequences(all_low_dim_ids, all_up_dim_ids);
+
+    using unique_sort_all_dim_ids = typename sequence_unique_sort<decltype(all_dim_ids),
+                                                                  math::less<index_t>,
+                                                                  math::equal<index_t>>::type;
+
+    return unique_sort_all_dim_ids::Size();
+}
+
+// Transforms: Tuple<transforms...>
+// LowerDimensionIdss : Tuple<Sequence<...>, ...>
+// UpperDimensionIdss : Tuple<Sequence<...>, ...>
+// VisibleDimensionIds> : Sequence<...>
+template <typename Transforms,
+          typename LowerDimensionIdss,
+          typename UpperDimensionIdss,
+          typename VisibleDimensionIds,
+          typename ElementSpaceSize>
+struct DynamicTensorDescriptor
+{
+    // TODO make these private
+    __host__ __device__ static constexpr index_t GetNumOfTransform() { return Transforms::Size(); }
+
+    __host__ __device__ static constexpr index_t GetNumOfVisibleDimension()
+    {
+        return VisibleDimensionIds::Size();
+    }
+
+    __host__ __device__ static constexpr index_t GetNumOfHiddenDimension()
+    {
+        constexpr auto all_low_dim_ids =
+            unpack([](auto&&... xs) constexpr { return merge_sequences(xs...); },
+                   LowerDimensionIdss{});
+
+        constexpr auto all_up_dim_ids =
+            unpack([](auto&&... xs) constexpr { return merge_sequences(xs...); },
+                   UpperDimensionIdss{});
+
+        constexpr auto all_dim_ids = merge_sequences(all_low_dim_ids, all_up_dim_ids);
+
+        using unique_sort_all_dim_ids = typename sequence_unique_sort<decltype(all_dim_ids),
+                                                                      math::less<index_t>,
+                                                                      math::equal<index_t>>::type;
+
+        return unique_sort_all_dim_ids::Size();
+    }
+
+    __host__ __device__ static constexpr auto InitializeElementSize(const Transforms& transforms)
+    {
+        const auto lengths = generate_tuple(
+            [&](auto idim_visible) {
+                constexpr auto tmp = GetTransformAndItsUpperDimension(idim_visible);
+
+                constexpr index_t itran   = tmp[Number<0>{}];
+                constexpr index_t idim_up = tmp[Number<1>{}];
+                constexpr bool found      = tmp[Number<2>{}];
+
+                static_assert(found == true,
+                              "wrong! not found matching transformation and upper-dimension");
+
+                const auto length =
+                    transforms[Number<itran>{}].GetUpperLengths()[Number<idim_up>{}];
+
+                return length;
+            },
+            Number<ndim_visible_>{});
+
+        // TODO: make container_reduce support tuple of Number and index_t
+        return container_reduce(lengths, math::multiplies_v2{}, Number<1>{});
+    }
+
+    template <index_t IDim>
+    __host__ __device__ static constexpr auto GetTransformAndItsUpperDimension(Number<IDim>)
+    {
+        constexpr auto idim_visible = Number<IDim>{};
+
+        constexpr index_t idim_hidden = VisibleDimensionIds::At(idim_visible);
+
+        index_t itran_found   = 0;
+        index_t idim_up_found = 0;
+        bool found            = false;
+
+        static_for<0, ntransform_, 1>{}([&](auto itran) {
+            constexpr auto up_dim_ids = UpperDimensionIdss{}[itran];
+
+            static_for<0, up_dim_ids.Size(), 1>{}([&](auto idim_up) {
+                if constexpr(up_dim_ids[idim_up] == idim_hidden)
+                {
+                    itran_found   = itran;
+                    idim_up_found = idim_up;
+                    found         = true;
+                }
+            });
+        });
+
+        return make_tuple(itran_found, idim_up_found, found);
+    }
+
+    constexpr static index_t ntransform_   = GetNumOfTransform();
+    constexpr static index_t ndim_visible_ = GetNumOfVisibleDimension();
+    constexpr static index_t ndim_hidden_  = GetNumOfHiddenDimension();
+
+    using VisibleIndex = MultiIndex<ndim_visible_>;
+    using HiddenIndex  = MultiIndex<ndim_hidden_>;
+    using Coordinate   = DynamicTensorCoordinate<ndim_hidden_, VisibleDimensionIds>;
+
+    // may be index_t or Number<>
+    using ElementSize = remove_cv_t<decltype(InitializeElementSize(Transforms{}))>;
+
+    public:
+    __host__ __device__ constexpr DynamicTensorDescriptor() = default;
+
+    __host__ __device__ constexpr DynamicTensorDescriptor(const Transforms& transforms,
+                                                          ElementSpaceSize element_space_size)
+        : transforms_{transforms},
+          element_size_{InitializeElementSize(transforms)},
+          element_space_size_{element_space_size}
+
+    {
+        static_assert(Transforms::Size() == ntransform_ &&
+                          LowerDimensionIdss::Size() == ntransform_ &&
+                          UpperDimensionIdss::Size() == ntransform_,
+                      "wrong! inconsistent # of transformations");
+
+        // TODO check dependency of dimensions is valid
+    }
+
+    __host__ __device__ static constexpr index_t GetNumOfDimension()
+    {
+        return GetNumOfVisibleDimension();
+    }
+
+    template <index_t IDim>
+    __host__ __device__ constexpr auto GetLength(Number<IDim>) const
+    {
+        static_assert(IDim >= 0 && IDim < ndim_visible_, "wrong! out of range");
+
+        constexpr auto tmp = GetTransformAndItsUpperDimension(Number<IDim>{});
+
+        constexpr index_t itran   = tmp[Number<0>{}];
+        constexpr index_t idim_up = tmp[Number<1>{}];
+        constexpr bool found      = tmp[Number<2>{}];
+
+        static_assert(found == true,
+                      "wrong! not found matching transformation and upper-dimension");
+
+        return transforms_[Number<itran>{}].GetUpperLengths()[Number<idim_up>{}];
+    }
+
+    __host__ __device__ constexpr auto GetElementSize() const { return element_size_; }
+
+    __host__ __device__ constexpr auto GetElementSpaceSize() const { return element_space_size_; }
+
+    template <typename Idx>
+    __host__ __device__ constexpr index_t CalculateOffset(const Idx& idx) const
+    {
+        static_assert(Idx::Size() == GetNumOfDimension(), "wrong! inconsistent # of dimension");
+
+        return make_dynamic_tensor_coordinate(*this, idx).GetOffset();
+    }
+
+    // TODO make these private
+    __host__ __device__ constexpr const auto& GetTransforms() const { return transforms_; }
+
+    __host__ __device__ static constexpr auto GetLowerDimensionIdss()
+    {
+        return LowerDimensionIdss{};
+    }
+
+    __host__ __device__ static constexpr auto GetUpperDimensionIdss()
+    {
+        return UpperDimensionIdss{};
+    }
+
+    __host__ __device__ static constexpr auto GetVisibleDimensionIds()
+    {
+        return VisibleDimensionIds{};
+    }
+
+    __host__ __device__ static constexpr bool IsKnownAtCompileTime()
+    {
+        bool is_known = true;
+
+        static_for<0, Transforms::Size(), 1>{}([&](auto i) {
+            is_known &=
+                remove_cv_t<remove_reference_t<decltype(Transforms{}[i])>>::IsKnownAtCompileTime();
+        });
+
+        return is_known && is_known_at_compile_time<ElementSize>::value &&
+               is_known_at_compile_time<ElementSpaceSize>::value;
+    }
+
+    __host__ __device__ void Print() const
+    {
+        printf("{");
+        printf("DynamicTensorDescriptor, ");
+        static_for<0, ntransform_, 1>{}([&](auto i) {
+            printf("transforms: ");
+            transforms_[i].Print();
+            printf("LowerDimensionIds:");
+            LowerDimensionIdss{}.At(i).Print();
+            printf("UpperDimensionIds:");
+            UpperDimensionIdss{}.At(i).Print();
+        });
+        printf("}");
+
+        VisibleDimensionIds::Print();
+    }
+
+    // TODO make these private
+    Transforms transforms_;
+    ElementSize element_size_;
+    ElementSpaceSize element_space_size_;
+};
+
+template <index_t NDimHidden, typename VisibleDimensionIds>
+struct DynamicTensorCoordinate
+{
+    // TODO make these private
+    static constexpr index_t ndim_visible_ = VisibleDimensionIds::Size();
+
+    using HiddenIndex  = MultiIndex<NDimHidden>;
+    using VisibleIndex = MultiIndex<ndim_visible_>;
+
+    public:
+    __host__ __device__ constexpr DynamicTensorCoordinate() = default;
+
+    __host__ __device__ constexpr DynamicTensorCoordinate(const HiddenIndex& idx_hidden)
+        : idx_hidden_{idx_hidden}
+    {
+    }
+
+    __host__ __device__ constexpr auto GetIndex() const { return GetVisibleIndex(); }
+
+    __host__ __device__ constexpr index_t GetOffset() const { return idx_hidden_[Number<0>{}]; }
+
+    // TODO make these private
+    __host__ __device__ constexpr const auto& GetHiddenIndex() const { return idx_hidden_; }
+
+    __host__ __device__ auto& GetHiddenIndex() { return idx_hidden_; }
+
+    __host__ __device__ constexpr auto GetVisibleIndex() const
+    {
+        return get_container_subset(idx_hidden_, VisibleDimensionIds{});
+    }
+
+    // TODO make these private
+    HiddenIndex idx_hidden_;
+};
+
+template <index_t NTransform, index_t NDimVisible, typename UpdateLowerIndexHack>
+struct DynamicTensorCoordinateIterator
+{
+    // TODO make these private
+    using VisibleIndex = MultiIndex<NDimVisible>;
+
+    public:
+    __host__ __device__ constexpr DynamicTensorCoordinateIterator() = default;
+
+    __host__ __device__ constexpr DynamicTensorCoordinateIterator(
+        const VisibleIndex& idx_diff_visible, const MultiIndex<NTransform>& do_transforms)
+        : idx_diff_visible_{idx_diff_visible}, do_transforms_{do_transforms}
+    {
+    }
+
+    __host__ __device__ constexpr const auto& GetIndexDiff() const { return GetVisibleIndexDiff(); }
+
+    // TODO make these private
+    __host__ __device__ constexpr const auto& GetVisibleIndexDiff() const
+    {
+        return idx_diff_visible_;
+    }
+
+    VisibleIndex idx_diff_visible_;
+    MultiIndex<NTransform> do_transforms_;
+
+    // HACK: control UpdateLowerIndex()
+    static constexpr UpdateLowerIndexHack update_lower_index_hack_;
+};
+
+// TODO: How to fix this? It uses an struct instead of lambda because lambda
+// doesn't have constructor, and to put it outside the scope where it is used
+// (transform_dynamic_tensor_descriptor) because template cannot be defined inside a function
+// template
+template <typename NewTransforms>
+struct lambda_get_up_dim_num
+{
+    template <typename I>
+    __host__ __device__ constexpr auto operator()(I) const
+    {
+        using Tran = remove_reference_t<decltype(NewTransforms{}.At(I{}))>;
+        return Number<Tran::GetNumOfUpperDimension()>{};
+    }
+};
+
+template <typename OldTensorDescriptor,
+          typename NewTransforms,
+          typename NewLowerDimensionOldVisibleIdss,
+          typename NewUpperDimensionNewVisibleIdss>
+__host__ __device__ constexpr auto
+transform_dynamic_tensor_descriptor(const OldTensorDescriptor& old_tensor_desc,
+                                    const NewTransforms& new_transforms,
+                                    NewLowerDimensionOldVisibleIdss,
+                                    NewUpperDimensionNewVisibleIdss)
+{
+    // lower dimension's hidden idss
+    // convert lower dimension visible idss (tuple of sequences) to hidden idss (tuple of
+    // sequences)
+    constexpr auto low_dim_hidden_idss = transform_tuples(
+        // convert lower dimension visible ids (a sequence) to hidden ids (a sequence)
+        [](auto low_dim_visible_ids) constexpr {
+            return transform_sequences(
+                // convert lower dimension visible id to hidden id
+                [](auto low_dim_visible_id) constexpr {
+                    return OldTensorDescriptor::GetVisibleDimensionIds()[low_dim_visible_id];
+                },
+                low_dim_visible_ids);
+        },
+        NewLowerDimensionOldVisibleIdss{});
+
+    constexpr index_t num_new_transform = NewTransforms::Size();
+
+    // upper dimension's hidden idss
+    constexpr index_t old_hidden_dim_number = OldTensorDescriptor::GetNumOfHiddenDimension();
+
+    constexpr auto up_dim_numbers =
+        generate_sequence(lambda_get_up_dim_num<NewTransforms>{}, Number<num_new_transform>{});
+
+    constexpr auto up_dim_numbers_scan = merge_sequences(
+        Sequence<0>{}, inclusive_scan_sequence(up_dim_numbers, math::plus<index_t>{}, Number<0>{}));
+
+    constexpr auto up_dim_hidden_idss =
+        generate_tuple([ old_hidden_dim_number, up_dim_numbers_scan ](auto i) constexpr {
+            return
+                typename arithmetic_sequence_gen<old_hidden_dim_number + up_dim_numbers_scan[i],
+                                                 old_hidden_dim_number + up_dim_numbers_scan[i + 1],
+                                                 1>::type{};
+        },
+                       Number<num_new_transform>{});
+
+    // new visible dimension's hidden ids
+    constexpr auto unordered_new_visible_dim_hidden_ids =
+        unpack([](auto... xs) constexpr { return merge_sequences(xs...); }, up_dim_hidden_idss);
+
+    constexpr auto new_visible_dim_unordered2ordered =
+        unpack([](auto... xs) constexpr { return merge_sequences(xs...); },
+               NewUpperDimensionNewVisibleIdss{});
+
+    constexpr auto new_visible_dim_hidden_ids =
+        unordered_new_visible_dim_hidden_ids.ReorderGivenOld2New(new_visible_dim_unordered2ordered);
+
+    // put everything together
+    const auto all_transforms = container_cat(old_tensor_desc.GetTransforms(), new_transforms);
+
+    constexpr auto all_low_dim_hidden_idss =
+        container_cat(OldTensorDescriptor::GetLowerDimensionIdss(), low_dim_hidden_idss);
+
+    constexpr auto all_up_dim_hidden_idss =
+        container_cat(OldTensorDescriptor::GetUpperDimensionIdss(), up_dim_hidden_idss);
+
+    const auto element_space_size = old_tensor_desc.GetElementSpaceSize();
+
+    return DynamicTensorDescriptor<remove_cv_t<decltype(all_transforms)>,
+                                   remove_cv_t<decltype(all_low_dim_hidden_idss)>,
+                                   remove_cv_t<decltype(all_up_dim_hidden_idss)>,
+                                   remove_cv_t<decltype(new_visible_dim_hidden_ids)>,
+                                   remove_cv_t<decltype(element_space_size)>>{all_transforms,
+                                                                              element_space_size};
+}
+
+template <typename TensorDesc, typename VisibleIndex>
+__host__ __device__ constexpr auto make_dynamic_tensor_coordinate(const TensorDesc& tensor_desc,
+                                                                  const VisibleIndex& idx_visible)
+{
+    static_assert(TensorDesc::GetNumOfDimension() == VisibleIndex::Size(),
+                  "wrong! # of dimension inconsistent");
+
+    constexpr index_t ntransform   = TensorDesc::GetNumOfTransform();
+    constexpr index_t ndim_hidden  = TensorDesc::GetNumOfHiddenDimension();
+    constexpr index_t ndim_visible = TensorDesc::GetNumOfVisibleDimension();
+    constexpr auto visible_dim_ids = TensorDesc::GetVisibleDimensionIds();
+
+    MultiIndex<ndim_hidden> idx_hidden;
+
+    // initialize visible index
+    set_container_subset(idx_hidden, visible_dim_ids, idx_visible);
+
+    // calculate hidden index
+    static_for<ntransform, 0, -1>{}([&tensor_desc, &idx_hidden](auto itran_p1) {
+        auto itran              = itran_p1 - Number<1>{};
+        const auto& tran        = tensor_desc.GetTransforms().At(itran);
+        constexpr auto dims_low = TensorDesc::GetLowerDimensionIdss().At(itran);
+        constexpr auto dims_up  = TensorDesc::GetUpperDimensionIdss().At(itran);
+
+        const auto idx_up = get_container_subset(idx_hidden, dims_up);
+
+        MultiIndex<dims_low.Size()> idx_low;
+
+        tran.CalculateLowerIndex(idx_low, idx_up);
+
+        set_container_subset(idx_hidden, dims_low, idx_low);
+    });
+
+    return DynamicTensorCoordinate<ndim_hidden, decltype(visible_dim_ids)>{idx_hidden};
+}
+
+// UpdateLowerIndexHack: Sequence<...>
+// HACK: control UpdateLowerIndex
+template <typename TensorDesc, typename VisibleIndex, typename UpdateLowerIndexHack>
+__host__ __device__ constexpr auto make_dynamic_tensor_coordinate_iterator(
+    const TensorDesc&, const VisibleIndex& idx_diff_visible, UpdateLowerIndexHack)
+{
+    static_assert(TensorDesc::GetNumOfDimension() == VisibleIndex::Size(),
+                  "wrong! # of dimension inconsistent");
+
+    constexpr index_t ntransform   = TensorDesc::GetNumOfTransform();
+    constexpr index_t ndim_hidden  = TensorDesc::GetNumOfHiddenDimension();
+    constexpr index_t ndim_visible = TensorDesc::GetNumOfVisibleDimension();
+    constexpr auto visible_dim_ids = TensorDesc::GetVisibleDimensionIds();
+
+    static_assert(UpdateLowerIndexHack::Size() == ntransform, "wrong!");
+
+    // use index_t for boolean type
+    auto do_transforms    = make_zero_multi_index<ntransform>();
+    auto is_non_zero_diff = make_zero_multi_index<ndim_hidden>();
+
+    // decide do_transform by checkout non-zero index diff components
+    MultiIndex<VisibleIndex::Size()> non_zero_diff_pick_visible;
+
+    static_for<0, ndim_visible, 1>{}(
+        [&](auto i) { non_zero_diff_pick_visible(i) = (idx_diff_visible[i] != 0); });
+
+    set_container_subset(is_non_zero_diff, visible_dim_ids, non_zero_diff_pick_visible);
+
+    static_for<ntransform - 1, -1, -1>{}([&](auto itran) {
+        constexpr auto dims_low = TensorDesc::GetLowerDimensionIdss().At(itran);
+        constexpr auto dims_up  = TensorDesc::GetUpperDimensionIdss().At(itran);
+
+        const auto non_zero_diff_pick_up = get_container_subset(is_non_zero_diff, dims_up);
+
+        MultiIndex<dims_low.Size()> non_zero_diff_pick_low;
+
+        // if any of upper index diff components is non-zero, then
+        //   1) Need to do this transform
+        //   2) all components of lower index diff will assume to be non-zero and need to be
+        //   computed
+        const bool idx_diff_up_has_non_zero = container_reduce(
+            non_zero_diff_pick_up, [](auto a, auto b) constexpr { return a or b; }, false);
+
+        do_transforms(itran) = idx_diff_up_has_non_zero;
+
+        static_for<0, dims_low.Size(), 1>{}(
+            [&](auto i) { non_zero_diff_pick_low(i) = idx_diff_up_has_non_zero; });
+
+        set_container_subset(is_non_zero_diff, dims_low, non_zero_diff_pick_low);
+    });
+
+    return DynamicTensorCoordinateIterator<ntransform, ndim_visible, UpdateLowerIndexHack>{
+        idx_diff_visible, do_transforms};
+}
+
+template <typename TensorDesc, typename VisibleIndex>
+__host__ __device__ constexpr auto
+make_dynamic_tensor_coordinate_iterator(const TensorDesc&, const VisibleIndex& idx_diff_visible)
+{
+    constexpr index_t ntransform = TensorDesc::GetNumOfTransform();
+
+    return make_dynamic_tensor_coordinate_iterator(
+        TensorDesc{}, idx_diff_visible, typename uniform_sequence_gen<ntransform, 0>::type{});
+}
+
+template <typename TensorDesc, typename TensorCoord, typename TensorCoordIterator>
+__host__ __device__ constexpr void move_dynamic_tensor_coordinate(
+    const TensorDesc& tensor_desc, TensorCoord& coord, const TensorCoordIterator& coord_iterator)
+{
+    constexpr index_t ndim_hidden  = TensorDesc::GetNumOfHiddenDimension();
+    constexpr index_t ndim_visible = TensorDesc::GetNumOfVisibleDimension();
+    constexpr index_t ntransform   = TensorDesc::GetNumOfTransform();
+
+    using HiddenIndex = MultiIndex<ndim_hidden>;
+
+    // this is what needs to be calculated
+    auto idx_diff_hidden = make_zero_multi_index<ndim_hidden>();
+
+    // initialize visible index diff
+    set_container_subset(idx_diff_hidden,
+                         TensorDesc::GetVisibleDimensionIds(),
+                         coord_iterator.GetVisibleIndexDiff());
+
+    // this is what needs to be updated
+    auto& idx_hidden = coord.GetHiddenIndex();
+
+    // update visible index
+    auto idx_hidden_pick_visible =
+        get_container_subset(idx_hidden, TensorDesc::GetVisibleDimensionIds());
+
+    idx_hidden_pick_visible += coord_iterator.GetIndexDiff();
+
+    set_container_subset(idx_hidden, TensorDesc::GetVisibleDimensionIds(), idx_hidden_pick_visible);
+
+    // update rest of hidden index
+    static_for<ntransform - 1, -1, -1>{}([&](auto itran) {
+        if(coord_iterator.do_transforms_[itran])
+        {
+            const auto& tran        = tensor_desc.GetTransforms().At(itran);
+            constexpr auto dims_low = TensorDesc::GetLowerDimensionIdss().At(itran);
+            constexpr auto dims_up  = TensorDesc::GetUpperDimensionIdss().At(itran);
+
+            const auto idx_up_new  = get_container_subset(idx_hidden, dims_up);
+            auto idx_low           = get_container_subset(idx_hidden, dims_low);
+            const auto idx_diff_up = get_container_subset(idx_diff_hidden, dims_up);
+
+            MultiIndex<dims_low.Size()> idx_diff_low;
+
+            // HACK: control UpdateLowerIndex for DynamicMerge using hack
+            constexpr index_t Hack = decltype(coord_iterator.update_lower_index_hack_)::At(itran);
+
+            tran.UpdateLowerIndex(idx_diff_low, idx_diff_up, idx_low, idx_up_new, Number<Hack>{});
+
+            set_container_subset(idx_diff_hidden, dims_low, idx_diff_low);
+            set_container_subset(idx_hidden, dims_low, idx_low);
+        }
+    });
+}
+
+template <typename TensorDesc, typename TensorCoord>
+__host__ __device__ constexpr bool
+coordinate_has_valid_offset_assuming_visible_index_is_valid(const TensorDesc& tensor_desc,
+                                                            const TensorCoord& coord)
+{
+    bool valid = true;
+
+    constexpr index_t ntransform = TensorDesc::GetNumOfTransform();
+
+    const auto& idx_hidden = coord.GetHiddenIndex();
+
+    static_for<ntransform - 1, -1, -1>{}([&tensor_desc, &idx_hidden, &valid](auto itran) {
+        const auto tran = tensor_desc.GetTransforms().At(itran);
+
+        // check validity, only if current transformation does not always has a valid mapping
+        if constexpr(!decltype(tran)::IsValidUpperIndexAlwaysMappedToValidLowerIndex())
+        {
+            const auto idx_up =
+                get_container_subset(idx_hidden, TensorDesc::GetUpperDimensionIdss().At(itran));
+
+            // Comment: using valid = valid && .. will result in weird control flow in ISA
+            valid &= tran.IsValidUpperIndexMappedToValidLowerIndex(idx_up);
+        }
+    });
+
+    return valid;
+}
+
+template <typename TensorDesc, typename TensorCoord>
+__host__ __device__ constexpr bool coordinate_has_valid_offset(const TensorDesc& tensor_desc,
+                                                               const TensorCoord& coord)
+{
+    // check visible index
+    const auto& idx_visible = coord.GetVisibleIndex();
+
+    bool is_visible_index_valid = true;
+
+    static_for<0, TensorDesc::GetNumOfDimension(), 1>{}(
+        [&is_visible_index_valid, &idx_visible, &tensor_desc](auto i) {
+            is_visible_index_valid =
+                is_visible_index_valid &&
+                (idx_visible[i] >= 0 && idx_visible[i] < tensor_desc.GetLength(i));
+        });
+
+    // check other hidden index
+    return is_visible_index_valid &&
+           coordinate_has_valid_offset_assuming_visible_index_is_valid(tensor_desc, coord);
+}
+
+template <typename TensorDesc>
+using DynamicTensorCoordinate_t = decltype(make_dynamic_tensor_coordinate(
+    TensorDesc{}, MultiIndex<remove_cv_t<remove_reference_t<TensorDesc>>::GetNumOfDimension()>{}));
+
+template <typename TensorDesc>
+using DynamicTensorCoordinateIterator_t = decltype(make_dynamic_tensor_coordinate_iterator(
+    TensorDesc{}, MultiIndex<remove_cv_t<remove_reference_t<TensorDesc>>::GetNumOfDimension()>{}));
+
+} // namespace ck
+#endif

composable_kernel/include/tensor_description/dynamic_tensor_descriptor_helper.hpp

@@ -0,0 +1,146 @@
+#ifndef CK_DYNAMIC_TENSOR_DESCRIPTOR_HELPER_HPP
+#define CK_DYNAMIC_TENSOR_DESCRIPTOR_HELPER_HPP
+
+#include "common_header.hpp"
+#include "dynamic_tensor_descriptor.hpp"
+#include "dynamic_multi_index_transform_helper.hpp"
+
+namespace ck {
+
+/*
+ * These functions create tensor descriptor at runtime. If they are not constexpr, you will
+ * likely see usage of scratch memory during construction of these tensor descriptors. So
+ * it's better to call these functions on host and then pass the constructed tensor descritpors
+ * to GPU. If the tensor descritpors being constructed are constexpr, then you can call these
+ * functions on GPU without worrying about scratch memory usage.
+ */
+
+#if CK_WORKAROUND_SWDEV_275126
+template <typename Lengths, typename Strides, index_t I, typename AccOld>
+__host__ __device__ constexpr auto calculate_element_space_size_impl(const Lengths& lengths,
+                                                                     const Strides& strides,
+                                                                     Number<I> i,
+                                                                     AccOld acc_old)
+{
+    auto acc_new = acc_old + (lengths[i] - Number<1>{}) * strides[i];
+
+    if constexpr(i.value < Lengths::Size() - 1)
+    {
+        return calculate_element_space_size_impl(lengths, strides, i + Number<1>{}, acc_new);
+    }
+    else
+    {
+        return acc_new;
+    }
+}
+#endif
+
+template <typename... Lengths,
+          typename... Strides,
+          typename std::enable_if<sizeof...(Lengths) == sizeof...(Strides), bool>::type = false>
+__host__ __device__ constexpr auto
+make_dynamic_naive_tensor_descriptor_v2(const Tuple<Lengths...>& lengths,
+                                        const Tuple<Strides...>& strides)
+{
+    constexpr index_t N = sizeof...(Lengths);
+
+    const auto transforms = make_tuple(make_embed_transform(lengths, strides));
+
+    constexpr auto low_dim_hidden_idss = make_tuple(Sequence<0>{});
+
+    constexpr auto up_dim_hidden_idss =
+        make_tuple(typename arithmetic_sequence_gen<1, N + 1, 1>::type{});
+
+    constexpr auto visible_dim_hidden_ids = typename arithmetic_sequence_gen<1, N + 1, 1>::type{};
+
+#if !CK_WORKAROUND_SWDEV_275126
+    // rocm-4.1 compiler would crash for recursive labmda
+    // recursive function for reduction
+    auto f = [&](auto fs, auto i, auto acc_old) {
+        auto acc_new = acc_old + (lengths[i] - Number<1>{}) * strides[i];
+
+        if constexpr(i.value < N - 1)
+        {
+            return fs(fs, i + Number<1>{}, acc_new);
+        }
+        else
+        {
+            return acc_new;
+        }
+    };
+
+    const auto element_space_size = f(f, Number<0>{}, Number<1>{});
+#else
+    const auto element_space_size =
+        calculate_element_space_size_impl(lengths, strides, Number<0>{}, Number<1>{});
+#endif
+
+    return DynamicTensorDescriptor<remove_cv_t<decltype(transforms)>,
+                                   remove_cv_t<decltype(low_dim_hidden_idss)>,
+                                   remove_cv_t<decltype(up_dim_hidden_idss)>,
+                                   remove_cv_t<decltype(visible_dim_hidden_ids)>,
+                                   remove_cv_t<decltype(element_space_size)>>{transforms,
+                                                                              element_space_size};
+}
+
+// Lengths... can be:
+//   1) index_t, which is known at run-time
+//   2) Number<>, which is known at compile-time
+template <typename... Lengths>
+__host__ __device__ constexpr auto
+make_dynamic_naive_tensor_descriptor_packed_v2(const Tuple<Lengths...>& lengths)
+{
+    constexpr index_t N = sizeof...(Lengths);
+
+    const auto transforms = make_tuple(make_unmerge_transform(lengths));
+
+    constexpr auto low_dim_hidden_idss = make_tuple(Sequence<0>{});
+
+    constexpr auto up_dim_hidden_idss =
+        make_tuple(typename arithmetic_sequence_gen<1, N + 1, 1>::type{});
+
+    constexpr auto visible_dim_hidden_ids = typename arithmetic_sequence_gen<1, N + 1, 1>::type{};
+
+    const auto element_space_size = container_reduce(lengths, math::multiplies_v2{}, Number<1>{});
+
+    return DynamicTensorDescriptor<remove_cv_t<decltype(transforms)>,
+                                   remove_cv_t<decltype(low_dim_hidden_idss)>,
+                                   remove_cv_t<decltype(up_dim_hidden_idss)>,
+                                   remove_cv_t<decltype(visible_dim_hidden_ids)>,
+                                   remove_cv_t<decltype(element_space_size)>>{transforms,
+                                                                              element_space_size};
+}
+
+template <typename... Lengths, typename Align>
+__host__ __device__ constexpr auto
+make_dynamic_naive_tensor_descriptor_aligned_v2(const Tuple<Lengths...>& lengths, Align align)
+{
+    constexpr index_t N = sizeof...(Lengths);
+
+    auto strides = generate_tuple(
+        [&](auto i) {
+            if constexpr(i.value == N - 1)
+            {
+                return Number<1>{};
+            }
+            else if constexpr(i.value == N - 2)
+            {
+                return math::lcm(lengths[Number<N - 1>{}], align);
+            }
+            else
+            {
+                return container_reduce(lengths,
+                                        math::multiplies_v2{},
+                                        math::lcm(lengths[Number<N - 1>{}], align),
+                                        i,
+                                        Number<N - 2>{},
+                                        Number<1>{});
+            }
+        },
+        Number<N>{});
+
+    return make_dynamic_naive_tensor_descriptor_v2(lengths, strides);
+}
+
+} // namespace ck
+#endif

composable_kernel/include/tensor_description/multi_index.hpp

@@ -0,0 +1,12 @@
+#ifndef CK_MULTI_INDEX_HPP
+#define CK_MULTI_INDEX_HPP
+
+#include "common_header.hpp"
+
+#if CK_USE_DYNAMICALLY_INDEXED_MULTI_INDEX
+#include "array_multi_index.hpp"
+#else
+#include "statically_indexed_array_multi_index.hpp"
+#endif
+
+#endif

composable_kernel/include/tensor_description/multi_index_transform.hpp

@@ -2,18 +2,10 @@
 #define CK_MULTI_INDEX_TRANSFORM_HPP
 
 #include "common_header.hpp"
+#include "multi_index.hpp"
 
 namespace ck {
 
-template <index_t N>
-using MultiIndex = Array<index_t, N>;
-
-template <typename... Xs>
-__host__ __device__ constexpr auto make_multi_index(Xs... xs)
-{
-    return MultiIndex<sizeof...(Xs)>(xs...);
-}
-
 template <index_t Length>
 struct PassThrough
 {
@@ -62,7 +54,7 @@ struct Pad
     using LowerIndex = MultiIndex<nDim>;
     using UpperIndex = MultiIndex<nDim>;
 
-    __host__ __device__ explicit constexpr Pad()
+    __host__ __device__ constexpr Pad()
     {
         static_assert(LowerLengths::GetSize() == nDim && LeftPads::GetSize() == nDim &&
                           RightPads::GetSize() == nDim,
@@ -123,7 +115,7 @@ struct Slice
     using LowerIndex = MultiIndex<nDim>;
     using UpperIndex = MultiIndex<nDim>;
 
-    __host__ __device__ explicit constexpr Slice()
+    __host__ __device__ constexpr Slice()
     {
         static_assert(LowerLengths::GetSize() == nDim && SliceBegins::GetSize() == nDim &&
                           SliceEnds::GetSize() == nDim,
@@ -197,8 +189,8 @@ struct Merge
         index_t& itmp;
         LowerIndex& idx_low;
 
-        __host__ __device__ explicit constexpr lambda_CalculateLowerIndex(index_t& itmp_,
-                                                                          LowerIndex& idx_low_)
+        __host__ __device__ constexpr lambda_CalculateLowerIndex(index_t& itmp_,
+                                                                 LowerIndex& idx_low_)
             : itmp(itmp_), idx_low(idx_low_)
         {
         }
@@ -216,7 +208,7 @@ struct Merge
     {
         LowerIndex idx_low;
 
-        index_t itmp = idx_up[0];
+        index_t itmp = idx_up[Number<0>{}];
 
         constexpr auto pseudo_low_strides =
             reverse_inclusive_scan_sequence(
@@ -226,7 +218,7 @@ struct Merge
         static_for<0, nDimLow - 1, 1>{}(
             lambda_CalculateLowerIndex<decltype(pseudo_low_strides)>(itmp, idx_low));
 
-        idx_low(nDimLow - 1) = itmp / pseudo_low_strides[nDimLow - 1];
+        idx_low(Number<nDimLow - 1>{}) = itmp / pseudo_low_strides[Number<nDimLow - 1>{}];
 
         return idx_low;
     }
@@ -240,9 +232,9 @@ struct Merge
                             const UpperIndex& /* idx_up_old */,
                             const LowerIndex& idx_low_old)
     {
-        if(idx_up_diff[0] == 0)
+        if(idx_up_diff[Number<0>{}] == 0)
         {
-            return make_zero_array<index_t, nDimLow>();
+            return make_zero_multi_index<nDimLow>();
         }
         else
         {
@@ -265,7 +257,7 @@ struct Merge
 
             LowerIndex idx_low_new = idx_low_old + idx_low_diff_tmp;
 
-            if(idx_up_diff[0] > 0)
+            if(idx_up_diff[Number<0>{}] > 0)
             {
                 // do carry check on each low dimension in reversed order
                 // starting from the first digit that changed
@@ -293,7 +285,7 @@ struct Merge
                 // highest dimension, no out-of-bound check
                 if(carry)
                 {
-                    ++idx_low_new(0);
+                    ++idx_low_new(Number<0>{});
                 }
             }
             else
@@ -324,7 +316,7 @@ struct Merge
                 // highest dimension, no out-of-bound check
                 if(borrow)
                 {
-                    --idx_low_new(0);
+                    --idx_low_new(Number<0>{});
                 }
             }
 
@@ -358,7 +350,7 @@ struct UnMerge
 
     __host__ __device__ static constexpr auto CalculateLowerIndex(const UpperIndex& idx_up)
     {
-        LowerIndex idx_low{0};
+        LowerIndex idx_low = make_multi_index(0);
 
         constexpr auto pseudo_up_strides =
             reverse_inclusive_scan_sequence(
@@ -366,7 +358,7 @@ struct UnMerge
                 .PushBack(Number<1>{});
 
         static_for<0, nDimUp, 1>{}(
-            [&](auto idim) { idx_low(0) += idx_up[idim] * pseudo_up_strides[idim]; });
+            [&](auto idim) { idx_low(Number<0>{}) += idx_up[idim] * pseudo_up_strides[idim]; });
 
         return idx_low;
     }
@@ -405,7 +397,7 @@ struct Embed
     using LowerIndex = MultiIndex<nDimLow>;
     using UpperIndex = MultiIndex<nDimUp>;
 
-    __host__ __device__ explicit constexpr Embed()
+    __host__ __device__ constexpr Embed()
     {
         static_assert(UpperLengths::GetSize() == nDimUp && Coefficients::GetSize() == nDimUp + 1,
                       "wrong! # of dimensions not consistent");
@@ -419,12 +411,10 @@ struct Embed
 
     __host__ __device__ static constexpr auto CalculateLowerIndex(const UpperIndex& idx_up)
     {
-        LowerIndex idx_low(Coefficients{}[nDimUp]);
+        LowerIndex idx_low = make_multi_index(Coefficients{}[Number<nDimUp>{}]);
 
-        for(index_t i = 0; i < nDimUp; ++i)
-        {
-            idx_low(0) += idx_up[i] * Coefficients{}[i];
-        }
+        static_for<0, nDimUp, 1>{}(
+            [&](auto i) { idx_low(Number<0>{}) += idx_up[i] * Coefficients{}[i]; });
 
         return idx_low;
     }
@@ -434,12 +424,10 @@ struct Embed
                             const UpperIndex& /* idx_up_old */,
                             const LowerIndex& /* idx_low_old */)
     {
-        LowerIndex idx_low_diff{0};
+        LowerIndex idx_low_diff = make_multi_index(0);
 
-        for(index_t i = 0; i < nDimUp; ++i)
-        {
-            idx_low_diff(0) += idx_up_diff[i] * Coefficients{}[i];
-        }
+        static_for<0, nDimUp, 1>{}(
+            [&](auto i) { idx_low_diff(Number<0>{}) += idx_up_diff[i] * Coefficients{}[i]; });
 
         return idx_low_diff;
     }
@@ -467,21 +455,21 @@ struct Embed
         for(index_t icorner = 0; icorner < ncorner; ++icorner)
         {
             // generate upper index for each corner
-            auto idx_up = make_zero_array<index_t, nDimUp>();
+            auto idx_up = make_zero_multi_index<nDimUp>();
 
             index_t itmp = icorner;
 
-            for(index_t idim = nDimUp - 1; idim >= 0; --idim)
-            {
-                idx_up(idim) = itmp % 2 == 0 ? 0 : UpperLengths::At(idim) - 1;
+            static_for<nDimUp, 0, -1>{}([&](auto idim) {
+                auto idim_m1    = idim - Number<1>{};
+                idx_up(idim_m1) = itmp % 2 == 0 ? 0 : UpperLengths::At(idim_m1) - 1;
                 itmp /= 2;
-            }
+            });
 
             // calculate lower index
             auto idx_low = CalculateLowerIndex(idx_up);
 
             // judge if lower index is valid
-            flag = flag && idx_low[0] >= 0 && idx_low[0] < LowerLength;
+            flag = flag && idx_low[Number<0>{}] >= 0 && idx_low[Number<0>{}] < LowerLength;
         }
 
         return flag;
@@ -499,7 +487,7 @@ struct Freeze
     using LowerIndex = MultiIndex<nDimLow>;
     using UpperIndex = MultiIndex<nDimUp>;
 
-    __host__ __device__ explicit constexpr Freeze()
+    __host__ __device__ constexpr Freeze()
     {
         // TODO: sanity check: LowerFreezePoint should be within range of LowerLengths
     }
@@ -512,7 +500,7 @@ struct Freeze
 
     __host__ __device__ static constexpr auto CalculateLowerIndex(const UpperIndex& /*idx_up*/)
     {
-        return to_array(LowerFreezePoint{});
+        return to_multi_index(LowerFreezePoint{});
     }
 
     __host__ __device__ static constexpr auto
@@ -520,49 +508,7 @@ struct Freeze
                             const UpperIndex& /* idx_up_old */,
                             const LowerIndex& /* idx_low_old */)
     {
-        return make_zero_array<index_t, nDimLow>();
-    }
-
-    __host__ __device__ static constexpr bool IsLinearTransform() { return true; }
-
-    __host__ __device__ static constexpr bool IsValidUpperIndexAlwaysMappedToValidLowerIndex()
-    {
-        return true;
-    }
-};
-
-template <index_t LowerLength, index_t VectorSize>
-struct Vectorize
-{
-    using LowerIndex = MultiIndex<1>;
-    using UpperIndex = MultiIndex<1>;
-
-    __host__ __device__ constexpr Vectorize()
-    {
-        static_assert(VectorSize > 0 && LowerLength % VectorSize == 0,
-                      "wrong! cannot evenly divide");
-    }
-
-    __host__ __device__ static constexpr auto GetNumOfLowerDimension() { return Number<1>{}; }
-
-    __host__ __device__ static constexpr auto GetNumOfUpperDimension() { return Number<1>{}; }
-
-    __host__ __device__ static constexpr auto GetUpperLengths()
-    {
-        return Sequence<LowerLength / VectorSize>{};
-    }
-
-    __host__ __device__ static constexpr auto CalculateLowerIndex(const UpperIndex& idx_up)
-    {
-        return VectorSize * idx_up;
-    }
-
-    __host__ __device__ static constexpr auto
-    CalculateLowerIndexDiff(const UpperIndex& idx_up_diff,
-                            const UpperIndex& /* idx_up_old */,
-                            const LowerIndex& /* idx_low_old */)
-    {
-        return VectorSize * idx_up_diff;
+        return make_zero_multi_index<nDimLow>();
     }
 
     __host__ __device__ static constexpr bool IsLinearTransform() { return true; }

composable_kernel/include/tensor_description/statically_indexed_array_multi_index.hpp

@@ -0,0 +1,107 @@
+#ifndef CK_STATICALLY_INDEXED_ARRAY_MULTI_INDEX_HPP
+#define CK_STATICALLY_INDEXED_ARRAY_MULTI_INDEX_HPP
+
+#include "common_header.hpp"
+
+namespace ck {
+
+template <index_t N>
+using MultiIndex = StaticallyIndexedArray<index_t, N>;
+
+template <typename... Xs>
+__host__ __device__ constexpr auto make_multi_index(Xs&&... xs)
+{
+    return make_statically_indexed_array<index_t>(index_t{xs}...);
+}
+
+template <index_t NSize>
+__host__ __device__ constexpr auto make_zero_multi_index()
+{
+    return unpack([](auto... xs) { return make_multi_index(xs...); },
+                  typename uniform_sequence_gen<NSize, 0>::type{});
+}
+
+template <typename T>
+__host__ __device__ constexpr auto to_multi_index(const T& x)
+{
+    return unpack([](auto... ys) { return make_multi_index(ys...); }, x);
+}
+
+// Here should use MultiIndex<NSize>, instead of Tuple<Ys...>, although the former
+// is the alias of the latter. This is because compiler cannot infer the NSize if
+// using MultiIndex<NSize>
+// TODO: how to fix this?
+template <typename... Ys, typename X>
+__host__ __device__ constexpr auto operator+=(Tuple<Ys...>& y, const X& x)
+{
+    static_assert(X::Size() == sizeof...(Ys), "wrong! size not the same");
+    constexpr index_t NSize = sizeof...(Ys);
+    static_for<0, NSize, 1>{}([&](auto i) { y(i) += x[i]; });
+    return y;
+}
+
+template <typename... Ys, typename X>
+__host__ __device__ constexpr auto operator-=(Tuple<Ys...>& y, const X& x)
+{
+    static_assert(X::Size() == sizeof...(Ys), "wrong! size not the same");
+    constexpr index_t NSize = sizeof...(Ys);
+    static_for<0, NSize, 1>{}([&](auto i) { y(i) -= x[i]; });
+    return y;
+}
+
+template <typename... Xs, typename Y>
+__host__ __device__ constexpr auto operator+(const Tuple<Xs...>& x, const Y& y)
+{
+    static_assert(Y::Size() == sizeof...(Xs), "wrong! size not the same");
+    constexpr index_t NSize = sizeof...(Xs);
+
+    Tuple<Xs...> r;
+    static_for<0, NSize, 1>{}([&](auto i) { r(i) = x[i] + y[i]; });
+    return r;
+}
+
+template <typename... Xs, typename Y>
+__host__ __device__ constexpr auto operator-(const Tuple<Xs...>& x, const Y& y)
+{
+    static_assert(Y::Size() == sizeof...(Xs), "wrong! size not the same");
+    constexpr index_t NSize = sizeof...(Xs);
+
+    Tuple<Xs...> r;
+    static_for<0, NSize, 1>{}([&](auto i) { r(i) = x[i] - y[i]; });
+    return r;
+}
+
+template <typename... Xs, typename Y>
+__host__ __device__ constexpr auto operator*(const Tuple<Xs...>& x, const Y& y)
+{
+    static_assert(Y::Size() == sizeof...(Xs), "wrong! size not the same");
+    constexpr index_t NSize = sizeof...(Xs);
+
+    Tuple<Xs...> r;
+    static_for<0, NSize, 1>{}([&](auto i) { r(i) = x[i] * y[i]; });
+    return r;
+}
+
+// MultiIndex = index_t * MultiIndex
+template <typename... Xs>
+__host__ __device__ constexpr auto operator*(index_t a, const Tuple<Xs...>& x)
+{
+    constexpr index_t NSize = sizeof...(Xs);
+
+    Tuple<Xs...> r;
+    static_for<0, NSize, 1>{}([&](auto i) { r(i) = a * x[i]; });
+    return r;
+}
+
+template <typename... Xs>
+__host__ __device__ void print_multi_index(const Tuple<Xs...>& x)
+{
+    printf("{");
+    printf("MultiIndex, ");
+    printf("size %d,", index_t{sizeof...(Xs)});
+    static_for<0, sizeof...(Xs), 1>{}([&](auto i) { printf("%d ", index_t{x.At(i)}); });
+    printf("}");
+}
+
+} // namespace ck
+#endif

composable_kernel/include/tensor_description/tensor_coordinate.hpp

@@ -41,13 +41,13 @@ struct NativeTensorCoordinate
 
     template <typename... Xs>
     __host__ __device__ constexpr NativeTensorCoordinate(Xs... xs)
-        : NativeTensorCoordinate(Index{xs...})
+        : NativeTensorCoordinate(make_multi_index(xs...))
     {
     }
 
     template <index_t... Xs>
     __host__ __device__ constexpr NativeTensorCoordinate(Sequence<Xs...>)
-        : NativeTensorCoordinate(Index{Xs...})
+        : NativeTensorCoordinate(make_mutli_index(Xs...))
     {
     }
 
@@ -267,18 +267,18 @@ struct TensorCoordinate
     private:
     template <typename... Ts>
     __host__ __device__ static constexpr auto
-        MakeDummyTensorCoordinate(NativeTensorDescriptor<Ts...>)
+    MakeDummyTensorCoordinate(NativeTensorDescriptor<Ts...>)
     {
         return NativeTensorCoordinate<NativeTensorDescriptor<Ts...>>(
-            make_zero_array<index_t, TensorDesc::GetNumOfDimension()>());
+            make_zero_multi_index<TensorDesc::GetNumOfDimension()>());
     }
 
     template <typename... Ts>
     __host__ __device__ static constexpr auto
-        MakeDummyTensorCoordinate(TransformedTensorDescriptor<Ts...>)
+    MakeDummyTensorCoordinate(TransformedTensorDescriptor<Ts...>)
     {
         return TransformedTensorCoordinate<TransformedTensorDescriptor<Ts...>>(
-            make_zero_array<index_t, TensorDesc::GetNumOfDimension()>());
+            make_zero_multi_index<TensorDesc::GetNumOfDimension()>());
     }
 
     public:

composable_kernel/include/tensor_description/tensor_descriptor.hpp

@@ -311,13 +311,13 @@ struct TransformedTensorDescriptor
         static_for<0, nTransform, 1>{}([&](auto itran) {
             constexpr auto tran = Transforms{}.At(itran);
 
-            const auto idx_up_part = pick_array_element(idx_up, UpDimensionIds{}.At(itran));
-            auto idx_low_part      = pick_array_element(idx_low, LowDimensionIds{}.At(itran));
+            const auto idx_up_part = pick_container_element(idx_up, UpDimensionIds{}.At(itran));
+            auto idx_low_part      = pick_container_element(idx_low, LowDimensionIds{}.At(itran));
 
             // this assume each lower (single) index is only assocaited with one transformation,
             //   which is required for index transformation, and has been checked during constructor
             //   of TransformedTensorDescriptor
-            idx_low_part = tran.CalculateLowerIndex(to_array(idx_up_part));
+            idx_low_part = tran.CalculateLowerIndex(to_multi_index(idx_up_part));
         });
 
         return idx_low;
@@ -333,20 +333,23 @@ struct TransformedTensorDescriptor
             constexpr auto tran = Transforms{}.At(itran);
 
             const auto idx_up_diff_part =
-                pick_array_element(idx_up_diff, UpDimensionIds{}.At(itran));
+                pick_container_element(idx_up_diff, UpDimensionIds{}.At(itran));
 
-            const auto idx_up_old_part = pick_array_element(idx_up_old, UpDimensionIds{}.At(itran));
+            const auto idx_up_old_part =
+                pick_container_element(idx_up_old, UpDimensionIds{}.At(itran));
 
             const auto idx_low_old_part =
-                pick_array_element(idx_low_old, LowDimensionIds{}.At(itran));
+                pick_container_element(idx_low_old, LowDimensionIds{}.At(itran));
 
-            auto idx_low_diff_part = pick_array_element(idx_low_diff, LowDimensionIds{}.At(itran));
+            auto idx_low_diff_part =
+                pick_container_element(idx_low_diff, LowDimensionIds{}.At(itran));
 
             // this assume each lower (single) index is associated with only one transformation,
             //   which is required for index transformation, and has been checked during constructor
             //   of TransformedTensorDescriptor
-            idx_low_diff_part = tran.CalculateLowerIndexDiff(
-                to_array(idx_up_diff_part), to_array(idx_up_old_part), to_array(idx_low_old_part));
+            idx_low_diff_part = tran.CalculateLowerIndexDiff(to_multi_index(idx_up_diff_part),
+                                                             to_multi_index(idx_up_old_part),
+                                                             to_multi_index(idx_low_old_part));
         });
 
         return idx_low_diff;
@@ -506,12 +509,12 @@ struct TransformedTensorDescriptor
                 constexpr auto low_dims_part = LowDimensionIds{}.At(itran);
                 constexpr auto low_lengths_part =
                     GetLowerTensorDescriptor().GetLengths(low_dims_part);
-                const auto idx_low_part = to_array(pick_array_element(idx_low, low_dims_part));
+                const auto idx_low_part =
+                    to_multi_index(pick_container_element(idx_low, low_dims_part));
 
-                for(index_t i = 0; i < low_dims_part.Size(); ++i)
-                {
+                static_for<0, decltype(low_dims_part)::Size(), 1>{}([&](auto i) {
                     flag = flag && idx_low_part[i] >= 0 && idx_low_part[i] < low_lengths_part[i];
-                }
+                });
             }
         });
 

composable_kernel/include/tensor_description/tensor_descriptor_helper.hpp

@@ -64,10 +64,10 @@ template <typename LowerTensorDescriptor,
           index_t... LowerDimensionIds,
           index_t... UpperDimensionIds>
 __host__ __device__ constexpr auto
-    reorder_transformed_tensor_descriptor_impl(LowerTensorDescriptor,
-                                               Sequence<LowerLengths...>,
-                                               Sequence<LowerDimensionIds...>,
-                                               Sequence<UpperDimensionIds...>)
+reorder_transformed_tensor_descriptor_impl(LowerTensorDescriptor,
+                                           Sequence<LowerLengths...>,
+                                           Sequence<LowerDimensionIds...>,
+                                           Sequence<UpperDimensionIds...>)
 {
     return TransformedTensorDescriptor<LowerTensorDescriptor,
                                        Tuple<PassThrough<LowerLengths>...>,
@@ -78,7 +78,7 @@ __host__ __device__ constexpr auto
 // reorder a NativeTensorDescriptor
 template <typename... Ts, typename MapLower2Upper>
 __host__ __device__ constexpr auto
-    reorder_tensor_descriptor_given_lower2upper(NativeTensorDescriptor<Ts...>, MapLower2Upper)
+reorder_tensor_descriptor_given_lower2upper(NativeTensorDescriptor<Ts...>, MapLower2Upper)
 {
     static_assert(is_valid_sequence_map<MapLower2Upper>{},
                   "wrong! MapLower2Upper is not a valid map");
@@ -96,7 +96,7 @@ __host__ __device__ constexpr auto
 // reorder a TransformedTensorDescriptor
 template <typename... Ts, typename MapLower2Upper>
 __host__ __device__ constexpr auto
-    reorder_tensor_descriptor_given_lower2upper(TransformedTensorDescriptor<Ts...>, MapLower2Upper)
+reorder_tensor_descriptor_given_lower2upper(TransformedTensorDescriptor<Ts...>, MapLower2Upper)
 {
     static_assert(is_valid_sequence_map<MapLower2Upper>{},
                   "wrong! MapLower2Upper is not a valid map");
@@ -172,41 +172,5 @@ __host__ __device__ constexpr auto unfold_tensor_descriptor(NativeTensorDescript
     return make_native_tensor_descriptor(new_lengths, new_strides);
 }
 
-// a cluster map 1d index to N-d index
-template <typename Lengths, typename ArrangeOrder>
-struct ClusterDescriptor
-{
-    static constexpr index_t nDim = Lengths::Size();
-
-    static constexpr auto mDesc = transform_tensor_descriptor(
-        make_native_tensor_descriptor_packed(Lengths{}),
-        make_tuple(Merge<decltype(Lengths::ReorderGivenNew2Old(ArrangeOrder{}))>{}),
-        make_tuple(ArrangeOrder{}),
-        make_tuple(Sequence<0>{}));
-
-    __host__ __device__ constexpr ClusterDescriptor()
-    {
-        static_assert(Lengths::Size() == nDim && ArrangeOrder::Size() == nDim,
-                      "wrong! size not the same");
-
-        static_assert(is_valid_sequence_map<ArrangeOrder>{}, "wrong! ArrangeOrder is wrong");
-    }
-
-    __host__ __device__ static constexpr index_t GetElementSize() { return mDesc.GetElementSize(); }
-
-    __host__ __device__ static constexpr auto CalculateClusterIndex(index_t idx_1d)
-    {
-        return mDesc.CalculateLowerIndex(MultiIndex<1>{idx_1d});
-    }
-};
-
-template <typename Lengths,
-          typename ArrangeOrder = typename arithmetic_sequence_gen<0, Lengths::Size(), 1>::type>
-__host__ __device__ constexpr auto make_cluster_descriptor(
-    Lengths, ArrangeOrder order = typename arithmetic_sequence_gen<0, Lengths::Size(), 1>::type{})
-{
-    return ClusterDescriptor<Lengths, decltype(order)>{};
-}
-
 } // namespace ck
 #endif

composable_kernel/include/tensor_operation/blockwise_batched_gemm.hpp

@@ -210,17 +210,16 @@ struct BlockwiseBatchGemmBlockABlockBThreadCTransANormalBNormalC_V2
 #pragma unroll
                     for(index_t m_repeat = 0; m_repeat < MRepeat; ++m_repeat)
                     {
-                        threadwise_matrix_copy(
-                            a_block_mtx,
-                            p_a_block +
-                                a_block_mtx.GetOffsetFromMultiIndex(k_begin,
-                                                                    m_repeat * MPerLevel1Cluster) +
-                                ib * BlockMatrixStrideA + mMyThreadOffsetA,
-                            a_thread_mtx,
-                            p_a_thread +
-                                a_thread_mtx.GetOffsetFromMultiIndex(0, m_repeat * MPerThreadSubC),
-                            a_thread_sub_mtx.GetLengths(),
-                            Number<DataPerReadA>{});
+                        threadwise_matrix_copy(a_block_mtx,
+                                               p_a_block +
+                                                   a_block_mtx.GetOffsetFromMultiIndex(
+                                                       k_begin, m_repeat * MPerLevel1Cluster) +
+                                                   ib * BlockMatrixStrideA + mMyThreadOffsetA,
+                                               a_thread_mtx,
+                                               p_a_thread + a_thread_mtx.GetOffsetFromMultiIndex(
+                                                                0, m_repeat * MPerThreadSubC),
+                                               a_thread_sub_mtx.GetLengths(),
+                                               Number<DataPerReadA>{});
                     }
                 }
 
@@ -229,17 +228,16 @@ struct BlockwiseBatchGemmBlockABlockBThreadCTransANormalBNormalC_V2
 #pragma unroll
                     for(index_t n_repeat = 0; n_repeat < NRepeat; ++n_repeat)
                     {
-                        threadwise_matrix_copy(
-                            b_block_mtx,
-                            p_b_block +
-                                b_block_mtx.GetOffsetFromMultiIndex(k_begin,
-                                                                    n_repeat * NPerLevel1Cluster) +
-                                ib * BlockMatrixStrideB + mMyThreadOffsetB,
-                            b_thread_mtx,
-                            p_b_thread +
-                                b_thread_mtx.GetOffsetFromMultiIndex(0, n_repeat * NPerThreadSubC),
-                            b_thread_sub_mtx.GetLengths(),
-                            Number<DataPerReadB>{});
+                        threadwise_matrix_copy(b_block_mtx,
+                                               p_b_block +
+                                                   b_block_mtx.GetOffsetFromMultiIndex(
+                                                       k_begin, n_repeat * NPerLevel1Cluster) +
+                                                   ib * BlockMatrixStrideB + mMyThreadOffsetB,
+                                               b_thread_mtx,
+                                               p_b_thread + b_thread_mtx.GetOffsetFromMultiIndex(
+                                                                0, n_repeat * NPerThreadSubC),
+                                               b_thread_sub_mtx.GetLengths(),
+                                               Number<DataPerReadB>{});
                     }
                 }
 
@@ -307,7 +305,7 @@ struct BlockwiseBatchGemmBlockABlockBThreadCTransANormalBNormalC_V2
                       "Run_amd_asm can only deal with BlockMatrixStrideA == 0 && BatchPerThread == "
                       "1 for now\n");
 
-        using Float4 = vector_type<float, 4>::MemoryType;
+        using Float4 = vector_type<float, 4>::type;
 
         Float4* reg_a = (Float4*)(p_a_thread);
         Float4* reg_b = (Float4*)(p_b_thread);
@@ -391,9 +389,8 @@ struct BlockwiseBatchGemmBlockABlockBThreadCTransANormalBNormalC_V2
             {
                 threadwise_matrix_copy(
                     c_thread_sub_mtx,
-                    p_c_thread +
-                        c_thread_sub_mtx.GetOffsetFromMultiIndex(m_repeat * MPerLevel1Cluster,
-                                                                 n_repeat * NPerLevel1Cluster),
+                    p_c_thread + c_thread_sub_mtx.GetOffsetFromMultiIndex(
+                                     m_repeat * MPerLevel1Cluster, n_repeat * NPerLevel1Cluster),
                     c_block_mtx,
                     p_c_block +
                         c_block_mtx.GetOffsetFromMultiIndex(m_repeat * MPerLevel1Cluster,
@@ -405,5 +402,5 @@ struct BlockwiseBatchGemmBlockABlockBThreadCTransANormalBNormalC_V2
     }
 };
 
-} // namespace
+} // namespace ck
 #endif

composable_kernel/include/tensor_operation/blockwise_dynamic_tensor_slice_transfer.hpp

@@ -0,0 +1,171 @@
+#ifndef CK_BLOCKWISE_DYNAMIC_TENSOR_SLICE_TRANSFER_HPP
+#define CK_BLOCKWISE_DYNAMIC_TENSOR_SLICE_TRANSFER_HPP
+
+#include "common_header.hpp"
+#include "dynamic_tensor_descriptor.hpp"
+#include "dynamic_tensor_descriptor_helper.hpp"
+#include "cluster_descriptor.hpp"
+#include "threadwise_dynamic_tensor_slice_transfer.hpp"
+
+namespace ck {
+
+// this version does following things to avoid scratch memory issue
+// 1. Use StaticallyIndexedArray instead of C array for thread buffer
+// 2. ThreadwiseDynamicTensorSliceTransfer_v3 does not keep reference to tensor descriptor
+// 3. ThreadwiseDynamicTensorSliceTransfer_v3::Run() does not construct new tensor coordinate
+template <index_t BlockSize,
+          InMemoryDataOperation DstInMemOp,
+          typename BlockSliceLengths,
+          typename ThreadSliceLengths,
+          typename ThreadClusterLengths,
+          typename ThreadClusterArrangeOrder,
+          typename SrcData,
+          typename DstData,
+          typename SrcDesc,
+          typename DstDesc,
+          typename SrcDimAccessOrder,
+          typename DstDimAccessOrder,
+          index_t SrcVectorDim,
+          index_t DstVectorDim,
+          index_t SrcScalarPerVector,
+          index_t DstScalarPerVector,
+          AddressSpace SrcAddressSpace,
+          AddressSpace DstAddressSpace,
+          index_t SrcScalarStrideInVector,
+          index_t DstScalarStrideInVector,
+          index_t ThreadTransferSrcResetCoordinateAfterRun,
+          index_t ThreadTransferDstResetCoordinateAfterRun>
+struct BlockwiseDynamicTensorSliceTransfer_v4
+{
+    static constexpr index_t nDim = remove_reference_t<SrcDesc>::GetNumOfDimension();
+
+    using Index = MultiIndex<nDim>;
+
+    __device__ constexpr BlockwiseDynamicTensorSliceTransfer_v4(const SrcDesc& src_desc,
+                                                                const Index& src_block_slice_origin,
+                                                                const DstDesc& dst_desc,
+                                                                const Index& dst_block_slice_origin)
+        : threadwise_transfer_(
+              src_desc, make_zero_multi_index<nDim>(), dst_desc, make_zero_multi_index<nDim>())
+
+    {
+        static_assert(nDim == remove_reference_t<remove_cv_t<SrcDesc>>::GetNumOfDimension() &&
+                          nDim == remove_reference_t<remove_cv_t<DstDesc>>::GetNumOfDimension() &&
+                          nDim == BlockSliceLengths::Size() && nDim == ThreadSliceLengths::Size() &&
+                          nDim == ThreadClusterLengths::Size() &&
+                          nDim == ThreadClusterArrangeOrder::Size() &&
+                          nDim == SrcDimAccessOrder::Size() && nDim == DstDimAccessOrder::Size(),
+                      "wrong! nDim not consistent");
+
+        static_assert(
+            is_same<BlockSliceLengths, decltype(ThreadSliceLengths{} * ThreadClusterLengths{})>{},
+            "wrong! threads should be mapped to cover entire slicing window");
+
+        static_assert(BlockSize >= thread_cluster_desc_.GetElementSize(),
+                      "wrong! BlockSize too small");
+
+        if(BlockSize == thread_cluster_desc_.GetElementSize() or
+           get_thread_local_1d_id() < thread_cluster_desc_.GetElementSize())
+        {
+            const auto thread_cluster_id =
+                thread_cluster_desc_.CalculateClusterIndex(get_thread_local_1d_id());
+
+            const auto thread_data_id_begin = thread_cluster_id * ThreadSliceLengths{};
+
+            threadwise_transfer_.SetSrcSliceOrigin(src_desc,
+                                                   src_block_slice_origin + thread_data_id_begin);
+            threadwise_transfer_.SetDstSliceOrigin(dst_desc,
+                                                   dst_block_slice_origin + thread_data_id_begin);
+        }
+    }
+
+    __device__ static constexpr auto CalculateThreadDataBegin()
+    {
+        const auto thread_cluster_id =
+            thread_cluster_desc_.CalculateClusterIndex(get_thread_local_1d_id());
+
+        return thread_cluster_id * ThreadSliceLengths{};
+    }
+
+    template <typename SrcIteratorHacks>
+    __device__ void RunRead(const SrcDesc& src_desc,
+                            const SrcData* p_src,
+                            const SrcIteratorHacks& src_iterator_hacks)
+    {
+        if(BlockSize == thread_cluster_desc_.GetElementSize() or
+           get_thread_local_1d_id() < thread_cluster_desc_.GetElementSize())
+        {
+            threadwise_transfer_.RunRead(src_desc, p_src, src_iterator_hacks);
+        }
+    }
+
+    __device__ void RunWrite(const DstDesc& dst_desc, DstData* p_dst)
+    {
+        if(BlockSize == thread_cluster_desc_.GetElementSize() or
+           get_thread_local_1d_id() < thread_cluster_desc_.GetElementSize())
+        {
+            threadwise_transfer_.RunWrite(dst_desc, p_dst);
+        }
+    }
+
+    __device__ void MoveSrcSliceWindow(const SrcDesc& src_desc, const Index& step)
+    {
+        if(BlockSize == thread_cluster_desc_.GetElementSize() or
+           get_thread_local_1d_id() < thread_cluster_desc_.GetElementSize())
+        {
+            threadwise_transfer_.MoveSrcSliceWindow(src_desc, step);
+        }
+    }
+
+    // SrcMoveSliceWindowIteratorHack to control index calculation move slice window
+    template <typename SrcMoveSliceWindowIteratorHack>
+    __device__ void
+    MoveSrcSliceWindow(const SrcDesc& src_desc,
+                       const Index& step,
+                       const SrcMoveSliceWindowIteratorHack& src_move_slice_window_iterator_hack)
+    {
+        if(BlockSize == thread_cluster_desc_.GetElementSize() or
+           get_thread_local_1d_id() < thread_cluster_desc_.GetElementSize())
+        {
+            threadwise_transfer_.MoveSrcSliceWindow(
+                src_desc, step, src_move_slice_window_iterator_hack);
+        }
+    }
+
+    __device__ void MoveDstSliceWindow(const DstDesc& dst_desc, const Index& step)
+    {
+        if(BlockSize == thread_cluster_desc_.GetElementSize() or
+           get_thread_local_1d_id() < thread_cluster_desc_.GetElementSize())
+        {
+            threadwise_transfer_.MoveDstSliceWindow(dst_desc, step);
+        }
+    }
+
+    static constexpr auto thread_cluster_desc_ =
+        make_cluster_descriptor(ThreadClusterLengths{}, ThreadClusterArrangeOrder{});
+
+    using ThreadwiseTransfer =
+        ThreadwiseDynamicTensorSliceTransfer_v3<ThreadSliceLengths,
+                                                DstInMemOp,
+                                                SrcData,
+                                                DstData,
+                                                SrcDesc,
+                                                DstDesc,
+                                                SrcDimAccessOrder,
+                                                DstDimAccessOrder,
+                                                SrcVectorDim,
+                                                DstVectorDim,
+                                                SrcScalarPerVector,
+                                                DstScalarPerVector,
+                                                SrcScalarStrideInVector,
+                                                DstScalarStrideInVector,
+                                                SrcAddressSpace,
+                                                DstAddressSpace,
+                                                ThreadTransferSrcResetCoordinateAfterRun,
+                                                ThreadTransferDstResetCoordinateAfterRun>;
+
+    ThreadwiseTransfer threadwise_transfer_;
+};
+
+} // namespace ck
+#endif

composable_kernel/include/tensor_operation/blockwise_gemm.hpp

@@ -95,26 +95,6 @@ struct BlockwiseGemmBlockABlockBThreadCTransANormalBNormalC_v2
                            level1_n_id * NPerLevel0Cluster + level0_n_id * NPerThreadSubC};
     }
 
-    __device__ static MatrixIndex GetDistanceFromBeginOfThreadMatrixC(index_t m_in_c,
-                                                                      index_t n_in_c)
-    {
-        constexpr auto c_thread_mtx = ThreadMatrixC{};
-
-        constexpr index_t MPerLevel1Cluster =
-            MPerThreadSubC * MLevel0ThreadCluster * MLevel1ThreadCluster;
-        constexpr index_t NPerLevel1Cluster =
-            NPerThreadSubC * NLevel0ThreadCluster * NLevel1ThreadCluster;
-
-        index_t m_repeat = m_in_c / MPerThreadSubC;
-        index_t n_repeat = n_in_c / NPerThreadSubC;
-
-        index_t m_in_sub_c = m_in_c % MPerThreadSubC;
-        index_t n_in_sub_c = n_in_c % NPerThreadSubC;
-
-        return MatrixIndex{m_repeat * MPerLevel1Cluster + m_in_sub_c,
-                           n_repeat * NPerLevel1Cluster + n_in_sub_c};
-    }
-
     template <typename FloatA, typename FloatB, typename FloatC>
     __device__ void
     Run_naive(const FloatA* p_a_block, const FloatB* p_b_block, FloatC* p_c_thread) const
@@ -336,9 +316,14 @@ struct BlockwiseGemmBlockABlockBThreadCTransANormalBNormalC_v2
         constexpr index_t MRepeat = MPerThread / MPerThreadSubC;
         constexpr index_t NRepeat = NPerThread / NPerThreadSubC;
 
-        static_if<MRepeat == 2 && NRepeat == 2>{}([&](auto) {
+        if constexpr(MRepeat == 2 && NRepeat == 2)
+        {
             Run_pipelined_2x2(p_a_block, p_b_block, p_c_thread);
-        }).Else([&](auto) { Run_naive(p_a_block, p_b_block, p_c_thread); });
+        }
+        else
+        {
+            Run_naive(p_a_block, p_b_block, p_c_thread);
+        }
 #else
         Run_naive(p_a_block, p_b_block, p_c_thread);
 #endif

composable_kernel/include/tensor_operation/blockwise_gemm_v2.hpp

@@ -0,0 +1,370 @@
+#ifndef CK_BLOCKWISE_GEMM_V2_HPP
+#define CK_BLOCKWISE_GEMM_V2_HPP
+
+#include "common_header.hpp"
+#include "threadwise_gemm_v2.hpp"
+
+namespace ck {
+
+// blockwise GEMM: C[M, N] += transpose(A[K, M]) * B[K, N]
+// A and B are visable to the whole block, C is distributed among each thread
+// If following number are power of 2, index calculation shall be greatly reduced:
+//    MPerThreadSubC, NPerThreadSubC, MLevel0ThreadCluster, NLevel0ThreadCluster,
+//    MLevel1ThreadCluster, NLevel1ThreadCluster
+template <index_t BlockSize,
+          typename BlockMatrixA,
+          typename BlockMatrixB,
+          typename ThreadMatrixC,
+          index_t MPerThreadSubC,
+          index_t NPerThreadSubC,
+          index_t KPerThreadLoop,
+          index_t MLevel0ThreadCluster,
+          index_t NLevel0ThreadCluster,
+          index_t MLevel1ThreadCluster,
+          index_t NLevel1ThreadCluster,
+          index_t ThreadGemmADataPerRead_M,
+          index_t ThreadGemmBDataPerRead_N>
+struct BlockwiseGemm_km_kn_m0m1n0n1_v1
+{
+    struct MatrixIndex
+    {
+        index_t row;
+        index_t col;
+    };
+
+    index_t mMyThreadOffsetA;
+    index_t mMyThreadOffsetB;
+
+    __device__ BlockwiseGemm_km_kn_m0m1n0n1_v1()
+    {
+        static_assert(BlockMatrixA::IsKnownAtCompileTime() &&
+                          BlockMatrixB::IsKnownAtCompileTime() &&
+                          ThreadMatrixC::IsKnownAtCompileTime(),
+                      "wrong! Desc should be known at compile-time");
+
+        constexpr auto I0 = Number<0>{};
+        constexpr auto I1 = Number<1>{};
+
+        constexpr index_t ThreadPerLevel1Cluster = MLevel0ThreadCluster * NLevel0ThreadCluster *
+                                                   MLevel1ThreadCluster * NLevel1ThreadCluster;
+
+        static_assert(BlockSize == ThreadPerLevel1Cluster, "wrong! wrong blocksize\n");
+
+        static_assert(BlockMatrixA{}.GetLength(I0) == BlockMatrixB{}.GetLength(I0),
+                      "wrong! K dimension not consistent\n");
+
+        constexpr index_t M = BlockMatrixA{}.GetLength(I1); // A is transposed
+        constexpr index_t N = BlockMatrixB{}.GetLength(I1);
+
+        static_assert(M % (MPerThreadSubC * MLevel0ThreadCluster * MLevel1ThreadCluster) == 0 &&
+                          N % (NPerThreadSubC * NLevel0ThreadCluster * NLevel1ThreadCluster) == 0,
+                      "wrong! Cannot evenly divide work among\n");
+
+        static_assert(ThreadMatrixC{}.GetLength(I0) == GetThreadMatrixCLengths()[I0] &&
+                          ThreadMatrixC{}.GetLength(I1) == GetThreadMatrixCLengths()[I1],
+                      "wrong! ThreadMatrixC lengths is wrong");
+
+        auto c_thread_mtx_index = GetBeginOfThreadMatrixC(get_thread_local_1d_id());
+
+        mMyThreadOffsetA = BlockMatrixA{}.CalculateOffset(make_tuple(0, c_thread_mtx_index.row));
+        mMyThreadOffsetB = BlockMatrixB{}.CalculateOffset(make_tuple(0, c_thread_mtx_index.col));
+    }
+
+    __device__ static constexpr auto GetThreadMatrixCLengths()
+    {
+        constexpr auto I1 = Number<1>{};
+
+        constexpr index_t M = BlockMatrixA{}.GetLength(I1); // A is transposed
+        constexpr index_t N = BlockMatrixB{}.GetLength(I1);
+
+        constexpr index_t MRepeat =
+            M / (MPerThreadSubC * MLevel0ThreadCluster * MLevel1ThreadCluster);
+        constexpr index_t NRepeat =
+            N / (NPerThreadSubC * NLevel0ThreadCluster * NLevel1ThreadCluster);
+
+        return Sequence<MRepeat * MPerThreadSubC, NRepeat * NPerThreadSubC>{};
+    }
+
+    __device__ static MatrixIndex GetBeginOfThreadMatrixC(index_t thread_id)
+    {
+        constexpr index_t ThreadPerLevel0Cluster = MLevel0ThreadCluster * NLevel0ThreadCluster;
+
+        index_t level1_id   = thread_id / ThreadPerLevel0Cluster;
+        index_t level1_m_id = level1_id / NLevel1ThreadCluster;
+        index_t level1_n_id = level1_id % NLevel1ThreadCluster;
+
+        index_t level0_id   = thread_id % ThreadPerLevel0Cluster;
+        index_t level0_m_id = level0_id / NLevel0ThreadCluster;
+        index_t level0_n_id = level0_id % NLevel0ThreadCluster;
+
+        constexpr index_t MPerLevel0Cluster = MPerThreadSubC * MLevel0ThreadCluster;
+        constexpr index_t NPerLevel0Cluster = NPerThreadSubC * NLevel0ThreadCluster;
+
+        return MatrixIndex{level1_m_id * MPerLevel0Cluster + level0_m_id * MPerThreadSubC,
+                           level1_n_id * NPerLevel0Cluster + level0_n_id * NPerThreadSubC};
+    }
+
+    template <typename FloatA, typename FloatB, typename FloatC>
+    __device__ void
+    Run_naive(const FloatA* p_a_block, const FloatB* p_b_block, FloatC* p_c_thread) const
+    {
+        constexpr auto I0 = Number<0>{};
+        constexpr auto I1 = Number<1>{};
+
+        constexpr auto a_block_mtx  = BlockMatrixA{};
+        constexpr auto b_block_mtx  = BlockMatrixB{};
+        constexpr auto c_thread_mtx = ThreadMatrixC{};
+
+        constexpr auto K = a_block_mtx.GetLength(I0);
+
+        constexpr auto MPerThread = c_thread_mtx.GetLength(I0);
+        constexpr auto NPerThread = c_thread_mtx.GetLength(I1);
+
+        constexpr index_t MPerLevel1Cluster =
+            MPerThreadSubC * MLevel0ThreadCluster * MLevel1ThreadCluster;
+        constexpr index_t NPerLevel1Cluster =
+            NPerThreadSubC * NLevel0ThreadCluster * NLevel1ThreadCluster;
+
+        constexpr index_t MRepeat = MPerThread / MPerThreadSubC;
+        constexpr index_t NRepeat = NPerThread / NPerThreadSubC;
+
+        // thread A, B for GEMM
+        constexpr auto a_thread_mtx = make_dynamic_naive_tensor_descriptor_packed_v2(
+            Number<KPerThreadLoop>{}, Number<MPerThread>{});
+
+        constexpr auto b_thread_mtx = make_dynamic_naive_tensor_descriptor_packed_v2(
+            Number<KPerThreadLoop>{}, Number<NPerThread>{});
+
+        FloatA p_a_thread[a_thread_mtx.GetElementSpace()];
+        FloatB p_b_thread[b_thread_mtx.GetElementSpace()];
+
+        constexpr auto a_thread_copy = ThreadwiseMatrixSliceCopy_v2<BlockMatrixA,
+                                                                    decltype(a_thread_mtx),
+                                                                    KPerThreadLoop,
+                                                                    MPerThreadSubC,
+                                                                    ThreadGemmADataPerRead_M>{};
+
+        constexpr auto b_thread_copy = ThreadwiseMatrixSliceCopy_v2<BlockMatrixB,
+                                                                    decltype(b_thread_mtx),
+                                                                    KPerThreadLoop,
+                                                                    NPerThreadSubC,
+                                                                    ThreadGemmBDataPerRead_N>{};
+
+        constexpr auto threadwise_gemm = ThreadwiseGemm_km_kn_mn_v1<decltype(a_thread_mtx),
+                                                                    decltype(b_thread_mtx),
+                                                                    decltype(c_thread_mtx)>{};
+#pragma unroll
+        // loop over k
+        for(index_t k_begin = 0; k_begin < K; k_begin += KPerThreadLoop)
+        {
+#pragma unroll
+            // read A
+            for(index_t m_repeat = 0; m_repeat < MRepeat; ++m_repeat)
+            {
+                a_thread_copy.Run(p_a_block +
+                                      a_block_mtx.CalculateOffset(
+                                          make_tuple(k_begin, m_repeat * MPerLevel1Cluster)) +
+                                      mMyThreadOffsetA,
+                                  p_a_thread + a_thread_mtx.CalculateOffset(
+                                                   make_tuple(0, m_repeat * MPerThreadSubC)));
+            }
+
+#pragma unroll
+            // read B
+            for(index_t n_repeat = 0; n_repeat < NRepeat; ++n_repeat)
+            {
+                b_thread_copy.Run(p_b_block +
+                                      b_block_mtx.CalculateOffset(
+                                          make_tuple(k_begin, n_repeat * NPerLevel1Cluster)) +
+                                      mMyThreadOffsetB,
+                                  p_b_thread + b_thread_mtx.CalculateOffset(
+                                                   make_tuple(0, n_repeat * NPerThreadSubC)));
+            }
+
+            // C += A * B
+            threadwise_gemm.Run(p_a_thread, p_b_thread, p_c_thread);
+        }
+    }
+
+    template <typename FloatA, typename FloatB, typename FloatC>
+    __device__ void
+    Run_pipelined_2x2(const FloatA* p_a_block, const FloatB* p_b_block, FloatC* p_c_thread) const
+    {
+        constexpr auto I0 = Number<0>{};
+        constexpr auto I1 = Number<1>{};
+
+        constexpr auto a_block_mtx  = BlockMatrixA{};
+        constexpr auto b_block_mtx  = BlockMatrixB{};
+        constexpr auto c_thread_mtx = ThreadMatrixC{};
+
+        constexpr auto K = a_block_mtx.GetLength(I0);
+
+        constexpr auto MPerThread = c_thread_mtx.GetLength(I0);
+        constexpr auto NPerThread = c_thread_mtx.GetLength(I1);
+
+        constexpr index_t MPerLevel1Cluster =
+            MPerThreadSubC * MLevel0ThreadCluster * MLevel1ThreadCluster;
+
+        constexpr index_t NPerLevel1Cluster =
+            NPerThreadSubC * NLevel0ThreadCluster * NLevel1ThreadCluster;
+
+        constexpr index_t MRepeat = MPerThread / MPerThreadSubC;
+        constexpr index_t NRepeat = NPerThread / NPerThreadSubC;
+
+        static_assert(MRepeat == 2 && NRepeat == 2,
+                      "wrong! inline asm cannot deal with this GEMM config yet");
+
+        // thread A, B
+        constexpr auto a_thread_mtx = make_dynamic_naive_tensor_descriptor_packed_v2(
+            make_tuple(Number<KPerThreadLoop>{}, Number<MPerThread>{}));
+
+        constexpr auto b_thread_mtx = make_dynamic_naive_tensor_descriptor_packed_v2(
+            make_tuple(Number<KPerThreadLoop>{}, Number<NPerThread>{}));
+
+        // thread A-sub, B-sub
+        constexpr auto a_thread_sub_mtx = make_dynamic_naive_tensor_descriptor_v2(
+            make_tuple(Number<KPerThreadLoop>{}, Number<MPerThreadSubC>{}),
+            make_tuple(Number<MPerThread>{}, Number<1>{}));
+
+        constexpr auto b_thread_sub_mtx = make_dynamic_naive_tensor_descriptor_v2(
+            make_tuple(Number<KPerThreadLoop>{}, Number<NPerThreadSubC>{}),
+            make_tuple(Number<NPerThread>{}, Number<1>{}));
+
+        constexpr auto c_thread_sub_mtx = make_dynamic_naive_tensor_descriptor_v2(
+            make_tuple(Number<MPerThreadSubC>{}, Number<NPerThreadSubC>{}),
+            make_tuple(Number<NPerThread>{}, Number<1>{}));
+
+        FloatA p_a_thread[a_thread_mtx.GetElementSpaceSize()];
+        FloatB p_b_thread[b_thread_mtx.GetElementSpaceSize()];
+
+        constexpr auto a_thread_copy = ThreadwiseMatrixSliceCopy_v2<BlockMatrixA,
+                                                                    decltype(a_thread_mtx),
+                                                                    KPerThreadLoop,
+                                                                    MPerThreadSubC,
+                                                                    ThreadGemmADataPerRead_M>{};
+
+        constexpr auto b_thread_copy = ThreadwiseMatrixSliceCopy_v2<BlockMatrixB,
+                                                                    decltype(b_thread_mtx),
+                                                                    KPerThreadLoop,
+                                                                    NPerThreadSubC,
+                                                                    ThreadGemmBDataPerRead_N>{};
+
+        constexpr auto threadwise_gemm = ThreadwiseGemm_km_kn_mn_v1<decltype(a_thread_sub_mtx),
+                                                                    decltype(b_thread_sub_mtx),
+                                                                    decltype(c_thread_sub_mtx)>{};
+
+        const FloatA* p_a_block_off = p_a_block + mMyThreadOffsetA;
+        const FloatB* p_b_block_off = p_b_block + mMyThreadOffsetB;
+
+        // read A_sub_0
+        a_thread_copy.Run(p_a_block_off, p_a_thread);
+
+        // read B_sub_0
+        b_thread_copy.Run(p_b_block_off, p_b_thread);
+
+        // read B_sub_1
+        b_thread_copy.Run(p_b_block_off +
+                              b_block_mtx.CalculateOffset(make_tuple(0, NPerLevel1Cluster)),
+                          p_b_thread + b_thread_mtx.CalculateOffset(make_tuple(0, NPerThreadSubC)));
+
+        // read A_sub_1
+        a_thread_copy.Run(p_a_block_off +
+                              a_block_mtx.CalculateOffset(make_tuple(0, MPerLevel1Cluster)),
+                          p_a_thread + a_thread_mtx.CalculateOffset(make_tuple(0, MPerThreadSubC)));
+
+        // C_sub_00 += transpose(A_sub_0) * B_sub_0
+        threadwise_gemm.Run(p_a_thread, p_b_thread, p_c_thread);
+
+        // C_sub_01 += transpose(A_sub_0) * B_sub_1
+        threadwise_gemm.Run(
+            p_a_thread,
+            p_b_thread + b_thread_mtx.CalculateOffset(make_tuple(0, NPerThreadSubC)),
+            p_c_thread + c_thread_mtx.CalculateOffset(make_tuple(0, NPerThreadSubC)));
+
+#pragma unroll
+        // loop over rest of k
+        for(index_t k = KPerThreadLoop; k < K; k += KPerThreadLoop)
+        {
+            // read A_sub_0
+            a_thread_copy.Run(p_a_block_off + a_block_mtx.CalculateOffset(make_tuple(k, 0)),
+                              p_a_thread);
+
+            // C_sub_10 += transpose(A_sub_1) * B_sub_0
+            threadwise_gemm.Run(
+                p_a_thread + a_thread_mtx.CalculateOffset(make_tuple(0, MPerThreadSubC)),
+                p_b_thread,
+                p_c_thread + c_thread_mtx.CalculateOffset(make_tuple(MPerThreadSubC, 0)));
+
+            // read B_sub_0
+            b_thread_copy.Run(p_b_block_off + b_block_mtx.CalculateOffset(make_tuple(k, 0)),
+                              p_b_thread);
+
+            // C_sub_11 += transpose(A_sub_1) * B_sub_1
+            threadwise_gemm.Run(
+                p_a_thread + a_thread_mtx.CalculateOffset(make_tuple(0, MPerThreadSubC)),
+                p_b_thread + b_thread_mtx.CalculateOffset(make_tuple(0, NPerThreadSubC)),
+                p_c_thread +
+                    c_thread_mtx.CalculateOffset(make_tuple(MPerThreadSubC, NPerThreadSubC)));
+
+            // read B_sub_1
+            b_thread_copy.Run(
+                p_b_block_off + b_block_mtx.CalculateOffset(make_tuple(k, NPerLevel1Cluster)),
+                p_b_thread + b_thread_mtx.CalculateOffset(make_tuple(0, NPerThreadSubC)));
+
+            // read A_sub_1
+            a_thread_copy.Run(
+                p_a_block_off + a_block_mtx.CalculateOffset(make_tuple(k, MPerLevel1Cluster)),
+                p_a_thread + a_thread_mtx.CalculateOffset(make_tuple(0, MPerThreadSubC)));
+
+            // C_sub_00 += transpose(A_sub_0) * B_sub_0
+            threadwise_gemm.Run(p_a_thread, p_b_thread, p_c_thread);
+
+            // C_sub_01 += transpose(A_sub_0) * B_sub_1
+            threadwise_gemm.Run(
+                p_a_thread,
+                p_b_thread + b_thread_mtx.CalculateOffset(make_tuple(0, NPerThreadSubC)),
+                p_c_thread + c_thread_mtx.CalculateOffset(make_tuple(0, NPerThreadSubC)));
+        }
+
+        // C_sub_10 += transpose(A_sub_1) * B_sub_0
+        threadwise_gemm.Run(
+            p_a_thread + a_thread_mtx.CalculateOffset(make_tuple(0, MPerThreadSubC)),
+            p_b_thread,
+            p_c_thread + c_thread_mtx.CalculateOffset(make_tuple(MPerThreadSubC, 0)));
+
+        // C_sub_11 += transpose(A_sub_1) * B_sub_1
+        threadwise_gemm.Run(
+            p_a_thread + a_thread_mtx.CalculateOffset(make_tuple(0, MPerThreadSubC)),
+            p_b_thread + b_thread_mtx.CalculateOffset(make_tuple(0, NPerThreadSubC)),
+            p_c_thread + c_thread_mtx.CalculateOffset(make_tuple(MPerThreadSubC, NPerThreadSubC)));
+    }
+
+    template <typename FloatA, typename FloatB, typename FloatC>
+    __device__ void Run(const FloatA* p_a_block, const FloatB* p_b_block, FloatC* p_c_thread) const
+    {
+#if CK_EXPERIMENTAL_BLOCKWISE_GEMM_USE_PIPELINE
+        constexpr auto I0 = Number<0>{};
+        constexpr auto I1 = Number<1>{};
+
+        constexpr index_t MPerThread = ThreadMatrixC{}.GetLength(I0);
+        constexpr index_t NPerThread = ThreadMatrixC{}.GetLength(I1);
+
+        constexpr index_t MRepeat = MPerThread / MPerThreadSubC;
+        constexpr index_t NRepeat = NPerThread / NPerThreadSubC;
+
+        if constexpr(MRepeat == 2 && NRepeat == 2)
+        {
+            Run_pipelined_2x2(p_a_block, p_b_block, p_c_thread);
+        }
+        else
+        {
+            Run_naive(p_a_block, p_b_block, p_c_thread);
+        }
+#else
+        Run_naive(p_a_block, p_b_block, p_c_thread);
+#endif
+    }
+};
+
+} // namespace ck
+#endif

composable_kernel/include/tensor_operation/blockwise_gemm_v3.hpp

@@ -0,0 +1,198 @@
+#ifndef CK_BLOCKWISE_GEMM_V3_HPP
+#define CK_BLOCKWISE_GEMM_V3_HPP
+
+#include "common_header.hpp"
+#include "threadwise_gemm_v3.hpp"
+
+namespace ck {
+
+// blockwise GEMM: C[M, N] += transpose(A[K, M]) * B[K, N]
+// A and B are visable to the whole block, C is distributed among each thread
+// If following number are power of 2, index calculation shall be greatly reduced:
+//    KPerThread, HPerThread, MLevel0ThreadCluster, NLevel0ThreadCluster,
+//    MLevel1ThreadCluster, NLevel1ThreadCluster
+template <index_t BlockSize,
+          typename BlockMatrixA,
+          typename BlockMatrixB,
+          typename ThreadMatrixC,
+          index_t KPerThread,
+          index_t HPerThread,
+          index_t WPerThread,
+          index_t EPerThreadLoop,
+          index_t ThreadGemmADataPerRead_K,
+          index_t ThreadGemmBDataPerRead_W>
+struct BlockwiseGemm_km_kn_m0m1n0n1_v3
+{
+    struct MatrixIndex
+    {
+        index_t k;
+        index_t h;
+        index_t w;
+    };
+
+    index_t mMyThreadOffsetA;
+
+    __device__ BlockwiseGemm_km_kn_m0m1n0n1_v3()
+    {
+        static_assert(BlockMatrixA::IsKnownAtCompileTime() &&
+                          BlockMatrixB::IsKnownAtCompileTime() &&
+                          ThreadMatrixC::IsKnownAtCompileTime(),
+                      "wrong! Desc should be known at compile-time");
+
+        constexpr auto I0 = Number<0>{};
+        constexpr auto I1 = Number<1>{};
+        constexpr auto I2 = Number<2>{};
+        constexpr auto I3 = Number<3>{};
+
+        static_assert(BlockMatrixA{}.GetLength(I0) == BlockMatrixB{}.GetLength(I0),
+                      "wrong! K dimension not consistent\n");
+
+        constexpr index_t K = BlockMatrixA{}.GetLength(I1); // A is transposed
+        constexpr index_t N = BlockMatrixB{}.GetLength(I1);
+        constexpr index_t H = BlockMatrixB{}.GetLength(I2);
+        constexpr index_t W = BlockMatrixB{}.GetLength(I3);
+
+        static_assert(K % KPerThread == 0 && H % HPerThread == 0 && W % WPerThread == 0,
+                      "wrong! Cannot evenly divide work among\n");
+
+        constexpr auto KThreadCluster = K / KPerThread;
+        constexpr auto HThreadCluster = H / HPerThread;
+        constexpr auto WThreadCluster = W / WPerThread;
+
+        static_assert(BlockSize == KThreadCluster * HThreadCluster * WThreadCluster,
+                      "wrong! wrong blocksize\n");
+
+        auto c_thread_mtx_index = GetBeginOfThreadMatrixC(get_thread_local_1d_id());
+
+        mMyThreadOffsetA =
+            BlockMatrixA{}.CalculateOffset(make_tuple(0, c_thread_mtx_index.k * KPerThread));
+    }
+
+    __device__ static constexpr auto GetThreadMatrixCLengths()
+    {
+        return Sequence<KPerThread, 1, HPerThread, WPerThread>{};
+    }
+
+    __device__ static MatrixIndex GetBeginOfThreadMatrixC(index_t thread_id)
+    {
+        constexpr index_t H = BlockMatrixB{}.GetLength(Number<2>{});
+        constexpr index_t W = BlockMatrixB{}.GetLength(Number<3>{});
+
+        constexpr auto num_w_threads  = W / WPerThread;
+        constexpr auto num_h_threads  = H / HPerThread;
+        constexpr auto num_hw_threads = num_w_threads * num_h_threads;
+
+        index_t k_thread_id  = thread_id / num_hw_threads;
+        index_t hw_thread_id = thread_id % num_hw_threads;
+
+        index_t h_thread_id = hw_thread_id / num_w_threads;
+        index_t w_thread_id = hw_thread_id % num_w_threads;
+
+        return MatrixIndex{k_thread_id, h_thread_id, w_thread_id};
+    }
+
+    template <typename SrcDesc,
+              typename DstDesc,
+              index_t NSliceRow,
+              index_t NSliceCol,
+              index_t DataPerAccess>
+    struct ThreadwiseSliceCopy_a
+    {
+        template <typename Data>
+        __device__ static void Run(const Data* p_src, Data* p_dst)
+        {
+            static_assert(SrcDesc::IsKnownAtCompileTime() && DstDesc::IsKnownAtCompileTime(),
+                          "wrong! Desc should be known at compile-time");
+
+            using vector_t = typename vector_type<Data, DataPerAccess>::type;
+
+            static_for<0, NSliceRow, 1>{}([&](auto i) {
+                static_for<0, NSliceCol, DataPerAccess>{}([&](auto j) {
+                    constexpr auto src_offset = SrcDesc{}.CalculateOffset(make_tuple(i, j));
+                    constexpr auto dst_offset = DstDesc{}.CalculateOffset(make_tuple(i, j));
+
+                    *reinterpret_cast<vector_t*>(&p_dst[dst_offset]) =
+                        *reinterpret_cast<const vector_t*>(&p_src[src_offset]);
+                });
+            });
+        }
+    };
+
+    template <typename FloatA, typename FloatB, typename FloatC>
+    __device__ void
+    Run_naive(const FloatA* p_a_block, const FloatB* p_b_thread, FloatC* p_c_thread) const
+    {
+
+        constexpr auto I0 = Number<0>{};
+        constexpr auto I1 = Number<1>{};
+        constexpr auto I2 = Number<2>{};
+        constexpr auto I3 = Number<3>{};
+
+        constexpr auto a_block_mtx = BlockMatrixA{};
+
+        constexpr auto EPerBlock = a_block_mtx.GetLength(I0);
+
+        constexpr auto KPerThreadSubC = 4;
+
+        static_assert(KPerThread % KPerThreadSubC == 0, "");
+        static_assert(HPerThread % 2 == 0, "");
+        static_assert(WPerThread % 2 == 0, "");
+
+        // thread A, B for GEMM
+        constexpr auto a_thread_mtx = make_dynamic_naive_tensor_descriptor_packed_v2(
+            make_tuple(Number<EPerThreadLoop>{}, Number<KPerThreadSubC>{}));
+
+        constexpr auto b_thread_mtx = make_dynamic_naive_tensor_descriptor_packed_v2(make_tuple(
+            Number<EPerThreadLoop>{}, Number<1>{}, Number<HPerThread>{}, Number<WPerThread>{}));
+
+        constexpr auto c_thread_mtx = make_dynamic_naive_tensor_descriptor_packed_v2(make_tuple(
+            Number<KPerThreadSubC>{}, Number<1>{}, Number<HPerThread>{}, Number<WPerThread>{}));
+
+        FloatA p_a_thread[a_thread_mtx.GetElementSpaceSize()];
+
+        constexpr auto a_thread_copy = ThreadwiseSliceCopy_a<BlockMatrixA,
+                                                             decltype(a_thread_mtx),
+                                                             EPerThreadLoop,
+                                                             KPerThreadSubC,
+                                                             ThreadGemmADataPerRead_K>{};
+
+        constexpr auto threadwise_gemm = ThreadwiseGemm_km_kn_mn_v3<decltype(a_thread_mtx),
+                                                                    decltype(b_thread_mtx),
+                                                                    decltype(c_thread_mtx)>{};
+        // loop over k
+#pragma unroll
+        for(index_t e_begin = 0; e_begin < EPerBlock; e_begin += EPerThreadLoop)
+        {
+#pragma unroll
+            for(index_t k_begin = 0; k_begin < KPerThread; k_begin += KPerThreadSubC)
+            {
+                a_thread_copy.Run(p_a_block +
+                                      a_block_mtx.CalculateOffset(make_tuple(e_begin, k_begin)) +
+                                      mMyThreadOffsetA,
+                                  p_a_thread);
+
+                for(index_t h_begin = 0; h_begin < HPerThread; h_begin += 2)
+                {
+
+                    for(index_t w_begin = 0; w_begin < WPerThread; w_begin += 2)
+                    {
+                        threadwise_gemm.Run(p_a_thread,
+                                            p_b_thread + b_thread_mtx.CalculateOffset(make_tuple(
+                                                             e_begin, 0, h_begin, w_begin)),
+                                            p_c_thread + c_thread_mtx.CalculateOffset(make_tuple(
+                                                             k_begin, 0, h_begin, w_begin)));
+                    }
+                }
+            }
+        }
+    }
+
+    template <typename FloatA, typename FloatB, typename FloatC>
+    __device__ void Run(const FloatA* p_a_block, const FloatB* p_b_thread, FloatC* p_c_thread) const
+    {
+        Run_naive(p_a_block, p_b_thread, p_c_thread);
+    }
+};
+
+} // namespace ck
+#endif

composable_kernel/include/tensor_operation/blockwise_generic_tensor_slice_copy.hpp

@@ -5,6 +5,7 @@
 #include "tensor_descriptor.hpp"
 #include "tensor_descriptor_helper.hpp"
 #include "tensor_coordinate.hpp"
+#include "cluster_descriptor.hpp"
 #include "threadwise_generic_tensor_slice_copy.hpp"
 
 namespace ck {
@@ -68,9 +69,9 @@ struct BlockwiseGenericTensorSliceCopy_v4
             const auto thread_data_id_begin = thread_cluster_id * ThreadSliceLengths{};
 
             mThreadwiseLoad.SetSrcSliceOrigin(src_block_slice_origin + thread_data_id_begin);
-            mThreadwiseLoad.SetDstSliceOrigin(make_zero_array<index_t, nDim>());
+            mThreadwiseLoad.SetDstSliceOrigin(make_zero_multi_index<nDim>());
 
-            mThreadwiseStore.SetSrcSliceOrigin(make_zero_array<index_t, nDim>());
+            mThreadwiseStore.SetSrcSliceOrigin(make_zero_multi_index<nDim>());
             mThreadwiseStore.SetDstSliceOrigin(dst_block_slice_origin + thread_data_id_begin);
         }
     }

composable_kernel/include/tensor_operation/gridwise_dynamic_gemm.hpp

@@ -0,0 +1,509 @@
+#ifndef CK_GRIDWISE_DYNAMIC_GEMM_HPP
+#define CK_GRIDWISE_DYNAMIC_GEMM_HPP
+
+#include "common_header.hpp"
+#include "dynamic_multi_index_transform_helper.hpp"
+#include "dynamic_tensor_descriptor.hpp"
+#include "dynamic_tensor_descriptor_helper.hpp"
+#include "blockwise_dynamic_tensor_slice_transfer.hpp"
+#include "threadwise_dynamic_tensor_slice_transfer.hpp"
+#include "blockwise_gemm_v2.hpp"
+
+namespace ck {
+
+template <index_t BlockSize,
+          typename FloatAB,
+          typename FloatAcc,
+          typename FloatC,
+          InMemoryDataOperation CGlobalMemoryDataOperation,
+          typename AGlobalDesc,
+          typename BGlobalDesc,
+          typename CGlobalDesc,
+          index_t MPerBlock,
+          index_t NPerBlock,
+          index_t KPerBlock,
+          index_t MPerThread,
+          index_t NPerThread,
+          index_t KPerThread,
+          index_t MLevel0Cluster,
+          index_t NLevel0Cluster,
+          index_t MLevel1Cluster,
+          index_t NLevel1Cluster,
+          typename ABlockTransferThreadSliceLengths_K_M,
+          typename ABlockTransferThreadClusterLengths_K_M,
+          typename ABlockTransferThreadClusterArrangeOrder,
+          typename ABlockTransferSrcAccessOrder,
+          index_t ABlockTransferSrcVectorDim,
+          index_t ABlockTransferSrcScalarPerVector,
+          index_t ABlockTransferDstScalarPerVector_M,
+          bool AThreadTransferSrcResetCoordinateAfterRun,
+          typename BBlockTransferThreadSliceLengths_K_N,
+          typename BBlockTransferThreadClusterLengths_K_N,
+          typename BBlockTransferThreadClusterArrangeOrder,
+          typename BBlockTransferSrcAccessOrder,
+          index_t BBlockTransferSrcVectorDim,
+          index_t BBlockTransferSrcScalarPerVector,
+          index_t BBlockTransferDstScalarPerVector_N,
+          bool BThreadTransferSrcResetCoordinateAfterRun,
+          typename CThreadTransferSrcDstAccessOrder,
+          index_t CThreadTransferSrcDstVectorDim,
+          index_t CThreadTransferDstScalarPerVector,
+          typename AGlobalIteratorHacks,
+          typename BGlobalIteratorHacks,
+          typename CGlobalIteratorHacks,
+          typename AGlobalMoveSliceWindowIteratorHacks,
+          typename BGlobalMoveSliceWindowIteratorHacks>
+struct GridwiseDynamicGemm_km_kn_m0m1n0n1_v1
+{
+    __host__ __device__ static constexpr index_t GetSharedMemoryNumberOfByte()
+    {
+        constexpr auto max_lds_align = math::lcm(Number<ABlockTransferDstScalarPerVector_M>{},
+                                                 Number<BBlockTransferDstScalarPerVector_N>{},
+                                                 Number<MPerThread>{},
+                                                 Number<NPerThread>{});
+
+        // A matrix in LDS memory, dst of blockwise copy
+        //   be careful of LDS alignment
+        constexpr auto a_k_m_block_desc = make_dynamic_naive_tensor_descriptor_aligned_v2(
+            make_tuple(Number<KPerBlock>{}, Number<MPerBlock>{}), max_lds_align);
+
+        // B matrix in LDS memory, dst of blockwise copy
+        //   be careful of LDS alignment
+        constexpr auto b_k_n_block_desc = make_dynamic_naive_tensor_descriptor_aligned_v2(
+            make_tuple(Number<KPerBlock>{}, Number<NPerBlock>{}), max_lds_align);
+
+        // LDS allocation for A and B: be careful of alignment
+        constexpr auto a_block_space_size =
+            math::integer_least_multiple(a_k_m_block_desc.GetElementSpaceSize(), max_lds_align);
+
+        constexpr auto b_block_space_size =
+            math::integer_least_multiple(b_k_n_block_desc.GetElementSpaceSize(), max_lds_align);
+
+        return 2 * (a_block_space_size + b_block_space_size) * sizeof(FloatAB);
+    }
+
+    template <bool HasMainKBlockLoop, bool HasDoubleTailKBlockLoop>
+    __device__ void Run(const AGlobalDesc& a_k_m_global_desc,
+                        const FloatAB* __restrict__ p_a_global,
+                        const BGlobalDesc& b_k_n_global_desc,
+                        const FloatAB* __restrict__ p_b_global,
+                        const CGlobalDesc& c_m0_m1_n0_n1_global_desc,
+                        FloatC* __restrict__ p_c_global,
+                        FloatAB* __restrict__ p_shared_block,
+                        integral_constant<bool, HasMainKBlockLoop>,
+                        integral_constant<bool, HasDoubleTailKBlockLoop>) const
+    {
+        constexpr auto I0 = Number<0>{};
+        constexpr auto I1 = Number<1>{};
+
+        const auto K = a_k_m_global_desc.GetLength(I0);
+        const auto M = a_k_m_global_desc.GetLength(I1);
+        const auto N = b_k_n_global_desc.GetLength(I1);
+
+        // divide block work by [M, N]
+#if 0
+        const auto m_block_work_num = M / Number<MPerBlock>{};
+        const auto n_block_work_num = N / Number<NPerBlock>{};
+
+        const index_t m_block_work_id = get_block_1d_id() / n_block_work_num;
+        const index_t n_block_work_id = get_block_1d_id() - m_block_work_id * n_block_work_num;
+
+#else
+        // Hack: this force result into SGPR
+        const index_t m_block_work_num = __builtin_amdgcn_readfirstlane(M / MPerBlock);
+        const index_t n_block_work_num = __builtin_amdgcn_readfirstlane(N / NPerBlock);
+
+        const index_t m_block_work_id =
+            __builtin_amdgcn_readfirstlane(get_block_1d_id() / n_block_work_num);
+        const index_t n_block_work_id = get_block_1d_id() - m_block_work_id * n_block_work_num;
+#endif
+
+        const index_t m_block_data_on_global = m_block_work_id * MPerBlock;
+        const index_t n_block_data_on_global = n_block_work_id * NPerBlock;
+
+        // lds max alignment
+        constexpr auto max_lds_align = math::lcm(Number<ABlockTransferDstScalarPerVector_M>{},
+                                                 Number<BBlockTransferDstScalarPerVector_N>{},
+                                                 Number<MPerThread>{},
+                                                 Number<NPerThread>{});
+
+        // A matrix in LDS memory, dst of blockwise copy
+        //   be careful of LDS alignment
+        constexpr auto a_k_m_block_desc = make_dynamic_naive_tensor_descriptor_aligned_v2(
+            make_tuple(Number<KPerBlock>{}, Number<MPerBlock>{}), max_lds_align);
+
+        // B matrix in LDS memory, dst of blockwise copy
+        //   be careful of LDS alignment
+        constexpr auto b_k_n_block_desc = make_dynamic_naive_tensor_descriptor_aligned_v2(
+            make_tuple(Number<KPerBlock>{}, Number<NPerBlock>{}), max_lds_align);
+
+        // A matrix blockwise copy
+        auto a_blockwise_copy =
+            BlockwiseDynamicTensorSliceTransfer_v4<BlockSize,
+                                                   InMemoryDataOperation::Set,
+                                                   Sequence<KPerBlock, MPerBlock>,
+                                                   ABlockTransferThreadSliceLengths_K_M,
+                                                   ABlockTransferThreadClusterLengths_K_M,
+                                                   ABlockTransferThreadClusterArrangeOrder,
+                                                   FloatAB,
+                                                   FloatAB,
+                                                   decltype(a_k_m_global_desc),
+                                                   decltype(a_k_m_block_desc),
+                                                   ABlockTransferSrcAccessOrder,
+                                                   Sequence<0, 1>,
+                                                   ABlockTransferSrcVectorDim,
+                                                   1,
+                                                   ABlockTransferSrcScalarPerVector,
+                                                   ABlockTransferDstScalarPerVector_M,
+                                                   AddressSpace::Global,
+                                                   AddressSpace::Lds,
+                                                   1,
+                                                   1,
+                                                   AThreadTransferSrcResetCoordinateAfterRun,
+                                                   true>(
+                a_k_m_global_desc,
+                make_multi_index(0, m_block_data_on_global),
+                a_k_m_block_desc,
+                make_multi_index(0, 0));
+
+        // B matrix blockwise copy
+        auto b_blockwise_copy =
+            BlockwiseDynamicTensorSliceTransfer_v4<BlockSize,
+                                                   InMemoryDataOperation::Set,
+                                                   Sequence<KPerBlock, NPerBlock>,
+                                                   BBlockTransferThreadSliceLengths_K_N,
+                                                   BBlockTransferThreadClusterLengths_K_N,
+                                                   BBlockTransferThreadClusterArrangeOrder,
+                                                   FloatAB,
+                                                   FloatAB,
+                                                   decltype(b_k_n_global_desc),
+                                                   decltype(b_k_n_block_desc),
+                                                   BBlockTransferSrcAccessOrder,
+                                                   Sequence<0, 1>,
+                                                   BBlockTransferSrcVectorDim,
+                                                   1,
+                                                   BBlockTransferSrcScalarPerVector,
+                                                   BBlockTransferDstScalarPerVector_N,
+                                                   AddressSpace::Global,
+                                                   AddressSpace::Lds,
+                                                   1,
+                                                   1,
+                                                   BThreadTransferSrcResetCoordinateAfterRun,
+                                                   true>(
+                b_k_n_global_desc,
+                make_multi_index(0, n_block_data_on_global),
+                b_k_n_block_desc,
+                make_multi_index(0, 0));
+
+        // GEMM definition
+        //   c_mtx += transpose(a_mtx) * b_mtx
+        //     a_mtx[KPerBlock, MPerBlock] is in LDS
+        //     b_mtx[KPerBlocl, NPerBlock] is in LDS
+        //     c_mtx[MPerBlock, NPerBlock] is distributed among threads, and saved in
+        //       register
+        // sanity check
+        static_assert(MPerBlock % (MPerThread * MLevel0Cluster * MLevel1Cluster) == 0 &&
+                          NPerBlock % (NPerThread * NLevel0Cluster * NLevel1Cluster) == 0,
+                      "wrong!");
+
+        constexpr index_t MRepeat = MPerBlock / (MPerThread * MLevel0Cluster * MLevel1Cluster);
+        constexpr index_t NRepeat = NPerBlock / (NPerThread * NLevel0Cluster * NLevel1Cluster);
+
+        // c_thread_mtx definition: this is a mess
+        // TODO:: more elegent way of defining c_thread_mtx
+        constexpr auto c_m0m1_n0n1_thread_desc = make_dynamic_naive_tensor_descriptor_packed_v2(
+            make_tuple(Number<MRepeat * MPerThread>{}, Number<NRepeat * NPerThread>{}));
+
+        const auto blockwise_gemm =
+            BlockwiseGemm_km_kn_m0m1n0n1_v1<BlockSize,
+                                            decltype(a_k_m_block_desc),
+                                            decltype(b_k_n_block_desc),
+                                            decltype(c_m0m1_n0n1_thread_desc),
+                                            MPerThread,
+                                            NPerThread,
+                                            KPerThread,
+                                            MLevel0Cluster,
+                                            NLevel0Cluster,
+                                            MLevel1Cluster,
+                                            NLevel1Cluster,
+                                            MPerThread,
+                                            NPerThread>{};
+
+        // LDS allocation for A and B: be careful of alignment
+        constexpr auto a_block_space_size =
+            math::integer_least_multiple(a_k_m_block_desc.GetElementSpaceSize(), max_lds_align);
+
+        constexpr auto b_block_space_size =
+            math::integer_least_multiple(b_k_n_block_desc.GetElementSpaceSize(), max_lds_align);
+
+        FloatAB* p_a_block_double = p_shared_block;
+        FloatAB* p_b_block_double = p_shared_block + 2 * a_block_space_size;
+
+        // register allocation for output
+        FloatAcc p_c_thread[c_m0m1_n0n1_thread_desc.GetElementSpaceSize()];
+
+        // zero out threadwise output
+        threadwise_matrix_set_zero_v2(c_m0m1_n0n1_thread_desc, p_c_thread);
+
+        constexpr auto a_block_slice_copy_step = make_multi_index(KPerBlock, 0);
+        constexpr auto b_block_slice_copy_step = make_multi_index(KPerBlock, 0);
+
+        // hack to control index calculation when iterating over A and B matrix for threadwise copy
+        constexpr auto a_k_m_global_iterator_hacks = AGlobalIteratorHacks{};
+        constexpr auto b_k_n_global_iterator_hacks = BGlobalIteratorHacks{};
+
+        // hack to control index calculation when move slice window for A and B matrix for
+        // threadwise copy
+        constexpr auto a_k_m_global_move_slice_window_iterator_hack =
+            AGlobalMoveSliceWindowIteratorHacks{};
+        constexpr auto b_k_n_global_move_slice_window_iterator_hack =
+            BGlobalMoveSliceWindowIteratorHacks{};
+
+        // LDS double buffer: preload data into LDS
+        {
+            a_blockwise_copy.RunRead(a_k_m_global_desc, p_a_global, a_k_m_global_iterator_hacks);
+            b_blockwise_copy.RunRead(b_k_n_global_desc, p_b_global, b_k_n_global_iterator_hacks);
+
+            a_blockwise_copy.RunWrite(a_k_m_block_desc, p_a_block_double);
+            b_blockwise_copy.RunWrite(b_k_n_block_desc, p_b_block_double);
+        }
+
+        if constexpr(HasMainKBlockLoop)
+        {
+            FloatAB* p_a_block_even = p_a_block_double;
+            FloatAB* p_b_block_even = p_b_block_double;
+
+            FloatAB* p_a_block_odd = p_a_block_double + a_block_space_size;
+            FloatAB* p_b_block_odd = p_b_block_double + b_block_space_size;
+
+            index_t k_block_data_begin = 0;
+
+            // LDS double buffer: main body
+            // use Do-While loop instead of For loop to simplify control flow
+            do
+            {
+                // even iteration
+                a_blockwise_copy.MoveSrcSliceWindow(a_k_m_global_desc,
+                                                    a_block_slice_copy_step,
+                                                    a_k_m_global_move_slice_window_iterator_hack);
+                b_blockwise_copy.MoveSrcSliceWindow(b_k_n_global_desc,
+                                                    b_block_slice_copy_step,
+                                                    b_k_n_global_move_slice_window_iterator_hack);
+
+                __syncthreads();
+
+                // LDS doubel buffer: load next data from device mem
+                a_blockwise_copy.RunRead(
+                    a_k_m_global_desc, p_a_global, a_k_m_global_iterator_hacks);
+                b_blockwise_copy.RunRead(
+                    b_k_n_global_desc, p_b_global, b_k_n_global_iterator_hacks);
+
+                // LDS double buffer: GEMM on current data
+                blockwise_gemm.Run(p_a_block_even, p_b_block_even, p_c_thread);
+
+                // LDS double buffer: store next data to LDS
+                a_blockwise_copy.RunWrite(a_k_m_block_desc, p_a_block_odd);
+                b_blockwise_copy.RunWrite(b_k_n_block_desc, p_b_block_odd);
+
+                // odd iteration
+                a_blockwise_copy.MoveSrcSliceWindow(a_k_m_global_desc,
+                                                    a_block_slice_copy_step,
+                                                    a_k_m_global_move_slice_window_iterator_hack);
+                b_blockwise_copy.MoveSrcSliceWindow(b_k_n_global_desc,
+                                                    b_block_slice_copy_step,
+                                                    b_k_n_global_move_slice_window_iterator_hack);
+
+                __syncthreads();
+
+                // LDS doubel buffer: load next data from device mem
+                a_blockwise_copy.RunRead(
+                    a_k_m_global_desc, p_a_global, a_k_m_global_iterator_hacks);
+                b_blockwise_copy.RunRead(
+                    b_k_n_global_desc, p_b_global, b_k_n_global_iterator_hacks);
+
+                // LDS double buffer: GEMM on current data
+                blockwise_gemm.Run(p_a_block_odd, p_b_block_odd, p_c_thread);
+
+                // LDS double buffer: store next data to LDS
+                a_blockwise_copy.RunWrite(a_k_m_block_desc, p_a_block_even);
+                b_blockwise_copy.RunWrite(b_k_n_block_desc, p_b_block_even);
+
+                k_block_data_begin += 2 * KPerBlock;
+            } while(k_block_data_begin < K - 2 * KPerBlock);
+        }
+
+        // LDS double buffer: tail
+        if constexpr(HasDoubleTailKBlockLoop) // if has 2 iteration left
+        {
+            a_blockwise_copy.MoveSrcSliceWindow(a_k_m_global_desc,
+                                                a_block_slice_copy_step,
+                                                a_k_m_global_move_slice_window_iterator_hack);
+            b_blockwise_copy.MoveSrcSliceWindow(b_k_n_global_desc,
+                                                b_block_slice_copy_step,
+                                                b_k_n_global_move_slice_window_iterator_hack);
+
+            __syncthreads();
+
+            // LDS double buffer: load last data from device mem
+            a_blockwise_copy.RunRead(a_k_m_global_desc, p_a_global, a_k_m_global_iterator_hacks);
+            b_blockwise_copy.RunRead(b_k_n_global_desc, p_b_global, b_k_n_global_iterator_hacks);
+
+            // LDS double buffer: GEMM on 2nd-last data
+            blockwise_gemm.Run(p_a_block_double, p_b_block_double, p_c_thread);
+
+            // LDS double buffer: store last data to LDS
+            a_blockwise_copy.RunWrite(a_k_m_block_desc, p_a_block_double + a_block_space_size);
+            b_blockwise_copy.RunWrite(b_k_n_block_desc, p_b_block_double + b_block_space_size);
+
+            __syncthreads();
+
+            // LDS double buffer: GEMM on last data
+            blockwise_gemm.Run(p_a_block_double + a_block_space_size,
+                               p_b_block_double + b_block_space_size,
+                               p_c_thread);
+        }
+        else // if has 1 iteration left
+        {
+            __syncthreads();
+
+            // LDS double buffer: GEMM on last data
+            blockwise_gemm.Run(p_a_block_double, p_b_block_double, p_c_thread);
+        }
+
+        // output: register to global memory
+        {
+            constexpr auto M1 = Number<MPerThread * MLevel0Cluster * MLevel1Cluster>{};
+            constexpr auto N1 = Number<NPerThread * NLevel0Cluster * NLevel1Cluster>{};
+
+            // define input tensor descriptor for threadwise copy
+            //     thread input tensor, src of threadwise copy
+            constexpr auto c_m0_m1_n0_n1_thread_desc =
+                make_dynamic_naive_tensor_descriptor_packed_v2(make_tuple(Number<MRepeat>{},
+                                                                          Number<MPerThread>{},
+                                                                          Number<NRepeat>{},
+                                                                          Number<NPerThread>{}));
+
+            // calculate origin of thread input tensor on global memory
+            //     blockwise GEMM c matrix starting index
+            const auto c_thread_mtx_on_block =
+                blockwise_gemm.GetBeginOfThreadMatrixC(get_thread_local_1d_id());
+
+            const index_t m_thread_data_on_global =
+                m_block_data_on_global + c_thread_mtx_on_block.row;
+
+            const index_t n_thread_data_on_global =
+                n_block_data_on_global + c_thread_mtx_on_block.col;
+
+            // hack to control index calculation when iterating over c_m0_m1_n0_n1_global tensor
+            constexpr auto c_m0_m1_n0_n1_global_tensor_iterator_hacks = CGlobalIteratorHacks{};
+
+            constexpr auto tmp = make_unmerge_transform(make_tuple(
+                Number<MRepeat>{}, Number<MPerThread>{}, Number<NRepeat>{}, Number<NPerThread>{}));
+
+            ThreadwiseDynamicTensorSliceTransfer_v1r3<
+                FloatAcc,
+                FloatC,
+                decltype(c_m0_m1_n0_n1_thread_desc),
+                decltype(c_m0_m1_n0_n1_global_desc),
+                Sequence<MRepeat, MPerThread, NRepeat, NPerThread>,
+                CThreadTransferSrcDstAccessOrder,
+                CThreadTransferSrcDstVectorDim,
+                CThreadTransferDstScalarPerVector,
+                AddressSpace::Vgpr,
+                AddressSpace::Global,
+                CGlobalMemoryDataOperation,
+                1,
+                true>(c_m0_m1_n0_n1_global_desc,
+                      make_multi_index(m_thread_data_on_global / M1,
+                                       m_thread_data_on_global % M1,
+                                       n_thread_data_on_global / N1,
+                                       n_thread_data_on_global % N1))
+                .Run(c_m0_m1_n0_n1_thread_desc,
+                     make_tuple(I0, I0, I0, I0),
+                     p_c_thread,
+                     c_m0_m1_n0_n1_global_desc,
+                     p_c_global,
+                     c_m0_m1_n0_n1_global_tensor_iterator_hacks);
+        }
+    }
+
+    // pass tensor descriptor by reference
+    template <bool HasMainKBlockLoop, bool HasDoubleTailKBlockLoop>
+    __device__ void Run(const AGlobalDesc& a_k_m_global_desc,
+                        const FloatAB* __restrict__ p_a_global,
+                        const BGlobalDesc& b_k_n_global_desc,
+                        const FloatAB* __restrict__ p_b_global,
+                        const CGlobalDesc& c_m0_m1_n0_n1_global_desc,
+                        FloatC* __restrict__ p_c_global,
+                        integral_constant<bool, HasMainKBlockLoop>,
+                        integral_constant<bool, HasDoubleTailKBlockLoop>) const
+    {
+        constexpr index_t shared_block_size = GetSharedMemoryNumberOfByte() / sizeof(FloatAB);
+
+        __shared__ FloatAB p_shared_block[shared_block_size];
+
+        Run(a_k_m_global_desc,
+            p_a_global,
+            b_k_n_global_desc,
+            p_b_global,
+            c_m0_m1_n0_n1_global_desc,
+            p_c_global,
+            p_shared_block,
+            integral_constant<bool, HasMainKBlockLoop>{},
+            integral_constant<bool, HasDoubleTailKBlockLoop>{});
+    }
+
+    // pass tensor descriptors by pointers
+    template <bool HasMainKBlockLoop, bool HasDoubleTailKBlockLoop>
+    __device__ void Run(const AGlobalDesc* p_a_k_m_global_desc,
+                        const FloatAB* __restrict__ p_a_global,
+                        const BGlobalDesc* p_b_k_n_global_desc,
+                        const FloatAB* __restrict__ p_b_global,
+                        const CGlobalDesc* p_c_m0_m1_n0_n1_global_desc,
+                        FloatC* __restrict__ p_c_global,
+                        integral_constant<bool, HasMainKBlockLoop>,
+                        integral_constant<bool, HasDoubleTailKBlockLoop>) const
+    {
+        const auto a_k_m_global_desc         = *p_a_k_m_global_desc;
+        const auto b_k_n_global_desc         = *p_b_k_n_global_desc;
+        const auto c_m0_m1_n0_n1_global_desc = *p_c_m0_m1_n0_n1_global_desc;
+
+        Run(a_k_m_global_desc,
+            p_a_global,
+            b_k_n_global_desc,
+            p_b_global,
+            c_m0_m1_n0_n1_global_desc,
+            p_c_global,
+            integral_constant<bool, HasMainKBlockLoop>{},
+            integral_constant<bool, HasDoubleTailKBlockLoop>{});
+    }
+
+    // pass tensor descriptors by void*
+    template <bool HasMainKBlockLoop, bool HasDoubleTailKBlockLoop>
+    __device__ void Run(const void* p_a_k_m_global_desc,
+                        const FloatAB* __restrict__ p_a_global,
+                        const void* p_b_k_n_global_desc,
+                        const FloatAB* __restrict__ p_b_global,
+                        const void* p_c_m0_m1_n0_n1_global_desc,
+                        FloatC* __restrict__ p_c_global,
+                        integral_constant<bool, HasMainKBlockLoop>,
+                        integral_constant<bool, HasDoubleTailKBlockLoop>) const
+    {
+        const auto a_k_m_global_desc = *reinterpret_cast<const AGlobalDesc*>(p_a_k_m_global_desc);
+        const auto b_k_n_global_desc = *reinterpret_cast<const BGlobalDesc*>(p_b_k_n_global_desc);
+        const auto c_m0_m1_n0_n1_global_desc =
+            *reinterpret_cast<const CGlobalDesc*>(p_c_m0_m1_n0_n1_global_desc);
+
+        Run(a_k_m_global_desc,
+            p_a_global,
+            b_k_n_global_desc,
+            p_b_global,
+            c_m0_m1_n0_n1_global_desc,
+            p_c_global,
+            integral_constant<bool, HasMainKBlockLoop>{},
+            integral_constant<bool, HasDoubleTailKBlockLoop>{});
+    }
+};
+
+} // namespace ck
+#endif

composable_kernel/include/tensor_operation/gridwise_dynamic_gemm_v2.hpp

@@ -0,0 +1,471 @@
+#ifndef CK_GRIDWISE_DYNAMIC_GEMM_V2_HPP
+#define CK_GRIDWISE_DYNAMIC_GEMM_V2_HPP
+
+#include "common_header.hpp"
+#include "dynamic_multi_index_transform_helper.hpp"
+#include "dynamic_tensor_descriptor.hpp"
+#include "dynamic_tensor_descriptor_helper.hpp"
+#include "blockwise_dynamic_tensor_slice_transfer.hpp"
+#include "threadwise_dynamic_tensor_slice_transfer.hpp"
+#include "blockwise_gemm_v3.hpp"
+
+namespace ck {
+
+template <index_t BlockSize,
+          typename FloatAB,
+          typename FloatAcc,
+          typename FloatC,
+          InMemoryDataOperation CGlobalMemoryDataOperation,
+          typename AGlobalDesc,
+          typename BGlobalDesc,
+          typename CGlobalDesc,
+          index_t KPerBlock,
+          index_t HoPerBlock,
+          index_t WoPerBlock,
+          index_t EPerBlock,
+          index_t KPerThread,
+          index_t HoPerThread,
+          index_t WoPerThread,
+          index_t EPerThread,
+          typename ABlockTransferThreadSliceLengths_E_K,
+          typename ABlockTransferThreadClusterLengths_E_K,
+          typename ABlockTransferThreadClusterArrangeOrder,
+          typename ABlockTransferSrcAccessOrder,
+          index_t ABlockTransferSrcVectorDim,
+          index_t ABlockTransferSrcScalarPerVector,
+          index_t ABlockTransferDstScalarPerVector_K,
+          bool AThreadTransferSrcResetCoordinateAfterRun,
+          typename BBlockTransferSrcAccessOrder,
+          index_t BBlockTransferSrcVectorDim,
+          index_t BBlockTransferSrcScalarPerVector,
+          bool BThreadTransferSrcResetCoordinateAfterRun,
+          typename CThreadTransferSrcDstAccessOrder,
+          index_t CThreadTransferSrcDstVectorDim,
+          index_t CThreadTransferDstScalarPerVector,
+          typename AGlobalIteratorHacks,
+          typename BGlobalIteratorHacks,
+          typename CGlobalIteratorHacks,
+          typename AGlobalMoveSliceWindowIteratorHacks,
+          typename BGlobalMoveSliceWindowIteratorHacks>
+struct GridwiseDynamicGemm_km_kn_mn_v3
+{
+    __host__ __device__ static constexpr index_t GetSharedMemoryNumberOfByte()
+    {
+        constexpr auto E = EPerBlock * 3 * 3;
+
+        constexpr auto max_lds_align =
+            math::lcm(Number<ABlockTransferDstScalarPerVector_K>{}, Number<KPerBlock>{});
+
+        // A matrix in LDS memory, dst of blockwise copy
+        //   be careful of LDS alignment
+        constexpr auto a_e_k_desc = make_dynamic_naive_tensor_descriptor_aligned_v2(
+            make_tuple(Number<E>{}, Number<KPerBlock>{}), max_lds_align);
+
+        // LDS allocation for A and B: be careful of alignment
+        constexpr auto a_block_space_size =
+            math::integer_least_multiple(a_e_k_desc.GetElementSpaceSize(), max_lds_align);
+
+        return a_block_space_size * sizeof(FloatAB);
+    }
+
+    template <bool HasMainKBlockLoop, bool HasDoubleTailKBlockLoop>
+    __device__ void Run(const AGlobalDesc& a_e_k_global_desc,
+                        const FloatAB* __restrict__ p_a_global,
+                        const BGlobalDesc& b_e_n_ho_wo_global_desc,
+                        const FloatAB* __restrict__ p_b_global,
+                        const CGlobalDesc& c_k_n_ho_wo_global_desc,
+                        FloatC* __restrict__ p_c_global,
+                        FloatAB* __restrict__ p_shared_block,
+                        integral_constant<bool, HasMainKBlockLoop>,
+                        integral_constant<bool, HasDoubleTailKBlockLoop>) const
+    {
+        constexpr auto I0 = Number<0>{};
+        constexpr auto I1 = Number<1>{};
+        constexpr auto I2 = Number<2>{};
+        constexpr auto I3 = Number<3>{};
+
+        constexpr auto E = EPerBlock * 3 * 3;
+
+        // const auto E = a_e_k_global_desc.GetLength(I0);
+        const auto K = a_e_k_global_desc.GetLength(I1);
+
+        const auto N  = b_e_n_ho_wo_global_desc.GetLength(I1);
+        const auto Ho = b_e_n_ho_wo_global_desc.GetLength(I2);
+        const auto Wo = b_e_n_ho_wo_global_desc.GetLength(I3);
+
+        // divide block work by [M, N]
+#if 0
+        const auto k_block_work_num   = K / Number<KPerBlock>{};
+        const auto ho_block_work_num  = Ho / Number<HoPerBlock>{};
+        const auto wo_block_work_num  = Wo / Number<WoPerBlock>{};
+        const auto hwo_block_work_num = ho_block_work_num * wo_block_work_num;
+
+        const index_t k_block_work_id   = get_block_1d_id() / hwo_block_work_num;
+        const index_t hwo_block_work_id = get_block_1d_id() - k_block_work_id * hwo_block_work_num;
+
+        const index_t ho_block_work_id = hwo_block_work_id / wo_block_work_num;
+        const index_t wo_block_work_id = hwo_block_work_id - ho_block_work_id * wo_block_work_num;
+#else
+        // Hack: this force result into SGPR
+        const index_t k_block_work_num   = __builtin_amdgcn_readfirstlane(K / KPerBlock);
+        const index_t ho_block_work_num  = __builtin_amdgcn_readfirstlane(Ho / HoPerBlock);
+        const index_t wo_block_work_num  = __builtin_amdgcn_readfirstlane(Wo / WoPerBlock);
+        const index_t hwo_block_work_num = ho_block_work_num * wo_block_work_num;
+
+        const index_t k_block_work_id =
+            __builtin_amdgcn_readfirstlane(get_block_1d_id() / hwo_block_work_num);
+        const index_t hwo_block_work_id = get_block_1d_id() - k_block_work_id * hwo_block_work_num;
+
+        const index_t ho_block_work_id =
+            __builtin_amdgcn_readfirstlane(hwo_block_work_id / wo_block_work_num);
+        const index_t wo_block_work_id = hwo_block_work_id - ho_block_work_id * wo_block_work_num;
+#endif
+
+        // lds max alignment
+        constexpr auto max_lds_align =
+            math::lcm(Number<ABlockTransferDstScalarPerVector_K>{}, Number<KPerBlock>{});
+
+        // A matrix in LDS memory, dst of blockwise copy
+        //   be careful of LDS alignment
+        constexpr auto a_e_k_block_desc = make_dynamic_naive_tensor_descriptor_aligned_v2(
+            make_tuple(Number<EPerBlock>{}, Number<KPerBlock>{}), max_lds_align);
+
+        constexpr auto a_e_k_desc = make_dynamic_naive_tensor_descriptor_aligned_v2(
+            make_tuple(Number<E>{}, Number<KPerBlock>{}), max_lds_align);
+
+        // B matrix in LDS memory, dst of blockwise copy
+        //   be careful of LDS alignment
+        constexpr auto b_e_n_ho_wo_block_desc =
+            make_dynamic_naive_tensor_descriptor_packed_v2(make_tuple(
+                Number<EPerBlock>{}, Number<1>{}, Number<HoPerBlock>{}, Number<WoPerBlock>{}));
+
+        // c_thread_mtx definition: this is a mess
+        // TODO:: more elegent way of defining c_thread_mtx
+        constexpr auto c_k_n_ho_wo_thread_desc =
+            make_dynamic_naive_tensor_descriptor_packed_v2(make_tuple(
+                Number<KPerThread>{}, Number<1>{}, Number<HoPerThread>{}, Number<WoPerThread>{}));
+
+        const auto blockwise_gemm =
+            BlockwiseGemm_km_kn_m0m1n0n1_v3<BlockSize,
+                                            decltype(a_e_k_block_desc),
+                                            decltype(b_e_n_ho_wo_block_desc),
+                                            decltype(c_k_n_ho_wo_thread_desc),
+                                            KPerThread,
+                                            HoPerThread,
+                                            WoPerThread,
+                                            EPerThread,
+                                            ABlockTransferSrcScalarPerVector,
+                                            ABlockTransferDstScalarPerVector_K>{};
+
+        auto c_thread_mtx_index = blockwise_gemm.GetBeginOfThreadMatrixC(get_thread_local_1d_id());
+
+        const auto k_thread_id  = c_thread_mtx_index.k;
+        const auto ho_thread_id = c_thread_mtx_index.h;
+        const auto wo_thread_id = c_thread_mtx_index.w;
+
+        const index_t k_block_data_on_global  = k_block_work_id * KPerBlock;
+        const index_t ho_block_data_on_global = ho_block_work_id * HoPerBlock;
+        const index_t wo_block_data_on_global = wo_block_work_id * WoPerBlock;
+
+        const index_t ho_thread_data_on_global =
+            ho_block_data_on_global + ho_thread_id * HoPerThread;
+        const index_t wo_thread_data_on_global =
+            wo_block_data_on_global + wo_thread_id * WoPerThread;
+
+        // A matrix blockwise copy
+        auto a_blockwise_copy =
+            BlockwiseDynamicTensorSliceTransfer_v4<BlockSize,
+                                                   InMemoryDataOperation::Set,
+                                                   Sequence<E, KPerBlock>,
+                                                   ABlockTransferThreadSliceLengths_E_K,
+                                                   ABlockTransferThreadClusterLengths_E_K,
+                                                   ABlockTransferThreadClusterArrangeOrder,
+                                                   FloatAB,
+                                                   FloatAB,
+                                                   decltype(a_e_k_global_desc),
+                                                   decltype(a_e_k_desc),
+                                                   ABlockTransferSrcAccessOrder,
+                                                   Sequence<0, 1>,
+                                                   ABlockTransferSrcVectorDim,
+                                                   1,
+                                                   ABlockTransferSrcScalarPerVector,
+                                                   ABlockTransferDstScalarPerVector_K,
+                                                   AddressSpace::Global,
+                                                   AddressSpace::Lds,
+                                                   1,
+                                                   1,
+                                                   AThreadTransferSrcResetCoordinateAfterRun,
+                                                   true>(
+                a_e_k_global_desc,
+                make_multi_index(0, k_block_data_on_global),
+                a_e_k_desc,
+                make_multi_index(0, 0));
+
+        constexpr auto b_e_n_ho_wo_thread_desc =
+            make_dynamic_naive_tensor_descriptor_packed_v2(make_tuple(
+                Number<EPerBlock>{}, Number<1>{}, Number<HoPerThread>{}, Number<WoPerThread>{}));
+
+        auto b_threadwise_transfer = ThreadwiseDynamicTensorSliceTransfer_v2<
+            FloatAB,
+            FloatAB,
+            decltype(b_e_n_ho_wo_global_desc),
+            decltype(b_e_n_ho_wo_thread_desc),
+            Sequence<EPerBlock, 1, HoPerThread, WoPerThread>,
+            BBlockTransferSrcAccessOrder,
+            BBlockTransferSrcVectorDim,
+            BBlockTransferSrcScalarPerVector,
+            AddressSpace::Global,
+            AddressSpace::Vgpr,
+            InMemoryDataOperation::Set,
+            1,
+            true>(b_e_n_ho_wo_global_desc,
+                  make_multi_index(0, 0, ho_thread_data_on_global, wo_thread_data_on_global));
+
+        FloatAB* p_a_block = p_shared_block;
+
+        // register allocation for output
+        FloatAcc p_c_thread[c_k_n_ho_wo_thread_desc.GetElementSpaceSize()];
+
+        // zero out threadwise output
+        threadwise_matrix_set_zero_v3(c_k_n_ho_wo_thread_desc, p_c_thread);
+
+        constexpr auto b_thread_slice_copy_step = make_multi_index(EPerBlock, 0, 0, 0);
+
+        // hack to control index calculation when iterating over A and B matrix for threadwise copy
+        constexpr auto a_e_k_global_iterator_hacks       = AGlobalIteratorHacks{};
+        constexpr auto b_e_n_ho_wo_global_iterator_hacks = BGlobalIteratorHacks{};
+
+        // hack to control index calculation when move slice window for A and B matrix for
+        // threadwise copy
+        constexpr auto a_e_k_global_move_slice_window_iterator_hack =
+            AGlobalMoveSliceWindowIteratorHacks{};
+        constexpr auto b_e_n_ho_wo_global_move_slice_window_iterator_hack =
+            BGlobalMoveSliceWindowIteratorHacks{};
+
+        constexpr auto b_thread_space_size = b_e_n_ho_wo_thread_desc.GetElementSpaceSize();
+        FloatAB p_b_thread[b_thread_space_size * 2];
+
+        FloatAB* p_b_thread_double = p_b_thread;
+
+        // LDS double buffer: preload data into LDS
+        {
+            a_blockwise_copy.RunRead(a_e_k_global_desc, p_a_global, a_e_k_global_iterator_hacks);
+
+            b_threadwise_transfer.Run(b_e_n_ho_wo_global_desc,
+                                      p_b_global,
+                                      b_e_n_ho_wo_thread_desc,
+                                      make_tuple(I0, I0, I0, I0),
+                                      p_b_thread_double,
+                                      b_e_n_ho_wo_global_iterator_hacks);
+
+            a_blockwise_copy.RunWrite(a_e_k_desc, p_a_block);
+        }
+
+        __syncthreads();
+
+        index_t b_block_data_begin = 0;
+
+#if 1
+        if constexpr(HasMainKBlockLoop)
+        {
+            FloatAB* p_b_thread_even = p_b_thread_double;
+            FloatAB* p_b_thread_odd  = p_b_thread_double + b_thread_space_size;
+
+            // LDS double buffer: main body
+            // use Do-While loop instead of For loop to simplify control flow
+            do
+            {
+                // even iteration
+                b_threadwise_transfer.MoveSrcSliceWindow(b_e_n_ho_wo_global_desc,
+                                                         b_thread_slice_copy_step);
+
+                b_threadwise_transfer.Run(b_e_n_ho_wo_global_desc,
+                                          p_b_global,
+                                          b_e_n_ho_wo_thread_desc,
+                                          make_tuple(I0, I0, I0, I0),
+                                          p_b_thread_odd,
+                                          b_e_n_ho_wo_global_iterator_hacks);
+
+                // LDS double buffer: GEMM on current data
+                blockwise_gemm.Run(
+                    p_a_block + a_e_k_block_desc.CalculateOffset(make_tuple(b_block_data_begin, 0)),
+                    p_b_thread_even,
+                    p_c_thread);
+
+                b_block_data_begin += EPerBlock;
+
+                b_threadwise_transfer.MoveSrcSliceWindow(b_e_n_ho_wo_global_desc,
+                                                         b_thread_slice_copy_step);
+
+                b_threadwise_transfer.Run(b_e_n_ho_wo_global_desc,
+                                          p_b_global,
+                                          b_e_n_ho_wo_thread_desc,
+                                          make_tuple(I0, I0, I0, I0),
+                                          p_b_thread_even,
+                                          b_e_n_ho_wo_global_iterator_hacks);
+
+                // LDS double buffer: GEMM on current data
+                blockwise_gemm.Run(
+                    p_a_block + a_e_k_block_desc.CalculateOffset(make_tuple(b_block_data_begin, 0)),
+                    p_b_thread_odd,
+                    p_c_thread);
+
+                b_block_data_begin += EPerBlock;
+
+            } while(b_block_data_begin < E - 2 * EPerBlock);
+        }
+
+        // LDS double buffer: tail
+        if constexpr(HasDoubleTailKBlockLoop) // if has 2 iteration left
+        {
+            b_threadwise_transfer.MoveSrcSliceWindow(b_e_n_ho_wo_global_desc,
+                                                     b_thread_slice_copy_step);
+
+            b_threadwise_transfer.Run(b_e_n_ho_wo_global_desc,
+                                      p_b_global,
+                                      b_e_n_ho_wo_thread_desc,
+                                      make_tuple(I0, I0, I0, I0),
+                                      p_b_thread_double + b_thread_space_size,
+                                      b_e_n_ho_wo_global_iterator_hacks);
+
+            // LDS double buffer: GEMM on 2nd-last data
+            blockwise_gemm.Run(
+                p_a_block + a_e_k_block_desc.CalculateOffset(make_tuple(b_block_data_begin, 0)),
+                p_b_thread_double,
+                p_c_thread);
+
+            b_block_data_begin += EPerBlock;
+
+            // LDS double buffer: GEMM on last data
+            blockwise_gemm.Run(
+                p_a_block + a_e_k_block_desc.CalculateOffset(make_tuple(b_block_data_begin, 0)),
+                p_b_thread_double + b_thread_space_size,
+                p_c_thread);
+        }
+        else // if has 1 iteration left
+        {
+            // LDS double buffer: GEMM on last data
+            blockwise_gemm.Run(
+                p_a_block + a_e_k_block_desc.CalculateOffset(make_tuple(b_block_data_begin, 0)),
+                p_b_thread_double,
+                p_c_thread);
+        }
+#endif
+
+#if 1
+        // output: register to global memory
+        {
+            // hack to control index calculation when iterating over c_k_n_ho_wo_global tensor
+            constexpr auto c_k_n_ho_wo_global_tensor_iterator_hacks = CGlobalIteratorHacks{};
+
+            const index_t k_thread_data_on_global =
+                k_block_data_on_global + k_thread_id * KPerThread;
+
+            ThreadwiseDynamicTensorSliceTransfer_v1r3<
+                FloatAcc,
+                FloatC,
+                decltype(c_k_n_ho_wo_thread_desc),
+                decltype(c_k_n_ho_wo_global_desc),
+                Sequence<KPerThread, 1, HoPerThread, WoPerThread>,
+                CThreadTransferSrcDstAccessOrder,
+                CThreadTransferSrcDstVectorDim,
+                CThreadTransferDstScalarPerVector,
+                AddressSpace::Vgpr,
+                AddressSpace::Global,
+                CGlobalMemoryDataOperation,
+                1,
+                true>(
+                c_k_n_ho_wo_global_desc,
+                make_multi_index(
+                    k_thread_data_on_global, 0, ho_thread_data_on_global, wo_thread_data_on_global))
+                .Run(c_k_n_ho_wo_thread_desc,
+                     make_tuple(I0, I0, I0, I0),
+                     p_c_thread,
+                     c_k_n_ho_wo_global_desc,
+                     p_c_global,
+                     c_k_n_ho_wo_global_tensor_iterator_hacks);
+        }
+#endif
+    }
+
+    // pass tensor descriptor by reference
+    template <bool HasMainKBlockLoop, bool HasDoubleTailKBlockLoop>
+    __device__ void Run(const AGlobalDesc& a_e_k_global_desc,
+                        const FloatAB* __restrict__ p_a_global,
+                        const BGlobalDesc& b_e_n_ho_wo_global_desc,
+                        const FloatAB* __restrict__ p_b_global,
+                        const CGlobalDesc& c_k_n_ho_wo_global_desc,
+                        FloatC* __restrict__ p_c_global,
+                        integral_constant<bool, HasMainKBlockLoop>,
+                        integral_constant<bool, HasDoubleTailKBlockLoop>) const
+    {
+        constexpr index_t shared_block_size = GetSharedMemoryNumberOfByte() / sizeof(FloatAB);
+
+        __shared__ FloatAB p_shared_block[shared_block_size];
+
+        Run(a_e_k_global_desc,
+            p_a_global,
+            b_e_n_ho_wo_global_desc,
+            p_b_global,
+            c_k_n_ho_wo_global_desc,
+            p_c_global,
+            p_shared_block,
+            integral_constant<bool, HasMainKBlockLoop>{},
+            integral_constant<bool, HasDoubleTailKBlockLoop>{});
+    }
+
+    // pass tensor descriptors by their pointers
+    template <bool HasMainKBlockLoop, bool HasDoubleTailKBlockLoop>
+    __device__ void Run(const AGlobalDesc* p_a_e_k_global_desc,
+                        const FloatAB* __restrict__ p_a_global,
+                        const BGlobalDesc* p_b_e_n_ho_wo_global_desc,
+                        const FloatAB* __restrict__ p_b_global,
+                        const CGlobalDesc* p_c_k_n_ho_wo_global_desc,
+                        FloatC* __restrict__ p_c_global,
+                        integral_constant<bool, HasMainKBlockLoop>,
+                        integral_constant<bool, HasDoubleTailKBlockLoop>) const
+    {
+        const auto a_e_k_global_desc       = *p_a_e_k_global_desc;
+        const auto b_e_n_ho_wo_global_desc = *p_b_e_n_ho_wo_global_desc;
+        const auto c_k_n_ho_wo_global_desc = *p_c_k_n_ho_wo_global_desc;
+
+        Run(a_e_k_global_desc,
+            p_a_global,
+            b_e_n_ho_wo_global_desc,
+            p_b_global,
+            c_k_n_ho_wo_global_desc,
+            p_c_global,
+            integral_constant<bool, HasMainKBlockLoop>{},
+            integral_constant<bool, HasDoubleTailKBlockLoop>{});
+    }
+
+    // pass tensor descriptors by void*
+    template <bool HasMainKBlockLoop, bool HasDoubleTailKBlockLoop>
+    __device__ void Run(const void* p_a_e_k_global_desc,
+                        const FloatAB* __restrict__ p_a_global,
+                        const void* p_b_e_n_ho_wo_global_desc,
+                        const FloatAB* __restrict__ p_b_global,
+                        const void* p_c_k_n_ho_wo_global_desc,
+                        FloatC* __restrict__ p_c_global,
+                        integral_constant<bool, HasMainKBlockLoop>,
+                        integral_constant<bool, HasDoubleTailKBlockLoop>) const
+    {
+        const auto a_e_k_global_desc = *reinterpret_cast<const AGlobalDesc*>(p_a_e_k_global_desc);
+        const auto b_e_n_ho_wo_global_desc =
+            *reinterpret_cast<const BGlobalDesc*>(p_b_e_n_ho_wo_global_desc);
+        const auto c_k_n_ho_wo_global_desc =
+            *reinterpret_cast<const CGlobalDesc*>(p_c_k_n_ho_wo_global_desc);
+
+        Run(a_e_k_global_desc,
+            p_a_global,
+            b_e_n_ho_wo_global_desc,
+            p_b_global,
+            c_k_n_ho_wo_global_desc,
+            p_c_global,
+            integral_constant<bool, HasMainKBlockLoop>{},
+            integral_constant<bool, HasDoubleTailKBlockLoop>{});
+    }
+};
+
+} // namespace ck
+#endif

composable_kernel/include/tensor_operation/gridwise_gemm.hpp

@@ -68,13 +68,13 @@ struct GridwiseGemmTransposedANormalBNormalC_v1
             Sequence<KPerBlock, NPerBlock>{}, Number<max_lds_align>{});
 
         // LDS allocation for A and B: be careful of alignment
-        constexpr index_t a_block_space =
+        constexpr index_t a_block_space_size =
             math::integer_least_multiple(a_k_m_block_desc.GetElementSpace(), max_lds_align);
 
-        constexpr index_t b_block_space =
+        constexpr index_t b_block_space_size =
             math::integer_least_multiple(b_k_n_block_desc.GetElementSpace(), max_lds_align);
 
-        return 2 * (a_block_space + b_block_space) * sizeof(Float);
+        return 2 * (a_block_space_size + b_block_space_size) * sizeof(Float);
     }
 
     __device__ void Run(const Float* __restrict__ p_a_global,
@@ -116,8 +116,8 @@ struct GridwiseGemmTransposedANormalBNormalC_v1
 
         const auto block_work_id = block_work_desc.CalculateClusterIndex(get_block_1d_id());
 
-        const index_t m_block_data_on_global = block_work_id[0] * MPerBlock;
-        const index_t n_block_data_on_global = block_work_id[1] * NPerBlock;
+        const index_t m_block_data_on_global = block_work_id[Number<0>{}] * MPerBlock;
+        const index_t n_block_data_on_global = block_work_id[Number<1>{}] * NPerBlock;
 
         // A matrix in LDS memory, dst of blockwise copy
         //   be careful of LDS alignment
@@ -143,7 +143,7 @@ struct GridwiseGemmTransposedANormalBNormalC_v1
                                                AddressSpace::Vgpr,
                                                AddressSpace::Lds,
                                                InMemoryDataOperation::Set>(
-                {0, m_block_data_on_global}, {0, 0});
+                make_multi_index(0, m_block_data_on_global), make_multi_index(0, 0));
 
         // B matrix in LDS memory, dst of blockwise copy
         //   be careful of LDS alignment
@@ -169,7 +169,7 @@ struct GridwiseGemmTransposedANormalBNormalC_v1
                                                AddressSpace::Vgpr,
                                                AddressSpace::Lds,
                                                InMemoryDataOperation::Set>(
-                {0, n_block_data_on_global}, {0, 0});
+                make_multi_index(0, n_block_data_on_global), make_multi_index(0, 0));
 
         // GEMM definition
         //   c_mtx += transpose(a_mtx) * b_mtx
@@ -209,14 +209,14 @@ struct GridwiseGemmTransposedANormalBNormalC_v1
             ThreadGemmBThreadCopySrcDataPerRead_N>{};
 
         // LDS allocation for A and B: be careful of alignment
-        constexpr index_t a_block_space =
+        constexpr index_t a_block_space_size =
             math::integer_least_multiple(a_k_m_block_desc.GetElementSpace(), max_lds_align);
 
-        constexpr index_t b_block_space =
+        constexpr index_t b_block_space_size =
             math::integer_least_multiple(b_k_n_block_desc.GetElementSpace(), max_lds_align);
 
         Float* p_a_block_double = p_shared_block;
-        Float* p_b_block_double = p_shared_block + 2 * a_block_space;
+        Float* p_b_block_double = p_shared_block + 2 * a_block_space_size;
 
         // register allocation for output
         AccFloat p_c_thread[c_m0m1_n0n1_thread_mtx_desc.GetElementSpace()];
@@ -230,47 +230,55 @@ struct GridwiseGemmTransposedANormalBNormalC_v1
             b_blockwise_copy.Run(p_b_global, p_b_block_double);
         }
 
-        constexpr auto a_block_slice_copy_steps = Sequence<KPerBlock, 0>{};
-        constexpr auto b_block_slice_copy_steps = Sequence<KPerBlock, 0>{};
+        constexpr auto a_block_slice_copy_step = Sequence<KPerBlock, 0>{};
+        constexpr auto b_block_slice_copy_step = Sequence<KPerBlock, 0>{};
+
+        Float* p_a_block_even = p_a_block_double;
+        Float* p_b_block_even = p_b_block_double;
+
+        Float* p_a_block_odd = p_a_block_double + a_block_space_size;
+        Float* p_b_block_odd = p_b_block_double + b_block_space_size;
 
         // LDS double buffer: main body
-        for(index_t k_block_data_begin = 0; k_block_data_begin + 2 * KPerBlock < K;
+        for(index_t k_block_data_begin = 0; k_block_data_begin < K - 2 * KPerBlock;
             k_block_data_begin += 2 * KPerBlock)
         {
-#pragma unroll
-            for(index_t iloop = 0; iloop < 2; ++iloop)
-            {
-                const bool even_loop = (iloop % 2 == 0);
+            Float p_a_thread_buffer[a_blockwise_copy.GetThreadBufferSize()];
+            Float p_b_thread_buffer[b_blockwise_copy.GetThreadBufferSize()];
 
-                Float* p_a_block_now =
-                    even_loop ? p_a_block_double : p_a_block_double + a_block_space;
-                Float* p_b_block_now =
-                    even_loop ? p_b_block_double : p_b_block_double + b_block_space;
+            // even iteration
+            a_blockwise_copy.MoveSrcSliceWindow(a_block_slice_copy_step, True);
+            b_blockwise_copy.MoveSrcSliceWindow(b_block_slice_copy_step, True);
 
-                Float* p_a_block_next =
-                    even_loop ? p_a_block_double + a_block_space : p_a_block_double;
-                Float* p_b_block_next =
-                    even_loop ? p_b_block_double + b_block_space : p_b_block_double;
+            __syncthreads();
 
-                Float p_a_thread_buffer[a_blockwise_copy.GetThreadBufferSize()];
-                Float p_b_thread_buffer[b_blockwise_copy.GetThreadBufferSize()];
+            // LDS doubel buffer: load next data from device mem
+            a_blockwise_copy.RunLoadThreadBuffer(p_a_global, p_a_thread_buffer);
+            b_blockwise_copy.RunLoadThreadBuffer(p_b_global, p_b_thread_buffer);
 
-                a_blockwise_copy.MoveSrcSliceWindow(a_block_slice_copy_steps, True);
-                b_blockwise_copy.MoveSrcSliceWindow(b_block_slice_copy_steps, True);
+            // LDS double buffer: GEMM on current data
+            blockwise_gemm.Run(p_a_block_even, p_b_block_even, p_c_thread);
 
-                __syncthreads();
+            // LDS double buffer: store next data to LDS
+            a_blockwise_copy.RunStoreThreadBuffer(p_a_thread_buffer, p_a_block_odd);
+            b_blockwise_copy.RunStoreThreadBuffer(p_b_thread_buffer, p_b_block_odd);
 
-                // LDS doubel buffer: load next data from device mem
-                a_blockwise_copy.RunLoadThreadBuffer(p_a_global, p_a_thread_buffer);
-                b_blockwise_copy.RunLoadThreadBuffer(p_b_global, p_b_thread_buffer);
+            // odd iteration
+            a_blockwise_copy.MoveSrcSliceWindow(a_block_slice_copy_step, True);
+            b_blockwise_copy.MoveSrcSliceWindow(b_block_slice_copy_step, True);
 
-                // LDS double buffer: GEMM on current data
-                blockwise_gemm.Run(p_a_block_now, p_b_block_now, p_c_thread);
+            __syncthreads();
 
-                // LDS double buffer: store next data to LDS
-                a_blockwise_copy.RunStoreThreadBuffer(p_a_thread_buffer, p_a_block_next);
-                b_blockwise_copy.RunStoreThreadBuffer(p_b_thread_buffer, p_b_block_next);
-            }
+            // LDS doubel buffer: load next data from device mem
+            a_blockwise_copy.RunLoadThreadBuffer(p_a_global, p_a_thread_buffer);
+            b_blockwise_copy.RunLoadThreadBuffer(p_b_global, p_b_thread_buffer);
+
+            // LDS double buffer: GEMM on current data
+            blockwise_gemm.Run(p_a_block_odd, p_b_block_odd, p_c_thread);
+
+            // LDS double buffer: store next data to LDS
+            a_blockwise_copy.RunStoreThreadBuffer(p_a_thread_buffer, p_a_block_even);
+            b_blockwise_copy.RunStoreThreadBuffer(p_b_thread_buffer, p_b_block_even);
         }
 
         // LDS double buffer: tail
@@ -282,8 +290,8 @@ struct GridwiseGemmTransposedANormalBNormalC_v1
                 Float p_a_thread_buffer[a_blockwise_copy.GetThreadBufferSize()];
                 Float p_b_thread_buffer[b_blockwise_copy.GetThreadBufferSize()];
 
-                a_blockwise_copy.MoveSrcSliceWindow(a_block_slice_copy_steps, True);
-                b_blockwise_copy.MoveSrcSliceWindow(b_block_slice_copy_steps, True);
+                a_blockwise_copy.MoveSrcSliceWindow(a_block_slice_copy_step, True);
+                b_blockwise_copy.MoveSrcSliceWindow(b_block_slice_copy_step, True);
 
                 __syncthreads();
 
@@ -296,15 +304,16 @@ struct GridwiseGemmTransposedANormalBNormalC_v1
 
                 // LDS double buffer: store last data to LDS
                 a_blockwise_copy.RunStoreThreadBuffer(p_a_thread_buffer,
-                                                      p_a_block_double + a_block_space);
+                                                      p_a_block_double + a_block_space_size);
                 b_blockwise_copy.RunStoreThreadBuffer(p_b_thread_buffer,
-                                                      p_b_block_double + b_block_space);
+                                                      p_b_block_double + b_block_space_size);
 
                 __syncthreads();
 
                 // LDS double buffer: GEMM on last data
-                blockwise_gemm.Run(
-                    p_a_block_double + a_block_space, p_b_block_double + b_block_space, p_c_thread);
+                blockwise_gemm.Run(p_a_block_double + a_block_space_size,
+                                   p_b_block_double + b_block_space_size,
+                                   p_c_thread);
             }
             else // if has 1 iteration left
             {
@@ -355,11 +364,11 @@ struct GridwiseGemmTransposedANormalBNormalC_v1
                                                   AddressSpace::Vgpr,
                                                   AddressSpace::Global,
                                                   CGlobalMemoryDataOperation>(
-                {0, 0, 0, 0},
-                {m_thread_data_on_global / M1,
-                 m_thread_data_on_global % M1,
-                 n_thread_data_on_global / N1,
-                 n_thread_data_on_global % N1})
+                make_multi_index(0, 0, 0, 0),
+                make_multi_index(m_thread_data_on_global / M1,
+                                 m_thread_data_on_global % M1,
+                                 n_thread_data_on_global / N1,
+                                 n_thread_data_on_global % N1))
                 .Run(p_c_thread, p_c_global);
         }
     }
@@ -433,13 +442,13 @@ struct GridwiseGemmTransposedANormalBNormalC_v2
             Sequence<KPerBlock, NPerBlock>{}, Number<max_lds_align>{});
 
         // LDS allocation for A and B: be careful of alignment
-        constexpr index_t a_block_space =
+        constexpr index_t a_block_space_size =
             math::integer_least_multiple(a_k_m_block_desc.GetElementSpace(), max_lds_align);
 
-        constexpr index_t b_block_space =
+        constexpr index_t b_block_space_size =
             math::integer_least_multiple(b_k_n_block_desc.GetElementSpace(), max_lds_align);
 
-        return 2 * (a_block_space + b_block_space) * sizeof(Float);
+        return 2 * (a_block_space_size + b_block_space_size) * sizeof(Float);
     }
 
     __device__ void Run(const Float* __restrict__ p_a_global,
@@ -447,21 +456,23 @@ struct GridwiseGemmTransposedANormalBNormalC_v2
                         Float* __restrict__ p_c_global,
                         Float* __restrict__ p_shared_block) const
     {
+        constexpr auto I0 = Number<0>{};
+        constexpr auto I1 = Number<1>{};
+        constexpr auto I2 = Number<2>{};
+        constexpr auto I3 = Number<3>{};
+
         constexpr auto True  = integral_constant<bool, true>{};
         constexpr auto False = integral_constant<bool, false>{};
 
-        constexpr auto I0 = Number<0>{};
-        constexpr auto I2 = Number<2>{};
-
         constexpr auto a_k0_k1_k2_m_global_desc = AGlobalDesc{};
         constexpr auto b_k0_k1_k2_n_global_desc = BGlobalDesc{};
         constexpr auto c_m_n_global_desc        = CGlobalDesc{};
 
-        constexpr auto K0 = a_k0_k1_k2_m_global_desc.GetLengths()[0];
-        constexpr auto K1 = a_k0_k1_k2_m_global_desc.GetLengths()[1];
-        constexpr auto K  = a_k0_k1_k2_m_global_desc.GetLengths()[2];
-        constexpr auto M  = c_m_n_global_desc.GetLengths()[0];
-        constexpr auto N  = c_m_n_global_desc.GetLengths()[1];
+        constexpr auto K0 = a_k0_k1_k2_m_global_desc.GetLengths()[I0];
+        constexpr auto K1 = a_k0_k1_k2_m_global_desc.GetLengths()[I1];
+        constexpr auto K  = a_k0_k1_k2_m_global_desc.GetLengths()[I2];
+        constexpr auto M  = c_m_n_global_desc.GetLengths()[I0];
+        constexpr auto N  = c_m_n_global_desc.GetLengths()[I1];
 
         // don't do anything if K == 0
         if(K == 0)
@@ -487,8 +498,8 @@ struct GridwiseGemmTransposedANormalBNormalC_v2
 
         const auto block_work_id = block_work_desc.CalculateClusterIndex(get_block_1d_id());
 
-        const index_t m_block_data_on_global = block_work_id[0] * MPerBlock;
-        const index_t n_block_data_on_global = block_work_id[1] * NPerBlock;
+        const index_t m_block_data_on_global = block_work_id[I0] * MPerBlock;
+        const index_t n_block_data_on_global = block_work_id[I1] * NPerBlock;
 
         // A matrix in LDS memory, dst of blockwise copy
         //   be careful of LDS alignment
@@ -514,7 +525,7 @@ struct GridwiseGemmTransposedANormalBNormalC_v2
                                                AddressSpace::Vgpr,
                                                AddressSpace::Lds,
                                                InMemoryDataOperation::Set>(
-                {0, 0, 0, m_block_data_on_global}, {0, 0, 0, 0});
+                make_multi_index(0, 0, 0, m_block_data_on_global), make_multi_index(0, 0, 0, 0));
 
         // B matrix in LDS memory, dst of blockwise copy
         //   be careful of LDS alignment
@@ -540,7 +551,7 @@ struct GridwiseGemmTransposedANormalBNormalC_v2
                                                AddressSpace::Vgpr,
                                                AddressSpace::Lds,
                                                InMemoryDataOperation::Set>(
-                {0, 0, 0, n_block_data_on_global}, {0, 0, 0, 0});
+                make_multi_index(0, 0, 0, n_block_data_on_global), make_multi_index(0, 0, 0, 0));
 
         // GEMM definition
         //   c_mtx += transpose(a_mtx) * b_mtx
@@ -582,14 +593,14 @@ struct GridwiseGemmTransposedANormalBNormalC_v2
             ThreadGemmBThreadCopySrcDataPerRead_N>{};
 
         // LDS allocation for A and B: be careful of alignment
-        constexpr index_t a_block_space =
+        constexpr index_t a_block_space_size =
             math::integer_least_multiple(a_k0_k1_k2_m_block_desc.GetElementSpace(), max_lds_align);
 
-        constexpr index_t b_block_space =
+        constexpr index_t b_block_space_size =
             math::integer_least_multiple(b_k0_k1_k2_n_block_desc.GetElementSpace(), max_lds_align);
 
         Float* p_a_block_double = p_shared_block;
-        Float* p_b_block_double = p_shared_block + 2 * a_block_space;
+        Float* p_b_block_double = p_shared_block + 2 * a_block_space_size;
 
         // register allocation for output
         AccFloat p_c_thread[c_m0m1_n0n1_thread_mtx_desc.GetElementSpace()];
@@ -601,15 +612,14 @@ struct GridwiseGemmTransposedANormalBNormalC_v2
         {
             for(index_t k1 = 0; k1 < K1; ++k1)
             {
-
                 // LDS double buffer: preload data into LDS
                 {
                     a_blockwise_copy.Run(p_a_global, p_a_block_double);
                     b_blockwise_copy.Run(p_b_global, p_b_block_double);
                 }
 
-                constexpr auto a_block_slice_copy_steps = Sequence<0, 0, KPerBlock, 0>{};
-                constexpr auto b_block_slice_copy_steps = Sequence<0, 0, KPerBlock, 0>{};
+                constexpr auto a_block_slice_copy_step = Sequence<0, 0, KPerBlock, 0>{};
+                constexpr auto b_block_slice_copy_step = Sequence<0, 0, KPerBlock, 0>{};
 
                 // LDS double buffer: main body
                 for(index_t k_block_data_begin = 0; k_block_data_begin + 2 * KPerBlock < K;
@@ -621,20 +631,20 @@ struct GridwiseGemmTransposedANormalBNormalC_v2
                         const bool even_loop = (iloop % 2 == 0);
 
                         Float* p_a_block_now =
-                            even_loop ? p_a_block_double : p_a_block_double + a_block_space;
+                            even_loop ? p_a_block_double : p_a_block_double + a_block_space_size;
                         Float* p_b_block_now =
-                            even_loop ? p_b_block_double : p_b_block_double + b_block_space;
+                            even_loop ? p_b_block_double : p_b_block_double + b_block_space_size;
 
                         Float* p_a_block_next =
-                            even_loop ? p_a_block_double + a_block_space : p_a_block_double;
+                            even_loop ? p_a_block_double + a_block_space_size : p_a_block_double;
                         Float* p_b_block_next =
-                            even_loop ? p_b_block_double + b_block_space : p_b_block_double;
+                            even_loop ? p_b_block_double + b_block_space_size : p_b_block_double;
 
                         Float p_a_thread_buffer[a_blockwise_copy.GetThreadBufferSize()];
                         Float p_b_thread_buffer[b_blockwise_copy.GetThreadBufferSize()];
 
-                        a_blockwise_copy.MoveSrcSliceWindow(a_block_slice_copy_steps, True);
-                        b_blockwise_copy.MoveSrcSliceWindow(b_block_slice_copy_steps, True);
+                        a_blockwise_copy.MoveSrcSliceWindow(a_block_slice_copy_step, True);
+                        b_blockwise_copy.MoveSrcSliceWindow(b_block_slice_copy_step, True);
 
                         __syncthreads();
 
@@ -660,8 +670,8 @@ struct GridwiseGemmTransposedANormalBNormalC_v2
                         Float p_a_thread_buffer[a_blockwise_copy.GetThreadBufferSize()];
                         Float p_b_thread_buffer[b_blockwise_copy.GetThreadBufferSize()];
 
-                        a_blockwise_copy.MoveSrcSliceWindow(a_block_slice_copy_steps, True);
-                        b_blockwise_copy.MoveSrcSliceWindow(b_block_slice_copy_steps, True);
+                        a_blockwise_copy.MoveSrcSliceWindow(a_block_slice_copy_step, True);
+                        b_blockwise_copy.MoveSrcSliceWindow(b_block_slice_copy_step, True);
 
                         __syncthreads();
 
@@ -673,16 +683,16 @@ struct GridwiseGemmTransposedANormalBNormalC_v2
                         blockwise_gemm.Run(p_a_block_double, p_b_block_double, p_c_thread);
 
                         // LDS double buffer: store last data to LDS
-                        a_blockwise_copy.RunStoreThreadBuffer(p_a_thread_buffer,
-                                                              p_a_block_double + a_block_space);
-                        b_blockwise_copy.RunStoreThreadBuffer(p_b_thread_buffer,
-                                                              p_b_block_double + b_block_space);
+                        a_blockwise_copy.RunStoreThreadBuffer(
+                            p_a_thread_buffer, p_a_block_double + a_block_space_size);
+                        b_blockwise_copy.RunStoreThreadBuffer(
+                            p_b_thread_buffer, p_b_block_double + b_block_space_size);
 
                         __syncthreads();
 
                         // LDS double buffer: GEMM on last data
-                        blockwise_gemm.Run(p_a_block_double + a_block_space,
-                                           p_b_block_double + b_block_space,
+                        blockwise_gemm.Run(p_a_block_double + a_block_space_size,
+                                           p_b_block_double + b_block_space_size,
                                            p_c_thread);
                     }
                     else // if has 1 iteration left
@@ -750,11 +760,11 @@ struct GridwiseGemmTransposedANormalBNormalC_v2
                                                   AddressSpace::Vgpr,
                                                   AddressSpace::Global,
                                                   CGlobalMemoryDataOperation>(
-                {0, 0, 0, 0},
-                {m_thread_data_on_global / M1,
-                 m_thread_data_on_global % M1,
-                 n_thread_data_on_global / N1,
-                 n_thread_data_on_global % N1})
+                make_multi_index(0, 0, 0, 0),
+                make_multi_index(m_thread_data_on_global / M1,
+                                 m_thread_data_on_global % M1,
+                                 n_thread_data_on_global / N1,
+                                 n_thread_data_on_global % N1))
                 .Run(p_c_thread, p_c_global);
         }
     }

composable_kernel/include/tensor_operation/gridwise_tensor_contraction.hpp

@@ -1,330 +0,0 @@
-#ifndef CK_GRIDWISE_TENSOR_CONTRACTION_HPP
-#define CK_GRIDWISE_TENSOR_CONTRACTION_HPP
-
-#include "common_header.hpp"
-#include "tensor_descriptor.hpp"
-#include "tensor_descriptor_helper.hpp"
-#include "ConstantMatrixDescriptor.hpp"
-#include "blockwise_generic_tensor_slice_copy.hpp"
-#include "threadwise_generic_tensor_slice_copy.hpp"
-#include "blockwise_gemm.hpp"
-
-namespace ck {
-
-template <index_t GridSize,
-          index_t BlockSize,
-          typename Float,
-          typename AccFloat,
-          typename AGlobalDesc,
-          typename BGlobalDesc,
-          typename CGlobalDesc,
-          typename CBlockLengths,
-          index_t KPerBlock,
-          InMemoryDataOperation CGlobalMemoryDataOperation>
-struct GridwiseTensorContraction_v1
-{
-    __host__ __device__ static constexpr index_t GetSharedMemoryNumberOfByte() {}
-
-    __device__ void Run(const Float* __restrict__ p_a_global,
-                        const Float* __restrict__ p_b_global,
-                        Float* __restrict__ p_c_global,
-                        Float* __restrict__ p_shared_block) const
-    {
-        /// \todo sanity-check on AGlobalDesc, BGlboalDesc, CGlobalDesc length consisitency
-        /// \todo santiy-check on CBlockLengtsh
-
-        constexpr auto True = integral_constant<bool, true>{};
-
-        constexpr auto a_global_desc = AGlobalDesc{};
-        constexpr auto b_global_desc = BGlobalDesc{};
-        constexpr auto c_global_desc = CGlobalDesc{};
-
-        constexpr auto K = a_global_desc.GetLengths()[0];
-
-        // don't do anything if K == 0
-        if(K == 0)
-        {
-            return;
-        }
-
-        // lds max alignment
-        constexpr index_t max_lds_align = math::lcm(ABlockCopyDstDataPerWrite_M,
-                                                    BBlockCopyDstDataPerWrite_N,
-                                                    ThreadGemmAThreadCopySrcDataPerRead_M,
-                                                    ThreadGemmBThreadCopySrcDataPerRead_N);
-
-        // divide block work by [M, N]
-        static_assert(M % MPerBlock == 0 && N % NPerBlock == 0 && K % KPerBlock == 0,
-                      "wrong! cannot divide work evenly among block");
-
-        constexpr index_t MBlockWork = M / MPerBlock;
-        constexpr index_t NBlockWork = N / NPerBlock;
-
-        constexpr auto block_work_desc =
-            make_cluster_descriptor(Sequence<MBlockWork, NBlockWork>{});
-
-        const auto block_work_id = block_work_desc.CalculateClusterIndex(get_block_1d_id());
-
-        const index_t m_block_data_on_global = block_work_id[0] * MPerBlock;
-        const index_t n_block_data_on_global = block_work_id[1] * NPerBlock;
-
-        // A matrix in LDS memory, dst of blockwise copy
-        //   be careful of LDS alignment
-        constexpr auto a_k_m_block_desc = make_native_tensor_descriptor_aligned(
-            Sequence<KPerBlock, MPerBlock>{}, Number<max_lds_align>{});
-
-        // A matrix blockwise copy
-        auto a_blockwise_copy =
-            BlockwiseGenericTensorSliceCopy_v4<BlockSize,
-                                               AGlobalDesc,
-                                               decltype(a_block_desc),
-                                               decltype(a_k_m_block_desc.GetLengths()),
-                                               ABlockCopyThreadSliceLengths_K_M,
-                                               ABlockCopyThreadClusterLengths_K_M,
-                                               ABlockCopyThreadClusterArrangeOrder,
-                                               ABlockCopySrcAccessOrder,
-                                               Sequence<0, 1>,
-                                               ABlockCopySrcVectorReadDim,
-                                               1,
-                                               ABlockCopySrcDataPerRead,
-                                               ABlockCopyDstDataPerWrite_M,
-                                               AddressSpace::Global,
-                                               AddressSpace::Vgpr,
-                                               AddressSpace::Lds,
-                                               InMemoryDataOperation::Set>(
-                {0, m_block_data_on_global}, {0, 0});
-
-        // B matrix in LDS memory, dst of blockwise copy
-        //   be careful of LDS alignment
-        constexpr auto b_k_n_block_desc = make_native_tensor_descriptor_aligned(
-            Sequence<KPerBlock, NPerBlock>{}, Number<max_lds_align>{});
-
-        // B matrix blockwise copy
-        auto b_blockwise_copy =
-            BlockwiseGenericTensorSliceCopy_v4<BlockSize,
-                                               decltype(b_k_n_global_desc),
-                                               decltype(b_k_n_block_desc),
-                                               decltype(b_k_n_block_desc.GetLengths()),
-                                               BBlockCopyThreadSliceLengths_K_N,
-                                               BBlockCopyThreadClusterLengths_K_N,
-                                               BBlockCopyThreadClusterArrangeOrder,
-                                               BBlockCopySrcAccessOrder,
-                                               Sequence<0, 1>,
-                                               BBlockCopySrcVectorReadDim,
-                                               1,
-                                               BBlockCopySrcDataPerRead,
-                                               BBlockCopyDstDataPerWrite_N,
-                                               AddressSpace::Global,
-                                               AddressSpace::Vgpr,
-                                               AddressSpace::Lds,
-                                               InMemoryDataOperation::Set>(
-                {0, n_block_data_on_global}, {0, 0});
-
-        // GEMM definition
-        //   c_mtx += transpose(a_mtx) * b_mtx
-        //     a_mtx[KPerBlock, MPerBlock] is in LDS
-        //     b_mtx[KPerBlocl, NPerBlock] is in LDS
-        //     c_mtx[MPerBlock, NPerBlock] is distributed among threads, and saved in
-        //       register
-        constexpr auto a_k_m_block_mtx_desc = make_ConstantMatrixDescriptor(a_k_m_block_desc);
-        constexpr auto b_k_n_block_mtx_desc = make_ConstantMatrixDescriptor(b_k_n_block_desc);
-
-        // sanity check
-        static_assert(MPerBlock % (MPerThread * MLevel0Cluster * MLevel1Cluster) == 0 &&
-                          NPerBlock % (NPerThread * NLevel0Cluster * NLevel1Cluster) == 0,
-                      "wrong!");
-
-        constexpr index_t GemmMRepeat = MPerBlock / (MPerThread * MLevel0Cluster * MLevel1Cluster);
-
-        constexpr index_t GemmNRepeat = NPerBlock / (NPerThread * NLevel0Cluster * NLevel1Cluster);
-
-        // c_thread_mtx definition: this is a mess
-        // TODO:: more elegent way of defining c_thread_mtx
-        constexpr auto c_m0m1_n0n1_thread_mtx_desc = make_ConstantMatrixDescriptor_packed(
-            Number<GemmMRepeat * MPerThread>{}, Number<GemmNRepeat * NPerThread>{});
-
-        const auto blockwise_gemm = BlockwiseGemmBlockABlockBThreadCTransANormalBNormalC_v2<
-            BlockSize,
-            decltype(a_k_m_block_mtx_desc),
-            decltype(b_k_n_block_mtx_desc),
-            decltype(c_m0m1_n0n1_thread_mtx_desc),
-            MPerThread,
-            NPerThread,
-            MLevel0Cluster,
-            NLevel0Cluster,
-            MLevel1Cluster,
-            NLevel1Cluster,
-            KPerThread,
-            ThreadGemmAThreadCopySrcDataPerRead_M,
-            ThreadGemmBThreadCopySrcDataPerRead_N>{};
-
-        // LDS allocation for A and B: be careful of alignment
-        constexpr index_t a_block_space =
-            math::integer_least_multiple(a_k_m_block_desc.GetElementSpace(), max_lds_align);
-
-        constexpr index_t b_block_space =
-            math::integer_least_multiple(b_k_n_block_desc.GetElementSpace(), max_lds_align);
-
-        Float* p_a_block_double = p_shared_block;
-        Float* p_b_block_double = p_shared_block + 2 * a_block_space;
-
-        // register allocation for output
-        AccFloat p_c_thread[c_m0m1_n0n1_thread_mtx_desc.GetElementSpace()];
-
-        // zero out threadwise output
-        threadwise_matrix_set_zero(c_m0m1_n0n1_thread_mtx_desc, p_c_thread);
-
-        // LDS double buffer: preload data into LDS
-        {
-            a_blockwise_copy.Run(p_a_global, p_a_block_double);
-            b_blockwise_copy.Run(p_b_global, p_b_block_double);
-        }
-
-        constexpr auto a_block_slice_copy_steps = Sequence<KPerBlock, 0>{};
-        constexpr auto b_block_slice_copy_steps = Sequence<KPerBlock, 0>{};
-
-        // LDS double buffer: main body
-        for(index_t k_block_data_begin = 0; k_block_data_begin + 2 * KPerBlock < K;
-            k_block_data_begin += 2 * KPerBlock)
-        {
-#pragma unroll
-            for(index_t iloop = 0; iloop < 2; ++iloop)
-            {
-                const bool even_loop = (iloop % 2 == 0);
-
-                Float* p_a_block_now =
-                    even_loop ? p_a_block_double : p_a_block_double + a_block_space;
-                Float* p_b_block_now =
-                    even_loop ? p_b_block_double : p_b_block_double + b_block_space;
-
-                Float* p_a_block_next =
-                    even_loop ? p_a_block_double + a_block_space : p_a_block_double;
-                Float* p_b_block_next =
-                    even_loop ? p_b_block_double + b_block_space : p_b_block_double;
-
-                Float p_a_thread_buffer[a_blockwise_copy.GetThreadBufferSize()];
-                Float p_b_thread_buffer[b_blockwise_copy.GetThreadBufferSize()];
-
-                a_blockwise_copy.MoveSrcSliceWindow(a_block_slice_copy_steps, True);
-                b_blockwise_copy.MoveSrcSliceWindow(b_block_slice_copy_steps, True);
-
-                __syncthreads();
-
-                // LDS doubel buffer: load next data from device mem
-                a_blockwise_copy.RunLoadThreadBuffer(p_a_global, p_a_thread_buffer);
-                b_blockwise_copy.RunLoadThreadBuffer(p_b_global, p_b_thread_buffer);
-
-                // LDS double buffer: GEMM on current data
-                blockwise_gemm.Run(p_a_block_now, p_b_block_now, p_c_thread);
-
-                // LDS double buffer: store next data to LDS
-                a_blockwise_copy.RunStoreThreadBuffer(p_a_thread_buffer, p_a_block_next);
-                b_blockwise_copy.RunStoreThreadBuffer(p_b_thread_buffer, p_b_block_next);
-            }
-        }
-
-        // LDS double buffer: tail
-        {
-            constexpr bool has_two_iteration_left = (K % (2 * KPerBlock) == 0);
-
-            if(has_two_iteration_left) // if has 2 iteration left
-            {
-                Float p_a_thread_buffer[a_blockwise_copy.GetThreadBufferSize()];
-                Float p_b_thread_buffer[b_blockwise_copy.GetThreadBufferSize()];
-
-                a_blockwise_copy.MoveSrcSliceWindow(a_block_slice_copy_steps, True);
-                b_blockwise_copy.MoveSrcSliceWindow(b_block_slice_copy_steps, True);
-
-                __syncthreads();
-
-                // LDS double buffer: load last data from device mem
-                a_blockwise_copy.RunLoadThreadBuffer(p_a_global, p_a_thread_buffer);
-                b_blockwise_copy.RunLoadThreadBuffer(p_b_global, p_b_thread_buffer);
-
-                // LDS double buffer: GEMM on 2nd-last data
-                blockwise_gemm.Run(p_a_block_double, p_b_block_double, p_c_thread);
-
-                // LDS double buffer: store last data to LDS
-                a_blockwise_copy.RunStoreThreadBuffer(p_a_thread_buffer,
-                                                      p_a_block_double + a_block_space);
-                b_blockwise_copy.RunStoreThreadBuffer(p_b_thread_buffer,
-                                                      p_b_block_double + b_block_space);
-
-                __syncthreads();
-
-                // LDS double buffer: GEMM on last data
-                blockwise_gemm.Run(
-                    p_a_block_double + a_block_space, p_b_block_double + b_block_space, p_c_thread);
-            }
-            else // if has 1 iteration left
-            {
-                __syncthreads();
-
-                // LDS double buffer: GEMM on last data
-                blockwise_gemm.Run(p_a_block_double, p_b_block_double, p_c_thread);
-            }
-        }
-
-        // input: register to global memory
-        {
-            constexpr index_t M1 = MPerThread * MLevel0Cluster * MLevel1Cluster;
-            constexpr index_t M0 = M / M1;
-
-            constexpr index_t N1 = NPerThread * NLevel0Cluster * NLevel1Cluster;
-            constexpr index_t N0 = N / N1;
-
-            // define input tensor descriptor for threadwise copy
-            //     thread input tensor, src of threadwise copy
-            constexpr auto c_m0_m1_n0_n1_thread_desc = make_native_tensor_descriptor_packed(
-                Sequence<GemmMRepeat, MPerThread, GemmNRepeat, NPerThread>{});
-
-            constexpr auto c_m0_m1_n0_n1_global_desc = transform_tensor_descriptor(
-                c_m_n_global_desc,
-                make_tuple(UnMerge<Sequence<M0, M1>>{}, UnMerge<Sequence<N0, N1>>{}),
-                make_tuple(Sequence<0>{}, Sequence<1>{}),
-                make_tuple(Sequence<0, 1>{}, Sequence<2, 3>{}));
-
-            // calculate origin of thread input tensor on global memory
-            //     blockwise GEMM c matrix starting index
-            const auto c_thread_mtx_on_block =
-                blockwise_gemm.GetBeginOfThreadMatrixC(get_thread_local_1d_id());
-
-            const index_t m_thread_data_on_global =
-                m_block_data_on_global + c_thread_mtx_on_block.row;
-
-            const index_t n_thread_data_on_global =
-                n_block_data_on_global + c_thread_mtx_on_block.col;
-
-            ThreadwiseGenericTensorSliceCopy_v4r2<decltype(c_m0_m1_n0_n1_thread_desc),
-                                                  decltype(c_m0_m1_n0_n1_global_desc),
-                                                  decltype(c_m0_m1_n0_n1_thread_desc.GetLengths()),
-                                                  CThreadCopySrcDstAccessOrder,
-                                                  CThreadCopySrcDstVectorReadWriteDim,
-                                                  1,
-                                                  CThreadCopyDstDataPerWrite,
-                                                  AddressSpace::Vgpr,
-                                                  AddressSpace::Global,
-                                                  CGlobalMemoryDataOperation>(
-                {0, 0, 0, 0},
-                {m_thread_data_on_global / M1,
-                 m_thread_data_on_global % M1,
-                 n_thread_data_on_global / N1,
-                 n_thread_data_on_global % N1})
-                .Run(p_c_thread, p_c_global);
-        }
-    }
-
-    __device__ void Run(const Float* __restrict__ p_a_global,
-                        const Float* __restrict__ p_b_global,
-                        Float* __restrict__ p_c_global) const
-    {
-        constexpr index_t shared_block_size = GetSharedMemoryNumberOfByte() / sizeof(Float);
-
-        __shared__ Float p_shared_block[shared_block_size];
-
-        Run(p_a_global, p_b_global, p_c_global, p_shared_block);
-    }
-};
-
-} // namespace ck
-#endif

composable_kernel/include/tensor_operation/threadwise_gemm.hpp

@@ -39,7 +39,7 @@ struct ThreadwiseMatrixSliceCopy
     template <typename Data>
     __device__ static void Run(const Data* p_src, Data* p_dst)
     {
-        using vector_t = typename vector_type<Data, DataPerAccess>::MemoryType;
+        using vector_t = typename vector_type<Data, DataPerAccess>::type;
 
         for(index_t i = 0; i < NSliceRow; ++i)
         {
@@ -153,9 +153,8 @@ struct ThreadwiseGemmTransANormalBNormalC
                                       (is_same<FloatA, half2_t>{} && is_same<FloatB, half2_t>{}) ||
                                       (is_same<FloatA, half4_t>{} && is_same<FloatB, half4_t>{}));
 
-        static_if<has_amd_asm>{}([&](auto fwd) {
-            Run_amd_asm(p_a, p_b, fwd(p_c));
-        }).Else([&](auto) { Run_source(p_a, p_b, p_c); });
+        static_if<has_amd_asm>{}([&](auto fwd) { Run_amd_asm(p_a, p_b, fwd(p_c)); })
+            .Else([&](auto) { Run_source(p_a, p_b, p_c); });
 #else
         Run_source(p_a, p_b, p_c);
 #endif

composable_kernel/include/tensor_operation/threadwise_gemm_v2.hpp

@@ -0,0 +1,172 @@
+#ifndef CK_THREADWISE_GEMM_V2_HPP
+#define CK_THREADWISE_GEMM_V2_HPP
+
+#include "common_header.hpp"
+#include "math.hpp"
+
+namespace ck {
+
+template <typename Float, typename Desc>
+__device__ void threadwise_matrix_set_zero_v2(Desc, Float* __restrict__ p_thread)
+{
+    static_assert(Desc::IsKnownAtCompileTime(), "wrong! Desc should be known at compile-time");
+
+    constexpr auto I0 = Number<0>{};
+    constexpr auto I1 = Number<1>{};
+
+    constexpr auto desc = Desc{};
+
+    constexpr auto M = desc.GetLength(I0);
+    constexpr auto N = desc.GetLength(I1);
+
+    static_for<0, M, 1>{}([&](auto i) {
+        static_for<0, N, 1>{}([&](auto j) {
+            constexpr auto offset = desc.CalculateOffset(make_tuple(i, j));
+
+            p_thread[offset] = Float(0);
+        });
+    });
+}
+
+template <typename SrcDesc,
+          typename DstDesc,
+          index_t NSliceRow,
+          index_t NSliceCol,
+          index_t DataPerAccess>
+struct ThreadwiseMatrixSliceCopy_v2
+{
+    template <typename Data>
+    __device__ static void Run(const Data* p_src, Data* p_dst)
+    {
+        static_assert(SrcDesc::IsKnownAtCompileTime() && DstDesc::IsKnownAtCompileTime(),
+                      "wrong! Desc should be known at compile-time");
+
+        using vector_t = typename vector_type<Data, DataPerAccess>::type;
+
+        static_for<0, NSliceRow, 1>{}([&](auto i) {
+            static_for<0, NSliceCol, DataPerAccess>{}([&](auto j) {
+                constexpr auto src_offset = SrcDesc{}.CalculateOffset(make_tuple(i, j));
+                constexpr auto dst_offset = DstDesc{}.CalculateOffset(make_tuple(i, j));
+
+                *reinterpret_cast<vector_t*>(&p_dst[dst_offset]) =
+                    *reinterpret_cast<const vector_t*>(&p_src[src_offset]);
+            });
+        });
+    }
+};
+
+// C[M, N] += transpose(A[K, M]) * B[K, N]
+//   Element of matrix can be vectorized data
+template <typename ADesc,
+          typename BDesc,
+          typename CDesc,
+          typename std::enable_if<ADesc::IsKnownAtCompileTime() && BDesc::IsKnownAtCompileTime() &&
+                                      CDesc::IsKnownAtCompileTime(),
+                                  bool>::type = false>
+struct ThreadwiseGemm_km_kn_mn_v1
+{
+    template <typename FloatA, typename FloatB, typename FloatC>
+    __device__ static void Run_source(const FloatA* p_a, const FloatB* p_b, FloatC* p_c)
+    {
+        static_assert(ADesc::IsKnownAtCompileTime() && BDesc::IsKnownAtCompileTime() &&
+                          CDesc::IsKnownAtCompileTime(),
+                      "wrong! Desc should be known at compile-time");
+
+        constexpr auto I0 = Number<0>{};
+        constexpr auto I1 = Number<1>{};
+
+        constexpr auto M = CDesc{}.GetLength(I0);
+        constexpr auto N = CDesc{}.GetLength(I1);
+        constexpr auto K = ADesc{}.GetLength(I0);
+
+        static_for<0, K, 1>{}([&](auto k) {
+            static_for<0, M, 1>{}([&](auto m) {
+                static_for<0, N, 1>{}([&](auto n) {
+                    constexpr auto a_offset = ADesc{}.CalculateOffset(make_tuple(k, m));
+                    constexpr auto b_offset = BDesc{}.CalculateOffset(make_tuple(k, n));
+                    constexpr auto c_offset = CDesc{}.CalculateOffset(make_tuple(m, n));
+
+                    p_c[c_offset] +=
+                        inner_product_with_conversion<FloatC>{}(p_a[a_offset], p_b[b_offset]);
+                });
+            });
+        });
+    }
+
+#if CK_THREADWISE_GEMM_USE_AMD_INLINE_ASM
+    template <typename FloatA, typename FloatB, typename FloatC>
+    __device__ static void Run_amd_asm(const FloatA* p_a, const FloatB* p_b, FloatC* p_c)
+    {
+        static_assert(ADesc::IsKnownAtCompileTime() && BDesc::IsKnownAtCompileTime() &&
+                          CDesc::IsKnownAtCompileTime(),
+                      "wrong! Desc should be known at compile-time");
+
+        constexpr auto I0 = Number<0>{};
+        constexpr auto I1 = Number<1>{};
+        constexpr auto I2 = Number<2>{};
+        constexpr auto I3 = Number<3>{};
+
+        constexpr auto M = CDesc{}.GetLength(I0);
+        constexpr auto N = CDesc{}.GetLength(I1);
+        constexpr auto K = ADesc{}.GetLength(I0);
+
+        static_assert(N == 4 || N == 2, "wrong! this config not supported by asm yet");
+
+        static_for<0, K, 1>{}([&](auto k) {
+            static_for<0, M, 1>{}([&](auto m) {
+                constexpr auto a_offset = ADesc{}.CalculateOffset(make_tuple(k, m));
+
+                if constexpr(N == 2)
+                {
+                    constexpr auto b_offset_0 = BDesc{}.CalculateOffset(make_tuple(k, I0));
+                    constexpr auto b_offset_1 = BDesc{}.CalculateOffset(make_tuple(k, I1));
+
+                    constexpr auto c_offset_0 = CDesc{}.CalculateOffset(make_tuple(m, I0));
+                    constexpr auto c_offset_1 = CDesc{}.CalculateOffset(make_tuple(m, I1));
+
+                    amd_assembly_outer_product_1x2(p_a[a_offset],
+                                                   p_b[b_offset_0],
+                                                   p_b[b_offset_1],
+                                                   p_c[c_offset_0],
+                                                   p_c[c_offset_1]);
+                }
+                else if constexpr(N == 4)
+                {
+                    constexpr auto b_offset_0 = BDesc{}.CalculateOffset(make_tuple(k, I0));
+                    constexpr auto b_offset_1 = BDesc{}.CalculateOffset(make_tuple(k, I1));
+                    constexpr auto b_offset_2 = BDesc{}.CalculateOffset(make_tuple(k, I2));
+                    constexpr auto b_offset_3 = BDesc{}.CalculateOffset(make_tuple(k, I3));
+
+                    constexpr auto c_offset_0 = CDesc{}.CalculateOffset(make_tuple(m, I0));
+                    constexpr auto c_offset_1 = CDesc{}.CalculateOffset(make_tuple(m, I1));
+                    constexpr auto c_offset_2 = CDesc{}.CalculateOffset(make_tuple(m, I2));
+                    constexpr auto c_offset_3 = CDesc{}.CalculateOffset(make_tuple(m, I3));
+
+                    amd_assembly_outer_product_1x4(p_a[a_offset],
+                                                   p_b[b_offset_0],
+                                                   p_b[b_offset_1],
+                                                   p_b[b_offset_2],
+                                                   p_b[b_offset_3],
+                                                   p_c[c_offset_0],
+                                                   p_c[c_offset_1],
+                                                   p_c[c_offset_2],
+                                                   p_c[c_offset_3]);
+                }
+            });
+        });
+    }
+#endif
+
+    template <typename FloatA, typename FloatB, typename FloatC>
+    __device__ static void Run(const FloatA* p_a, const FloatB* p_b, FloatC* p_c)
+    {
+#if CK_THREADWISE_GEMM_USE_AMD_INLINE_ASM
+        Run_amd_asm(p_a, p_b, p_c);
+#else
+        Run_source(p_a, p_b, p_c);
+#endif
+    }
+};
+
+} // namespace ck
+#endif

composable_kernel/include/tensor_operation/threadwise_gemm_v3.hpp

@@ -0,0 +1,141 @@
+#ifndef CK_THREADWISE_GEMM_V3_HPP
+#define CK_THREADWISE_GEMM_V3_HPP
+
+#include "common_header.hpp"
+#include "math.hpp"
+
+namespace ck {
+
+template <typename Float, typename Desc>
+__device__ void threadwise_matrix_set_zero_v3(Desc, Float* __restrict__ p_thread)
+{
+    static_assert(Desc::IsKnownAtCompileTime(), "wrong! Desc should be known at compile-time");
+
+    constexpr auto I0 = Number<0>{};
+    constexpr auto I2 = Number<2>{};
+    constexpr auto I3 = Number<3>{};
+
+    constexpr auto desc = Desc{};
+
+    constexpr auto K = desc.GetLength(I0);
+    constexpr auto H = desc.GetLength(I2);
+    constexpr auto W = desc.GetLength(I3);
+
+    static_for<0, K, 1>{}([&](auto i) {
+        static_for<0, H, 1>{}([&](auto j) {
+            static_for<0, W, 1>{}([&](auto k) {
+                constexpr auto offset = desc.CalculateOffset(make_tuple(i, 0, j, k));
+
+                p_thread[offset] = Float(0);
+            });
+        });
+    });
+}
+
+// C[M, N] += transpose(A[K, M]) * B[K, N]
+//   Element of matrix can be vectorized data
+template <typename ADesc,
+          typename BDesc,
+          typename CDesc,
+          typename std::enable_if<ADesc::IsKnownAtCompileTime() && BDesc::IsKnownAtCompileTime() &&
+                                      CDesc::IsKnownAtCompileTime(),
+                                  bool>::type = false>
+struct ThreadwiseGemm_km_kn_mn_v3
+{
+    template <typename FloatA, typename FloatB, typename FloatC>
+    __device__ static void Run_source(const FloatA* p_a, const FloatB* p_b, FloatC* p_c)
+    {
+        static_assert(ADesc::IsKnownAtCompileTime() && BDesc::IsKnownAtCompileTime() &&
+                          CDesc::IsKnownAtCompileTime(),
+                      "wrong! Desc should be known at compile-time");
+
+        constexpr auto I0 = Number<0>{};
+        constexpr auto I1 = Number<1>{};
+        constexpr auto I2 = Number<2>{};
+        constexpr auto I3 = Number<3>{};
+
+        // constexpr auto H = BDesc{}.GetLength(I2);
+        // constexpr auto W = BDesc{}.GetLength(I3);
+        constexpr auto H = 2;
+        constexpr auto W = 2;
+
+        constexpr auto E = ADesc{}.GetLength(I0);
+        constexpr auto K = ADesc{}.GetLength(I1);
+
+        static_for<0, E, 1>{}([&](auto e) {
+            static_for<0, K, 1>{}([&](auto k) {
+                constexpr auto a_offset = ADesc{}.CalculateOffset(make_tuple(e, k));
+
+                if constexpr(H == 2 && W == 2)
+                {
+
+                    constexpr auto b_offset_0 = BDesc{}.CalculateOffset(make_tuple(e, 0, 0, 0));
+                    constexpr auto b_offset_1 = BDesc{}.CalculateOffset(make_tuple(e, 0, 0, 1));
+                    constexpr auto b_offset_2 = BDesc{}.CalculateOffset(make_tuple(e, 0, 1, 0));
+                    constexpr auto b_offset_3 = BDesc{}.CalculateOffset(make_tuple(e, 0, 1, 1));
+
+                    constexpr auto c_offset_0 = CDesc{}.CalculateOffset(make_tuple(k, 0, 0, 0));
+                    constexpr auto c_offset_1 = CDesc{}.CalculateOffset(make_tuple(k, 0, 0, 1));
+                    constexpr auto c_offset_2 = CDesc{}.CalculateOffset(make_tuple(k, 0, 1, 0));
+                    constexpr auto c_offset_3 = CDesc{}.CalculateOffset(make_tuple(k, 0, 1, 1));
+
+                    amd_assembly_outer_product_1x4(p_a[a_offset],
+                                                   p_b[b_offset_0],
+                                                   p_b[b_offset_1],
+                                                   p_b[b_offset_2],
+                                                   p_b[b_offset_3],
+                                                   p_c[c_offset_0],
+                                                   p_c[c_offset_1],
+                                                   p_c[c_offset_2],
+                                                   p_c[c_offset_3]);
+                }
+                else if constexpr(H == 4 && W == 1)
+                {
+
+                    constexpr auto b_offset_0 = BDesc{}.CalculateOffset(make_tuple(e, 0, 0, 0));
+                    constexpr auto b_offset_1 = BDesc{}.CalculateOffset(make_tuple(e, 0, 1, 0));
+                    constexpr auto b_offset_2 = BDesc{}.CalculateOffset(make_tuple(e, 0, 2, 0));
+                    constexpr auto b_offset_3 = BDesc{}.CalculateOffset(make_tuple(e, 0, 3, 0));
+
+                    constexpr auto c_offset_0 = CDesc{}.CalculateOffset(make_tuple(k, 0, 0, 0));
+                    constexpr auto c_offset_1 = CDesc{}.CalculateOffset(make_tuple(k, 0, 1, 0));
+                    constexpr auto c_offset_2 = CDesc{}.CalculateOffset(make_tuple(k, 0, 2, 0));
+                    constexpr auto c_offset_3 = CDesc{}.CalculateOffset(make_tuple(k, 0, 3, 0));
+
+                    amd_assembly_outer_product_1x4(p_a[a_offset],
+                                                   p_b[b_offset_0],
+                                                   p_b[b_offset_1],
+                                                   p_b[b_offset_2],
+                                                   p_b[b_offset_3],
+                                                   p_c[c_offset_0],
+                                                   p_c[c_offset_1],
+                                                   p_c[c_offset_2],
+                                                   p_c[c_offset_3]);
+                }
+                else
+                {
+                    static_for<0, H, 1>{}([&](auto h) {
+                        static_for<0, W, 1>{}([&](auto w) {
+                            constexpr auto b_offset =
+                                BDesc{}.CalculateOffset(make_tuple(e, 0, h, w));
+                            constexpr auto c_offset =
+                                CDesc{}.CalculateOffset(make_tuple(k, 0, h, w));
+
+                            p_c[c_offset] += inner_product_with_conversion<FloatC>{}(p_a[a_offset],
+                                                                                     p_b[b_offset]);
+                        });
+                    });
+                }
+            });
+        });
+    }
+
+    template <typename FloatA, typename FloatB, typename FloatC>
+    __device__ static void Run(const FloatA* p_a, const FloatB* p_b, FloatC* p_c)
+    {
+        Run_source(p_a, p_b, p_c);
+    }
+};
+
+} // namespace ck
+#endif

composable_kernel/include/tensor_operation/threadwise_generic_tensor_slice_copy.hpp

@@ -54,8 +54,8 @@ struct ThreadwiseGenericTensorSliceCopy_v4r2
     }
 
     __device__ constexpr ThreadwiseGenericTensorSliceCopy_v4r2()
-        : ThreadwiseGenericTensorSliceCopy_v4r2(make_zero_array<index_t, nDim>(),
-                                                make_zero_array<index_t, nDim>())
+        : ThreadwiseGenericTensorSliceCopy_v4r2(make_zero_multi_index<nDim>(),
+                                                make_zero_multi_index<nDim>())
     {
     }
 
@@ -82,113 +82,104 @@ struct ThreadwiseGenericTensorSliceCopy_v4r2
         constexpr auto long_vector_access_lengths = SliceLengths::Modify(
             vector_access_dim, SliceLengths::Get(vector_access_dim) / long_vector_size);
 
-        ford<decltype(long_vector_access_lengths), SrcDstDimAccessOrder>{}([&](
-            auto long_vector_access_id) {
+        ford<decltype(long_vector_access_lengths), SrcDstDimAccessOrder>{}(
+            [&](auto long_vector_access_id) {
 
-            // data id w.r.t slicing-window
-            auto long_vector_data_begin_id = long_vector_access_id;
-            long_vector_data_begin_id(vector_access_dim) =
-                long_vector_size * long_vector_access_id[vector_access_dim];
+                // data id w.r.t slicing-window
+                auto long_vector_data_begin_id = long_vector_access_id;
+                long_vector_data_begin_id(vector_access_dim) =
+                    long_vector_size * long_vector_access_id[vector_access_dim];
 
-            // buffer to hold a src long-vector
-            SrcData p_src_long_vector[long_vector_size];
+                // buffer to hold a src long-vector
+                SrcData p_src_long_vector[long_vector_size];
 
-#if 1
-            // zero out buffer
-            for(index_t i = 0; i < long_vector_size; ++i)
-            {
-                p_src_long_vector[i] = 0;
-            }
-#endif
+                // load data from src to the long-vector buffer
+                static_for<0, long_vector_size / src_data_per_access, 1>{}([&](auto i) {
+                    auto scalar_id               = make_zero_multi_index<nDim>();
+                    scalar_id(vector_access_dim) = i * src_data_per_access;
 
-            // load data from src to the long-vector buffer
-            for(index_t i = 0; i < long_vector_size / src_data_per_access; ++i)
-            {
-                auto scalar_id               = make_zero_array<index_t, nDim>();
-                scalar_id(vector_access_dim) = i * src_data_per_access;
+                    const index_t buffer_offset = i * src_data_per_access;
 
-                const index_t buffer_offset = i * src_data_per_access;
+                    const auto src_coord =
+                        mSrcSliceOrigin + (long_vector_data_begin_id + scalar_id);
 
-                const auto src_coord = mSrcSliceOrigin + (long_vector_data_begin_id + scalar_id);
+                    // Check src data's valid mapping situation, only check the first data in this
+                    // src
+                    //   vector. It's user's responsiblity to make sure all data in the src vector
+                    //   has the valid/invalid mapping situation
+                    transfer_data<SrcData,
+                                  SrcDataPerRead,
+                                  SrcAddressSpace,
+                                  AddressSpace::Vgpr,
+                                  InMemoryDataOperation::Set,
+                                  SrcDataStride,
+                                  1>(p_src,
+                                     src_coord.GetOffset(),
+                                     src_coord.IsOffsetValidAssumingUpperIndexIsValid(),
+                                     SrcDesc::GetElementSpace(),
+                                     p_src_long_vector,
+                                     buffer_offset,
+                                     true,
+                                     long_vector_size);
+                });
 
-                // Check src data's valid mapping situation, only check the first data in this src
-                //   vector. It's user's responsiblity to make sure all data in the src vector
-                //   has the valid/invalid mapping situation
-                transfer_data<SrcData,
-                              SrcDataPerRead,
-                              SrcAddressSpace,
-                              AddressSpace::Vgpr,
-                              InMemoryDataOperation::Set,
-                              SrcDataStride,
-                              1>(p_src,
-                                 src_coord.GetOffset(),
-                                 src_coord.IsOffsetValidAssumingUpperIndexIsValid(),
-                                 SrcDesc::GetElementSpace(),
-                                 p_src_long_vector,
-                                 buffer_offset,
-                                 true,
-                                 long_vector_size);
-            }
+                // SrcData to DstData conversion
+                DstData p_dst_long_vector[long_vector_size];
 
-            // SrcData to DstData conversion
-            DstData p_dst_long_vector[long_vector_size];
+                static_for<0, long_vector_size, 1>{}([&](auto i) {
+                    p_dst_long_vector[i] = type_convert<DstData>{}(p_src_long_vector[i]);
+                });
 
-            for(index_t i = 0; i < long_vector_size; ++i)
-            {
-                p_dst_long_vector[i] = type_convert<DstData>{}(p_src_long_vector[i]);
-            }
+                // store data from the long-vector buffer to dst
+                static_for<0, long_vector_size / dst_data_per_access, 1>{}([&](auto i) {
+                    auto scalar_id               = make_zero_multi_index<nDim>();
+                    scalar_id(vector_access_dim) = i * dst_data_per_access;
 
-            // store data from the long-vector buffer to dst
-            for(index_t i = 0; i < long_vector_size / dst_data_per_access; ++i)
-            {
-                auto scalar_id               = make_zero_array<index_t, nDim>();
-                scalar_id(vector_access_dim) = i * dst_data_per_access;
+                    const index_t buffer_offset = i * dst_data_per_access;
 
-                const index_t buffer_offset = i * dst_data_per_access;
+                    const auto dst_coord =
+                        mDstSliceOrigin + (long_vector_data_begin_id + scalar_id);
 
-                const auto dst_coord = mDstSliceOrigin + (long_vector_data_begin_id + scalar_id);
-
-                // Check dst data's valid mapping situation, only check the first data in this dst
-                //   vector. It's user's responsiblity to make sure all data in the dst vector
-                //   has the valid/invalid mapping situation
-                transfer_data<DstData,
-                              DstDataPerWrite,
-                              AddressSpace::Vgpr,
-                              DstAddressSpace,
-                              DstInMemOp,
-                              1,
-                              DstDataStride>(p_dst_long_vector,
-                                             buffer_offset,
-                                             true,
-                                             long_vector_size,
-                                             p_dst,
-                                             dst_coord.GetOffset(),
-                                             dst_coord.IsOffsetValidAssumingUpperIndexIsValid(),
-                                             DstDesc::GetElementSpace());
-            }
-        });
+                    // Check dst data's valid mapping situation, only check the first data in this
+                    // dst
+                    //   vector. It's user's responsiblity to make sure all data in the dst vector
+                    //   has the valid/invalid mapping situation
+                    transfer_data<DstData,
+                                  DstDataPerWrite,
+                                  AddressSpace::Vgpr,
+                                  DstAddressSpace,
+                                  DstInMemOp,
+                                  1,
+                                  DstDataStride>(p_dst_long_vector,
+                                                 buffer_offset,
+                                                 true,
+                                                 long_vector_size,
+                                                 p_dst,
+                                                 dst_coord.GetOffset(),
+                                                 dst_coord.IsOffsetValidAssumingUpperIndexIsValid(),
+                                                 DstDesc::GetElementSpace());
+                });
+            });
     }
 
     template <typename T, bool PositiveDirection>
     __device__ void MoveSrcSliceWindow(const T& step_sizes_,
                                        integral_constant<bool, PositiveDirection>)
     {
-        const auto step_sizes = to_array(step_sizes_);
+        const auto step_sizes = to_multi_index(step_sizes_);
 
-        static_if<PositiveDirection>{}([&](auto) {
-            mSrcSliceOrigin += to_array(step_sizes);
-        }).Else([&](auto) { mSrcSliceOrigin -= step_sizes; });
+        static_if<PositiveDirection>{}([&](auto) { mSrcSliceOrigin += to_multi_index(step_sizes); })
+            .Else([&](auto) { mSrcSliceOrigin -= step_sizes; });
     }
 
     template <typename T, bool PositiveDirection>
     __device__ void MoveDstSliceWindow(const T& step_sizes_,
                                        integral_constant<bool, PositiveDirection>)
     {
-        const auto step_sizes = to_array(step_sizes_);
+        const auto step_sizes = to_multi_index(step_sizes_);
 
-        static_if<PositiveDirection>{}([&](auto) {
-            mDstSliceOrigin += step_sizes;
-        }).Else([&](auto) { mDstSliceOrigin -= step_sizes; });
+        static_if<PositiveDirection>{}([&](auto) { mDstSliceOrigin += step_sizes; })
+            .Else([&](auto) { mDstSliceOrigin -= step_sizes; });
     }
 
     private:

composable_kernel/include/utility/amd_buffer_addressing_v2.hpp

@@ -0,0 +1,365 @@
+#ifndef CK_AMD_BUFFER_ADDRESSING_V2_HPP
+#define CK_AMD_BUFFER_ADDRESSING_V2_HPP
+
+#include "float_type.hpp"
+
+namespace ck {
+
+template <typename T>
+union BufferResource
+{
+    // 128 bit SGPRs to supply buffer resource in buffer instructions
+    // https://rocm-documentation.readthedocs.io/en/latest/GCN_ISA_Manuals/testdocbook.html#vector-memory-buffer-instructions
+    int32x4_t data;
+    T* address[2];
+    int32_t range[4];
+    int32_t config[4];
+};
+
+template <typename T>
+__device__ int32x4_t make_wave_buffer_resource(T* p_wave, index_t data_space_size)
+{
+    BufferResource<T> wave_buffer_resource;
+
+    // wavewise base address (64 bit)
+    wave_buffer_resource.address[0] = const_cast<remove_cv_t<T>*>(p_wave);
+    // wavewise range (32 bit)
+    wave_buffer_resource.range[2] = data_space_size * sizeof(T);
+    // wavewise setting (32 bit)
+    wave_buffer_resource.config[3] = CK_BUFFER_RESOURCE_3RD_DWORD;
+
+    return wave_buffer_resource.data;
+}
+
+// load
+__device__ int8_t
+__llvm_amdgcn_raw_buffer_load_i8(int32x4_t srsrc,
+                                 index_t voffset,
+                                 index_t soffset,
+                                 index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.i8");
+__device__ int16_t
+__llvm_amdgcn_raw_buffer_load_i16(int32x4_t srsrc,
+                                  index_t voffset,
+                                  index_t soffset,
+                                  index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.i32");
+__device__ int32_t
+__llvm_amdgcn_raw_buffer_load_i32(int32x4_t srsrc,
+                                  index_t voffset,
+                                  index_t soffset,
+                                  index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.i32");
+
+__device__ int32x2_t
+__llvm_amdgcn_raw_buffer_load_i32x2(int32x4_t srsrc,
+                                    index_t voffset,
+                                    index_t soffset,
+                                    index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.v2i32");
+
+__device__ int32x4_t
+__llvm_amdgcn_raw_buffer_load_i32x4(int32x4_t srsrc,
+                                    index_t voffset,
+                                    index_t soffset,
+                                    index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.v4i32");
+
+__device__ float
+__llvm_amdgcn_raw_buffer_load_fp32(int32x4_t srsrc,
+                                   index_t voffset,
+                                   index_t soffset,
+                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.f32");
+
+__device__ float2_t
+__llvm_amdgcn_raw_buffer_load_fp32x2(int32x4_t srsrc,
+                                     index_t voffset,
+                                     index_t soffset,
+                                     index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.v2f32");
+
+__device__ float4_t
+__llvm_amdgcn_raw_buffer_load_fp32x4(int32x4_t srsrc,
+                                     index_t voffset,
+                                     index_t soffset,
+                                     index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.load.v4f32");
+
+// store
+__device__ void
+__llvm_amdgcn_raw_buffer_store_i8(int8_t vdata,
+                                  int32x4_t rsrc,
+                                  index_t voffset,
+                                  index_t soffset,
+                                  index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.i8");
+
+__device__ void
+__llvm_amdgcn_raw_buffer_store_i16(int16_t vdata,
+                                   int32x4_t rsrc,
+                                   index_t voffset,
+                                   index_t soffset,
+                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.i16");
+
+__device__ void
+__llvm_amdgcn_raw_buffer_store_i32(int32_t vdata,
+                                   int32x4_t rsrc,
+                                   index_t voffset,
+                                   index_t soffset,
+                                   index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.i32");
+
+__device__ void
+__llvm_amdgcn_raw_buffer_store_i32x2(int32x2_t vdata,
+                                     int32x4_t rsrc,
+                                     index_t voffset,
+                                     index_t soffset,
+                                     index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v2i32");
+
+__device__ void
+__llvm_amdgcn_raw_buffer_store_i32x4(int32x4_t vdata,
+                                     int32x4_t rsrc,
+                                     index_t voffset,
+                                     index_t soffset,
+                                     index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v4i32");
+
+__device__ void
+__llvm_amdgcn_raw_buffer_store_fp32(float vdata,
+                                    int32x4_t rsrc,
+                                    index_t voffset,
+                                    index_t soffset,
+                                    index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.f32");
+
+__device__ void
+__llvm_amdgcn_raw_buffer_store_fp32x2(float2_t vdata,
+                                      int32x4_t rsrc,
+                                      index_t voffset,
+                                      index_t soffset,
+                                      index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v2f32");
+
+__device__ void
+__llvm_amdgcn_raw_buffer_store_fp32x4(float4_t vdata,
+                                      int32x4_t rsrc,
+                                      index_t voffset,
+                                      index_t soffset,
+                                      index_t glc_slc) __asm("llvm.amdgcn.raw.buffer.store.v4f32");
+
+template <typename T, index_t N>
+__device__ typename vector_type<T, N>::type
+amd_buffer_load_impl_v2(int32x4_t src_wave_buffer_resource,
+                        index_t src_thread_addr_offset,
+                        index_t src_wave_addr_offset)
+{
+    static_assert((is_same<T, float>::value && (N == 1 || N == 2 || N == 4 || N == 8)) ||
+                      (is_same<T, int32_t>::value && (N == 1 || N == 2 || N == 4 || N == 8)),
+                  "wrong! not implemented");
+
+    if constexpr(is_same<T, float>::value)
+    {
+        if constexpr(N == 1)
+        {
+            return __llvm_amdgcn_raw_buffer_load_fp32(
+                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
+        }
+        else if constexpr(N == 2)
+        {
+            return __llvm_amdgcn_raw_buffer_load_fp32x2(
+                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
+        }
+        else if constexpr(N == 4)
+        {
+            return __llvm_amdgcn_raw_buffer_load_fp32x4(
+                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
+        }
+        else if constexpr(N == 8)
+        {
+            vector_type<float, 8> tmp;
+
+            tmp.Vectors(Number<4>{})(Number<0>{}) = __llvm_amdgcn_raw_buffer_load_fp32x4(
+                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
+
+            tmp.Vectors(Number<4>{})(Number<1>{}) = __llvm_amdgcn_raw_buffer_load_fp32x4(
+                src_wave_buffer_resource, src_thread_addr_offset, 4 * sizeof(float), 0);
+
+            return tmp.Vector();
+        }
+    }
+    else if constexpr(is_same<T, int32_t>::value)
+    {
+        if constexpr(N == 1)
+        {
+            return __llvm_amdgcn_raw_buffer_load_i32(
+                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
+        }
+        else if constexpr(N == 2)
+        {
+            return __llvm_amdgcn_raw_buffer_load_i32x2(
+                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
+        }
+        else if constexpr(N == 4)
+        {
+            return __llvm_amdgcn_raw_buffer_load_i32x4(
+                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
+        }
+        else if constexpr(N == 8)
+        {
+            vector_type<int32_t, 8> tmp;
+
+            tmp.Vectors(Number<4>{})(Number<0>{}) = __llvm_amdgcn_raw_buffer_load_i32x4(
+                src_wave_buffer_resource, src_thread_addr_offset, src_wave_addr_offset, 0);
+
+            tmp.Vectors(Number<4>{})(Number<1>{}) = __llvm_amdgcn_raw_buffer_load_i32x4(
+                src_wave_buffer_resource, src_thread_addr_offset, 4 * sizeof(int32_t), 0);
+
+            return tmp.Vector();
+        }
+    }
+}
+
+template <typename T, index_t N>
+__device__ void amd_buffer_store_impl_v2(const typename vector_type<T, N>::type src_thread_data,
+                                         int32x4_t dst_wave_buffer_resource,
+                                         index_t dst_thread_addr_offset,
+                                         index_t dst_wave_addr_offset)
+{
+    static_assert((is_same<T, float>::value && (N == 1 || N == 2 || N == 4)) ||
+                      (is_same<T, int32_t>::value && (N == 1 || N == 2 || N == 4)) ||
+                      (is_same<T, int8_t>::value && (N == 1 || N == 2 || N == 4)),
+                  "wrong! not implemented");
+
+    if constexpr(is_same<T, float>::value)
+    {
+        if constexpr(N == 1)
+        {
+            __llvm_amdgcn_raw_buffer_store_fp32(src_thread_data,
+                                                dst_wave_buffer_resource,
+                                                dst_thread_addr_offset,
+                                                dst_wave_addr_offset,
+                                                0);
+        }
+        else if constexpr(N == 2)
+        {
+            __llvm_amdgcn_raw_buffer_store_fp32x2(src_thread_data,
+                                                  dst_wave_buffer_resource,
+                                                  dst_thread_addr_offset,
+                                                  dst_wave_addr_offset,
+                                                  0);
+        }
+        else if constexpr(N == 4)
+        {
+            __llvm_amdgcn_raw_buffer_store_fp32x4(src_thread_data,
+                                                  dst_wave_buffer_resource,
+                                                  dst_thread_addr_offset,
+                                                  dst_wave_addr_offset,
+                                                  0);
+        }
+    }
+    else if constexpr(is_same<T, int32_t>::value)
+    {
+        if constexpr(N == 1)
+        {
+            __llvm_amdgcn_raw_buffer_store_i32(src_thread_data,
+                                               dst_wave_buffer_resource,
+                                               dst_thread_addr_offset,
+                                               dst_wave_addr_offset,
+                                               0);
+        }
+        else if constexpr(N == 2)
+        {
+            __llvm_amdgcn_raw_buffer_store_i32x2(src_thread_data,
+                                                 dst_wave_buffer_resource,
+                                                 dst_thread_addr_offset,
+                                                 dst_wave_addr_offset,
+                                                 0);
+        }
+        else if constexpr(N == 4)
+        {
+            __llvm_amdgcn_raw_buffer_store_i32x4(src_thread_data,
+                                                 dst_wave_buffer_resource,
+                                                 dst_thread_addr_offset,
+                                                 dst_wave_addr_offset,
+                                                 0);
+        }
+    }
+    else if constexpr(is_same<T, int8_t>::value)
+    {
+        if constexpr(N == 1)
+        {
+            __llvm_amdgcn_raw_buffer_store_i8(src_thread_data,
+                                              dst_wave_buffer_resource,
+                                              dst_thread_addr_offset,
+                                              dst_wave_addr_offset,
+                                              0);
+        }
+        else if constexpr(N == 2)
+        {
+            __llvm_amdgcn_raw_buffer_store_i16(src_thread_data,
+                                               dst_wave_buffer_resource,
+                                               dst_thread_addr_offset,
+                                               dst_wave_addr_offset,
+                                               0);
+        }
+        else if constexpr(N == 4)
+        {
+            __llvm_amdgcn_raw_buffer_store_i32(src_thread_data,
+                                               dst_wave_buffer_resource,
+                                               dst_thread_addr_offset,
+                                               dst_wave_addr_offset,
+                                               0);
+        }
+    }
+}
+
+// buffer_load requires:
+//   1) p_src_wave must be in global memory space
+//   2) p_src_wave to be a wavewise pointer.
+// It is user's responsibility to make sure that is true.
+template <typename T, index_t N>
+__device__ typename vector_type<T, N>::type amd_buffer_load_v2(const T* p_src_wave,
+                                                               index_t src_thread_data_offset,
+                                                               bool src_thread_data_valid,
+                                                               index_t src_element_space)
+{
+    const int32x4_t src_wave_buffer_resource =
+        make_wave_buffer_resource(p_src_wave, src_element_space);
+
+    index_t src_thread_addr_offset = src_thread_data_offset * sizeof(T);
+
+#if CK_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK
+    uint32_t src_addr_shift = src_thread_data_valid ? 0 : 0x7fffffff;
+
+    return amd_buffer_load_impl_v2<T, N>(
+        src_wave_buffer_resource, src_addr_shift + src_thread_addr_offset, 0);
+#else
+    using vector_t = typename vector_type<T, N>::type;
+
+    vector_t tmp =
+        amd_buffer_load_impl_v2<T, N>(src_wave_buffer_resource, src_thread_addr_offset, 0);
+
+    return src_thread_data_valid ? tmp : vector_t(0);
+#endif
+}
+
+// buffer_store requires:
+//   1) p_dst_wave must be global memory
+//   2) p_dst_wave to be a wavewise pointer.
+// It is user's responsibility to make sure that is true.
+template <typename T, index_t N>
+__device__ void amd_buffer_store_v2(const typename vector_type<T, N>::type src_thread_data,
+                                    T* p_dst_wave,
+                                    const index_t dst_thread_data_offset,
+                                    const bool dst_thread_data_valid,
+                                    const index_t dst_element_space)
+{
+    const int32x4_t dst_wave_buffer_resource =
+        make_wave_buffer_resource(p_dst_wave, dst_element_space);
+
+    index_t dst_thread_addr_offset = dst_thread_data_offset * sizeof(T);
+
+#if CK_EXPERIMENTAL_USE_BUFFER_STORE_OOB_CHECK_OFFSET_TRICK
+    uint32_t dst_addr_shift = dst_thread_data_valid ? 0 : 0x7fffffff;
+
+    amd_buffer_store_impl_v2<T, N>(
+        src_thread_data, dst_wave_buffer_resource, dst_addr_shift + dst_thread_addr_offset, 0);
+#else
+    if(dst_thread_data_valid)
+    {
+        amd_buffer_store_impl_v2<T, N>(
+            src_thread_data, dst_wave_buffer_resource, dst_thread_addr_offset, 0);
+    }
+#endif
+}
+
+} // namespace ck
+#endif

composable_kernel/include/utility/amd_inline_asm.hpp

@@ -5,21 +5,44 @@
 
 namespace ck {
 
-// outer-product: c[i,j] += inner_product(a[i], b[j])
+// c0 += inner_product(a, b0)
+// c1 += inner_product(a, b1)
 __device__ void amd_assembly_outer_product_1x2(float a, float b0, float b1, float& c0, float& c1)
 {
+#if CK_USE_AMD_V_FMAC_F32
+    asm volatile("\n \
+            v_fmac_f32 %0, %2, %3 \n \
+            v_fmac_f32 %1, %2, %4 \n \
+            "
+                 : "=v"(c0), "=v"(c1)
+                 : "v"(a), "v"(b0), "v"(b1), "0"(c0), "1"(c1));
+#else
     asm volatile("\n \
             v_mac_f32 %0, %2, %3 \n \
             v_mac_f32 %1, %2, %4 \n \
             "
                  : "=v"(c0), "=v"(c1)
                  : "v"(a), "v"(b0), "v"(b1), "0"(c0), "1"(c1));
+#endif
 }
 
-// outer-product: c[i,j] += inner_product(a[i], b[j])
+// c0 += inner_product(a, b0)
+// c1 += inner_product(a, b1)
+// c2 += inner_product(a, b2)
+// c3 += inner_product(a, b3)
 __device__ void amd_assembly_outer_product_1x4(
     float a, float b0, float b1, float b2, float b3, float& c0, float& c1, float& c2, float& c3)
 {
+#if CK_USE_AMD_V_FMAC_F32
+    asm volatile("\n \
+            v_fmac_f32 %0, %4, %5 \n \
+            v_fmac_f32 %1, %4, %6 \n \
+            v_fmac_f32 %2, %4, %7 \n \
+            v_fmac_f32 %3, %4, %8 \n \
+            "
+                 : "=v"(c0), "=v"(c1), "=v"(c2), "=v"(c3)
+                 : "v"(a), "v"(b0), "v"(b1), "v"(b2), "v"(b3), "0"(c0), "1"(c1), "2"(c2), "3"(c3));
+#else
     asm volatile("\n \
             v_mac_f32 %0, %4, %5 \n \
             v_mac_f32 %1, %4, %6 \n \
@@ -28,9 +51,11 @@ __device__ void amd_assembly_outer_product_1x4(
             "
                  : "=v"(c0), "=v"(c1), "=v"(c2), "=v"(c3)
                  : "v"(a), "v"(b0), "v"(b1), "v"(b2), "v"(b3), "0"(c0), "1"(c1), "2"(c2), "3"(c3));
+#endif
 }
 
-// outer-product: c[i,j] += inner_product(a[i], b[j])
+// c0 += inner_product(a, b0)
+// c1 += inner_product(a, b1)
 __device__ void
 amd_assembly_outer_product_1x2(half2_t a, half2_t b0, half2_t b1, float& c0, float& c1)
 {
@@ -38,15 +63,12 @@ amd_assembly_outer_product_1x2(half2_t a, half2_t b0, half2_t b1, float& c0, flo
             v_dot2_f32_f16 %0, %2, %3, %0\n \
             v_dot2_f32_f16 %1, %2, %4, %1\n \
             "
-                 : "=v"(c0), "=v"(c1) // Dest registers
-                 : "v"(a),            // 1st Src register for 1 half2 registers
-                   "v"(b0),           // 2nd Src register
-                   "v"(b1),
-                   "0"(c0), // 3rd Src register
-                   "1"(c1));
+                 : "=v"(c0), "=v"(c1)
+                 : "v"(a), "v"(b0), "v"(b1), "0"(c0), "1"(c1));
 }
 
-// outer-product: c[i,j] += inner_product(a[i], b[j])
+// c0 += inner_product(a, b0)
+// c1 += inner_product(a, b1)
 __device__ void
 amd_assembly_outer_product_1x2(half4_t a, half4_t b0, half4_t b1, float& c0, float& c1)
 {
@@ -61,18 +83,21 @@ amd_assembly_outer_product_1x2(half4_t a, half4_t b0, half4_t b1, float& c0, flo
             v_dot2_f32_f16 %0, %3, %5, %0\n \
             v_dot2_f32_f16 %1, %3, %7, %1\n \
             "
-                 : "=v"(c0), "=v"(c1) // Dest registers
+                 : "=v"(c0), "=v"(c1)
                  : "v"(p_a_half2[0]),
-                   "v"(p_a_half2[1]), // 1st Src registers for 2 half2 registers
+                   "v"(p_a_half2[1]),
                    "v"(p_b0_half2[0]),
                    "v"(p_b0_half2[1]),
                    "v"(p_b1_half2[0]),
-                   "v"(p_b1_half2[1]), // 2nd Src registers for 2 half2 registers
+                   "v"(p_b1_half2[1]),
                    "0"(c0),
-                   "1"(c1)); // 3rd Src Acc registers for 2 half2 registers
+                   "1"(c1));
 }
 
-// outer-product: c[i,j] += inner_product(a[i], b[j])
+// c0 += inner_product(a, b0)
+// c1 += inner_product(a, b1)
+// c2 += inner_product(a, b2)
+// c3 += inner_product(a, b3)
 __device__ void amd_assembly_outer_product_1x4(half2_t a,
                                                half2_t b0,
                                                half2_t b1,
@@ -89,19 +114,14 @@ __device__ void amd_assembly_outer_product_1x4(half2_t a,
             v_dot2_f32_f16 %2, %4, %7, %2\n \
             v_dot2_f32_f16 %3, %4, %8, %3\n \
             "
-                 : "=v"(c0), "=v"(c1), "=v"(c2), "=v"(c3) // Dest registers
-                 : "v"(a),                                // 1st Src register for 1 half2 registers
-                   "v"(b0),                               // 2nd Src register
-                   "v"(b1),
-                   "v"(b2),
-                   "v"(b3),
-                   "0"(c0), // 3rd Src register
-                   "1"(c1),
-                   "2"(c2),
-                   "3"(c3));
+                 : "=v"(c0), "=v"(c1), "=v"(c2), "=v"(c3)
+                 : "v"(a), "v"(b0), "v"(b1), "v"(b2), "v"(b3), "0"(c0), "1"(c1), "2"(c2), "3"(c3));
 }
 
-// outer-product: c[i,j] += inner_product(a[i], b[j])
+// c0 += inner_product(a, b0)
+// c1 += inner_product(a, b1)
+// c2 += inner_product(a, b2)
+// c3 += inner_product(a, b3)
 __device__ void amd_assembly_outer_product_1x4(half4_t a,
                                                half4_t b0,
                                                half4_t b1,
@@ -129,21 +149,70 @@ __device__ void amd_assembly_outer_product_1x4(half4_t a,
             v_dot2_f32_f16 %2, %5, %11, %2\n \
             v_dot2_f32_f16 %3, %5, %13, %3\n \
             "
-                 : "=v"(c0), "=v"(c1), "=v"(c2), "=v"(c3) // Dest registers
+                 : "=v"(c0), "=v"(c1), "=v"(c2), "=v"(c3)
                  : "v"(p_a_half2[0]),
-                   "v"(p_a_half2[1]), // 1st Src registers for 2 half2 registers
+                   "v"(p_a_half2[1]),
                    "v"(p_b0_half2[0]),
                    "v"(p_b0_half2[1]),
                    "v"(p_b1_half2[0]),
-                   "v"(p_b1_half2[1]), // 2nd Src registers for 2 half2 registers
+                   "v"(p_b1_half2[1]),
                    "v"(p_b2_half2[0]),
                    "v"(p_b2_half2[1]),
                    "v"(p_b3_half2[0]),
-                   "v"(p_b3_half2[1]), // 2nd Src registers for 2 half2 registers
+                   "v"(p_b3_half2[1]),
                    "0"(c0),
                    "1"(c1),
                    "2"(c2),
-                   "3"(c3)); // 3rd Src Acc registers for 2 half2 registers
+                   "3"(c3));
+}
+
+// c0 += inner_product(a, b0)
+// c1 += inner_product(a, b1)
+__device__ void
+amd_assembly_outer_product_1x2(int8x4_t a, int8x4_t b0, int8x4_t b1, int32_t& c0, int32_t& c1)
+{
+#if 1
+    asm volatile("\n \
+            v_dot4_i32_i8 %0, %2, %3, %0\n \
+            v_dot4_i32_i8 %1, %2, %4, %1\n \
+            "
+                 : "=v"(c0), "=v"(c1)
+                 : "v"(a), "v"(b0), "v"(b1), "0"(c0), "1"(c1));
+#else
+    c0 = __builtin_amdgcn_sdot4(a, b0, c0, false);
+    c1 = __builtin_amdgcn_sdot4(a, b1, c1, false);
+#endif
+}
+
+// c0 += inner_product(a, b0)
+// c1 += inner_product(a, b1)
+// c2 += inner_product(a, b2)
+// c3 += inner_product(a, b3)
+__device__ void amd_assembly_outer_product_1x4(int8x4_t a,
+                                               int8x4_t b0,
+                                               int8x4_t b1,
+                                               int8x4_t b2,
+                                               int8x4_t b3,
+                                               int32_t& c0,
+                                               int32_t& c1,
+                                               int32_t& c2,
+                                               int32_t& c3)
+{
+#if 1
+    asm volatile("\n \
+            v_dot4_i32_i8 %0, %4, %5, %0\n \
+            v_dot4_i32_i8 %1, %4, %6, %1\n \
+            v_dot4_i32_i8 %2, %4, %7, %2\n \
+            v_dot4_i32_i8 %3, %4, %8, %3\n \
+            "
+                 : "=v"(c0), "=v"(c1), "=v"(c2), "=v"(c3)
+                 : "v"(a), "v"(b0), "v"(b1), "v"(b2), "v"(b3), "0"(c0), "1"(c1), "2"(c2), "3"(c3));
+#else
+    c0 = __builtin_amdgcn_sdot4(a, b0, c0, false);
+    c1 = __builtin_amdgcn_sdot4(a, b1, c1, false);
+    c2 = __builtin_amdgcn_sdot4(a, b2, c2, false);
+    c3 = __builtin_amdgcn_sdot4(a, b3, c3, false);
+#endif
 }
 
 } // namespace ck

composable_kernel/include/utility/amd_llvm_intrinsic.hpp

@@ -0,0 +1,11 @@
+#ifndef CK_AMD_LLVM_INTRINSIC_HPP
+#define CK_AMD_LLVM_INTRINSIC_HPP
+
+#include "float_type.hpp"
+
+namespace ck {
+
+__device__ int32_t __llvm_amdgcn_readfirstlane_i32(int32_t i) __asm("llvm.amdgcn.readfirstlane");
+
+} // namespace ck
+#endif

composable_kernel/include/utility/array.hpp

@@ -1,419 +1,62 @@
 #ifndef CK_ARRAY_HPP
 #define CK_ARRAY_HPP
 
-#include "sequence.hpp"
 #include "functional2.hpp"
+#include "sequence.hpp"
 
 namespace ck {
 
 template <typename TData, index_t NSize>
 struct Array
 {
-    using type      = Array<TData, NSize>;
+    using type      = Array;
     using data_type = TData;
 
-    // TODO: implement empty Array
-    index_t mData[NSize];
-
-    __host__ __device__ explicit constexpr Array() {}
-
-    template <typename X, typename... Xs>
-    __host__ __device__ constexpr Array(X x, Xs... xs)
-        : mData{static_cast<TData>(x), static_cast<TData>(xs)...}
-    {
-        static_assert(sizeof...(Xs) + 1 == NSize, "wrong! size");
-    }
+    TData mData[NSize];
 
     __host__ __device__ static constexpr index_t Size() { return NSize; }
 
-    // TODO: remove
-    __host__ __device__ static constexpr index_t GetSize() { return Size(); }
-
-    template <index_t I>
-    __host__ __device__ constexpr const TData& At(Number<I>) const
-    {
-        static_assert(I < NSize, "wrong!");
-
-        return mData[I];
-    }
-
-    template <index_t I>
-    __host__ __device__ constexpr TData& At(Number<I>)
-    {
-        static_assert(I < NSize, "wrong!");
-
-        return mData[I];
-    }
-
     __host__ __device__ constexpr const TData& At(index_t i) const { return mData[i]; }
 
     __host__ __device__ constexpr TData& At(index_t i) { return mData[i]; }
 
-    template <typename I>
-    __host__ __device__ constexpr const TData& operator[](I i) const
-    {
-        return At(i);
-    }
+    __host__ __device__ constexpr const TData& operator[](index_t i) const { return At(i); }
 
-    template <typename I>
-    __host__ __device__ constexpr TData& operator()(I i)
-    {
-        return At(i);
-    }
+    __host__ __device__ constexpr TData& operator()(index_t i) { return At(i); }
 
     template <typename T>
-    __host__ __device__ constexpr type& operator=(const T& x)
+    __host__ __device__ constexpr auto operator=(const T& a)
     {
-        static_for<0, Size(), 1>{}([&](auto i) { operator()(i) = x[i]; });
+        static_assert(T::Size() == Size(), "wrong! size not the same");
 
-        return *this;
-    }
-
-    struct lambda_PushBack // emulate constexpr lambda
-    {
-        const Array<TData, NSize>& old_array;
-        Array<TData, NSize + 1>& new_array;
-
-        __host__ __device__ constexpr lambda_PushBack(const Array<TData, NSize>& old_array_,
-                                                      Array<TData, NSize + 1>& new_array_)
-            : old_array(old_array_), new_array(new_array_)
-        {
-        }
-
-        template <index_t I>
-        __host__ __device__ constexpr void operator()(Number<I>) const
-        {
-            new_array(Number<I>{}) = old_array[I];
-        }
-    };
-
-    __host__ __device__ constexpr auto PushBack(TData x) const
-    {
-        Array<TData, NSize + 1> new_array;
-
-        static_for<0, NSize, 1>{}(lambda_PushBack(*this, new_array));
-
-        new_array(Number<NSize>{}) = x;
-
-        return new_array;
-    }
-};
-
-// Arr: Array
-// Picks: Sequence<...>
-template <typename Arr, typename Picks>
-struct ArrayElementPicker
-{
-    using type      = ArrayElementPicker;
-    using data_type = typename Arr::data_type;
-
-    __host__ __device__ constexpr ArrayElementPicker() = delete;
-
-    __host__ __device__ explicit constexpr ArrayElementPicker(Arr& array) : mArray{array}
-    {
-        constexpr index_t imax = reduce_on_sequence(Picks{}, math::maxer<index_t>{}, Number<0>{});
-
-        static_assert(imax < Arr::Size(), "wrong! exceeding # array element");
-    }
-
-    __host__ __device__ static constexpr auto Size() { return Picks::Size(); }
-
-    template <index_t I>
-    __host__ __device__ constexpr const data_type& At(Number<I>) const
-    {
-        static_assert(I < Size(), "wrong!");
-
-        constexpr auto IP = Picks{}[I];
-        return mArray[IP];
-    }
-
-    template <index_t I>
-    __host__ __device__ constexpr data_type& At(Number<I>)
-    {
-        static_assert(I < Size(), "wrong!");
-
-        constexpr auto IP = Picks{}[I];
-        return mArray(IP);
-    }
-
-    template <typename I>
-    __host__ __device__ constexpr const data_type& operator[](I i) const
-    {
-        return At(i);
-    }
-
-    template <typename I>
-    __host__ __device__ constexpr data_type& operator()(I i)
-    {
-        return At(i);
-    }
-
-    template <typename T>
-    __host__ __device__ constexpr type& operator=(const T& a)
-    {
         static_for<0, Size(), 1>{}([&](auto i) { operator()(i) = a[i]; });
 
         return *this;
     }
-
-    Arr& mArray;
 };
 
-template <typename Arr, typename Picks>
-__host__ __device__ constexpr auto pick_array_element(Arr& a, Picks)
+// empty Array
+template <typename TData>
+struct Array<TData, 0>
 {
-    return ArrayElementPicker<Arr, Picks>(a);
-}
+    using type      = Array;
+    using data_type = TData;
 
-template <typename T>
-__host__ __device__ constexpr auto to_array(const T& x)
-{
-    Array<typename T::data_type, T::Size()> y;
-
-    static_for<0, T::Size(), 1>{}([&](auto i) { y.At(i) = x.At(i); });
-
-    return y;
-}
-
-// TODO: remove this
-template <index_t... Is>
-__host__ __device__ constexpr auto sequence2array(Sequence<Is...>)
-{
-    return Array<index_t, sizeof...(Is)>{Is...};
-}
-
-template <typename TData, index_t NSize>
-__host__ __device__ constexpr auto make_zero_array()
-{
-    constexpr auto zero_sequence = typename uniform_sequence_gen<NSize, 0>::type{};
-    constexpr auto zero_array    = sequence2array(zero_sequence);
-    return zero_array;
-}
-
-template <typename TData, index_t NSize, index_t... IRs>
-__host__ __device__ constexpr auto reorder_array_given_new2old(const Array<TData, NSize>& old_array,
-                                                               Sequence<IRs...> /*new2old*/)
-{
-    static_assert(NSize == sizeof...(IRs), "NSize not consistent");
-
-    static_assert(is_valid_sequence_map<Sequence<IRs...>>{}, "wrong! invalid reorder map");
-
-    return Array<TData, NSize>{old_array[IRs]...};
-}
-
-template <typename TData, index_t NSize, typename MapOld2New>
-struct lambda_reorder_array_given_old2new
-{
-    const Array<TData, NSize>& old_array;
-    Array<TData, NSize>& new_array;
-
-    __host__ __device__ constexpr lambda_reorder_array_given_old2new(
-        const Array<TData, NSize>& old_array_, Array<TData, NSize>& new_array_)
-        : old_array(old_array_), new_array(new_array_)
-    {
-    }
-
-    template <index_t IOldDim>
-    __host__ __device__ constexpr void operator()(Number<IOldDim>) const
-    {
-        TData old_data = old_array[IOldDim];
-
-        constexpr index_t INewDim = MapOld2New::At(Number<IOldDim>{});
-
-        new_array(Number<INewDim>{}) = old_data;
-    }
+    __host__ __device__ static constexpr index_t Size() { return 0; }
 };
 
-template <typename TData, index_t NSize, index_t... IRs>
-__host__ __device__ constexpr auto reorder_array_given_old2new(const Array<TData, NSize>& old_array,
-                                                               Sequence<IRs...> /*old2new*/)
+template <typename X, typename... Xs>
+__host__ __device__ constexpr auto make_array(X&& x, Xs&&... xs)
 {
-    Array<TData, NSize> new_array;
-
-    static_assert(NSize == sizeof...(IRs), "NSize not consistent");
-
-    static_assert(is_valid_sequence_map<Sequence<IRs...>>::value, "wrong! invalid reorder map");
-
-    static_for<0, NSize, 1>{}(
-        lambda_reorder_array_given_old2new<TData, NSize, Sequence<IRs...>>(old_array, new_array));
-
-    return new_array;
+    using data_type = remove_cv_t<remove_reference_t<X>>;
+    return Array<data_type, sizeof...(Xs) + 1>{{std::forward<X>(x), std::forward<Xs>(xs)...}};
 }
 
-template <typename TData, index_t NSize, typename ExtractSeq>
-__host__ __device__ constexpr auto extract_array(const Array<TData, NSize>& old_array, ExtractSeq)
+// make empty array
+template <typename X>
+__host__ __device__ constexpr auto make_array()
 {
-    Array<TData, ExtractSeq::GetSize()> new_array;
-
-    constexpr index_t new_size = ExtractSeq::GetSize();
-
-    static_assert(new_size <= NSize, "wrong! too many extract");
-
-    static_for<0, new_size, 1>{}([&](auto I) { new_array(I) = old_array[ExtractSeq::At(I)]; });
-
-    return new_array;
-}
-
-// emulate constepxr lambda for array
-template <typename F, typename X, typename Y, typename Z>
-struct lambda_array_math
-{
-    const F& f;
-    const X& x;
-    const Y& y;
-    Z& z;
-
-    __host__ __device__ constexpr lambda_array_math(const F& f_, const X& x_, const Y& y_, Z& z_)
-        : f(f_), x(x_), y(y_), z(z_)
-    {
-    }
-
-    template <index_t IDim_>
-    __host__ __device__ constexpr void operator()(Number<IDim_>) const
-    {
-        constexpr auto IDim = Number<IDim_>{};
-        z(IDim)             = f(x[IDim], y[IDim]);
-    }
-};
-
-// Array = Array + Array
-template <typename TData, index_t NSize>
-__host__ __device__ constexpr auto operator+(Array<TData, NSize> a, Array<TData, NSize> b)
-{
-    Array<TData, NSize> result;
-
-    auto f = math::plus<index_t>{};
-
-    static_for<0, NSize, 1>{}(
-        lambda_array_math<decltype(f), decltype(a), decltype(b), decltype(result)>(
-            f, a, b, result));
-
-    return result;
-}
-
-// Array = Array - Array
-template <typename TData, index_t NSize>
-__host__ __device__ constexpr auto operator-(Array<TData, NSize> a, Array<TData, NSize> b)
-{
-    Array<TData, NSize> result;
-
-    auto f = math::minus<index_t>{};
-
-    static_for<0, NSize, 1>{}(
-        lambda_array_math<decltype(f), decltype(a), decltype(b), decltype(result)>(
-            f, a, b, result));
-
-    return result;
-}
-
-// Array += Array
-template <typename TData, index_t NSize>
-__host__ __device__ constexpr auto operator+=(Array<TData, NSize>& a, Array<TData, NSize> b)
-{
-    a = a + b;
-    return a;
-}
-
-// Array -= Array
-template <typename TData, index_t NSize>
-__host__ __device__ constexpr auto operator-=(Array<TData, NSize>& a, Array<TData, NSize> b)
-{
-    a = a - b;
-    return a;
-}
-// Array = Array + Sequence
-template <typename TData, index_t NSize, index_t... Is>
-__host__ __device__ constexpr auto operator+(Array<TData, NSize> a, Sequence<Is...> b)
-{
-    static_assert(sizeof...(Is) == NSize, "wrong! size not the same");
-
-    Array<TData, NSize> result;
-
-    auto f = math::plus<index_t>{};
-
-    static_for<0, NSize, 1>{}(
-        lambda_array_math<decltype(f), decltype(a), decltype(b), decltype(result)>(
-            f, a, b, result));
-
-    return result;
-}
-
-// Array = Array - Sequence
-template <typename TData, index_t NSize, index_t... Is>
-__host__ __device__ constexpr auto operator-(Array<TData, NSize> a, Sequence<Is...> b)
-{
-    static_assert(sizeof...(Is) == NSize, "wrong! size not the same");
-
-    Array<TData, NSize> result;
-
-    auto f = math::minus<index_t>{};
-
-    static_for<0, NSize, 1>{}(
-        lambda_array_math<decltype(f), decltype(a), decltype(b), decltype(result)>(
-            f, a, b, result));
-
-    return result;
-}
-
-// Array = Array * Sequence
-template <typename TData, index_t NSize, index_t... Is>
-__host__ __device__ constexpr auto operator*(Array<TData, NSize> a, Sequence<Is...> b)
-{
-    static_assert(sizeof...(Is) == NSize, "wrong! size not the same");
-
-    Array<TData, NSize> result;
-
-    auto f = math::multiplies<index_t>{};
-
-    static_for<0, NSize, 1>{}(
-        lambda_array_math<decltype(f), decltype(a), decltype(b), decltype(result)>(
-            f, a, b, result));
-
-    return result;
-}
-
-// Array = Sequence - Array
-template <typename TData, index_t NSize, index_t... Is>
-__host__ __device__ constexpr auto operator-(Sequence<Is...> a, Array<TData, NSize> b)
-{
-    static_assert(sizeof...(Is) == NSize, "wrong! size not the same");
-
-    Array<TData, NSize> result;
-
-    auto f = math::minus<index_t>{};
-
-    static_for<0, NSize, 1>{}(
-        lambda_array_math<decltype(f), decltype(a), decltype(b), decltype(result)>(
-            f, a, b, result));
-
-    return result;
-}
-
-// Array = Array * TData
-template <typename TData, index_t NSize>
-__host__ __device__ constexpr auto operator*(TData v, Array<TData, NSize> a)
-{
-    Array<TData, NSize> result;
-
-    for(index_t i = 0; i < NSize; ++i)
-    {
-        result(i) = a[i] * v;
-    }
-
-    return result;
-}
-
-template <typename TData, index_t NSize, typename Reduce>
-__host__ __device__ constexpr TData
-accumulate_on_array(const Array<TData, NSize>& a, Reduce f, TData init)
-{
-    TData result = init;
-
-    static_assert(NSize > 0, "wrong");
-
-    static_for<0, NSize, 1>{}([&](auto I) { result = f(result, a[I]); });
-
-    return result;
+    return Array<X, 0>{};
 }
 
 } // namespace ck

composable_kernel/include/utility/common_header.hpp

@@ -2,21 +2,26 @@
 #define CK_COMMON_HEADER_HPP
 
 #include "config.hpp"
-#include "utility.hpp"
-#include "integral_constant.hpp"
-#include "number.hpp"
-#include "float_type.hpp"
-#include "type.hpp"
-#include "tuple.hpp"
-#include "math.hpp"
-#include "sequence.hpp"
 #include "array.hpp"
+#include "container_helper.hpp"
+#include "statically_indexed_array.hpp"
+#include "container_element_picker.hpp"
+#include "float_type.hpp"
 #include "functional.hpp"
 #include "functional2.hpp"
 #include "functional3.hpp"
 #include "functional4.hpp"
 #include "in_memory_operation.hpp"
+#include "integral_constant.hpp"
+#include "math.hpp"
+#include "number.hpp"
+#include "sequence.hpp"
+#include "sequence_helper.hpp"
 #include "synchronization.hpp"
+#include "tuple.hpp"
+#include "tuple_helper.hpp"
+#include "type.hpp"
+#include "utility.hpp"
 
 #if CK_USE_AMD_INLINE_ASM
 #include "amd_inline_asm.hpp"

composable_kernel/include/utility/config.amd.hpp.in

@@ -1,16 +1,43 @@
 #ifndef CK_CONFIG_AMD_HPP
 #define CK_CONFIG_AMD_HPP
 
+#ifndef MIOPEN_DONT_USE_HIP_RUNTIME_HEADERS
 #include "hip/hip_runtime.h"
 #include "hip/hip_fp16.h"
+#endif
 #include "bfloat16_dev.hpp"
 
-// index type: unsigned or signed
-#define CK_UNSIGNED_INDEX_TYPE 0
-
 // device backend
 #define CK_DEVICE_BACKEND_AMD 1
 
+// GPU ID
+#define CK_AMD_GPU_GFX906 1
+#define CK_AMD_GPU_GFX908 0
+#define CK_AMD_GPU_GFX1030 0
+
+// HIP version
+#ifndef CK_HIP_VERSION_FLAT
+#define CK_HIP_VERSION_FLAT 0
+#endif
+
+// launch bounds
+#define CK_USE_LAUNCH_BOUNDS 0
+
+#ifdef CK_USE_LAUNCH_BOUNDS
+#define CK_MAX_THREAD_PER_BLOCK 256
+#define CK_MIN_BLOCK_PER_CU 1
+#endif
+
+// buffer resourse
+#if defined(CK_AMD_GPU_GFX906) || defined(CK_AMD_GPU_GFX908)
+#define CK_BUFFER_RESOURCE_3RD_DWORD 0x00020000
+#elif defined(CK_AMD_GPU_GFX1030)
+#define CK_BUFFER_RESOURCE_3RD_DWORD 0x31014000
+#endif
+
+// multi index
+#define CK_USE_DYNAMICALLY_INDEXED_MULTI_INDEX 0
+
 // AMD inline asm
 #ifndef CK_USE_AMD_INLINE_ASM
 #define CK_USE_AMD_INLINE_ASM 1
@@ -20,14 +47,18 @@
 #define CK_THREADWISE_GEMM_USE_AMD_INLINE_ASM 1
 #endif
 
+#ifndef CK_USE_AMD_V_FMAC_F32
+#define CK_USE_AMD_V_FMAC_F32 1
+#endif
+
 // AMD buffer addressing
 #ifndef CK_USE_AMD_BUFFER_ADDRESSING
 #define CK_USE_AMD_BUFFER_ADDRESSING 1
 #endif
 
 // only gfx908 support native floating point atomic add
-#ifndef CK_USE_AMD_BUFFER_ATOMIC_ADD
-#define CK_USE_AMD_BUFFER_ATOMIC_ADD 0
+#ifndef CK_USE_AMD_BUFFER_ATOMIC_FADD
+#define CK_USE_AMD_BUFFER_ATOMIC_FADD 0
 #endif
 
 // AMD XDLOPS
@@ -49,8 +80,16 @@
 #endif
 
 // experimental implementation
-#ifndef CK_EXPERIMENTAL_AMD_BUFFER_ADDRESSING_USE_OFFSET_TRICK
-#define CK_EXPERIMENTAL_AMD_BUFFER_ADDRESSING_USE_OFFSET_TRICK 1
+#ifndef CK_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK
+#define CK_EXPERIMENTAL_USE_BUFFER_LOAD_OOB_CHECK_OFFSET_TRICK 1
+#endif
+
+#ifndef CK_EXPERIMENTAL_USE_BUFFER_STORE_OOB_CHECK_OFFSET_TRICK
+#define CK_EXPERIMENTAL_USE_BUFFER_STORE_OOB_CHECK_OFFSET_TRICK 1
+#endif
+
+#ifndef CK_EXPERIMENTAL_USE_BUFFER_ATOMIC_OOB_CHECK_OFFSET_TRICK
+#define CK_EXPERIMENTAL_USE_BUFFER_ATOMIC_OOB_CHECK_OFFSET_TRICK 1
 #endif
 
 #ifndef CK_EXPERIMENTAL_BLOCKWISE_GEMM_USE_PIPELINE
@@ -65,14 +104,33 @@
 #define CK_EXPERIMENTAL_IMPLICIT_GEMM_BACKWARD_DATA_V4R1_INPUT_SKIP_OUT_OF_BOUND_CHECK 0
 #endif
 
+// pass tensor descriptor by value, pointer or void*
+#define CK_EXPERIMENTAL_PASS_TENSOR_DESCRIPTOR_BY_VALUE 1
+#define CK_EXPERIMENTAL_PASS_TENSOR_DESCRIPTOR_BY_POINTER 0
+#define CK_EXPERIMENTAL_PASS_TENSOR_DESCRIPTOR_BY_VOID_POINTER 0
+
+// hack: have underlying assumption that need to be satsified, otherwise it's a bug
+// hack for forcing register to keep idx_diff_low_const in SGPR. idx_diff_low_const must be
+// thread-invariant, otherwise it's a bug
+// TODO: separate index calculation into "compile-time", "global", "block", "wave", "thread"
+#ifndef CK_HACK_DYNAMIC_MERGE_CALCULATE_IDX_DIFF_LOW_CONST_USE_AMD_GCN_READ_FIRST_LANE
+#define CK_HACK_DYNAMIC_MERGE_CALCULATE_IDX_DIFF_LOW_CONST_USE_AMD_GCN_READ_FIRST_LANE 0
+#endif
+
 // workaround: put all workaround here
-// workaround for unnecessary VGPA <--> AGRP data movement when using mfma LLVM intrinsic
+// workaround for unnecessary VGPR <--> AGPR data movement when using mfma LLVM intrinsic
 #ifndef CK_WORKAROUND_SWDEV_229564
 #define CK_WORKAROUND_SWDEV_229564 1
 #endif
-// workaround for buffer load/store fp16/bfp16 intrinsic bug
-#ifndef CK_WORKAROUND_SWDEV_231101
-#define CK_WORKAROUND_SWDEV_231101 1
+
+// workaround for accvgpr over-allocation
+#ifndef CK_WORKAROUND_SWDEV_241664
+#define CK_WORKAROUND_SWDEV_241664 1
+#endif
+
+// workaround for compiler crash when compiling recursive lambda
+#ifndef CK_WORKAROUND_SWDEV_275126
+#define CK_WORKAROUND_SWDEV_275126 1
 #endif
 
 namespace ck {
@@ -91,14 +149,8 @@ enum InMemoryDataOperation
     AtomicAdd
 };
 
-#if CK_UNSIGNED_INDEX_TYPE
-using index_t = uint32_t;
-#else
+// index type
 using index_t = int32_t;
-#endif
-
-// int32x4_t use by buffer_load and buffer_store llvm intrinsic
-typedef int32_t int32x4_t __attribute__((ext_vector_type(4)));
 
 } // namespace ck
 #endif

composable_kernel/include/utility/container_element_picker.hpp

@@ -0,0 +1,153 @@
+#ifndef CK_CONTAINER_ELEMENT_PICKER_HPP
+#define CK_CONTAINER_ELEMENT_PICKER_HPP
+
+#include "functional2.hpp"
+#include "sequence.hpp"
+
+namespace ck {
+
+// Arr: Array or StaticallyIndexedArray
+// Picks: Sequence<...>
+template <typename Arr, typename Picks>
+struct ContainerElementPicker
+{
+    using type = ContainerElementPicker;
+#if 0
+    using data_type = typename Arr::data_type;
+#endif
+
+    __host__ __device__ constexpr ContainerElementPicker() = delete;
+
+    __host__ __device__ constexpr ContainerElementPicker(Arr& array) : mArray{array}
+    {
+        constexpr index_t imax = reduce_on_sequence(Picks{}, math::maxer<index_t>{}, Number<0>{});
+
+        static_assert(imax < Arr::Size(), "wrong! exceeding # array element");
+    }
+
+    __host__ __device__ static constexpr auto Size() { return Picks::Size(); }
+
+    template <index_t I>
+    __host__ __device__ constexpr const auto& At(Number<I> i) const
+    {
+        static_assert(I < Size(), "wrong!");
+
+        constexpr auto IP = Picks{}[i];
+        return mArray[IP];
+    }
+
+    template <index_t I>
+    __host__ __device__ constexpr auto& At(Number<I> i)
+    {
+        static_assert(I < Size(), "wrong!");
+
+        constexpr auto IP = Picks{}[i];
+        return mArray(IP);
+    }
+
+    template <index_t I>
+    __host__ __device__ constexpr const auto& operator[](Number<I> i) const
+    {
+        return At(i);
+    }
+
+    template <index_t I>
+    __host__ __device__ constexpr auto& operator()(Number<I> i)
+    {
+        return At(i);
+    }
+
+    template <typename T>
+    __host__ __device__ constexpr auto operator=(const T& a)
+    {
+        static_assert(T::Size() == Size(), "wrong! size not the same");
+
+        static_for<0, Size(), 1>{}([&](auto i) { operator()(i) = a[i]; });
+
+        return *this;
+    }
+
+    private:
+    Arr& mArray;
+};
+
+// Arr: Array or StaticallyIndexedArray
+// Picks: Sequence<...>
+template <typename Arr, typename Picks>
+struct ConstantContainerElementPicker
+{
+    using type = ConstantContainerElementPicker;
+#if 0
+    using data_type = typename Arr::data_type;
+#endif
+
+    __host__ __device__ constexpr ConstantContainerElementPicker() = delete;
+
+    __host__ __device__ constexpr ConstantContainerElementPicker(const Arr& array) : mArray{array}
+    {
+        constexpr index_t imax = reduce_on_sequence(Picks{}, math::maxer<index_t>{}, Number<0>{});
+
+        static_assert(imax < Arr::Size(), "wrong! exceeding # array element");
+    }
+
+    __host__ __device__ static constexpr auto Size() { return Picks::Size(); }
+
+    template <index_t I>
+    __host__ __device__ constexpr const auto& At(Number<I> i) const
+    {
+        static_assert(I < Size(), "wrong!");
+
+        constexpr auto IP = Picks{}[i];
+        return mArray[IP];
+    }
+
+    template <index_t I>
+    __host__ __device__ constexpr const auto& operator[](Number<I> i) const
+    {
+        return At(i);
+    }
+
+    private:
+    const Arr& mArray;
+};
+
+template <typename Arr, typename Picks, typename X>
+__host__ __device__ constexpr auto operator+=(ContainerElementPicker<Arr, Picks>& y, const X& x)
+{
+    using Y                 = ContainerElementPicker<Arr, Picks>;
+    constexpr index_t nsize = Y::Size();
+
+    static_assert(nsize == X::Size(), "wrong! size not the same");
+
+    static_for<0, nsize, 1>{}([&](auto i) { y(i) += x[i]; });
+
+    return y;
+}
+
+template <typename Arr, typename Picks, typename X>
+__host__ __device__ constexpr auto operator-=(ContainerElementPicker<Arr, Picks>& y, const X& x)
+{
+    using Y                 = ContainerElementPicker<Arr, Picks>;
+    constexpr index_t nsize = Y::Size();
+
+    static_assert(nsize == X::Size(), "wrong! size not the same");
+
+    static_for<0, nsize, 1>{}([&](auto i) { y(i) -= x[i]; });
+
+    return y;
+}
+
+template <typename Arr, typename Picks>
+__host__ __device__ constexpr auto pick_container_element(Arr& a, Picks)
+{
+    return ContainerElementPicker<Arr, Picks>(a);
+}
+
+template <typename Arr, typename Picks>
+__host__ __device__ constexpr auto pick_container_element(const Arr& a, Picks)
+{
+    return ConstantContainerElementPicker<Arr, Picks>(a);
+}
+
+} // namespace ck
+#endif

composable_kernel/include/utility/container_helper.hpp

@@ -0,0 +1,375 @@
+#ifndef CK_CONTAINER_HELPER_HPP
+#define CK_CONTAINER_HELPER_HPP
+
+#include "sequence.hpp"
+#include "sequence_helper.hpp"
+#include "array.hpp"
+#include "tuple.hpp"
+#include "tuple_helper.hpp"
+#include "statically_indexed_array.hpp"
+#include "container_element_picker.hpp"
+
+namespace ck {
+
+template <typename TData, index_t NSize>
+__host__ __device__ constexpr auto container_push_back(const Array<TData, NSize>& a, const TData& x)
+{
+    Array<TData, NSize + 1> r;
+
+    static_for<0, NSize, 1>{}([&r, &a ](auto i) constexpr { r(i) = a[i]; });
+
+    r(Number<NSize>{}) = x;
+
+    return r;
+}
+
+template <typename... Ts, typename T>
+__host__ __device__ constexpr auto container_push_front(const Tuple<Ts...>& a, const T& x)
+{
+    return container_cat(make_tuple(x), a);
+}
+
+template <typename... Ts, typename T>
+__host__ __device__ constexpr auto container_push_back(const Tuple<Ts...>& a, const T& x)
+{
+    return container_cat(a, make_tuple(x));
+}
+
+template <typename TData, index_t NSize, index_t... IRs>
+__host__ __device__ constexpr auto
+container_reorder_given_new2old(const Array<TData, NSize>& old_array, Sequence<IRs...> /*new2old*/)
+{
+    static_assert(NSize == sizeof...(IRs), "wrong! size not consistent");
+
+    static_assert(is_valid_sequence_map<Sequence<IRs...>>{}, "wrong! invalid reorder map");
+
+    return make_array(old_array[Number<IRs>{}]...);
+}
+
+template <typename TData, index_t NSize, index_t... IRs>
+__host__ __device__ constexpr auto
+container_reorder_given_old2new(const Array<TData, NSize>& old_array, Sequence<IRs...> old2new)
+{
+    return container_reorder_given_new2old(
+        old_array, typename sequence_map_inverse<decltype(old2new)>::type{});
+}
+
+template <typename... Ts, index_t... IRs>
+__host__ __device__ constexpr auto container_reorder_given_new2old(const Tuple<Ts...>& old_tuple,
+                                                                   Sequence<IRs...> /*new2old*/)
+{
+    static_assert(sizeof...(Ts) == sizeof...(IRs), "wrong! size not consistent");
+
+    static_assert(is_valid_sequence_map<Sequence<IRs...>>{}, "wrong! invalid reorder map");
+
+    return make_tuple(old_tuple[Number<IRs>{}]...);
+}
+
+template <typename... Ts, index_t... IRs>
+__host__ __device__ constexpr auto container_reorder_given_old2new(const Tuple<Ts...>& old_tuple,
+                                                                   Sequence<IRs...> old2new)
+{
+    return container_reorder_given_new2old(
+        old_tuple, typename sequence_map_inverse<decltype(old2new)>::type{});
+}
+
+template <index_t... Is, index_t... IRs>
+__host__ __device__ constexpr auto container_reorder_given_new2old(Sequence<Is...> /* old_seq */,
+                                                                   Sequence<IRs...> /*new2old*/)
+{
+    static_assert(sizeof...(Is) == sizeof...(IRs), "wrong! size not consistent");
+
+    static_assert(is_valid_sequence_map<Sequence<IRs...>>{}, "wrong! invalid reorder map");
+
+    return Sequence<Sequence<Is...>::At(Number<IRs>{})...>{};
+}
+
+template <index_t... Is, index_t... IRs>
+__host__ __device__ constexpr auto container_reorder_given_old2new(Sequence<Is...> old_seq,
+                                                                   Sequence<IRs...> /* old2new */)
+{
+    static_assert(sizeof...(Is) == sizeof...(IRs), "wrong! size not consistent");
+
+    static_assert(is_valid_sequence_map<Sequence<IRs...>>{}, "wrong! invalid reorder map");
+
+    constexpr auto new2old = typename sequence_map_inverse<Sequence<IRs...>>::type{};
+
+    return container_reorder_give_new2old(old_seq, new2old);
+}
+
+#if !CK_WORKAROUND_SWDEV_275126
+// rocm-4.1 compiler would crash for recursive lambda
+template <typename Container,
+          typename Reduce,
+          typename Init,
+          index_t IBegin = 0,
+          index_t IEnd   = Container::Size(),
+          index_t IStep  = 1>
+__host__ __device__ constexpr auto container_reduce(const Container& x,
+                                                    Reduce reduce,
+                                                    Init init,
+                                                    Number<IBegin> = Number<0>{},
+                                                    Number<IEnd>   = Number<Container::Size()>{},
+                                                    Number<IStep>  = Number<1>{})
+{
+    static_assert((IEnd - IBegin) % IStep == 0, "wrong!");
+
+    // f is recursive function, fs is a dummy of f
+    // i is index, y_old is current scan, r_old is current reduction
+    auto f = [&](auto fs, auto i, auto r_old) {
+        auto r_new = reduce(x[i], r_old);
+
+        if constexpr(i.value < IEnd - IStep)
+        {
+            // recursively call f/fs
+            return fs(fs, i + Number<IStep>{}, r_new);
+        }
+        else
+        {
+            return r_new;
+        }
+    };
+
+    // start recursion
+    return f(f, Number<IBegin>{}, init);
+}
+#else
+// i is index, y_old is current scan, r_old is current reduction
+template <typename Container,
+          typename Reduce,
+          typename ROld,
+          index_t I,
+          index_t IEnd,
+          index_t IStep>
+__host__ __device__ constexpr auto container_reduce_impl(
+    const Container& x, Reduce reduce, ROld r_old, Number<I> i, Number<IEnd>, Number<IStep>)
+{
+    auto r_new = reduce(x[i], r_old);
+
+    if constexpr(i.value < IEnd - IStep)
+    {
+        return container_reduce_impl(
+            x, reduce, r_new, i + Number<IStep>{}, Number<IEnd>{}, Number<IStep>{});
+    }
+    else
+    {
+        return r_new;
+    }
+}
+
+// rocm-4.1 compiler would crash for recursive lambda
+template <typename Container,
+          typename Reduce,
+          typename Init,
+          index_t IBegin = 0,
+          index_t IEnd   = Container::Size(),
+          index_t IStep  = 1>
+__host__ __device__ constexpr auto container_reduce(const Container& x,
+                                                    Reduce reduce,
+                                                    Init init,
+                                                    Number<IBegin> = Number<0>{},
+                                                    Number<IEnd>   = Number<Container::Size()>{},
+                                                    Number<IStep>  = Number<1>{})
+{
+    static_assert((IEnd - IBegin) % IStep == 0, "wrong!");
+
+    return container_reduce_impl(
+        x, reduce, init, Number<IBegin>{}, Number<IEnd>{}, Number<IStep>{});
+}
+#endif
+
+template <typename TData, index_t NSize, typename Reduce>
+__host__ __device__ constexpr auto
+container_reverse_inclusive_scan(const Array<TData, NSize>& x, Reduce f, TData init)
+{
+    Array<TData, NSize> y;
+
+    TData r = init;
+
+    static_for<NSize - 1, 0, -1>{}([&](auto i) {
+        r    = f(r, x[i]);
+        y(i) = r;
+    });
+
+    r              = f(r, x[Number<0>{}]);
+    y(Number<0>{}) = r;
+
+    return y;
+}
+
+template <typename TData, index_t NSize, typename Reduce>
+__host__ __device__ constexpr auto
+container_reverse_exclusive_scan(const Array<TData, NSize>& x, Reduce f, TData init)
+{
+    Array<TData, NSize> y;
+
+    TData r = init;
+
+    static_for<NSize - 1, 0, -1>{}([&](auto i) {
+        y(i) = r;
+        r    = f(r, x[i]);
+    });
+
+    y(Number<0>{}) = r;
+
+    return y;
+}
+
+#if !CK_WORKAROUND_SWDEV_275126
+// rocm4.1 compiler would crash with recursive lambda
+template <typename... Xs, typename Reduce, typename Init>
+__host__ __device__ constexpr auto
+container_reverse_exclusive_scan(const Tuple<Xs...>& x, Reduce reduce, Init init)
+{
+    constexpr index_t NSize = sizeof...(Xs);
+
+    // f is recursive function, fs is a dummy of f
+    // i is index, y_old is current scan, r_old is current reduction
+    auto f = [&](auto fs, auto i, auto y_old, auto r_old) {
+        auto r_new = reduce(x[i], r_old);
+
+        auto y_new = container_push_front(y_old, r_new);
+
+        if constexpr(i.value > 1)
+        {
+            // recursively call f/fs
+            return fs(fs, i - Number<1>{}, y_new, r_new);
+        }
+        else
+        {
+            return y_new;
+        }
+    };
+
+    // start recursion
+    return f(f, Number<NSize - 1>{}, make_tuple(init), init);
+}
+#else
+// i is index, y_old is current scan, r_old is current reduction
+template <typename... Xs, typename Reduce, index_t I, typename YOld, typename ROld>
+__host__ __device__ constexpr auto container_reverse_exclusive_scan_impl(
+    const Tuple<Xs...>& x, Reduce reduce, Number<I> i, YOld y_old, ROld r_old)
+{
+    auto r_new = reduce(x[i], r_old);
+
+    auto y_new = container_push_front(y_old, r_new);
+
+    if constexpr(i.value > 1)
+    {
+        // recursively call f/fs
+        return container_reverse_exclusive_scan_impl(x, reduce, i - Number<1>{}, y_new, r_new);
+    }
+    else
+    {
+        return y_new;
+    }
+}
+
+template <typename... Xs, typename Reduce, typename Init>
+__host__ __device__ constexpr auto
+container_reverse_exclusive_scan(const Tuple<Xs...>& x, Reduce reduce, Init init)
+{
+    constexpr index_t NSize = sizeof...(Xs);
+
+    return container_reverse_exclusive_scan_impl(
+        x, reduce, Number<NSize - 1>{}, make_tuple(init), init);
+}
+#endif
+
+// TODO: update to like container_reverse_exclusive_scan to deal with Tuple of Numebr<>
+template <typename... Xs, typename Reduce, typename TData>
+__host__ __device__ constexpr auto
+container_reverse_inclusive_scan(const Tuple<Xs...>& x, Reduce f, TData init)
+{
+    constexpr index_t NSize = sizeof...(Xs);
+
+    Tuple<Xs...> y;
+
+    TData r = init;
+
+    static_for<NSize - 1, 0, -1>{}([&](auto i) {
+        r    = f(r, x[i]);
+        y(i) = r;
+    });
+
+    r              = f(r, x[Number<0>{}]);
+    y(Number<0>{}) = r;
+
+    return y;
+}
+
+template <typename X, typename... Ys>
+__host__ __device__ constexpr auto container_cat(const X& x, const Ys&... ys)
+{
+    return container_cat(x, container_cat(ys...));
+}
+
+template <typename T, index_t NX, index_t NY>
+__host__ __device__ constexpr auto container_cat(const Array<T, NX>& ax, const Array<T, NY>& ay)
+{
+    return unpack2(
+        [&](auto&&... zs) { return make_array(std::forward<decltype(zs)>(zs)...); }, ax, ay);
+}
+
+template <typename... X, typename... Y>
+__host__ __device__ constexpr auto container_cat(const Tuple<X...>& tx, const Tuple<Y...>& ty)
+{
+    return unpack2(
+        [&](auto&&... zs) { return make_tuple(std::forward<decltype(zs)>(zs)...); }, tx, ty);
+}
+
+template <typename Container>
+__host__ __device__ constexpr auto container_cat(const Container& x)
+{
+    return x;
+}
+
+template <typename T, index_t N, index_t... Is>
+__host__ __device__ constexpr auto get_container_subset(const Array<T, N>& arr, Sequence<Is...>)
+{
+    static_assert(N >= sizeof...(Is), "wrong! size");
+
+    return make_array(arr[Number<Is>{}]...);
+}
+
+template <typename... Ts, index_t... Is>
+__host__ __device__ constexpr auto get_container_subset(const Tuple<Ts...>& tup, Sequence<Is...>)
+{
+    static_assert(sizeof...(Ts) >= sizeof...(Is), "wrong! size");
+
+    return make_tuple(tup[Number<Is>{}]...);
+}
+
+template <typename T, index_t N, index_t... Is>
+__host__ __device__ constexpr void
+set_container_subset(Array<T, N>& y, Sequence<Is...> picks, const Array<T, sizeof...(Is)>& x)
+{
+    static_assert(N >= sizeof...(Is), "wrong! size");
+
+    static_for<0, sizeof...(Is), 1>{}([&](auto i) { y(picks[i]) = x[i]; });
+}
+
+template <typename... Ys, index_t... Is, typename... Xs>
+__host__ __device__ constexpr void
+set_container_subset(Tuple<Ys...>& y, Sequence<Is...> picks, const Tuple<Xs...>& x)
+{
+    static_assert(sizeof...(Ys) >= sizeof...(Is) && sizeof...(Is) == sizeof...(Xs), "wrong! size");
+
+    static_for<0, sizeof...(Is), 1>{}([&](auto i) { y(picks[i]) = x[i]; });
+}
+
+template <index_t... Is>
+__host__ __device__ constexpr auto sequence_to_tuple_of_number(Sequence<Is...>)
+{
+    using Seq = Sequence<Is...>;
+
+    return generate_tuple(
+        [&](auto i) {
+            constexpr index_t tmp = Seq::At(i);
+            return Number<tmp>{};
+        },
+        Seq::Size());
+}
+
+} // namespace ck
+#endif

composable_kernel/include/utility/float_type.amd.hpp.in

@@ -3,263 +3,278 @@
 
 namespace ck {
 
-// For some reason, HIP compiler need this definition to generate optimal ISA
-// float
-typedef float float2_t __attribute__((ext_vector_type(2)));
-typedef float float4_t __attribute__((ext_vector_type(4)));
-typedef float float16_t __attribute__((ext_vector_type(16)));
-typedef float float32_t __attribute__((ext_vector_type(32)));
+template <typename T, index_t N>
+struct vector_type;
 
-// float16
-typedef _Float16 half_t;
-typedef _Float16 half2_t __attribute__((ext_vector_type(2)));
-typedef _Float16 half4_t __attribute__((ext_vector_type(4)));
-typedef _Float16 half8_t __attribute__((ext_vector_type(8)));
-
-// bfloat16
-typedef ushort ushort2_t __attribute__((ext_vector_type(2)));
-typedef ushort ushort4_t __attribute__((ext_vector_type(4)));
-typedef ushort ushort8_t __attribute__((ext_vector_type(8)));
-
-template <class T, index_t N>
-struct vector_type
+template <typename T>
+struct vector_type<T, 1>
 {
-    typedef struct
+    using type = T;
+
+    union
     {
-        T scalar[N];
-    } MemoryType;
+        T d1_;
+        StaticallyIndexedArray<T, 1> d1x1_;
+    } data_;
+
+    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
+
+    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
+
+    __host__ __device__ static constexpr index_t Size() { return 1; }
+
+    __host__ __device__ constexpr const auto& Vector() const { return data_.d1_; }
+
+    __host__ __device__ constexpr auto& Vector() { return data_.d1_; }
+
+    __host__ __device__ constexpr const auto& Scalars() const { return data_.d1x1_; }
+
+    __host__ __device__ constexpr auto& Scalars() { return data_.d1x1_; }
+
+    __host__ __device__ constexpr const auto& Vectors(Number<1>) const { return data_.d1x1_; }
+
+    __host__ __device__ constexpr auto& Vectors(Number<1>) { return data_.d1x1_; }
+};
+
+template <typename T>
+struct vector_type<T, 2>
+{
+    using d1_t = T;
+    typedef T d2_t __attribute__((ext_vector_type(2)));
+
+    using type = d2_t;
+
+    union
+    {
+        d2_t d2_;
+        StaticallyIndexedArray<d1_t, 2> d1x2_;
+        StaticallyIndexedArray<d2_t, 1> d2x1_;
+    } data_;
+
+    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
+
+    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
+
+    __host__ __device__ static constexpr index_t Size() { return 2; }
+
+    __host__ __device__ constexpr const auto& Vector() const { return data_.d2_; }
+
+    __host__ __device__ constexpr auto& Vector() { return data_.d2_; }
+
+    __host__ __device__ constexpr const auto& Scalars() const { return data_.d1x2_; }
+
+    __host__ __device__ constexpr auto& Scalars() { return data_.d1x2_; }
+
+    __host__ __device__ constexpr const auto& Vectors(Number<1>) const { return data_.d1x2_; }
+
+    __host__ __device__ constexpr const auto& Vectors(Number<2>) const { return data_.d2x1_; }
+
+    __host__ __device__ constexpr auto& Vectors(Number<1>) { return data_.d1x2_; }
+
+    __host__ __device__ constexpr auto& Vectors(Number<2>) { return data_.d2x1_; }
+};
+
+template <typename T>
+struct vector_type<T, 4>
+{
+    using d1_t = T;
+    typedef T d2_t __attribute__((ext_vector_type(2)));
+    typedef T d4_t __attribute__((ext_vector_type(4)));
+
+    using type = d4_t;
+
+    union
+    {
+        d4_t d4_;
+        StaticallyIndexedArray<d1_t, 4> d1x4_;
+        StaticallyIndexedArray<d2_t, 2> d2x2_;
+        StaticallyIndexedArray<d4_t, 1> d4x1_;
+    } data_;
+
+    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
+
+    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
+
+    __host__ __device__ static constexpr index_t Size() { return 4; }
+
+    __host__ __device__ constexpr const auto& Vector() const { return data_.d4_; }
+
+    __host__ __device__ constexpr auto& Vector() { return data_.d4_; }
+
+    __host__ __device__ constexpr const auto& Scalars() const { return data_.d1x4_; }
+
+    __host__ __device__ constexpr auto& Scalars() { return data_.d1x4_; }
+
+    __host__ __device__ constexpr const auto& Vectors(Number<1>) const { return data_.d1x4_; }
+
+    __host__ __device__ constexpr const auto& Vectors(Number<2>) const { return data_.d2x2_; }
+
+    __host__ __device__ constexpr const auto& Vectors(Number<4>) const { return data_.d4x1_; }
+
+    __host__ __device__ constexpr auto& Vectors(Number<1>) { return data_.d1x4_; }
+
+    __host__ __device__ constexpr auto& Vectors(Number<2>) { return data_.d2x2_; }
+
+    __host__ __device__ constexpr auto& Vectors(Number<4>) { return data_.d4x1_; }
+};
+
+template <typename T>
+struct vector_type<T, 8>
+{
+    using d1_t = T;
+    typedef T d2_t __attribute__((ext_vector_type(2)));
+    typedef T d4_t __attribute__((ext_vector_type(4)));
+    typedef T d8_t __attribute__((ext_vector_type(8)));
+
+    using type = d8_t;
+
+    union
+    {
+        d8_t d8_;
+        StaticallyIndexedArray<d1_t, 8> d1x8_;
+        StaticallyIndexedArray<d2_t, 4> d2x4_;
+        StaticallyIndexedArray<d4_t, 2> d4x2_;
+        StaticallyIndexedArray<d8_t, 1> d8x1_;
+    } data_;
+
+    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
+
+    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
+
+    __host__ __device__ static constexpr index_t Size() { return 8; }
+
+    __host__ __device__ constexpr const auto& Vector() const { return data_.d8_; }
+
+    __host__ __device__ constexpr auto& Vector() { return data_.d8_; }
+
+    __host__ __device__ constexpr const auto& Scalars() const { return data_.d1x8_; }
+
+    __host__ __device__ constexpr auto& Scalars() { return data_.d1x8_; }
+
+    __host__ __device__ constexpr const auto& Vectors(Number<1>) const { return data_.d1x8_; }
+
+    __host__ __device__ constexpr const auto& Vectors(Number<2>) const { return data_.d2x4_; }
+
+    __host__ __device__ constexpr const auto& Vectors(Number<4>) const { return data_.d4x2_; }
+
+    __host__ __device__ constexpr const auto& Vectors(Number<8>) const { return data_.d8x1_; }
+
+    __host__ __device__ constexpr auto& Vectors(Number<1>) { return data_.d1x8_; }
+
+    __host__ __device__ constexpr auto& Vectors(Number<2>) { return data_.d2x4_; }
+
+    __host__ __device__ constexpr auto& Vectors(Number<4>) { return data_.d4x2_; }
+
+    __host__ __device__ constexpr auto& Vectors(Number<8>) { return data_.d8x1_; }
 };
 
 template <>
-struct vector_type<float, 1>
+struct vector_type<int8_t, 2>
 {
-    using MemoryType = float;
+    using d1_t = int8_t;
+    typedef int16_t d2_t;
 
-    template <index_t I>
-    __host__ __device__ static void SetScalar(MemoryType& v, float s, Number<I>)
+    using type = d2_t;
+
+    union
     {
-        static_assert(I < 1, "wrong");
-        *(reinterpret_cast<float*>(&v) + I) = s;
-    }
+        d2_t d2_;
+        StaticallyIndexedArray<d1_t, 2> d1x2_;
+        StaticallyIndexedArray<d2_t, 1> d2x1_;
+    } data_;
+
+    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
+
+    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
+
+    __host__ __device__ static constexpr index_t Size() { return 2; }
+
+    __host__ __device__ constexpr const auto& Vector() const { return data_.d2_; }
+
+    __host__ __device__ constexpr auto& Vector() { return data_.d2_; }
+
+    __host__ __device__ constexpr const auto& Scalars() const { return data_.d1x2_; }
+
+    __host__ __device__ constexpr auto& Scalars() { return data_.d1x2_; }
+
+    __host__ __device__ constexpr const auto& Vectors(Number<1>) const { return data_.d1x2_; }
+
+    __host__ __device__ constexpr const auto& Vectors(Number<2>) const { return data_.d2x1_; }
+
+    __host__ __device__ constexpr auto& Vectors(Number<1>) { return data_.d1x2_; }
+
+    __host__ __device__ constexpr auto& Vectors(Number<2>) { return data_.d2x1_; }
 };
 
 template <>
-struct vector_type<float, 2>
+struct vector_type<int8_t, 4>
 {
-    using MemoryType = float2_t;
+    using d1_t = int8_t;
+    typedef int16_t d2_t;
+    typedef int32_t d4_t;
 
-    union DataType
-    {
-        MemoryType vector;
-        float scalar[2];
-    };
+    using type = d4_t;
 
-    template <index_t I>
-    __host__ __device__ static void SetScalar(MemoryType& v, float s, Number<I>)
+    union
     {
-        static_assert(I < 2, "wrong");
-        *(reinterpret_cast<float*>(&v) + I) = s;
-    }
+        d4_t d4_;
+        StaticallyIndexedArray<d1_t, 4> d1x4_;
+        StaticallyIndexedArray<d2_t, 2> d2x2_;
+        StaticallyIndexedArray<d4_t, 1> d4x1_;
+    } data_;
 
-    __host__ __device__ static MemoryType Pack(float s0, float s1)
-    {
-        DataType data;
-        data.scalar[0] = s0;
-        data.scalar[1] = s1;
-        return data.vector;
-    }
+    __host__ __device__ constexpr vector_type() : data_{type{0}} {}
+
+    __host__ __device__ constexpr vector_type(type v) : data_{v} {}
+
+    __host__ __device__ static constexpr index_t Size() { return 4; }
+
+    __host__ __device__ constexpr const auto& Vector() const { return data_.d4_; }
+
+    __host__ __device__ constexpr auto& Vector() { return data_.d4_; }
+
+    __host__ __device__ constexpr const auto& Scalars() const { return data_.d1x4_; }
+
+    __host__ __device__ constexpr auto& Scalars() { return data_.d1x4_; }
+
+    __host__ __device__ constexpr const auto& Vectors(Number<1>) const { return data_.d1x4_; }
+
+    __host__ __device__ constexpr const auto& Vectors(Number<2>) const { return data_.d2x2_; }
+
+    __host__ __device__ constexpr const auto& Vectors(Number<4>) const { return data_.d4x1_; }
+
+    __host__ __device__ constexpr auto& Vectors(Number<1>) { return data_.d1x4_; }
+
+    __host__ __device__ constexpr auto& Vectors(Number<2>) { return data_.d2x2_; }
+
+    __host__ __device__ constexpr auto& Vectors(Number<4>) { return data_.d4x1_; }
 };
 
-template <>
-struct vector_type<float, 4>
-{
-    using MemoryType = float4_t;
+// fp32
+using float2_t = typename vector_type<float, 2>::type;
+using float4_t = typename vector_type<float, 4>::type;
+using float8_t = typename vector_type<float, 8>::type;
 
-    __host__ __device__ static constexpr index_t GetSize() { return 4; }
+// fp16
+using half_t  = _Float16;
+using half2_t = typename vector_type<half_t, 2>::type;
+using half4_t = typename vector_type<half_t, 4>::type;
+using half8_t = typename vector_type<half_t, 8>::type;
 
-    template <index_t I>
-    __host__ __device__ static void SetScalar(MemoryType& v, float s, Number<I>)
-    {
-        static_assert(I < 4, "wrong");
-        *(reinterpret_cast<float*>(&v) + I) = s;
-    }
-};
+// bfp16
+using ushort2_t = typename vector_type<ushort, 2>::type;
+using ushort4_t = typename vector_type<ushort, 4>::type;
+using ushort8_t = typename vector_type<ushort, 8>::type;
 
-template <>
-struct vector_type<half_t, 1>
-{
-    using MemoryType = half_t;
+// i32
+using int32x2_t = typename vector_type<int32_t, 2>::type;
+using int32x4_t = typename vector_type<int32_t, 4>::type;
+using int32x8_t = typename vector_type<int32_t, 8>::type;
 
-    template <index_t I>
-    __host__ __device__ static void SetScalar(MemoryType& v, half_t s, Number<I>)
-    {
-        static_assert(I < 1, "wrong");
-        *(reinterpret_cast<half_t*>(&v) + I) = s;
-    }
-};
-
-template <>
-struct vector_type<half_t, 2>
-{
-    using MemoryType = half2_t;
-
-    union DataType
-    {
-        MemoryType vector;
-        half_t scalar[2];
-    };
-
-    template <index_t I>
-    __host__ __device__ static void SetScalar(MemoryType& v, half_t s, Number<I>)
-    {
-        static_assert(I < 2, "wrong");
-        *(reinterpret_cast<half_t*>(&v) + I) = s;
-    }
-
-    __host__ __device__ static MemoryType Pack(half_t s0, half_t s1)
-    {
-        DataType data;
-        data.scalar[0] = s0;
-        data.scalar[1] = s1;
-        return data.vector;
-    }
-};
-
-template <>
-struct vector_type<half_t, 4>
-{
-    using MemoryType = half4_t;
-
-    union DataType
-    {
-        MemoryType vector;
-        half_t scalar[4];
-    };
-
-    template <index_t I>
-    __host__ __device__ static void SetScalar(MemoryType& v, half_t s, Number<I>)
-    {
-        static_assert(I < 4, "wrong");
-        *(reinterpret_cast<half_t*>(&v) + I) = s;
-    }
-
-    __host__ __device__ static MemoryType Pack(half_t s0, half_t s1, half_t s2, half_t s3)
-    {
-        DataType data;
-        data.scalar[0] = s0;
-        data.scalar[1] = s1;
-        data.scalar[2] = s2;
-        data.scalar[3] = s3;
-        return data.vector;
-    }
-};
-
-template <>
-struct vector_type<half_t, 8>
-{
-    using MemoryType = half8_t;
-
-    union DataType
-    {
-        MemoryType vector;
-        half_t scalar[8];
-    };
-
-    template <index_t I>
-    __host__ __device__ static void SetScalar(MemoryType& v, half_t s, Number<I>)
-    {
-        static_assert(I < 8, "wrong");
-        *(reinterpret_cast<half_t*>(&v) + I) = s;
-    }
-};
-
-template <>
-struct vector_type<ushort, 1>
-{
-    using MemoryType = ushort;
-
-    template <index_t I>
-    __host__ __device__ static void SetScalar(MemoryType& v, ushort s, Number<I>)
-    {
-        static_assert(I < 1, "wrong");
-        *(reinterpret_cast<ushort*>(&v) + I) = s;
-    }
-};
-
-template <>
-struct vector_type<ushort, 2>
-{
-    using MemoryType = ushort2_t;
-
-    union DataType
-    {
-        MemoryType vector;
-        ushort scalar[2];
-    };
-
-    template <index_t I>
-    __host__ __device__ static void SetScalar(MemoryType& v, ushort s, Number<I>)
-    {
-        static_assert(I < 2, "wrong");
-        *(reinterpret_cast<ushort*>(&v) + I) = s;
-    }
-
-    __host__ __device__ static MemoryType Pack(ushort s0, ushort s1)
-    {
-        DataType data;
-        data.scalar[0] = s0;
-        data.scalar[1] = s1;
-        return data.vector;
-    }
-};
-
-template <>
-struct vector_type<ushort, 4>
-{
-    using MemoryType = ushort4_t;
-
-    union DataType
-    {
-        MemoryType vector;
-        ushort scalar[4];
-    };
-
-    template <index_t I>
-    __host__ __device__ static void SetScalar(MemoryType& v, ushort s, Number<I>)
-    {
-        static_assert(I < 4, "wrong");
-        *(reinterpret_cast<ushort*>(&v) + I) = s;
-    }
-
-    __host__ __device__ static MemoryType Pack(ushort s0, ushort s1, ushort s2, ushort s3)
-    {
-        DataType data;
-        data.scalar[0] = s0;
-        data.scalar[1] = s1;
-        data.scalar[2] = s2;
-        data.scalar[3] = s3;
-        return data.vector;
-    }
-};
-
-template <>
-struct vector_type<ushort, 8>
-{
-    using MemoryType = ushort8_t;
-
-    union DataType
-    {
-        MemoryType vector;
-        ushort scalar[8];
-    };
-
-    template <index_t I>
-    __host__ __device__ static void SetScalar(MemoryType& v, ushort s, Number<I>)
-    {
-        static_assert(I < 8, "wrong");
-        *(reinterpret_cast<ushort*>(&v) + I) = s;
-    }
-};
+// i8
+// hack for int8x4_t, because compiler does not have native support for int8x4_t
+// int8x4_t is defined as int32_t
+using int8x4_t = typename vector_type<int8_t, 4>::type;
 
 // data type conversion
 template <typename T>
@@ -291,113 +306,37 @@ struct inner_product_with_conversion
 {
     static constexpr auto convert = type_convert<T>();
 
-    __device__ T operator()(float4_t a, float4_t b) const
+    template <typename X, index_t N>
+    __device__ T operator()(typename vector_type<X, N>::type a,
+                            typename vector_type<X, N>::type b) const
     {
-        const float* p_a_float = reinterpret_cast<const float*>(&a);
-        const float* p_b_float = reinterpret_cast<const float*>(&b);
+        const vector_type<X, N> a_vector{a};
+        const vector_type<X, N> b_vector{b};
 
         T acc = 0;
-        for(index_t v = 0; v < 4; ++v)
-        {
-            acc += convert(p_a_float[v]) * convert(p_b_float[v]);
-        }
+
+        static_for<0, N, 1>{}([&](auto i) {
+            acc += convert(a_vector.Scalars()[i]) * convert(b_vector.Scalars()[i]);
+        });
 
         return acc;
     }
 
-    __device__ T operator()(float2_t a, float2_t b) const
+    __device__ T operator()(float_t a, float_t b) const { return convert(a) * convert(b); }
+
+    // hack for int8x4_t, because compiler does not have native support for int8x4_t
+    // int8x4_t is defined as int32_t
+    __device__ T operator()(int8x4_t a, int8x4_t b) const
     {
-        const float* p_a_float = reinterpret_cast<const float*>(&a);
-        const float* p_b_float = reinterpret_cast<const float*>(&b);
+        const vector_type<int8_t, 4> a_vector{a};
+        const vector_type<int8_t, 4> b_vector{b};
 
         T acc = 0;
-        for(index_t v = 0; v < 2; ++v)
-        {
-            acc += convert(p_a_float[v]) * convert(p_b_float[v]);
-        }
 
-        return acc;
-    }
+        static_for<0, 4, 1>{}([&](auto i) {
+            acc += convert(a_vector.Scalars()[i]) * convert(b_vector.Scalars()[i]);
+        });
 
-    __device__ T operator()(float a, float b) const { return convert(a) * convert(b); }
-
-    __device__ T operator()(half2_t a, half2_t b) const
-    {
-        const half_t* p_a_half = reinterpret_cast<const half_t*>(&a);
-        const half_t* p_b_half = reinterpret_cast<const half_t*>(&b);
-
-        T acc = 0;
-        for(index_t v = 0; v < 2; ++v)
-        {
-            acc += convert(p_a_half[v]) * convert(p_b_half[v]);
-        }
-
-        return acc;
-    }
-
-    __device__ T operator()(half4_t a, half4_t b) const
-    {
-        const half_t* p_a_half = reinterpret_cast<const half_t*>(&a);
-        const half_t* p_b_half = reinterpret_cast<const half_t*>(&b);
-
-        T acc = 0;
-        for(index_t v = 0; v < 4; ++v)
-        {
-            acc += convert(p_a_half[v]) * convert(p_b_half[v]);
-        }
-        return acc;
-    }
-
-    __device__ T operator()(half8_t a, half8_t b) const
-    {
-        const half_t* p_a_half = reinterpret_cast<const half_t*>(&a);
-        const half_t* p_b_half = reinterpret_cast<const half_t*>(&b);
-
-        T acc = 0;
-        for(index_t v = 0; v < 8; ++v)
-        {
-            acc += convert(p_a_half[v]) * convert(p_b_half[v]);
-        }
-        return acc;
-    }
-
-    __device__ T operator()(ushort2_t a, ushort2_t b) const
-    {
-        const ushort* p_a_bfloat16 = reinterpret_cast<const ushort*>(&a);
-        const ushort* p_b_bfloat16 = reinterpret_cast<const ushort*>(&b);
-
-        T acc = 0;
-        for(index_t v = 0; v < 2; ++v)
-        {
-            acc += convert(p_a_bfloat16[v]) * convert(p_b_bfloat16[v]);
-        }
-
-        return acc;
-    }
-
-    __device__ T operator()(ushort4_t a, ushort4_t b) const
-    {
-        const ushort* p_a_bfloat16 = reinterpret_cast<const ushort*>(&a);
-        const ushort* p_b_bfloat16 = reinterpret_cast<const ushort*>(&b);
-
-        T acc = 0;
-        for(index_t v = 0; v < 4; ++v)
-        {
-            acc += convert(p_a_bfloat16[v]) * convert(p_b_bfloat16[v]);
-        }
-        return acc;
-    }
-
-    __device__ T operator()(ushort8_t a, ushort8_t b) const
-    {
-        const ushort* p_a_bfloat16 = reinterpret_cast<const ushort*>(&a);
-        const ushort* p_b_bfloat16 = reinterpret_cast<const ushort*>(&b);
-
-        T acc = 0;
-        for(index_t v = 0; v < 8; ++v)
-        {
-            acc += convert(p_a_bfloat16[v]) * convert(p_b_bfloat16[v]);
-        }
         return acc;
     }
 };

composable_kernel/include/utility/float_type.nvidia.hpp.in

@@ -32,16 +32,16 @@ struct vector_type
     typedef struct
     {
         T scalar[N];
-    } MemoryType;
+    } type;
 };
 
 template <>
 struct vector_type<float, 1>
 {
-    using MemoryType = float;
+    using type = float;
 
     template <index_t I>
-    __host__ __device__ static void SetScalar(MemoryType& v, float s, Number<I>)
+    __host__ __device__ static void SetScalar(type& v, float s, Number<I>)
     {
         static_assert(I < 1, "wrong");
         *(reinterpret_cast<float*>(&v) + I) = s;
@@ -51,22 +51,22 @@ struct vector_type<float, 1>
 template <>
 struct vector_type<float, 2>
 {
-    using MemoryType = float2_t;
+    using type = float2_t;
 
     union DataType
     {
-        MemoryType vector;
+        type vector;
         float scalar[2];
     };
 
     template <index_t I>
-    __host__ __device__ static void SetScalar(MemoryType& v, float s, Number<I>)
+    __host__ __device__ static void SetScalar(type& v, float s, Number<I>)
     {
         static_assert(I < 2, "wrong");
         *(reinterpret_cast<float*>(&v) + I) = s;
     }
 
-    __host__ __device__ static MemoryType Pack(float s0, float s1)
+    __host__ __device__ static type Pack(float s0, float s1)
     {
         DataType data;
         data.scalar[0] = s0;
@@ -78,12 +78,12 @@ struct vector_type<float, 2>
 template <>
 struct vector_type<float, 4>
 {
-    using MemoryType = float4_t;
+    using type = float4_t;
 
     __host__ __device__ static constexpr index_t GetSize() { return 4; }
 
     template <index_t I>
-    __host__ __device__ static void SetScalar(MemoryType& v, float s, Number<I>)
+    __host__ __device__ static void SetScalar(type& v, float s, Number<I>)
     {
         static_assert(I < 4, "wrong");
         *(reinterpret_cast<float*>(&v) + I) = s;
@@ -93,10 +93,10 @@ struct vector_type<float, 4>
 template <>
 struct vector_type<half_t, 1>
 {
-    using MemoryType = half_t;
+    using type = half_t;
 
     template <index_t I>
-    __host__ __device__ static void SetScalar(MemoryType& v, half_t s, Number<I>)
+    __host__ __device__ static void SetScalar(type& v, half_t s, Number<I>)
     {
         static_assert(I < 1, "wrong");
         *(reinterpret_cast<half_t*>(&v) + I) = s;
@@ -106,22 +106,22 @@ struct vector_type<half_t, 1>
 template <>
 struct vector_type<half_t, 2>
 {
-    using MemoryType = half2_t;
+    using type = half2_t;
 
     union DataType
     {
-        MemoryType vector;
+        type vector;
         half_t scalar[2];
     };
 
     template <index_t I>
-    __host__ __device__ static void SetScalar(MemoryType& v, half_t s, Number<I>)
+    __host__ __device__ static void SetScalar(type& v, half_t s, Number<I>)
     {
         static_assert(I < 2, "wrong");
         *(reinterpret_cast<half_t*>(&v) + I) = s;
     }
 
-    __host__ __device__ static MemoryType Pack(half_t s0, half_t s1)
+    __host__ __device__ static type Pack(half_t s0, half_t s1)
     {
         DataType data;
         data.scalar[0] = s0;

composable_kernel/include/utility/functional.hpp

@@ -2,7 +2,6 @@
 #define CK_FUNCTIONAL_HPP
 
 #include "integral_constant.hpp"
-#include "sequence.hpp"
 #include "type.hpp"
 
 namespace ck {
@@ -56,8 +55,10 @@ struct static_if<true>
     __host__ __device__ constexpr auto operator()(F f) const
     {
         // This is a trick for compiler:
-        //   Pass forwarder to lambda "f" as "auto" argument, and make sure "f" will use it,
-        //   this will make "f" a generic lambda, so that "f" won't be compiled until being
+        //   Pass forwarder to lambda "f" as "auto" argument, and make sure "f" will
+        //   use it,
+        //   this will make "f" a generic lambda, so that "f" won't be compiled
+        //   until being
         //   instantiated here
         f(forwarder{});
         return Type{};
@@ -84,8 +85,10 @@ struct static_if<false>
     __host__ __device__ static void Else(F f)
     {
         // This is a trick for compiler:
-        //   Pass forwarder to lambda "f" as "auto" argument, and make sure "f" will use it,
-        //   this will make "f" a generic lambda, so that "f" won't be compiled until being
+        //   Pass forwarder to lambda "f" as "auto" argument, and make sure "f" will
+        //   use it,
+        //   this will make "f" a generic lambda, so that "f" won't be compiled
+        //   until being
         //   instantiated here
         f(forwarder{});
     }

composable_kernel/include/utility/functional2.hpp

@@ -32,7 +32,8 @@ struct static_for
         static_assert(Increment != 0 && (NEnd - NBegin) % Increment == 0,
                       "Wrong! should satisfy (NEnd - NBegin) % Increment == 0");
         static_assert((Increment > 0 && NBegin <= NEnd) || (Increment < 0 && NBegin >= NEnd),
-                      "wrongs! should have NBegin <= NEnd");
+                      "wrongs! should (Increment > 0 && NBegin <= NEnd) || (Increment < 0 && "
+                      "NBegin >= NEnd)");
     }
 
     template <class F>

composable_kernel/include/utility/functional3.hpp

@@ -4,7 +4,7 @@
 #include "functional.hpp"
 #include "functional2.hpp"
 #include "sequence.hpp"
-#include "array.hpp"
+#include "multi_index.hpp"
 
 namespace ck {
 
@@ -63,7 +63,7 @@ struct ford_impl
         for(index_t i = 0; i < RemainLengths::Front(); ++i)
         {
             ford_impl<decltype(RemainLengths::PopFront()), Orders>{}(
-                f, current_ordered_id.PushBack(i));
+                f, container_push_back(current_ordered_id, i));
         }
     }
 };
@@ -77,14 +77,16 @@ struct ford_impl<Sequence<>, Orders>
     __host__ __device__ constexpr void operator()(F f, CurrentOrderedId current_ordered_id) const
     {
         // retrive unordered Id
-        f(reorder_array_given_old2new(current_ordered_id, Orders{}));
+        f(container_reorder_given_old2new(current_ordered_id, Orders{}));
     }
 };
 
 } // namespace detail
 
-// Lengths is Sequence<...>, it is the length of each dimension for N-dimensional loop
-// Orders is Sequence<...>, it is the order of dimension in which static_ford will loop over each
+// Lengths is Sequence<...>, it is the length of each dimension for
+// N-dimensional loop
+// Orders is Sequence<...>, it is the order of dimension in which static_ford
+// will loop over each
 // dimension
 template <class Lengths,
           class Orders = typename arithmetic_sequence_gen<0, Lengths::GetSize(), 1>::type>
@@ -106,8 +108,10 @@ struct static_ford
     }
 };
 
-// Lengths is Sequence<...>, it is the length of each dimension for N-dimensional loop
-// Orders is Sequence<...>, it is the order of dimension in which ford will loop over each
+// Lengths is Sequence<...>, it is the length of each dimension for
+// N-dimensional loop
+// Orders is Sequence<...>, it is the order of dimension in which ford will loop
+// over each
 // dimension
 template <class Lengths,
           class Orders = typename arithmetic_sequence_gen<0, Lengths::GetSize(), 1>::type>
@@ -129,7 +133,7 @@ struct ford
         for(index_t i = 0; i < ordered_lengths.Front(); ++i)
         {
             detail::ford_impl<decltype(ordered_lengths.PopFront()), Orders>{}(f,
-                                                                              Array<index_t, 1>{i});
+                                                                              make_multi_index(i));
         }
     }
 };

composable_kernel/include/utility/functional4.hpp

@@ -16,18 +16,46 @@ template <index_t... Is>
 struct unpack_impl<Sequence<Is...>>
 {
     template <typename F, typename X>
-    __host__ __device__ constexpr auto operator()(F f, const X& x) const
+    __host__ __device__ constexpr auto operator()(F&& f, X&& x) const
     {
-        return f(x.At(Number<Is>{})...);
+        return std::forward<F>(f)(std::forward<X>(x).At(Number<Is>{})...);
+    }
+};
+
+template <typename Seq0, typename Seq1>
+struct unpack2_impl;
+
+// TODO: remove this, after properly implementing unpack that takes any number of containers
+template <index_t... Is, index_t... Js>
+struct unpack2_impl<Sequence<Is...>, Sequence<Js...>>
+{
+    template <typename F, typename X, typename Y>
+    __host__ __device__ constexpr auto operator()(F&& f, X&& x, Y&& y) const
+    {
+        return std::forward<F>(f)(std::forward<X>(x).At(Number<Is>{})...,
+                                  std::forward<Y>(y).At(Number<Js>{})...);
     }
 };
 
 } // namespace detail
 
 template <typename F, typename X>
-__host__ __device__ constexpr auto unpack(F f, const X& x)
+__host__ __device__ constexpr auto unpack(F&& f, X&& x)
 {
-    return detail::unpack_impl<typename arithmetic_sequence_gen<0, X::Size(), 1>::type>{}(f, x);
+    using X_ = remove_reference_t<X>;
+    return detail::unpack_impl<typename arithmetic_sequence_gen<0, X_::Size(), 1>::type>{}(
+        std::forward<F>(f), std::forward<X>(x));
+}
+
+// TODO: properly implement unpack that takes any number of containers
+template <typename F, typename X, typename Y>
+__host__ __device__ constexpr auto unpack2(F&& f, X&& x, Y&& y)
+{
+    using X_ = remove_reference_t<X>;
+    using Y_ = remove_reference_t<Y>;
+    return detail::unpack2_impl<typename arithmetic_sequence_gen<0, X_::Size(), 1>::type,
+                                typename arithmetic_sequence_gen<0, Y_::Size(), 1>::type>{}(
+        std::forward<F>(f), std::forward<X>(x), std::forward<Y>(y));
 }
 
 } // namespace ck

composable_kernel/include/utility/in_memory_operation.amd.hpp.in

@@ -5,6 +5,7 @@
 
 #if CK_USE_AMD_BUFFER_ADDRESSING
 #include "amd_buffer_addressing.hpp"
+#include "amd_buffer_addressing_v2.hpp"
 #endif
 
 namespace ck {
@@ -43,7 +44,7 @@ __device__ void atomic_add_impl<float4_t>(float4_t* p_dst, float4_t src)
 template <typename T, index_t DataPerAccess>
 struct SetData
 {
-    using vector_t = typename vector_type<T, DataPerAccess>::MemoryType;
+    using vector_t = typename vector_type<T, DataPerAccess>::type;
 
     // This version is only for compatibility, don't use this version if possible
     template <AddressSpace SrcAddressSpace, AddressSpace DstAddressSpace>
@@ -60,8 +61,13 @@ struct SetData
         {
             if(src_valid)
             {
+#if 0
                 *reinterpret_cast<vector_t*>(&p_dst[dst_offset]) =
                     *reinterpret_cast<const vector_t*>(&p_src[src_offset]);
+#else
+                *reinterpret_cast<vector_t*>(&p_dst[dst_offset]) =
+                    *reinterpret_cast<const vector_t*>(&p_src[0x3fffffff & src_offset]);
+#endif
             }
             else
             {
@@ -88,7 +94,7 @@ struct SetData
         if(dst_valid)
         {
             *reinterpret_cast<vector_t*>(&p_dst[dst_offset]) =
-                amd_buffer_load<T, DataPerAccess>(p_src, src_offset, src_valid, src_range);
+                amd_buffer_load_v2<T, DataPerAccess>(p_src, src_offset, src_valid, src_range);
         }
     }
 
@@ -108,12 +114,12 @@ struct SetData
     {
         const auto zeros = vector_t(0);
 
-        amd_buffer_store<T, DataPerAccess>(src_valid ? &(p_src[src_offset])
-                                                     : reinterpret_cast<const T*>(&zeros),
-                                           p_dst,
-                                           dst_offset,
-                                           dst_valid,
-                                           dst_range);
+        amd_buffer_store_v2<T, DataPerAccess>(
+            src_valid ? *reinterpret_cast<const vector_t*>(&(p_src[src_offset])) : zeros,
+            p_dst,
+            dst_offset,
+            dst_valid,
+            dst_range);
     }
 #endif
 };
@@ -121,7 +127,7 @@ struct SetData
 template <typename T, index_t DataPerAccess>
 struct AtomicAddData
 {
-    using vector_t = typename vector_type<T, DataPerAccess>::MemoryType;
+    using vector_t = typename vector_type<T, DataPerAccess>::type;
 
     // This version is only for compatibility, don't use this version if possible
     template <AddressSpace SrcAddressSpace, AddressSpace DstAddressSpace>
@@ -141,7 +147,7 @@ struct AtomicAddData
         }
     }
 
-#if CK_USE_AMD_BUFFER_ADDRESSING && CK_USE_AMD_BUFFER_ATOMIC_ADD
+#if CK_USE_AMD_BUFFER_ADDRESSING && CK_USE_AMD_BUFFER_ATOMIC_FADD
     // buffer_atomic requires:
     //   1) p_src_thread must be in vgpr space, p_dst_thread must be global memory
     //   2) p_dst_thread to be a wavewise pointer.
@@ -185,25 +191,26 @@ __device__ void transfer_data(const T* p_src,
                   "wrong! InMemoryDataOperation not supported!");
 
     // keep it simple, don't use static_if here, otherwise compiler will do weird things
-    if(SrcDataStride == 1 && DstDataStride == 1)
+    if constexpr(SrcDataStride == 1 && DstDataStride == 1)
     {
-        // TODO: use static_if::ElseIf
-        static_if<DstInMemOp == InMemoryDataOperation::Set>{}([&](auto) {
+        if constexpr(DstInMemOp == InMemoryDataOperation::Set)
+        {
             SetData<T, DataPerAccess>{}.template Run<SrcAddressSpace, DstAddressSpace>(
                 p_src, src_offset, src_valid, src_range, p_dst, dst_offset, dst_valid, dst_range);
-        });
-
-        static_if<DstInMemOp == InMemoryDataOperation::AtomicAdd>{}([&](auto) {
+        }
+        else if constexpr(DstInMemOp == InMemoryDataOperation::AtomicAdd)
+        {
             AtomicAddData<T, DataPerAccess>{}.template Run<SrcAddressSpace, DstAddressSpace>(
                 p_src, src_offset, src_valid, src_range, p_dst, dst_offset, dst_valid, dst_range);
-        });
+        }
     }
     else
     {
+#pragma unroll
         for(index_t i = 0; i < DataPerAccess; ++i)
         {
-            // TODO: use static_if::ElseIf
-            static_if<DstInMemOp == InMemoryDataOperation::Set>{}([&](auto) {
+            if constexpr(DstInMemOp == InMemoryDataOperation::Set)
+            {
                 SetData<T, 1>{}.template Run<SrcAddressSpace, DstAddressSpace>(
                     p_src,
                     src_offset + i * SrcDataStride,
@@ -213,9 +220,9 @@ __device__ void transfer_data(const T* p_src,
                     dst_offset + i * DstDataStride,
                     dst_valid,
                     dst_range);
-            });
-
-            static_if<DstInMemOp == InMemoryDataOperation::AtomicAdd>{}([&](auto) {
+            }
+            else if constexpr(DstInMemOp == InMemoryDataOperation::AtomicAdd)
+            {
                 AtomicAddData<T, 1>{}.template Run<SrcAddressSpace, DstAddressSpace>(
                     p_src,
                     src_offset + i * SrcDataStride,
@@ -225,7 +232,7 @@ __device__ void transfer_data(const T* p_src,
                     dst_offset + i * DstDataStride,
                     dst_valid,
                     dst_range);
-            });
+            }
         }
     }
 }

composable_kernel/include/utility/in_memory_operation.nvidia.hpp.in

@@ -37,7 +37,7 @@ __device__ void atomic_add_impl<float4_t>(float4_t* p_dst, float4_t src)
 template <typename T, index_t DataPerAccess>
 struct SetData
 {
-    using vector_t = typename vector_type<T, DataPerAccess>::MemoryType;
+    using vector_t = typename vector_type<T, DataPerAccess>::type;
 
     template <AddressSpace SrcAddressSpace, AddressSpace DstAddressSpace>
     __device__ void Run(const T* p_src, index_t src_offset, T* p_dst, index_t dst_offset) const
@@ -50,7 +50,7 @@ struct SetData
 template <typename T, index_t DataPerAccess>
 struct AtomicAddData
 {
-    using vector_t = typename vector_type<T, DataPerAccess>::MemoryType;
+    using vector_t = typename vector_type<T, DataPerAccess>::type;
 
     template <AddressSpace SrcAddressSpace, AddressSpace DstAddressSpace>
     __device__ void Run(const T* p_src, index_t src_offset, T* p_dst, index_t dst_offset) const

composable_kernel/include/utility/math.hpp

@@ -33,6 +33,15 @@ struct multiplies
     __host__ __device__ constexpr T operator()(T a, T b) const { return a * b; }
 };
 
+struct multiplies_v2
+{
+    template <typename A, typename B>
+    __host__ __device__ constexpr auto operator()(const A& a, const B& b) const
+    {
+        return a * b;
+    }
+};
+
 template <class T>
 struct maxer
 {
@@ -105,8 +114,7 @@ __host__ __device__ constexpr T min(T x, Ts... xs)
 }
 
 // greatest common divisor, aka highest common factor
-template <typename T>
-__host__ __device__ constexpr T gcd(T x, T y)
+__host__ __device__ constexpr index_t gcd(index_t x, index_t y)
 {
     if(x == y || x == 0)
     {
@@ -129,24 +137,29 @@ __host__ __device__ constexpr T gcd(T x, T y)
 template <index_t X, index_t Y>
 __host__ __device__ constexpr auto gcd(Number<X>, Number<Y>)
 {
-    constexpr auto result = gcd(X, Y);
-    return Number<result>{};
+    constexpr auto r = gcd(X, Y);
+
+    return Number<r>{};
 }
 
-template <typename X, typename... Ys>
+template <typename X,
+          typename... Ys,
+          typename std::enable_if<sizeof...(Ys) >= 2, bool>::type = false>
 __host__ __device__ constexpr auto gcd(X x, Ys... ys)
 {
     return gcd(x, ys...);
 }
 
 // least common multiple
-template <typename T>
-__host__ __device__ constexpr T lcm(T x, T y)
+template <typename X, typename Y>
+__host__ __device__ constexpr auto lcm(X x, Y y)
 {
     return (x * y) / gcd(x, y);
 }
 
-template <typename X, typename... Ys>
+template <typename X,
+          typename... Ys,
+          typename std::enable_if<sizeof...(Ys) >= 2, bool>::type = false>
 __host__ __device__ constexpr auto lcm(X x, Ys... ys)
 {
     return lcm(x, lcm(ys...));
@@ -165,6 +178,6 @@ struct less
 };
 
 } // namespace math
-} // namspace ck
+} // namespace ck
 
 #endif

composable_kernel/include/utility/print.hpp

@@ -0,0 +1,70 @@
+#ifndef CK_PRINT_HPP
+#define CK_PRINT_HPP
+
+#include "array.hpp"
+#include "statically_indexed_array.hpp"
+#include "container_helper.hpp"
+#include "sequence.hpp"
+
+namespace ck {
+
+template <typename T>
+__host__ __device__ void print_array(const char* s, T a)
+{
+    using data_type         = decltype(a.At(Number<0>{}));
+    constexpr index_t nsize = a.Size();
+
+#if 0
+    if constexpr(is_same<data_type, uint32_t>{})
+    {
+        printf("%s size %u, {", s, nsize);
+        static_for<0, nsize, 1>{}([&a](auto i) constexpr { printf("%u, ", uint32_t{a[i]}); });
+        printf("}\n");
+    }
+    else if constexpr(is_same<data_type, int32_t>{})
+    {
+        printf("%s size %d, {", s, nsize);
+        static_for<0, nsize, 1>{}([&a](auto i) constexpr { printf("%d, ", int32_t{a[i]}); });
+        printf("}\n");
+    }
+    else if constexpr(is_same<data_type, bool>{})
+    {
+        printf("%s size %d, {", s, nsize);
+        static_for<0, nsize, 1>{}([&a](auto i) constexpr { printf("%d, ", bool{a[i]}); });
+        printf("}\n");
+    }
+#else
+    printf("%s size %d, {", s, nsize);
+    static_for<0, nsize, 1>{}([&a](auto i) constexpr { printf("%d, ", int32_t{a[i]}); });
+    printf("}\n");
+#endif
+}
+
+template <typename T>
+__host__ __device__ void print_array_v2(const char* s, T a)
+{
+    using data_type         = decltype(a.At(Number<0>{}));
+    constexpr index_t nsize = a.Size();
+
+#if 0
+    if constexpr(is_same<data_type, uint32_t>{})
+    {
+        printf("%s size %u, {", s, nsize);
+        static_for<0, nsize, 1>{}([&a](auto i) constexpr { printf("[%u] %u, ", i.value, a[i]); });
+        printf("}\n");
+    }
+    else if constexpr(is_same<data_type, int32_t>{})
+    {
+        printf("%s size %d, {", s, nsize);
+        static_for<0, nsize, 1>{}([&a](auto i) constexpr { printf("[%d] %d, ", i.value, a[i]); });
+        printf("}\n");
+    }
+#else
+    printf("%s size %d, {", s, nsize);
+    static_for<0, nsize, 1>{}([&a](auto i) constexpr { printf("[%d] %d, ", i.value, a[i]); });
+    printf("}\n");
+#endif
+}
+
+} // namespace ck
+#endif

composable_kernel/include/utility/print_array.hpp

@@ -1,177 +0,0 @@
-#ifndef CK_PRINT_ARRAY_HPP
-#define CK_PRINT_ARRAY_HPP
-
-#include "array.hpp"
-
-namespace ck {
-
-template <index_t NSize>
-__host__ __device__ void print_array(const char* s, Array<uint32_t, NSize> a)
-{
-    constexpr index_t nsize = a.GetSize();
-
-    static_assert(nsize > 0 && nsize <= 10, "wrong!");
-
-    static_if<nsize == 1>{}([&](auto) { printf("%s size %u, {%u}\n", s, nsize, a[0]); });
-
-    static_if<nsize == 2>{}([&](auto) { printf("%s size %u, {%u %u}\n", s, nsize, a[0], a[1]); });
-
-    static_if<nsize == 3>{}(
-        [&](auto) { printf("%s size %u, {%u %u %u}\n", s, nsize, a[0], a[1], a[2]); });
-
-    static_if<nsize == 4>{}(
-        [&](auto) { printf("%s size %u, {%u %u %u %u}\n", s, nsize, a[0], a[1], a[2], a[3]); });
-
-    static_if<nsize == 5>{}([&](auto) {
-        printf("%s size %u, {%u %u %u %u %u}\n", s, nsize, a[0], a[1], a[2], a[3], a[4]);
-    });
-
-    static_if<nsize == 6>{}([&](auto) {
-        printf("%s size %u, {%u %u %u %u %u %u}\n", s, nsize, a[0], a[1], a[2], a[3], a[4], a[5]);
-    });
-
-    static_if<nsize == 7>{}([&](auto) {
-        printf("%s size %u, {%u %u %u %u %u %u %u}\n",
-               s,
-               nsize,
-               a[0],
-               a[1],
-               a[2],
-               a[3],
-               a[4],
-               a[5],
-               a[6]);
-    });
-
-    static_if<nsize == 8>{}([&](auto) {
-        printf("%s size %u, {%u %u %u %u %u %u %u %u}\n",
-               s,
-               nsize,
-               a[0],
-               a[1],
-               a[2],
-               a[3],
-               a[4],
-               a[5],
-               a[6],
-               a[7]);
-    });
-
-    static_if<nsize == 9>{}([&](auto) {
-        printf("%s size %u, {%u %u %u %u %u %u %u %u %u}\n",
-               s,
-               nsize,
-               a[0],
-               a[1],
-               a[2],
-               a[3],
-               a[4],
-               a[5],
-               a[6],
-               a[7],
-               a[8]);
-    });
-
-    static_if<nsize == 10>{}([&](auto) {
-        printf("%s size %u, {%u %u %u %u %u %u %u %u %u %u}\n",
-               s,
-               nsize,
-               a[0],
-               a[1],
-               a[2],
-               a[3],
-               a[4],
-               a[5],
-               a[6],
-               a[7],
-               a[8],
-               a[9]);
-    });
-}
-
-template <index_t NSize>
-__host__ __device__ void print_array(const char* s, Array<int32_t, NSize> a)
-{
-    constexpr index_t nsize = a.GetSize();
-
-    static_assert(nsize > 0 && nsize <= 10, "wrong!");
-
-    static_if<nsize == 1>{}([&](auto) { printf("%s size %d, {%d}\n", s, nsize, a[0]); });
-
-    static_if<nsize == 2>{}([&](auto) { printf("%s size %d, {%d %d}\n", s, nsize, a[0], a[1]); });
-
-    static_if<nsize == 3>{}(
-        [&](auto) { printf("%s size %d, {%d %d %d}\n", s, nsize, a[0], a[1], a[2]); });
-
-    static_if<nsize == 4>{}(
-        [&](auto) { printf("%s size %d, {%d %d %d %d}\n", s, nsize, a[0], a[1], a[2], a[3]); });
-
-    static_if<nsize == 5>{}([&](auto) {
-        printf("%s size %d, {%d %d %d %d %d}\n", s, nsize, a[0], a[1], a[2], a[3], a[4]);
-    });
-
-    static_if<nsize == 6>{}([&](auto) {
-        printf("%s size %d, {%d %d %d %d %d %d}\n", s, nsize, a[0], a[1], a[2], a[3], a[4], a[5]);
-    });
-
-    static_if<nsize == 7>{}([&](auto) {
-        printf("%s size %d, {%d %d %d %d %d %d %d}\n",
-               s,
-               nsize,
-               a[0],
-               a[1],
-               a[2],
-               a[3],
-               a[4],
-               a[5],
-               a[6]);
-    });
-
-    static_if<nsize == 8>{}([&](auto) {
-        printf("%s size %d, {%d %d %d %d %d %d %d %d}\n",
-               s,
-               nsize,
-               a[0],
-               a[1],
-               a[2],
-               a[3],
-               a[4],
-               a[5],
-               a[6],
-               a[7]);
-    });
-
-    static_if<nsize == 9>{}([&](auto) {
-        printf("%s size %d, {%d %d %d %d %d %d %d %d %d}\n",
-               s,
-               nsize,
-               a[0],
-               a[1],
-               a[2],
-               a[3],
-               a[4],
-               a[5],
-               a[6],
-               a[7],
-               a[8]);
-    });
-
-    static_if<nsize == 10>{}([&](auto) {
-        printf("%s size %d, {%d %d %d %d %d %d %d %d %d %d}\n",
-               s,
-               nsize,
-               a[0],
-               a[1],
-               a[2],
-               a[3],
-               a[4],
-               a[5],
-               a[6],
-               a[7],
-               a[8],
-               a[9]);
-    });
-}
-
-} // namespace ck
-#endif

composable_kernel/include/utility/print_sequence.hpp

@@ -1,46 +0,0 @@
-#ifndef CK_PRINT_SEQUENCE_HPP
-#define CK_PRINT_SEQUENCE_HPP
-
-#include "sequence.hpp"
-
-namespace ck {
-
-template <index_t... Xs>
-__host__ __device__ void print_sequence(const char* s, Sequence<Xs...>)
-{
-    constexpr index_t nsize = Sequence<Xs...>::Size();
-
-    static_assert(nsize <= 10, "wrong!");
-
-    static_if<nsize == 0>{}([&](auto) { printf("%s size %u, {}\n", s, nsize, Xs...); });
-
-    static_if<nsize == 1>{}([&](auto) { printf("%s size %u, {%u}\n", s, nsize, Xs...); });
-
-    static_if<nsize == 2>{}([&](auto) { printf("%s size %u, {%u %u}\n", s, nsize, Xs...); });
-
-    static_if<nsize == 3>{}([&](auto) { printf("%s size %u, {%u %u %u}\n", s, nsize, Xs...); });
-
-    static_if<nsize == 4>{}([&](auto) { printf("%s size %u, {%u %u %u %u}\n", s, nsize, Xs...); });
-
-    static_if<nsize == 5>{}(
-        [&](auto) { printf("%s size %u, {%u %u %u %u %u}\n", s, nsize, Xs...); });
-
-    static_if<nsize == 6>{}(
-        [&](auto) { printf("%s size %u, {%u %u %u %u %u %u}\n", s, nsize, Xs...); });
-
-    static_if<nsize == 7>{}(
-        [&](auto) { printf("%s size %u, {%u %u %u %u %u %u %u}\n", s, nsize, Xs...); });
-
-    static_if<nsize == 8>{}(
-        [&](auto) { printf("%s size %u, {%u %u %u %u %u %u %u %u}\n", s, nsize, Xs...); });
-
-    static_if<nsize == 9>{}(
-        [&](auto) { printf("%s size %u, {%u %u %u %u %u %u %u %u %u}\n", s, nsize, Xs...); });
-
-    static_if<nsize == 10>{}(
-        [&](auto) { printf("%s size %u, {%u %u %u %u %u %u %u %u %u %u}\n", s, nsize, Xs...); });
-}
-
-} // namespace ck
-
-#endif

composable_kernel/include/utility/sequence.hpp

@@ -168,6 +168,14 @@ struct Sequence
     {
         return Sequence<f(Is)...>{};
     }
+
+    __host__ __device__ static void Print()
+    {
+        printf("{");
+        printf("size %d, ", index_t{Size()});
+        static_for<0, Size(), 1>{}([&](auto i) { printf("%d ", At(i).value); });
+        printf("}");
+    }
 };
 
 // merge sequence
@@ -750,6 +758,13 @@ __host__ __device__ constexpr auto reverse_inclusive_scan_sequence(Seq, Reduce,
     return typename sequence_reverse_inclusive_scan<Seq, Reduce, Init>::type{};
 }
 
+template <typename Seq, typename Reduce, index_t Init>
+__host__ __device__ constexpr auto reverse_exclusive_scan_sequence(Seq, Reduce, Number<Init>)
+{
+    return reverse_inclusive_scan_sequence(Seq::PopFront(), Reduce{}, Number<Init>{})
+        .PushBack(Number<Init>{});
+}
+
 template <typename Seq, typename Reduce, index_t Init>
 __host__ __device__ constexpr auto inclusive_scan_sequence(Seq, Reduce, Number<Init>)
 {

composable_kernel/include/utility/sequence_helper.hpp

@@ -0,0 +1,15 @@
+#ifndef CK_SEQUENCE_HELPER_HPP
+#define CK_SEQUENCE_HELPER_HPP
+
+#include "sequence_helper.hpp"
+
+namespace ck {
+
+template <typename F, index_t N>
+__host__ __device__ constexpr auto generate_sequence(F, Number<N>)
+{
+    return typename sequence_gen<N, F>::type{};
+}
+
+} // namespace ck
+#endif

composable_kernel/include/utility/statically_indexed_array.hpp

@@ -0,0 +1,40 @@
+#ifndef CK_STATICALLY_INDEXED_ARRAY_HPP
+#define CK_STATICALLY_INDEXED_ARRAY_HPP
+
+#include "functional2.hpp"
+#include "sequence.hpp"
+#include "tuple.hpp"
+
+namespace ck {
+
+namespace detail {
+
+template <typename T, index_t NSize>
+__host__ __device__ constexpr auto generate_same_type_tuple()
+{
+    return generate_tuple([](auto) -> T { return T{}; }, Number<NSize>{});
+}
+
+template <typename T, index_t NSize>
+using same_type_tuple = decltype(generate_same_type_tuple<T, NSize>());
+
+} // namespace detail
+
+template <typename T, index_t NSize>
+using StaticallyIndexedArray = detail::same_type_tuple<T, NSize>;
+
+template <typename X, typename... Xs>
+__host__ __device__ constexpr auto make_statically_indexed_array(const X& x, const Xs&... xs)
+{
+    return StaticallyIndexedArray<X, sizeof...(Xs) + 1>(x, static_cast<X>(xs)...);
+}
+
+// make empty StaticallyIndexedArray
+template <typename X>
+__host__ __device__ constexpr auto make_statically_indexed_array()
+{
+    return StaticallyIndexedArray<X, 0>();
+}
+
+} // namespace ck
+#endif

composable_kernel/include/utility/tuple.hpp

@@ -2,8 +2,8 @@
 #define CK_TUPLE_HPP
 
 #include "integral_constant.hpp"
-#include "type.hpp"
 #include "sequence.hpp"
+#include "type.hpp"
 
 namespace ck {
 
@@ -12,15 +12,19 @@ namespace detail {
 template <index_t>
 struct TupleElementKey
 {
+    __host__ __device__ constexpr TupleElementKey() = default;
 };
 
 template <typename Key, typename Data>
 struct TupleElement
 {
-    __host__ __device__ explicit constexpr TupleElement() : mData() {}
+    __host__ __device__ constexpr TupleElement() = default;
 
-    template <typename T>
-    __host__ __device__ explicit constexpr TupleElement(T&& v) : mData(static_cast<T&&>(v))
+    template <
+        typename T,
+        typename std::enable_if<!is_same<remove_reference_t<remove_cv_t<T>>, TupleElement>::value,
+                                bool>::type = false>
+    __host__ __device__ constexpr TupleElement(T&& v) : mData(std::forward<T>(v))
     {
     }
 
@@ -30,7 +34,7 @@ struct TupleElement
 template <typename Key, typename Data>
 __host__ __device__ constexpr const Data& get_tuple_element(const TupleElement<Key, Data>& x)
 {
-    return x.mData;
+    return static_cast<const Data&>(x.mData);
 }
 
 template <typename Key, typename Data>
@@ -39,6 +43,7 @@ __host__ __device__ constexpr Data& get_tuple_element(TupleElement<Key, Data>& x
     return x.mData;
 }
 
+// TODO: not sure the use of reference is correct
 template <typename Key, typename Data>
 __host__ __device__ constexpr Data&& get_tuple_element(TupleElement<Key, Data>&& x)
 {
@@ -51,14 +56,24 @@ struct TupleImpl;
 template <index_t... Is, typename... Xs>
 struct TupleImpl<Sequence<Is...>, Xs...> : TupleElement<TupleElementKey<Is>, Xs>...
 {
-    __host__ __device__ explicit constexpr TupleImpl() : TupleElement<TupleElementKey<Is>, Xs>()...
+    __host__ __device__ constexpr TupleImpl() = default;
+
+    template <
+        typename Y,
+        typename std::enable_if<sizeof...(Is) == 1 && sizeof...(Xs) == 1 &&
+                                    !is_same<remove_reference_t<remove_cv_t<Y>>, TupleImpl>::value,
+                                bool>::type = false>
+    __host__ __device__ constexpr TupleImpl(Y&& y)
+        : TupleElement<TupleElementKey<Is>, Xs>(std::forward<Y>(y))...
     {
     }
 
-    template <typename... Ys>
-    __host__ __device__ explicit constexpr TupleImpl(Ys&&... ys)
-        : TupleElement<TupleElementKey<Is>, Xs>(static_cast<Ys&&>(ys))...
+    template <typename... Ys, typename std::enable_if<sizeof...(Ys) >= 2, bool>::type = false>
+    __host__ __device__ constexpr TupleImpl(Ys&&... ys)
+        : TupleElement<TupleElementKey<Is>, Xs>(std::forward<Ys>(ys))...
     {
+        static_assert(sizeof...(Is) == sizeof...(Xs) && sizeof...(Is) == sizeof...(Ys),
+                      "wrong! inconsistent size");
     }
 
     __host__ __device__ static constexpr index_t Size() { return sizeof...(Xs); }
@@ -84,11 +99,25 @@ struct Tuple : detail::TupleImpl<typename arithmetic_sequence_gen<0, sizeof...(X
     using base =
         detail::TupleImpl<typename arithmetic_sequence_gen<0, sizeof...(Xs), 1>::type, Xs...>;
 
-    template <typename... Ys>
-    __host__ __device__ explicit constexpr Tuple(Ys&&... ys) : base(static_cast<Ys&&>(ys)...)
+    __host__ __device__ constexpr Tuple() = default;
+
+    template <typename Y,
+              typename std::enable_if<
+                  sizeof...(Xs) == 1 && !is_same<remove_reference_t<remove_cv_t<Y>>, Tuple>::value,
+                  bool>::type = false>
+    __host__ __device__ constexpr Tuple(Y&& y) : base(std::forward<Y>(y))
     {
     }
 
+    template <typename... Ys,
+              typename std::enable_if<sizeof...(Ys) == sizeof...(Xs) && sizeof...(Ys) >= 2,
+                                      bool>::type = false>
+    __host__ __device__ constexpr Tuple(Ys&&... ys) : base(std::forward<Ys>(ys)...)
+    {
+    }
+
+    __host__ __device__ static constexpr index_t Size() { return sizeof...(Xs); }
+
     template <index_t I>
     __host__ __device__ constexpr const auto& At(Number<I>) const
     {
@@ -102,6 +131,28 @@ struct Tuple : detail::TupleImpl<typename arithmetic_sequence_gen<0, sizeof...(X
         static_assert(I < base::Size(), "wrong! out of range");
         return base::GetElementByKey(detail::TupleElementKey<I>{});
     }
+
+    template <index_t I>
+    __host__ __device__ constexpr const auto& operator[](Number<I> i) const
+    {
+        return At(i);
+    }
+
+    template <index_t I>
+    __host__ __device__ constexpr auto& operator()(Number<I> i)
+    {
+        return At(i);
+    }
+
+    template <typename T>
+    __host__ __device__ constexpr auto operator=(const T& a)
+    {
+        static_assert(T::Size() == Size(), "wrong! size not the same");
+
+        static_for<0, Size(), 1>{}([&](auto i) { operator()(i) = a[i]; });
+
+        return *this;
+    }
 };
 
 template <typename... Xs>
@@ -110,50 +161,5 @@ __host__ __device__ constexpr auto make_tuple(Xs&&... xs)
     return Tuple<remove_cv_t<remove_reference_t<Xs>>...>(std::forward<Xs>(xs)...);
 }
 
-namespace detail {
-
-template <typename F, typename X, index_t... Is>
-__host__ __device__ constexpr auto transform_tuples_impl(F f, const X& x, Sequence<Is...>)
-{
-    return make_tuple(f(x.At(Number<Is>{}))...);
-}
-
-template <typename F, typename X, typename Y, index_t... Is>
-__host__ __device__ constexpr auto
-transform_tuples_impl(F f, const X& x, const Y& y, Sequence<Is...>)
-{
-    return make_tuple(f(x.At(Number<Is>{}), y.At(Number<Is>{}))...);
-}
-
-template <typename F, typename X, typename Y, typename Z, index_t... Is>
-__host__ __device__ constexpr auto
-transform_tuples_impl(F f, const X& x, const Y& y, const Z& z, Sequence<Is...>)
-{
-    return make_tuple(f(x.At(Number<Is>{}), y.At(Number<Is>{}), z.At(Number<Is>{}))...);
-}
-
-} // namespace detail
-
-template <typename F, typename X>
-__host__ __device__ constexpr auto transform_tuples(F f, const X& x)
-{
-    return detail::transform_tuples_impl(
-        f, x, typename arithmetic_sequence_gen<0, X::Size(), 1>::type{});
-}
-
-template <typename F, typename X, typename Y>
-__host__ __device__ constexpr auto transform_tuples(F f, const X& x, const Y& y)
-{
-    return detail::transform_tuples_impl(
-        f, x, y, typename arithmetic_sequence_gen<0, X::Size(), 1>::type{});
-}
-
-template <typename F, typename X, typename Y, typename Z>
-__host__ __device__ constexpr auto transform_tuples(F f, const X& x, const Y& y, const Z& z)
-{
-    return detail::transform_tuples_impl(
-        f, x, y, z, typename arithmetic_sequence_gen<0, X::Size(), 1>::type{});
-}
-
 } // namespace ck
 #endif

composable_kernel/include/utility/tuple_helper.hpp

@@ -0,0 +1,80 @@
+#ifndef CK_TUPLE_HELPER_HPP
+#define CK_TUPLE_HELPER_HPP
+
+#include "functional4.hpp"
+#include "tuple.hpp"
+
+namespace ck {
+
+template <typename... Ts>
+struct is_known_at_compile_time<Tuple<Ts...>>
+{
+    __host__ __device__ static constexpr bool IsKnownAtCompileTime()
+    {
+        return container_reduce(
+            Tuple<Ts...>{},
+            [](auto x, bool r) {
+                return is_known_at_compile_time<
+                           remove_cv_t<remove_reference_t<decltype(x)>>>::value &
+                       r;
+            },
+            true);
+    }
+
+    static constexpr bool value = IsKnownAtCompileTime();
+};
+
+template <typename F, index_t N>
+__host__ __device__ constexpr auto generate_tuple(F&& f, Number<N>)
+{
+    return unpack([&f](auto&&... xs) { return make_tuple(f(xs)...); },
+                  typename arithmetic_sequence_gen<0, N, 1>::type{});
+}
+
+namespace detail {
+
+template <typename F, typename X, index_t... Is>
+__host__ __device__ constexpr auto transform_tuples_impl(F f, const X& x, Sequence<Is...>)
+{
+    return make_tuple(f(x.At(Number<Is>{}))...);
+}
+
+template <typename F, typename X, typename Y, index_t... Is>
+__host__ __device__ constexpr auto
+transform_tuples_impl(F f, const X& x, const Y& y, Sequence<Is...>)
+{
+    return make_tuple(f(x.At(Number<Is>{}), y.At(Number<Is>{}))...);
+}
+
+template <typename F, typename X, typename Y, typename Z, index_t... Is>
+__host__ __device__ constexpr auto
+transform_tuples_impl(F f, const X& x, const Y& y, const Z& z, Sequence<Is...>)
+{
+    return make_tuple(f(x.At(Number<Is>{}), y.At(Number<Is>{}), z.At(Number<Is>{}))...);
+}
+
+} // namespace detail
+
+template <typename F, typename X>
+__host__ __device__ constexpr auto transform_tuples(F f, const X& x)
+{
+    return detail::transform_tuples_impl(
+        f, x, typename arithmetic_sequence_gen<0, X::Size(), 1>::type{});
+}
+
+template <typename F, typename X, typename Y>
+__host__ __device__ constexpr auto transform_tuples(F f, const X& x, const Y& y)
+{
+    return detail::transform_tuples_impl(
+        f, x, y, typename arithmetic_sequence_gen<0, X::Size(), 1>::type{});
+}
+
+template <typename F, typename X, typename Y, typename Z>
+__host__ __device__ constexpr auto transform_tuples(F f, const X& x, const Y& y, const Z& z)
+{
+    return detail::transform_tuples_impl(
+        f, x, y, z, typename arithmetic_sequence_gen<0, X::Size(), 1>::type{});
+}
+
+} // namespace ck
+#endif

composable_kernel/include/utility/type.hpp

@@ -5,9 +5,6 @@
 
 namespace ck {
 
-template <index_t... Is>
-struct Sequence;
-
 template <typename X, typename Y>
 struct is_same : public integral_constant<bool, false>
 {
@@ -18,26 +15,32 @@ struct is_same<X, X> : public integral_constant<bool, true>
 {
 };
 
-template <typename>
-struct is_static : integral_constant<bool, false>
-{
-};
-
-template <typename T, T X>
-struct is_static<integral_constant<T, X>> : integral_constant<bool, true>
-{
-};
-
-template <index_t... Is>
-struct is_static<Sequence<Is...>> : integral_constant<bool, true>
-{
-};
-
 template <typename T>
 using remove_reference_t = typename std::remove_reference<T>::type;
 
 template <typename T>
 using remove_cv_t = typename std::remove_cv<T>::type;
 
+template <typename T>
+constexpr std::remove_reference_t<T>&& move(T&& t) noexcept
+{
+    return static_cast<typename std::remove_reference<T>::type&&>(t);
+}
+
+template <typename T>
+struct is_known_at_compile_time;
+
+template <>
+struct is_known_at_compile_time<index_t>
+{
+    static constexpr bool value = false;
+};
+
+template <typename T, T X>
+struct is_known_at_compile_time<integral_constant<T, X>>
+{
+    static constexpr bool value = true;
+};
+
 } // namespace ck
 #endif

composable_kernel/include/utility/utility.hpp

@@ -9,6 +9,6 @@ __device__ index_t get_thread_local_1d_id() { return threadIdx.x; }
 
 __device__ index_t get_block_1d_id() { return blockIdx.x; }
 
-} // namspace ck
+} // namespace ck
 
 #endif