driver variale name

composable_kernel/include/tensor_operation/blockwise_gemm_xdlops.hpp

@@ -168,8 +168,8 @@ struct BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1
             make_tuple(make_pass_through_transform(G),
                        make_unmerge_transform(make_tuple(MRepeat, MWaves, MPerXDL)),
                        make_unmerge_transform(make_tuple(NRepeat, NWaves, NPerXDL))),
-            make_tuple(Sequence<0>{}, Sequence<1>{}),
-            make_tuple(Sequence<0, 2, 4>{}, Sequence<1, 3, 5>{}));
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}),
+            make_tuple(Sequence<0>{}, Sequence<1, 3, 5>{}, Sequence<2, 4, 6>{}));
 
         return xdlops_gemm.MakeCGM0N0M1N1M2M3M4N2Descriptor(c_g_m0_n0_m1_n1_m2_n2_grid_desc);
     }

composable_kernel/include/tensor_operation/gridwise_gemm_xdlops_v3r1.hpp

+#ifndef CK_GRIDWISE_GEMM_XDLOPS_V3R1_HPP
+#define CK_GRIDWISE_GEMM_XDLOPS_V3R1_HPP
+
+#include "common_header.hpp"
+#include "multi_index_transform_helper.hpp"
+#include "tensor_descriptor.hpp"
+#include "tensor_descriptor_helper.hpp"
+#include "blockwise_gemm_xdlops.hpp"
+#include "blockwise_tensor_slice_transfer.hpp"
+#include "threadwise_tensor_slice_transfer.hpp"
+#include "threadwise_tensor_slice_set.hpp"
+
+namespace ck {
+
+#if CK_EXPERIMENTAL_PASS_TENSOR_DESCRIPTOR_BY_VALUE
+template <typename GridwiseGemm,
+          typename FloatAB,
+          typename FloatC,
+          typename AGK0MK1GridDesc,
+          typename BGK0NK1GridDesc,
+          typename CM0N0M1N1M2M3M4N2GridDesc,
+          typename CBlockClusterAdaptor>
+__global__ void
+#if CK_USE_LAUNCH_BOUNDS
+    __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, CK_MIN_BLOCK_PER_CU)
+#endif
+        kernel_gemm_xdlops_v3r1(const FloatAB* __restrict__ p_a_grid,
+                                const FloatAB* __restrict__ p_b_grid,
+                                FloatC* __restrict__ p_c_grid,
+                                const AGK0MK1GridDesc a_k0_m_k1_grid_desc,
+                                const BGK0NK1GridDesc b_k0_n_k1_grid_desc,
+                                const CM0N0M1N1M2M3M4N2GridDesc c_m0_m1_m2_n_grid_desc,
+                                const CBlockClusterAdaptor c_block_cluster_adaptor)
+{
+    constexpr index_t shared_block_size =
+        GridwiseGemm::GetSharedMemoryNumberOfByte() / sizeof(FloatAB);
+
+    __shared__ FloatAB p_shared_block[shared_block_size];
+
+    GridwiseGemm::Run(p_a_grid,
+                      p_b_grid,
+                      p_c_grid,
+                      p_shared_block,
+                      a_k0_m_k1_grid_desc,
+                      b_k0_n_k1_grid_desc,
+                      c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                      c_block_cluster_adaptor);
+}
+#elif CK_EXPERIMENTAL_PASS_TENSOR_DESCRIPTOR_BY_VOID_POINTER
+template <typename GridwiseGemm,
+          typename FloatAB,
+          typename FloatC,
+          typename AGK0MK1GridDesc,
+          typename BGK0NK1GridDesc,
+          typename CM0N0M1N1M2M3M4N2GridDesc,
+          typename CBlockClusterAdaptor>
+__global__ void
+#if CK_USE_LAUNCH_BOUNDS
+    __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, CK_MIN_BLOCK_PER_CU)
+#endif
+        kernel_gemm_xdlops_v3r1(const FloatAB* __restrict__ p_a_grid,
+                                const FloatAB* __restrict__ p_b_grid,
+                                FloatC* __restrict__ p_c_grid,
+                                const void CONSTANT* p_a_k0_m_k1_grid_desc,
+                                const void CONSTANT* p_b_k0_n_k1_grid_desc,
+                                const void CONSTANT* p_c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                                const void CONSTANT* p_c_block_cluster_adaptor)
+{
+    constexpr index_t shared_block_size =
+        GridwiseGemm::GetSharedMemoryNumberOfByte() / sizeof(FloatAB);
+
+    const auto a_k0_m_k1_grid_desc = *reinterpret_cast<const AGK0MK1GridDesc*>(
+        cast_pointer_to_generic_address_space(p_a_k0_m_k1_grid_desc));
+    const auto b_k0_n_k1_grid_desc = *reinterpret_cast<const BGK0NK1GridDesc*>(
+        cast_pointer_to_generic_address_space(p_b_k0_n_k1_grid_desc));
+    const auto c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc =
+        *reinterpret_cast<const CM0N0M1N1M2M3M4N2GridDesc*>(
+            cast_pointer_to_generic_address_space(p_c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc));
+    const auto c_block_cluster_adaptor = *reinterpret_cast<const CBlockClusterAdaptor*>(
+        cast_pointer_to_generic_address_space(p_c_block_cluster_adaptor));
+
+    __shared__ FloatAB p_shared_block[shared_block_size];
+
+    GridwiseGemm::Run(p_a_grid,
+                      p_b_grid,
+                      p_c_grid,
+                      p_shared_block,
+                      a_k0_m_k1_grid_desc,
+                      b_k0_n_k1_grid_desc,
+                      c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                      c_block_cluster_adaptor);
+}
+#endif
+
+template <index_t BlockSize,
+          typename FloatAB,
+          typename FloatAcc,
+          typename FloatC,
+          InMemoryDataOperationEnum_t CGlobalMemoryDataOperation,
+          typename AGK0MK1GridDesc,
+          typename BGK0NK1GridDesc,
+          typename CGMNGridDesc,
+          index_t MPerBlock,
+          index_t NPerBlock,
+          index_t KPerBlock,
+          index_t MPerXDL,
+          index_t NPerXDL,
+          index_t K1Value,
+          index_t MRepeat,
+          index_t NRepeat,
+          typename ABlockTransferThreadSliceLengths_G_K0_M_K1,
+          typename ABlockTransferThreadClusterLengths_G_K0_M_K1,
+          typename ABlockTransferThreadClusterArrangeOrder,
+          typename ABlockTransferSrcAccessOrder,
+          index_t ABlockTransferSrcVectorDim,
+          index_t ABlockTransferSrcScalarPerVector,
+          index_t ABlockTransferDstScalarPerVector_K1,
+          bool AThreadTransferSrcResetCoordinateAfterRun,
+          typename BBlockTransferThreadSliceLengths_G_K0_N_K1,
+          typename BBlockTransferThreadClusterLengths_G_K0_N_K1,
+          typename BBlockTransferThreadClusterArrangeOrder,
+          typename BBlockTransferSrcAccessOrder,
+          index_t BBlockTransferSrcVectorDim,
+          index_t BBlockTransferSrcScalarPerVector,
+          index_t BBlockTransferDstScalarPerVector_K1,
+          bool BThreadTransferSrcResetCoordinateAfterRun,
+          typename CThreadTransferSrcDstAccessOrder,
+          index_t CThreadTransferSrcDstVectorDim,
+          index_t CThreadTransferDstScalarPerVector,
+          typename AGridStepHacks,
+          typename BGridStepHacks,
+          typename CGridStepHacks,
+          typename AGridMoveSliceWindowStepHacks,
+          typename BGridMoveSliceWindowStepHacks,
+          bool CAccessOrderMRepeatNRepeat>
+struct GridwiseGemm_gk0mk1_gk0nk1_gmn_xdlops_v3r1
+{
+    static constexpr auto I0 = Number<0>{};
+    static constexpr auto I1 = Number<1>{};
+    static constexpr auto I2 = Number<2>{};
+    static constexpr auto I3 = Number<3>{};
+    static constexpr auto I4 = Number<4>{};
+    static constexpr auto I5 = Number<5>{};
+    static constexpr auto I6 = Number<6>{};
+
+    // K1 should be Number<...>
+    static constexpr auto K1 = Number<K1Value>{};
+
+    __host__ __device__ static constexpr index_t GetSharedMemoryNumberOfByte()
+    {
+        constexpr auto max_lds_align = K1;
+
+        // A matrix in LDS memory, dst of blockwise copy
+        //   be careful of LDS alignment
+        constexpr auto a_k0_m_k1_block_desc = make_naive_tensor_descriptor_aligned(
+            make_tuple(Number<KPerBlock>{}, Number<MPerBlock>{}, K1), max_lds_align);
+
+        // B matrix in LDS memory, dst of blockwise copy
+        //   be careful of LDS alignment
+        constexpr auto b_k0_n_k1_block_desc = make_naive_tensor_descriptor_aligned(
+            make_tuple(Number<KPerBlock>{}, Number<NPerBlock>{}, K1), max_lds_align);
+
+        // LDS allocation for A and B: be careful of alignment
+        constexpr auto a_block_space_size =
+            math::integer_least_multiple(a_k0_m_k1_block_desc.GetElementSpaceSize(), max_lds_align);
+
+        constexpr auto b_block_space_size =
+            math::integer_least_multiple(b_k0_n_k1_block_desc.GetElementSpaceSize(), max_lds_align);
+
+        return (a_block_space_size + b_block_space_size) * sizeof(FloatAB);
+    }
+
+    __host__ __device__ static constexpr bool
+    CheckValidity(const AGK0MK1GridDesc& a_g_k0_m_k1_grid_desc,
+                  const BGK0NK1GridDesc& b_g_k0_n_k1_grid_desc,
+                  const CGMNGridDesc& c_g_m_n_grid_desc)
+    {
+        // TODO: turn on this
+        static_assert(is_known_at_compile_time<remove_cv_t<decltype(K1)>>::value,
+                      "wrong! K1 need to be known at compile-time");
+
+        const auto M  = a_g_k0_m_k1_grid_desc.GetLength(I2);
+        const auto N  = b_g_k0_n_k1_grid_desc.GetLength(I2);
+        const auto K0 = a_g_k0_m_k1_grid_desc.GetLength(I1);
+
+        static_assert((MPerBlock % (MPerXDL * MRepeat) == 0) &&
+                          (NPerBlock % (NRepeat * NPerXDL)) == 0,
+                      "Invalid tuning param!");
+
+        // TODO: also check validity of all components (blockwise-copy, threadwise-copy, etc)
+        return (M == c_g_m_n_grid_desc.GetLength(I1) && N == c_g_m_n_grid_desc.GetLength(I2) &&
+                K0 == b_g_k0_n_k1_grid_desc.GetLength(I1) &&
+                K1 == a_g_k0_m_k1_grid_desc.GetLength(I3) &&
+                K1 == b_g_k0_n_k1_grid_desc.GetLength(I3)) &&
+               (M % MPerBlock == 0 && N % NPerBlock == 0 && K0 % KPerBlock == 0);
+    }
+
+    __host__ __device__ static constexpr index_t
+    CalculateGridSize(const CGMNGridDesc& c_g_m_n_grid_desc)
+    {
+        const auto G = c_g_m_n_grid_desc.GetLength(I0);
+        const auto M = c_g_m_n_grid_desc.GetLength(I1);
+        const auto N = c_g_m_n_grid_desc.GetLength(I2);
+
+        const index_t grid_size = G * (M / MPerBlock) * (N / NPerBlock);
+
+        return grid_size;
+    }
+
+    __host__ __device__ static constexpr auto
+    MakeCGM0N0M1N1M2M3M4N2GridDescriptor(const CGMNGridDesc& c_g_m_n_grid_desc)
+    {
+        constexpr auto max_lds_align = K1;
+
+        constexpr auto a_k0_m_k1_block_desc = make_naive_tensor_descriptor_aligned(
+            make_tuple(Number<KPerBlock>{}, Number<MPerBlock>{}, K1), max_lds_align);
+
+        constexpr auto b_k0_n_k1_block_desc = make_naive_tensor_descriptor_aligned(
+            make_tuple(Number<KPerBlock>{}, Number<NPerBlock>{}, K1), max_lds_align);
+
+        using BlockwiseGemm =
+            BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1<BlockSize,
+                                                                FloatAB,
+                                                                decltype(a_k0_m_k1_block_desc),
+                                                                decltype(b_k0_n_k1_block_desc),
+                                                                MPerXDL,
+                                                                NPerXDL,
+                                                                MRepeat,
+                                                                NRepeat,
+                                                                K1>;
+
+        return BlockwiseGemm::MakeCGM0N0M1N1M2M3M4N2GridDescriptor(c_g_m_n_grid_desc);
+    }
+
+    __host__ __device__ static constexpr auto
+    MakeCBlockClusterAdaptor(const CGMNGridDesc& c_g_m_n_grid_desc)
+    {
+        const auto G = c_g_m_n_grid_desc.GetLength(I0);
+        const auto M = c_g_m_n_grid_desc.GetLength(I1);
+        const auto N = c_g_m_n_grid_desc.GetLength(I2);
+
+        constexpr auto M1 = Number<MPerBlock>{};
+        constexpr auto N1 = Number<NPerBlock>{};
+
+        const auto M0 = M / M1;
+        const auto N0 = N / N1;
+
+#if 1
+        const auto c_blockid_to_m0_n0_block_cluster_adaptor =
+            make_single_stage_tensor_adaptor(make_tuple(make_merge_transform(make_tuple(G, M0, N0))),
+                                             make_tuple(Sequence<0, 1, 2>{}),
+                                             make_tuple(Sequence<0>{}));
+#elif 1
+        const auto c_blockid_to_m0_n0_block_cluster_adaptor =
+            make_single_stage_tensor_adaptor(make_tuple(make_merge_transform(make_tuple(G, N0, M0))),
+                                             make_tuple(Sequence<0, 2, 1>{}),
+                                             make_tuple(Sequence<0>{}));
+#endif
+
+        return c_blockid_to_m0_n0_block_cluster_adaptor;
+    }
+
+    using CM0N0M1N1M2M3M4N2GridDesc = decltype(MakeCGM0N0M1N1M2M3M4N2GridDescriptor(CGMNGridDesc{}));
+    using CBlockClusterAdaptor      = decltype(MakeCBlockClusterAdaptor(CGMNGridDesc{}));
+
+    __device__ static void Run(const FloatAB* __restrict__ p_a_grid,
+                               const FloatAB* __restrict__ p_b_grid,
+                               FloatC* __restrict__ p_c_grid,
+                               FloatAB* __restrict__ p_shared_block,
+                               const AGK0MK1GridDesc& a_k0_m_k1_grid_desc,
+                               const BGK0NK1GridDesc& b_k0_n_k1_grid_desc,
+                               const CM0N0M1N1M2M3M4N2GridDesc& c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                               const CBlockClusterAdaptor& c_block_cluster_adaptor)
+    {
+        /*    const auto a_grid_buf = make_dynamic_buffer<AddressSpaceEnum_t::Global>(
+                p_a_grid, a_k0_m_k1_grid_desc.GetElementSpaceSize());
+            const auto b_grid_buf = make_dynamic_buffer<AddressSpaceEnum_t::Global>(
+                p_b_grid, b_k0_n_k1_grid_desc.GetElementSpaceSize());
+            auto c_grid_buf = make_dynamic_buffer<AddressSpaceEnum_t::Global>(
+                p_c_grid, c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc.GetElementSpaceSize());
+
+            const auto K0 = a_k0_m_k1_grid_desc.GetLength(I0);
+
+            // divide block work by [M, N]
+            const auto block_work_idx =
+                c_block_cluster_adaptor.CalculateBottomIndex(make_multi_index(get_block_1d_id()));
+
+            // HACK: this force m/n_block_data_idx_on_grid into SGPR
+            const index_t m_block_data_idx_on_grid =
+                __builtin_amdgcn_readfirstlane(block_work_idx[I0] * MPerBlock);
+
+            const index_t n_block_data_idx_on_grid =
+                __builtin_amdgcn_readfirstlane(block_work_idx[I1] * NPerBlock);
+
+            // lds max alignment
+            constexpr auto max_lds_align = K1;
+
+            // A matrix in LDS memory, dst of blockwise copy
+            //   be careful of LDS alignment
+            constexpr auto a_k0_m_k1_block_desc = make_naive_tensor_descriptor_aligned(
+                make_tuple(Number<KPerBlock>{}, Number<MPerBlock>{}, K1), max_lds_align);
+
+            // B matrix in LDS memory, dst of blockwise copy
+            //   be careful of LDS alignment
+            constexpr auto b_k0_n_k1_block_desc = make_naive_tensor_descriptor_aligned(
+                make_tuple(Number<KPerBlock>{}, Number<NPerBlock>{}, K1), max_lds_align);
+
+            // A matrix blockwise copy
+            auto a_blockwise_copy =
+                BlockwiseTensorSliceTransfer_v4<BlockSize,
+                                                InMemoryDataOperationEnum_t::Set,
+                                                Sequence<KPerBlock, MPerBlock, K1>,
+                                                ABlockTransferThreadSliceLengths_K0_M_K1,
+                                                ABlockTransferThreadClusterLengths_K0_M_K1,
+                                                ABlockTransferThreadClusterArrangeOrder,
+                                                FloatAB,
+                                                FloatAB,
+                                                decltype(a_k0_m_k1_grid_desc),
+                                                decltype(a_k0_m_k1_block_desc),
+                                                ABlockTransferSrcAccessOrder,
+                                                Sequence<1, 0, 2>,
+                                                ABlockTransferSrcVectorDim,
+                                                2,
+                                                ABlockTransferSrcScalarPerVector,
+                                                ABlockTransferDstScalarPerVector_K1,
+                                                1,
+                                                1,
+                                                AThreadTransferSrcResetCoordinateAfterRun,
+                                                true>(a_k0_m_k1_grid_desc,
+                                                      make_multi_index(0, m_block_data_idx_on_grid,
+           0), a_k0_m_k1_block_desc, make_multi_index(0, 0, 0));
+
+            // B matrix blockwise copy
+            auto b_blockwise_copy =
+                BlockwiseTensorSliceTransfer_v4<BlockSize,
+                                                InMemoryDataOperationEnum_t::Set,
+                                                Sequence<KPerBlock, NPerBlock, K1>,
+                                                BBlockTransferThreadSliceLengths_K0_N_K1,
+                                                BBlockTransferThreadClusterLengths_K0_N_K1,
+                                                BBlockTransferThreadClusterArrangeOrder,
+                                                FloatAB,
+                                                FloatAB,
+                                                decltype(b_k0_n_k1_grid_desc),
+                                                decltype(b_k0_n_k1_block_desc),
+                                                BBlockTransferSrcAccessOrder,
+                                                Sequence<1, 0, 2>,
+                                                BBlockTransferSrcVectorDim,
+                                                2,
+                                                BBlockTransferSrcScalarPerVector,
+                                                BBlockTransferDstScalarPerVector_K1,
+                                                1,
+                                                1,
+                                                BThreadTransferSrcResetCoordinateAfterRun,
+                                                true>(b_k0_n_k1_grid_desc,
+                                                      make_multi_index(0, n_block_data_idx_on_grid,
+           0), b_k0_n_k1_block_desc, make_multi_index(0, 0, 0));
+
+            // GEMM definition
+            //   c_mtx += transpose(a_mtx) * b_mtx
+            //     a_mtx[KPerBlock, MPerBlock] is in LDS
+            //     b_mtx[KPerBlock, NPerBlock] is in LDS
+            //     c_mtx[MPerBlock, NPerBlock] is distributed among threads, and saved in
+            //       register
+            // sanity check
+
+            const auto blockwise_gemm =
+                BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1<BlockSize,
+                                                                    FloatAB,
+                                                                    decltype(a_k0_m_k1_block_desc),
+                                                                    decltype(b_k0_n_k1_block_desc),
+                                                                    MPerXDL,
+                                                                    NPerXDL,
+                                                                    MRepeat,
+                                                                    NRepeat,
+                                                                    K1>{};
+
+            constexpr auto c_mr_nr_blk_desc =
+                make_naive_tensor_descriptor_packed(make_tuple(Number<MRepeat>{},
+           Number<NRepeat>{}));
+
+            constexpr auto c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc =
+                blockwise_gemm.GetCM0N0M1N1M2M3M4N2ThreadDescriptor();
+            constexpr auto CBlkSize = c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc.GetElementSpaceSize();
+
+            StaticBuffer<AddressSpaceEnum_t::Vgpr,
+                         vector_type<FloatAcc, CBlkSize>,
+                         c_mr_nr_blk_desc.GetElementSpaceSize(),
+                         true>
+                c_thread_buf;
+
+            // LDS allocation for A and B: be careful of alignment
+            constexpr auto a_block_space_size =
+                math::integer_least_multiple(a_k0_m_k1_block_desc.GetElementSpaceSize(),
+           max_lds_align);
+
+            FloatAB* p_a_block = p_shared_block;
+            FloatAB* p_b_block = p_shared_block + a_block_space_size;
+
+            constexpr auto a_block_slice_copy_step = make_multi_index(KPerBlock, 0, 0);
+            constexpr auto b_block_slice_copy_step = make_multi_index(KPerBlock, 0, 0);
+
+            // hack to control index calculation when iterating over A and B matrix for threadwise
+           copy constexpr auto a_k0_m_k1_grid_step_hacks = AGridStepHacks{}; constexpr auto
+           b_k0_n_k1_grid_step_hacks = BGridStepHacks{};
+
+            // hack to control index calculation when move slice window for A and B matrix for
+            // threadwise copy
+            constexpr auto a_k0_m_k1_grid_move_slice_window_step_hack =
+           AGridMoveSliceWindowStepHacks{}; constexpr auto
+           b_k0_n_k1_grid_move_slice_window_step_hack = BGridMoveSliceWindowStepHacks{};
+
+            auto a_block_buf = make_dynamic_buffer<AddressSpaceEnum_t::Lds>(
+                p_a_block, a_k0_m_k1_block_desc.GetElementSpaceSize());
+            auto b_block_buf = make_dynamic_buffer<AddressSpaceEnum_t::Lds>(
+                p_b_block, b_k0_n_k1_block_desc.GetElementSpaceSize());
+
+            // preload data into LDS
+            {
+                a_blockwise_copy.RunRead(a_k0_m_k1_grid_desc, a_grid_buf,
+           a_k0_m_k1_grid_step_hacks); b_blockwise_copy.RunRead(b_k0_n_k1_grid_desc, b_grid_buf,
+           b_k0_n_k1_grid_step_hacks);
+
+                a_blockwise_copy.RunWrite(a_k0_m_k1_block_desc, a_block_buf);
+                b_blockwise_copy.RunWrite(b_k0_n_k1_block_desc, b_block_buf);
+            }
+
+            // main body
+            index_t k_block_data_begin = 0;
+
+            do
+            {
+                a_blockwise_copy.MoveSrcSliceWindow(a_k0_m_k1_grid_desc,
+                                                    a_block_slice_copy_step,
+                                                    a_k0_m_k1_grid_move_slice_window_step_hack);
+                b_blockwise_copy.MoveSrcSliceWindow(b_k0_n_k1_grid_desc,
+                                                    b_block_slice_copy_step,
+                                                    b_k0_n_k1_grid_move_slice_window_step_hack);
+
+                a_blockwise_copy.RunRead(a_k0_m_k1_grid_desc, a_grid_buf,
+           a_k0_m_k1_grid_step_hacks);
+
+                block_sync_lds();
+
+                b_blockwise_copy.RunRead(b_k0_n_k1_grid_desc, b_grid_buf,
+           b_k0_n_k1_grid_step_hacks);
+
+                blockwise_gemm.Run(a_block_buf, b_block_buf, c_thread_buf);
+
+                block_sync_lds();
+
+                a_blockwise_copy.RunWrite(a_k0_m_k1_block_desc, a_block_buf);
+                b_blockwise_copy.RunWrite(b_k0_n_k1_block_desc, b_block_buf);
+
+                k_block_data_begin += KPerBlock;
+            } while(k_block_data_begin < (K0 - KPerBlock));
+
+            // tail
+            {
+                block_sync_lds();
+
+                blockwise_gemm.Run(a_block_buf, b_block_buf, c_thread_buf);
+            }
+
+            // output: register to global memory
+            {
+                constexpr auto c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc =
+                    blockwise_gemm.GetCM0N0M1N1M2M3M4N2BlockDescriptor();
+
+                constexpr auto M2 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I4);
+                constexpr auto M3 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I5);
+                constexpr auto M4 = c_m0_n0_m1_n1_m2_m3_m4_n2_block_desc.GetLength(I6);
+
+                // calculate origin of thread output tensor on global memory
+                //     blockwise GEMM c matrix starting index
+                const auto c_thread_mtx_on_block =
+                    blockwise_gemm.CalculateCThreadOriginDataIndex(I0, I0, I0, I0);
+
+                const index_t m_thread_data_on_grid =
+                    m_block_data_idx_on_grid + c_thread_mtx_on_block[I0];
+
+                const index_t n_thread_data_on_grid =
+                    n_block_data_idx_on_grid + c_thread_mtx_on_block[I1];
+
+                constexpr auto c_m0_n0_m1_n1_m2_m3_m4_n2_grid_tensor_step_hacks = CGridStepHacks{};
+
+                auto c_thread_copy =
+                    ThreadwiseTensorSliceTransfer_v1r3<FloatC,
+                                                       FloatC,
+                                                       decltype(c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc),
+                                                       decltype(c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc),
+                                                       Sequence<I1, I1, I1, I1, M2, I1, M4, I1>,
+                                                       CThreadTransferSrcDstAccessOrder,
+                                                       CThreadTransferSrcDstVectorDim,
+                                                       CThreadTransferDstScalarPerVector,
+                                                       CGlobalMemoryDataOperation,
+                                                       1,
+                                                       true>{
+                        c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                        make_multi_index(0,
+                                         0,
+                                         0,
+                                         0,
+                                         m_thread_data_on_grid / (M3 * M4),
+                                         m_thread_data_on_grid % (M3 * M4) / M4,
+                                         m_thread_data_on_grid % M4,
+                                         n_thread_data_on_grid)};
+
+                auto init_copy = [&](auto c_thread_idx_) {
+                    constexpr auto blk_off = c_mr_nr_blk_desc.CalculateOffset(c_thread_idx_);
+                    c_thread_copy.Run(c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc,
+                                      make_tuple(I0, I0, I0, I0, I0, I0, I0, I0),
+                                      c_thread_buf[Number<blk_off>{}].template AsType<FloatAcc>(),
+                                      c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                                      c_grid_buf,
+                                      c_m0_n0_m1_n1_m2_m3_m4_n2_grid_tensor_step_hacks);
+
+                    return c_thread_idx_;
+                };
+
+                auto mrepeat_plus_copy = [&](auto c_thread_idx_) {
+                    constexpr auto mrepeat_step_plus = make_multi_index(1, 0, 0, 0, 0, 0, 0, 0);
+                    c_thread_copy.MoveDstSliceWindow(c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                                                     mrepeat_step_plus);
+
+                    constexpr auto blk_off = c_mr_nr_blk_desc.CalculateOffset(c_thread_idx_);
+                    c_thread_copy.Run(c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc,
+                                      make_tuple(I0, I0, I0, I0, I0, I0, I0, I0),
+                                      c_thread_buf[Number<blk_off>{}].template AsType<FloatAcc>(),
+                                      c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                                      c_grid_buf,
+                                      c_m0_n0_m1_n1_m2_m3_m4_n2_grid_tensor_step_hacks);
+                };
+
+                auto nrepeat_plus_copy = [&](auto c_thread_idx_) {
+                    constexpr auto nrepeat_step_plus = make_multi_index(0, 1, 0, 0, 0, 0, 0, 0);
+                    c_thread_copy.MoveDstSliceWindow(c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                                                     nrepeat_step_plus);
+
+                    constexpr auto blk_off = c_mr_nr_blk_desc.CalculateOffset(c_thread_idx_);
+                    c_thread_copy.Run(c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc,
+                                      make_tuple(I0, I0, I0, I0, I0, I0, I0, I0),
+                                      c_thread_buf[Number<blk_off>{}].template AsType<FloatAcc>(),
+                                      c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                                      c_grid_buf,
+                                      c_m0_n0_m1_n1_m2_m3_m4_n2_grid_tensor_step_hacks);
+                };
+
+                auto mrepeat_minus_copy = [&](auto c_thread_idx_) {
+                    constexpr auto mrepeat_step_plus = make_multi_index(-1, 0, 0, 0, 0, 0, 0, 0);
+                    c_thread_copy.MoveDstSliceWindow(c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                                                     mrepeat_step_plus);
+
+                    constexpr auto blk_off = c_mr_nr_blk_desc.CalculateOffset(c_thread_idx_);
+                    c_thread_copy.Run(c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc,
+                                      make_tuple(I0, I0, I0, I0, I0, I0, I0, I0),
+                                      c_thread_buf[Number<blk_off>{}].template AsType<FloatAcc>(),
+                                      c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                                      c_grid_buf,
+                                      c_m0_n0_m1_n1_m2_m3_m4_n2_grid_tensor_step_hacks);
+                };
+
+                auto nrepeat_minus_copy = [&](auto c_thread_idx_) {
+                    constexpr auto nrepeat_step_minus = make_multi_index(0, -1, 0, 0, 0, 0, 0, 0);
+                    c_thread_copy.MoveDstSliceWindow(c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                                                     nrepeat_step_minus);
+
+                    constexpr auto blk_off = c_mr_nr_blk_desc.CalculateOffset(c_thread_idx_);
+                    c_thread_copy.Run(c_m0_n0_m1_n1_m2_m3_m4_n2_thread_desc,
+                                      make_tuple(I0, I0, I0, I0, I0, I0, I0, I0),
+                                      c_thread_buf[Number<blk_off>{}].template AsType<FloatAcc>(),
+                                      c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc,
+                                      c_grid_buf,
+                                      c_m0_n0_m1_n1_m2_m3_m4_n2_grid_tensor_step_hacks);
+                };
+
+                static_assert((MRepeat == 4 && NRepeat == 4) or (MRepeat == 4 && NRepeat == 2) or
+                                  (MRepeat == 2 && NRepeat == 4) or (MRepeat == 2 && NRepeat == 2)
+           or (MRepeat == 2 && NRepeat == 1) or (MRepeat == 1 && NRepeat == 2) or (MRepeat == 1 &&
+           NRepeat == 1), "wrong");
+
+                if constexpr(MRepeat == 4 && NRepeat == 4)
+                {
+                    init_copy(make_tuple(I0, I0));
+
+                    if constexpr(CAccessOrderMRepeatNRepeat)
+                    {
+                        nrepeat_plus_copy(make_tuple(I0, I1));
+                        nrepeat_plus_copy(make_tuple(I0, I2));
+                        nrepeat_plus_copy(make_tuple(I0, I3));
+                        mrepeat_plus_copy(make_tuple(I1, I3));
+                        nrepeat_minus_copy(make_tuple(I1, I2));
+                        nrepeat_minus_copy(make_tuple(I1, I1));
+                        nrepeat_minus_copy(make_tuple(I1, I0));
+                        mrepeat_plus_copy(make_tuple(I2, I0));
+                        nrepeat_plus_copy(make_tuple(I2, I1));
+                        nrepeat_plus_copy(make_tuple(I2, I2));
+                        nrepeat_plus_copy(make_tuple(I2, I3));
+                        mrepeat_plus_copy(make_tuple(I3, I3));
+                        nrepeat_minus_copy(make_tuple(I3, I2));
+                        nrepeat_minus_copy(make_tuple(I3, I1));
+                        nrepeat_minus_copy(make_tuple(I3, I0));
+                    }
+                    else
+                    {
+                        mrepeat_plus_copy(make_tuple(I1, I0));
+                        mrepeat_plus_copy(make_tuple(I2, I0));
+                        mrepeat_plus_copy(make_tuple(I3, I0));
+                        nrepeat_plus_copy(make_tuple(I3, I1));
+                        mrepeat_minus_copy(make_tuple(I2, I1));
+                        mrepeat_minus_copy(make_tuple(I1, I1));
+                        mrepeat_minus_copy(make_tuple(I0, I1));
+                        nrepeat_plus_copy(make_tuple(I0, I2));
+                        mrepeat_plus_copy(make_tuple(I1, I2));
+                        mrepeat_plus_copy(make_tuple(I2, I2));
+                        mrepeat_plus_copy(make_tuple(I3, I2));
+                        nrepeat_plus_copy(make_tuple(I3, I3));
+                        mrepeat_minus_copy(make_tuple(I2, I3));
+                        mrepeat_minus_copy(make_tuple(I1, I3));
+                        mrepeat_minus_copy(make_tuple(I0, I3));
+                    }
+                }
+                else if constexpr(MRepeat == 4 && NRepeat == 2)
+                {
+                    init_copy(make_tuple(I0, I0));
+
+                    if constexpr(CAccessOrderMRepeatNRepeat)
+                    {
+                        nrepeat_plus_copy(make_tuple(I0, I1));
+                        mrepeat_plus_copy(make_tuple(I1, I1));
+                        nrepeat_minus_copy(make_tuple(I1, I0));
+                        mrepeat_plus_copy(make_tuple(I2, I0));
+                        nrepeat_plus_copy(make_tuple(I2, I1));
+                        mrepeat_plus_copy(make_tuple(I3, I1));
+                        nrepeat_minus_copy(make_tuple(I3, I0));
+                    }
+                    else
+                    {
+                        mrepeat_plus_copy(make_tuple(I1, I0));
+                        mrepeat_plus_copy(make_tuple(I2, I0));
+                        mrepeat_plus_copy(make_tuple(I3, I0));
+                        nrepeat_plus_copy(make_tuple(I3, I1));
+                        mrepeat_minus_copy(make_tuple(I2, I1));
+                        mrepeat_minus_copy(make_tuple(I1, I1));
+                        mrepeat_minus_copy(make_tuple(I0, I1));
+                    }
+                }
+                else if constexpr(MRepeat == 2 && NRepeat == 4)
+                {
+                    init_copy(make_tuple(I0, I0));
+
+                    if constexpr(CAccessOrderMRepeatNRepeat)
+                    {
+                        nrepeat_plus_copy(make_tuple(I0, I1));
+                        nrepeat_plus_copy(make_tuple(I0, I2));
+                        nrepeat_plus_copy(make_tuple(I0, I3));
+                        mrepeat_plus_copy(make_tuple(I1, I3));
+                        nrepeat_minus_copy(make_tuple(I1, I2));
+                        nrepeat_minus_copy(make_tuple(I1, I1));
+                        nrepeat_minus_copy(make_tuple(I1, I0));
+                    }
+                    else
+                    {
+                        mrepeat_plus_copy(make_tuple(I1, I0));
+                        nrepeat_plus_copy(make_tuple(I1, I1));
+                        mrepeat_minus_copy(make_tuple(I0, I1));
+                        nrepeat_plus_copy(make_tuple(I0, I2));
+                        mrepeat_plus_copy(make_tuple(I1, I2));
+                        nrepeat_plus_copy(make_tuple(I1, I3));
+                        mrepeat_minus_copy(make_tuple(I0, I3));
+                    }
+                }
+                else if constexpr(MRepeat == 2 && NRepeat == 2)
+                {
+                    init_copy(make_tuple(I0, I0));
+
+                    if constexpr(CAccessOrderMRepeatNRepeat)
+                    {
+                        nrepeat_plus_copy(make_tuple(I0, I1));
+                        mrepeat_plus_copy(make_tuple(I1, I1));
+                        nrepeat_minus_copy(make_tuple(I1, I0));
+                    }
+                    else
+                    {
+                        mrepeat_plus_copy(make_tuple(I1, I0));
+                        nrepeat_plus_copy(make_tuple(I1, I1));
+                        mrepeat_minus_copy(make_tuple(I0, I1));
+                    }
+                }
+                else if constexpr(MRepeat == 2 && NRepeat == 1)
+                {
+                    init_copy(make_tuple(I0, I0));
+                    mrepeat_plus_copy(make_tuple(I1, I0));
+                }
+                else if constexpr(MRepeat == 1 && NRepeat == 2)
+                {
+                    init_copy(make_tuple(I0, I0));
+                    nrepeat_plus_copy(make_tuple(I0, I1));
+                }
+                else if constexpr(MRepeat == 1 && NRepeat == 1)
+                {
+                    init_copy(make_tuple(I0, I0));
+                }
+            }*/
+    }
+}; // namespace ck
+
+} // namespace ck
+#endif

host/driver_offline/include/driver_gemm_xdlops_v3r1.hpp

@@ -11,9 +11,9 @@ template <ck::index_t BlockSize,
           typename FloatAcc,
           typename FloatC,
           ck::InMemoryDataOperationEnum_t CGlobalMemoryDataOperation,
-          typename AK0MK1GridDesc,
-          typename BK0NK1GridDesc,
-          typename CMNGridDesc,
+          typename AGK0MK1GridDesc,
+          typename BGK0NK1GridDesc,
+          typename CGMNGridDesc,
           ck::index_t MPerBlock,
           ck::index_t NPerBlock,
           ck::index_t KPerBlock,
@@ -50,9 +50,9 @@ template <ck::index_t BlockSize,
 __host__ float driver_gemm_xdlops_v3r1(const FloatAB* p_a_grid,
                                        const FloatAB* p_b_grid,
                                        FloatC* p_c_grid,
-                                       const AK0MK1GridDesc& a_g_k0_m_k1_grid_desc,
-                                       const BK0NK1GridDesc& b_g_k0_n_k1_grid_desc,
-                                       const CMNGridDesc& c_g_m_n_grid_desc,
+                                       const AGK0MK1GridDesc& a_g_k0_m_k1_grid_desc,
+                                       const BGK0NK1GridDesc& b_g_k0_n_k1_grid_desc,
+                                       const CGMNGridDesc& c_g_m_n_grid_desc,
                                        AGridStepHacks,
                                        BGridStepHacks,
                                        CGridStepHacks,
@@ -69,14 +69,14 @@ __host__ float driver_gemm_xdlops_v3r1(const FloatAB* p_a_grid,
     constexpr auto I3 = Number<3>{};
 
     using GridwiseGemm =
-        GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v3r1<BlockSize,
+        GridwiseGemm_gk0mk1_gk0nk1_gmn_xdlops_v3r1<BlockSize,
                                                 FloatAB,
                                                 FloatAcc,
                                                 FloatC,
                                                 CGlobalMemoryDataOperation,
-                                                AK0MK1GridDesc,
-                                                BK0NK1GridDesc,
-                                                CMNGridDesc,
+                                                AGK0MK1GridDesc,
+                                                BGK0NK1GridDesc,
+                                                CGMNGridDesc,
                                                 MPerBlock,
                                                 NPerBlock,
                                                 KPerBlock,
@@ -134,26 +134,26 @@ __host__ float driver_gemm_xdlops_v3r1(const FloatAB* p_a_grid,
             "wrong! GridwiseGemm_km_kn_m0m1n0n1_xdlops_v2r3 has invalid setting");
     }
 
-       const auto c_gemmg_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc =
-           GridwiseGemm::MakeCM0N0M1N1M2M3M4N2GridDescriptor(c_g_m_n_grid_desc);
+       const auto c_g_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc =
+           GridwiseGemm::MakeCGM0N0M1N1M2M3M4N2GridDescriptor(c_g_m_n_grid_desc);
 
-/*       using CM0N0M1N1M2M3M4N2GridDesc = decltype(c_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc);
+       using CGM0N0M1N1M2M3M4N2GridDesc = decltype(c_g_m0_n0_m1_n1_m2_m3_m4_n2_grid_desc);
 
-       const auto c_block_cluster_adaptor = GridwiseGemm::MakeCBlockClusterAdaptor(c_m_n_grid_desc);
+       const auto c_block_cluster_adaptor = GridwiseGemm::MakeCBlockClusterAdaptor(c_g_m_n_grid_desc);
 
        using CBlockClusterAdaptor = decltype(c_block_cluster_adaptor);
 
-       const index_t grid_size = GridwiseGemm::CalculateGridSize(c_m_n_grid_desc);
+       const index_t grid_size = GridwiseGemm::CalculateGridSize(c_g_m_n_grid_desc);
 
-       const auto kernel = kernel_gemm_xdlops_v2r3<GridwiseGemm,
+       const auto kernel = kernel_gemm_xdlops_v3r1<GridwiseGemm,
                                                    FloatAB,
                                                    FloatC,
-                                                   remove_reference_t<AK0MK1GridDesc>,
-                                                   remove_reference_t<BK0NK1GridDesc>,
-                                                   remove_reference_t<CM0N0M1N1M2M3M4N2GridDesc>,
+                                                   remove_reference_t<AGK0MK1GridDesc>,
+                                                   remove_reference_t<BGK0NK1GridDesc>,
+                                                   remove_reference_t<CGM0N0M1N1M2M3M4N2GridDesc>,
                                                    remove_reference_t<CBlockClusterAdaptor>>;
 
-   #if CK_EXPERIMENTAL_PASS_TENSOR_DESCRIPTOR_BY_VALUE
+/*   #if CK_EXPERIMENTAL_PASS_TENSOR_DESCRIPTOR_BY_VALUE
        float ave_time = launch_and_time_kernel(kernel,
                                                nrepeat,
                                                dim3(grid_size),