From 98dd392f53c07d13fafafbf884b483bb3b98ba6b Mon Sep 17 00:00:00 2001 From: Bartlomiej Kocot Date: Fri, 5 Sep 2025 18:24:59 -0400 Subject: [PATCH] tmp --- .../blockwise_gemm_xdlops_skip_b_lds.hpp | 76 +++--- ...nv_bwd_data_multiple_d_xdl_cshuffle_v1.hpp | 139 ++++++---- ..._xdlops_skip_b_lds_multiple_d_cshuffle.hpp | 251 +++++++++--------- 3 files changed, 259 insertions(+), 207 deletions(-) diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_xdlops_skip_b_lds.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_xdlops_skip_b_lds.hpp index 84ee096cba..8e6bc1b4f3 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_xdlops_skip_b_lds.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_xdlops_skip_b_lds.hpp @@ -13,8 +13,8 @@ namespace ck { template {}; @@ -119,8 +119,8 @@ struct BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1r1 __host__ __device__ BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1r1() { - static_assert(AK0MK1BlockDesc::IsKnownAtCompileTime() && - BK0K0BN0N1N2N3K1BlockDesc::IsKnownAtCompileTime(), + static_assert(BK0NK1BlockDesc::IsKnownAtCompileTime() && + AK0K0BM0M1M2M3K1BlockDesc::IsKnownAtCompileTime(), "wrong! Desc should be known at compile-time"); static_assert(BlockSize == MWaves * NWaves * WaveSize, @@ -221,58 +221,58 @@ struct BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1r1 c_grid_desc_g_m0_n0_m1_n1_m2_n2); } - __host__ __device__ static constexpr auto MakeABlockDescriptor_M0_M1_M2_K() + __host__ __device__ static constexpr auto MakeBBlockDescriptor_N0_N1_N2_K() { return transform_tensor_descriptor( - AK0MK1BlockDesc{}, + BK0NK1BlockDesc{}, make_tuple( - make_merge_transform_v3_division_mod(make_tuple(Number{}, Number{})), + make_merge_transform_v3_division_mod(make_tuple(Number{}, Number{})), make_unmerge_transform( - make_tuple(Number{}, Number{}, Number{}))), + make_tuple(Number{}, Number{}, Number{}))), make_tuple(Sequence<0, 2>{}, Sequence<1>{}), make_tuple(Sequence<3>{}, Sequence<0, 1, 2>{})); } - __device__ void MoveABlockSliceWindow() + __device__ void MoveBBlockSliceWindow() { - a_thread_copy_.MoveSrcSliceWindow(a_block_desc_m0_m1_m2_k, + b_thread_copy_.MoveSrcSliceWindow(b_block_desc_n0_n1_n2_k, make_multi_index(0, 0, 0, K0PerBlock * KPack)); } - __device__ void ResetABlockStartWindow() + __device__ void ResetBBlockStartWindow() { - a_thread_copy_.SetSrcCoord(CalculateAThreadOriginDataIndex()); + b_thread_copy_.SetSrcCoord(CalculateBThreadOriginDataIndex()); } - static constexpr auto a_block_desc_m0_m1_m2_k = MakeABlockDescriptor_M0_M1_M2_K(); + static constexpr auto b_block_desc_n0_n1_n2_k = MakeBBlockDescriptor_N0_N1_N2_K(); template - __device__ void Run(const ABlockBuffer& a_block_buf, - const BBlockBuffer& b_thread_buf, + __device__ void Run(const ABlockBuffer& a_thread_buf, + const BBlockBuffer& b_block_buf, CThreadBuffer& c_thread_buf) const { - auto a_thread_buf = make_static_buffer( - a_thread_desc_.GetElementSpaceSize()); + auto b_thread_buf = make_static_buffer( + b_thread_desc_.GetElementSpaceSize()); - static_for<0, MRepeat, 1>{}([&](auto m0) { + static_for<0, NRepeat, 1>{}([&](auto n0) { // read A - a_thread_copy_.Run(a_block_desc_m0_m1_m2_k, - make_tuple(m0, I0, I0, I0), - a_block_buf, - a_thread_desc_, + b_thread_copy_.Run(b_block_desc_n0_n1_n2_k, + make_tuple(n0, I0, I0, I0), + b_block_buf, + b_thread_desc_, make_tuple(I0, I0, I0, I0), - a_thread_buf); + b_thread_buf); - static_for<0, NRepeat, 1>{}([&](auto n0) { + static_for<0, MRepeat, 1>{}([&](auto m0) { // read B static_for<0, KPerThread, KPack>{}([&](auto k) { vector_type a_thread_vec; vector_type b_thread_vec; constexpr index_t k0 = k / KPack; static_for<0, KPack, 1>{}([&](auto i) { - a_thread_vec.template AsType()(i) = a_thread_buf - [Number{}]; b_thread_vec.template AsType()(i) = b_thread_buf - [Number{}]; + [Number{}]; + a_thread_vec.template AsType()(i) = a_thread_buf + [Number{}]; }); using mfma_input_type = @@ -291,30 +291,30 @@ struct BlockwiseGemmXdlops_k0mk1_k0nk1_m0n0m1n1m2m3m4n2_v1r1 private: // A[M0, M1, M2, KPerThread] - static constexpr auto a_thread_desc_ = + static constexpr auto b_thread_desc_ = make_naive_tensor_descriptor_packed(make_tuple(I1, I1, I1, Number{})); // B[N0, N1, N2, KPerThread] - static constexpr auto b_thread_desc_ = + static constexpr auto a_thread_desc_ = make_naive_tensor_descriptor_packed(make_tuple(Number{}, // KPerThread - Number{}, // repeat + Number{}, // repeat Number{})); // C[M, N, NumRegXdlops] static constexpr auto c_thread_desc_ = make_naive_tensor_descriptor_packed( make_tuple(Number{}, Number{}, xdlops_gemm.GetRegSizePerXdlops())); - using AThreadCopy = ThreadwiseTensorSliceTransfer_v4, Sequence<0, 1, 2, 3>, 3, - A_K1, - A_K1>; + B_K1, + B_K1>; - AThreadCopy a_thread_copy_{CalculateAThreadOriginDataIndex()}; + BThreadCopy b_thread_copy_{CalculateBThreadOriginDataIndex()}; }; } // namespace ck diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_xdl_cshuffle_v1.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_xdl_cshuffle_v1.hpp index 90690c15dc..57f3812314 100644 --- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_xdl_cshuffle_v1.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_xdl_cshuffle_v1.hpp @@ -79,7 +79,7 @@ template + bool SkipALds> __global__ void #if CK_USE_LAUNCH_BOUNDS __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, CK_MIN_BLOCK_PER_CU) @@ -151,7 +151,7 @@ __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, CK_MIN_BLOCK_PER_CU) } // If constexpr to be compatible with skip LDS gridwise gemm - if constexpr(SkipBLds) + if constexpr(SkipALds) { if constexpr(HasMainKBlockLoopInAllGemm || NoMainKBlockLoopInAllGemm) { @@ -348,7 +348,8 @@ template + bool SkipALds = false, + index_t ABlockBufferSize = 1> struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1 : public DeviceGroupedConvBwdDataMultipleD, element_wise::PassThrough> && !SkipBLds; + std::is_same_v, element_wise::PassThrough> && !SkipALds; // TODO: Add support for different A and B data types. using ABDataType = ADataType; @@ -411,7 +412,7 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1 static constexpr bool CTranspose = (NeedTransposeKernel == false) && (is_same_v || is_same_v) && - !SkipBLds; + !SkipALds; using ALayoutAfterTranspose = std::conditional_t< is_NGCHW_NGKHW() && NeedTransposeKernel, @@ -530,26 +531,53 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1 CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, \ CDEBlockTransferScalarPerVector_NPerBlock, LoopSched, PipelineVersion::v1, BComputeType - static constexpr index_t BBlockBufferSize = 1; + // static constexpr index_t BBlockBufferSize = 1; // Force to 1, due to KN layout for GKYXC - static constexpr index_t BScalarPerVectorSkipLds = 1; + // static constexpr index_t BScalarPerVectorSkipLds = 1; -#define GridwiseGemmMultiDSkipBLdsTemplateParams \ + + static constexpr index_t SharedMemoryABlockBufferRequiredSize() + { + constexpr auto max_lds_align = AK1; + + // A matrix in LDS memory, dst of blockwise copy + constexpr auto b_block_desc_k0_n_k1 = [&]() { + if constexpr(BBlockLdsExtraN) + { + return make_naive_tensor_descriptor( + make_tuple(Number<(KPerBlock / AK1) * ABlockBufferSize>{}, Number{}, AK1), + make_tuple(Number{} * AK1, AK1, I1)); + } + else + { + return make_naive_tensor_descriptor_packed( + make_tuple(Number<(KPerBlock / AK1) * ABlockBufferSize>{}, Number{}, AK1)); + } + }(); + + constexpr auto b_block_space_size_aligned = + math::integer_least_multiple(b_block_desc_k0_n_k1.GetElementSpaceSize(), max_lds_align); + + return b_block_space_size_aligned * sizeof(ABDataType); + } + + static constexpr index_t ABlockBufferSizeAligned = std::max(1, std::min(65536 / SharedMemoryABlockBufferRequiredSize(), ABlockBufferSize)); + +#define GridwiseGemmMultiDSkipALdsTemplateParams \ BlockSize, ABDataType, AccDataType, CShuffleDataType, DsDataType, EDataType, \ InMemoryDataOperationEnum::Set, element_wise::PassThrough, element_wise::PassThrough, \ element_wise::PassThrough, MPerBlock, NPerBlock, KPerBlock / AK1, MPerXDL, NPerXDL, AK1, \ - MXdlPerWave, NXdlPerWave, ABlockTransferThreadClusterLengths_AK0_M_AK1, \ - ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, \ - ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, \ - ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, BScalarPerVectorSkipLds, \ - false, BBlockBufferSize, CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, \ + MXdlPerWave, NXdlPerWave, \ + ABlockTransferSrcScalarPerVector, false, ABlockBufferSizeAligned, \ + BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, \ + CShuffleMXdlPerWavePerShuffle, CShuffleNXdlPerWavePerShuffle, \ CDEBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, \ CDEBlockTransferScalarPerVector_NPerBlock using GridwiseGemm = - std::conditional_t, + GridwiseGemmMultiDSkipALdsTemplateParams>, GridwiseGemmMultipleD_xdl_cshuffle>; using GridwiseGemmCTranspose = std::conditional_t< CTranspose, @@ -1005,7 +1033,7 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1 const auto GemmK = a_grid_desc_m_k.GetLength(I1); const auto GemmK0 = a_grid_desc_ak0_m_ak1.GetLength(I0); bool HasMainKBlockLoop = true; - if constexpr(SkipBLds) + if constexpr(SkipALds) { HasMainKBlockLoop = GridwiseGemmCTranspose::CalculateHasMainK0BlockLoop(GemmK0); @@ -1091,6 +1119,33 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1 compute_ptr_offset_of_workspace_n_.BatchStrideE_ = e_g_n_c_wis_strides[1] * conv_N_per_block_; } + + has_loop_in_all_gemm_.resize(gemm_kernel_args_.size()); + no_loop_in_all_gemm_.resize(gemm_kernel_args_.size()); + for(std::size_t gemm_set_id = 0; gemm_set_id < gemm_kernel_args_.size(); + gemm_set_id++) + { + const std::array& gemm_kernel_args = + gemm_kernel_args_[gemm_set_id]; + + const index_t gemms_count_for_set = + gemm_set_id == gemm_kernel_args_.size() - 1 + ? gemms_count_ - MaxGroupedGemmGroupsNum * gemm_set_id + : MaxGroupedGemmGroupsNum; + + bool has_loop_in_all_gemm = true; + bool no_loop_in_all_gemm = true; + for(auto i = 0; i < gemms_count_for_set; i++) + { + has_loop_in_all_gemm &= gemm_kernel_args[i].HasMainKBlockLoop_; + no_loop_in_all_gemm &= !gemm_kernel_args[i].HasMainKBlockLoop_; + } + has_loop_in_all_gemm_[gemm_set_id] =has_loop_in_all_gemm; + no_loop_in_all_gemm_[gemm_set_id] =no_loop_in_all_gemm; + } + + gdy_ = num_group_; + gdz_ = num_workgroups_per_Conv_N_ * k_batch_; } std::size_t GetWorkspaceATensorSizeBytes() const @@ -1209,9 +1264,12 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1 std::vector gemms_grid_size_; index_t gemms_count_ = 0; std::vector> gemm_kernel_args_; + std::vector has_loop_in_all_gemm_; + std::vector no_loop_in_all_gemm_; bool bwd_needs_zero_out; long_index_t e_space_size_bytes; + index_t gdy_, gdz_; }; // Invoker @@ -1224,9 +1282,6 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1 { float ave_time = 0; - const index_t gdy = arg.num_group_; - const index_t gdz = arg.num_workgroups_per_Conv_N_ * arg.k_batch_; - const ADataType* p_a_grid = arg.p_a_grid_; const BDataType* p_b_grid = arg.p_b_grid_; EDataType* p_e_grid = arg.p_e_grid_; @@ -1257,8 +1312,6 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1 gemm_set_id == arg.gemm_kernel_args_.size() - 1 ? arg.gemms_count_ - MaxGroupedGemmGroupsNum * gemm_set_id : MaxGroupedGemmGroupsNum; - const std::array& gemm_kernel_args = - arg.gemm_kernel_args_[gemm_set_id]; const auto clear_workspace = [&]() { if(arg.bwd_needs_zero_out && gemm_set_id == 0) @@ -1268,14 +1321,6 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1 } }; - bool has_loop_in_all_gemm = true; - bool no_loop_in_all_gemm = true; - for(auto i = 0; i < gemms_count_for_set; i++) - { - has_loop_in_all_gemm &= gemm_kernel_args[i].HasMainKBlockLoop_; - no_loop_in_all_gemm &= !gemm_kernel_args[i].HasMainKBlockLoop_; - } - auto launch_kernel = [&](auto has_main_k_block_loop, auto no_main_k_block_loop) { constexpr bool has_main_loop = has_main_k_block_loop.value; constexpr bool no_main_loop = no_main_k_block_loop.value; @@ -1297,20 +1342,20 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1 has_main_loop, no_main_loop, CTranspose, - SkipBLds>; + SkipALds>; return launch_and_time_kernel_with_preprocess( stream_config, clear_workspace, kernel, - dim3(gdx, gdy, gdz), + dim3(gdx, arg.gdy_, arg.gdz_), dim3(BlockSize), 0, p_b_grid, p_a_grid, arg.p_ds_grid_, p_e_grid, - gemm_kernel_args, + arg.gemm_kernel_args_[gemm_set_id], gemms_count_for_set, arg.b_element_op_, arg.a_element_op_, @@ -1337,20 +1382,20 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1 has_main_loop, no_main_loop, CTranspose, - SkipBLds>; + SkipALds>; return launch_and_time_kernel_with_preprocess( stream_config, clear_workspace, kernel, - dim3(gdx, gdy, gdz), + dim3(gdx, arg.gdy_, arg.gdz_), dim3(BlockSize), 0, p_a_grid, p_b_grid, arg.p_ds_grid_, p_e_grid, - gemm_kernel_args, + arg.gemm_kernel_args_[gemm_set_id], gemms_count_for_set, arg.a_element_op_, arg.b_element_op_, @@ -1360,12 +1405,12 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1 arg.k_batch_); } }; - if(has_loop_in_all_gemm) + if(arg.has_loop_in_all_gemm_[gemm_set_id]) { ave_time += launch_kernel(integral_constant{}, integral_constant{}); } - else if(no_loop_in_all_gemm) + else if(arg.no_loop_in_all_gemm_[gemm_set_id]) { ave_time += launch_kernel(integral_constant{}, integral_constant{}); @@ -1384,10 +1429,10 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1 { float ave_time = 0; - if(stream_config.log_level_ > 0) - { - arg.Print(); - } + // if(stream_config.log_level_ > 0) + // { + // arg.Print(); + // } // Transpose from NGKHW to NHWGK if constexpr(NeedTransposeKernel) @@ -1699,7 +1744,7 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1 // Gridwise GEMM size for(std::size_t i = 0; i < arg.a_grid_desc_m_k_container_.size(); i++) { - if constexpr(SkipBLds) + if constexpr(SkipALds) { if(!GridwiseGemmCTranspose::CheckValidity( arg.gemm_kernel_args_[i / MaxGroupedGemmGroupsNum] @@ -1884,9 +1929,9 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1 auto str = std::stringstream(); str << "DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1"; - if constexpr(SkipBLds) + if constexpr(SkipALds) { - str << "_SkipBLds"; + str << "_SkipALds"; } // clang-format off @@ -1914,6 +1959,10 @@ struct DeviceGroupedConvBwdDataMultipleD_Xdl_CShuffle_v1 << "TransposeTransferOutScalarPerVectorAligned: " << TransposeTransferOutScalarPerVectorAligned; } + if constexpr(SkipALds) + { + str << ", " << ABlockBufferSize; + } str << ">"; diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_skip_b_lds_multiple_d_cshuffle.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_skip_b_lds_multiple_d_cshuffle.hpp index fa2f28cdc9..942c78343e 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_skip_b_lds_multiple_d_cshuffle.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdlops_skip_b_lds_multiple_d_cshuffle.hpp @@ -35,17 +35,20 @@ template {}, Number{}, K1), - make_tuple(Number{} * K1, K1, I1)); + make_tuple(Number{}, Number{}, K1), + make_tuple(Number{} * K1, K1, I1)); } else { return make_naive_tensor_descriptor_aligned( - make_tuple(Number{}, Number{}, K1), + make_tuple(Number{}, Number{}, K1), max_lds_align); } }(); - return a_block_desc_k0_m_k1; + return b_block_desc_k0_n_k1; } template @@ -155,7 +158,7 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle // 2-stage prefetch currently only support even number of K0 loop // TODO: add support for odd number of K0 loop - if(!((K0 / K0PerBlock) % BBlockBufferSize == 0)) + if(!((K0 / K0PerBlock) % ABlockBufferSize == 0)) { return false; } @@ -167,28 +170,28 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle // TODO move this function into GEMM-pipeline class __host__ __device__ static constexpr bool CalculateHasMainK0BlockLoop(index_t K0) { - const bool has_main_k0_block_loop = (K0 / (BBlockBufferSize * K0PerBlock)) > 1; + const bool has_main_k0_block_loop = (K0 / (ABlockBufferSize * K0PerBlock)) > 1; return has_main_k0_block_loop; } - template + template __host__ __device__ static constexpr auto - MakeBGridDescriptor_K0_K1_K2_N0_N1_N2_N3_K3(const BGridDesc_K0_N_K1& b_grid_desc_k0_n_k1) + MakeAGridDescriptor_K0_K1_K2_M0_M1_M2_M3_K3(const AGridDesc_K0_M_K1& a_grid_desc_k0_m_k1) { - const auto K0 = b_grid_desc_k0_n_k1.GetLength(I0); - const auto N = b_grid_desc_k0_n_k1.GetLength(I1); + const auto K0 = a_grid_desc_k0_m_k1.GetLength(I0); + const auto M = a_grid_desc_k0_m_k1.GetLength(I1); - const auto b_griddesc_k0_nblockid_nrepeat_waves_nperxdlops_k1 = transform_tensor_descriptor( - b_grid_desc_k0_n_k1, + const auto a_griddesc_k0_mblockid_mrepeat_waves_mperxdlops_k1 = transform_tensor_descriptor( + a_grid_desc_k0_m_k1, make_tuple(make_unmerge_transform( make_tuple(K0 / K0PerBlock, xdlops_gemm.K0PerXdlops, K0PerThread)), make_unmerge_transform(make_tuple( - N / (NXdlPerWave * NWaves * NPerXdl), NXdlPerWave, NWaves, NPerXdl)), + M / (MXdlPerWave * MWaves * MPerXdl), MXdlPerWave, MWaves, MPerXdl)), make_pass_through_transform(K1)), make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}), make_tuple(Sequence<0, 1, 2>{}, Sequence<3, 4, 5, 6>{}, Sequence<7>{})); - return b_griddesc_k0_nblockid_nrepeat_waves_nperxdlops_k1; + return a_griddesc_k0_mblockid_mrepeat_waves_mperxdlops_k1; } template @@ -235,14 +238,14 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle return threadid_to_wave_idx_adaptor.CalculateBottomIndex(make_multi_index(thread_id)); } - __device__ static auto GetWaveKNIdx(const index_t thread_id) + __device__ static auto GetWaveKMIdx(const index_t thread_id) { - constexpr auto wave_threadid_to_nk_idx_adaptor = make_single_stage_tensor_adaptor( - make_tuple(make_merge_transform(make_tuple(xdlops_gemm.K0PerXdlops, NPerXdl))), + constexpr auto wave_threadid_to_mk_idx_adaptor = make_single_stage_tensor_adaptor( + make_tuple(make_merge_transform(make_tuple(xdlops_gemm.K0PerXdlops, MPerXdl))), make_tuple(Sequence<0, 1>{}), make_tuple(Sequence<0>{})); - return wave_threadid_to_nk_idx_adaptor.CalculateBottomIndex(make_multi_index(thread_id)); + return wave_threadid_to_mk_idx_adaptor.CalculateBottomIndex(make_multi_index(thread_id)); } __device__ static constexpr auto GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock() @@ -263,12 +266,12 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle __host__ __device__ static constexpr index_t GetSharedMemoryNumberOfByte() { // LDS allocation for A and B: be careful of alignment - constexpr auto a_block_desc_k0_m_k1 = GetABlockDescriptor_K0PerBlock_MPerBlock_K1(); + constexpr auto b_block_desc_k0_n_k1 = GetBBlockDescriptor_K0PerBlock_NPerBlock_K1(); constexpr auto max_lds_align = K1; - constexpr auto a_block_space_size_aligned = - math::integer_least_multiple(a_block_desc_k0_m_k1.GetElementSpaceSize(), max_lds_align); + constexpr auto b_block_space_size_aligned = + math::integer_least_multiple(b_block_desc_k0_n_k1.GetElementSpaceSize(), max_lds_align); constexpr auto c_shuffle_block_desc_mblock_mperblock_nblock_nperblock = GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(); @@ -276,7 +279,7 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle constexpr auto c_block_size = c_shuffle_block_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize(); - return math::max((a_block_space_size_aligned) * sizeof(ABDataType), + return math::max((b_block_space_size_aligned) * sizeof(ABDataType), c_block_size * sizeof(CShuffleDataType)); } @@ -303,13 +306,13 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle const Block2CTileMap& block_2_ctile_map) { constexpr index_t NumDTensor = DsGridDesc_MBlock_MPerBlock_NBlock_NPerBlock::Size(); - const auto b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3 = - MakeBGridDescriptor_K0_K1_K2_N0_N1_N2_N3_K3(b_grid_desc_k0_n_k1); + const auto a_grid_desc_k0_k1_k2_m0_m1_m2_m3_k3 = + MakeAGridDescriptor_K0_K1_K2_M0_M1_M2_M3_K3(a_grid_desc_k0_m_k1); const auto a_grid_buf = make_dynamic_buffer( - p_a_grid, a_grid_desc_k0_m_k1.GetElementSpaceSize()); + p_a_grid, a_grid_desc_k0_k1_k2_m0_m1_m2_m3_k3.GetElementSpaceSize()); const auto b_grid_buf = make_dynamic_buffer( - p_b_grid, b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3.GetElementSpaceSize()); + p_b_grid, b_grid_desc_k0_n_k1.GetElementSpaceSize()); const auto ds_grid_buf = generate_tuple( [&](auto i) { return make_dynamic_buffer( @@ -320,113 +323,113 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle auto c_grid_buf = make_dynamic_buffer( p_c_grid, c_grid_desc_mblock_mperblock_nblock_nperblock.GetElementSpaceSize()); - const auto K0 = a_grid_desc_k0_m_k1.GetLength(I0); + const auto K0 = b_grid_desc_k0_n_k1.GetLength(I0); // divide block work by [M, N] const auto block_work_idx = block_2_ctile_map.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); // A matrix in LDS memory, dst of blockwise copy - constexpr auto a_block_desc_k0_m_k1 = GetABlockDescriptor_K0PerBlock_MPerBlock_K1(); + constexpr auto b_block_desc_k0_n_k1 = GetBBlockDescriptor_K0PerBlock_NPerBlock_K1(); // A matrix blockwise copy - auto a_blockwise_copy = ThreadGroupTensorSliceTransfer_v4r1< + auto b_blockwise_copy = ThreadGroupTensorSliceTransfer_v4r1< ThisThreadBlock, - AElementwiseOperation, + BElementwiseOperation, ck::tensor_operation::element_wise::PassThrough, InMemoryDataOperationEnum::Set, - Sequence, - ABlockTransferThreadClusterLengths_K0_M_K1, - ABlockTransferThreadClusterArrangeOrder, + Sequence, + BBlockTransferThreadClusterLengths_K0_N_K1, + BBlockTransferThreadClusterArrangeOrder, ABDataType, ABDataType, - decltype(a_grid_desc_k0_m_k1), - decltype(a_block_desc_k0_m_k1), - ABlockTransferSrcAccessOrder, + decltype(b_grid_desc_k0_n_k1), + decltype(b_block_desc_k0_n_k1), + BBlockTransferSrcAccessOrder, Sequence<1, 0, 2>, - ABlockTransferSrcVectorDim, + BBlockTransferSrcVectorDim, 2, - ABlockTransferSrcScalarPerVector, - ABlockTransferDstScalarPerVector_K1, + BBlockTransferSrcScalarPerVector, + BBlockTransferDstScalarPerVector_K1, 1, 1, - AThreadTransferSrcResetCoordinateAfterRun, + BThreadTransferSrcResetCoordinateAfterRun, true, - 1>(a_grid_desc_k0_m_k1, - make_multi_index(0, m_block_data_idx_on_grid, 0), - a_element_op, - a_block_desc_k0_m_k1, + 1>(b_grid_desc_k0_n_k1, + make_multi_index(0, n_block_data_idx_on_grid, 0), + b_element_op, + b_block_desc_k0_n_k1, make_multi_index(0, 0, 0), ck::tensor_operation::element_wise::PassThrough{}); - ignore = b_element_op; + ignore = a_element_op; // B matrix threadwise copy - constexpr auto b_thread_desc_k0_k1_k2_n0_n1_n2_n3_k3 = + constexpr auto a_thread_desc_k0_k1_k2_m0_m1_m2_m3_k3 = make_naive_tensor_descriptor_packed(make_tuple(I1, I1, Number{}, // K0PerThread I1, // NBlockId - Number{}, // repeat + Number{}, // repeat I1, // waves I1, // NPerXdlops Number{})); - auto b_thread_buf = generate_tuple( + auto a_thread_buf = generate_tuple( [&](auto i) { ignore = i; return StaticBuffer{}; }, - Number{}); + Number{}); const auto wave_id = GetWaveIdx(); - const auto wave_k_n_id = GetWaveKNIdx(wave_id[I2]); + const auto wave_k_m_id = GetWaveKMIdx(wave_id[I2]); -#if 0 - const index_t block_id = get_block_1d_id(); - const index_t thread_id = get_thread_local_1d_id(); - printf("block id: %d m blockid: %d n block id: %d ,thread id: %d, wave id :{%d %d %d} " - "kn id: {%d %d}\n", - block_id, - block_work_idx[I0], - block_work_idx[I1], - thread_id, - wave_id[I0], - wave_id[I1], - wave_id[I2], - wave_k_n_id[I0], - wave_k_n_id[I1]); - printf("mfma thread k per xdlops: %d K0PerThread: %d HasMainK0BlockLoop: %d K0: %d \t", - xdlops_gemm.K0PerXdlops, K0PerThread, HasMainK0BlockLoop, b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3.GetLength(I0)); -#endif +// #if 0 +// const index_t block_id = get_block_1d_id(); +// const index_t thread_id = get_thread_local_1d_id(); +// printf("block id: %d m blockid: %d n block id: %d ,thread id: %d, wave id :{%d %d %d} " +// "kn id: {%d %d}\n", +// block_id, +// block_work_idx[I0], +// block_work_idx[I1], +// thread_id, +// wave_id[I0], +// wave_id[I1], +// wave_id[I2], +// wave_k_n_id[I0], +// wave_k_n_id[I1]); +// printf("mfma thread k per xdlops: %d K0PerThread: %d HasMainK0BlockLoop: %d K0: %d \t", +// xdlops_gemm.K0PerXdlops, K0PerThread, HasMainK0BlockLoop, b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3.GetLength(I0)); +// #endif - auto b_threadwise_copy = + auto a_threadwise_copy = ThreadwiseTensorSliceTransfer_v2{}, I1, - Number{}, + Number{}, I1, I1, Number{}>, Sequence<0, 1, 2, 3, 4, 5, 6, 7>, 7, - BBlockTransferSrcScalarPerVector, - BThreadTransferSrcResetCoordinateAfterRun, + ABlockTransferSrcScalarPerVector, + AThreadTransferSrcResetCoordinateAfterRun, true>( - b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3, + a_grid_desc_k0_k1_k2_m0_m1_m2_m3_k3, make_multi_index( - 0, wave_k_n_id[I0], 0, block_work_idx[I1], 0, wave_id[I1], wave_k_n_id[I1], 0)); + 0, wave_k_m_id[I0], 0, block_work_idx[I0], 0, wave_id[I0], wave_k_m_id[I1], 0)); // GEMM definition // c_mtx += transpose(a_mtx) * b_mtx @@ -439,8 +442,8 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle BlockSize, ABDataType, AccDataType, - decltype(a_block_desc_k0_m_k1), - decltype(b_thread_desc_k0_k1_k2_n0_n1_n2_n3_k3), + decltype(a_thread_desc_k0_k1_k2_m0_m1_m2_m3_k3), + decltype(b_block_desc_k0_n_k1), MPerBlock, NPerBlock, K0PerBlock, @@ -453,64 +456,64 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle auto c_thread_buf = blockwise_gemm.GetCThreadBuffer(); // LDS allocation for A - auto a_block_buf = make_dynamic_buffer( - static_cast(p_shared), a_block_desc_k0_m_k1.GetElementSpaceSize()); + auto b_block_buf = make_dynamic_buffer( + static_cast(p_shared), b_block_desc_k0_n_k1.GetElementSpaceSize()); // gridwise GEMM pipeline - constexpr auto a_block_slice_copy_step = - make_multi_index(K0PerBlock * BBlockBufferSize, 0, 0); - constexpr auto b_thread_slice_copy_step = make_multi_index(1, 0, 0, 0, 0, 0, 0, 0); + constexpr auto a_thread_slice_copy_step = make_multi_index(1, 0, 0, 0, 0, 0, 0, 0); + constexpr auto b_block_slice_copy_step = + make_multi_index(K0PerBlock * ABlockBufferSize, 0, 0); // preload data to regiester and LDS { // Read - a_blockwise_copy.RunRead(a_grid_desc_k0_m_k1, a_grid_buf); - a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc_k0_m_k1, a_block_slice_copy_step); + b_blockwise_copy.RunRead(b_grid_desc_k0_n_k1, b_grid_buf); + b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc_k0_n_k1, b_block_slice_copy_step); - static_for<0, BBlockBufferSize, 1>{}([&](auto ii) { - b_threadwise_copy.Run(b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3, - b_grid_buf, - b_thread_desc_k0_k1_k2_n0_n1_n2_n3_k3, + static_for<0, ABlockBufferSize, 1>{}([&](auto ii) { + a_threadwise_copy.Run(a_grid_desc_k0_k1_k2_m0_m1_m2_m3_k3, + a_grid_buf, + a_thread_desc_k0_k1_k2_m0_m1_m2_m3_k3, make_tuple(I0, I0, I0, I0, I0, I0, I0, I0), - b_thread_buf(Number{})); - b_threadwise_copy.MoveSrcSliceWindow(b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3, - b_thread_slice_copy_step); + a_thread_buf(Number{})); + a_threadwise_copy.MoveSrcSliceWindow(a_grid_desc_k0_k1_k2_m0_m1_m2_m3_k3, + a_thread_slice_copy_step); }); // Initialize C c_thread_buf.Clear(); // a data write to lds - a_blockwise_copy.RunWrite(a_block_desc_k0_m_k1, a_block_buf); + b_blockwise_copy.RunWrite(b_block_desc_k0_n_k1, b_block_buf); // main body if constexpr(HasMainK0BlockLoop) { index_t K0BlockMainLoop = - __builtin_amdgcn_readfirstlane(K0 / (BBlockBufferSize * K0PerBlock)); + __builtin_amdgcn_readfirstlane(K0 / (ABlockBufferSize * K0PerBlock)); index_t i = 0; do { - a_blockwise_copy.RunRead(a_grid_desc_k0_m_k1, a_grid_buf); - blockwise_gemm.ResetABlockStartWindow(); + b_blockwise_copy.RunRead(b_grid_desc_k0_n_k1, b_grid_buf); + blockwise_gemm.ResetBBlockStartWindow(); block_sync_lds(); - static_for<0, BBlockBufferSize, 1>{}([&](auto ii) { - blockwise_gemm.Run(a_block_buf, b_thread_buf(Number{}), c_thread_buf); - blockwise_gemm.MoveABlockSliceWindow(); + static_for<0, ABlockBufferSize, 1>{}([&](auto ii) { + blockwise_gemm.Run(a_thread_buf(Number{}), b_block_buf, c_thread_buf); + blockwise_gemm.MoveBBlockSliceWindow(); s_nop(); - b_threadwise_copy.Run(b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3, - b_grid_buf, - b_thread_desc_k0_k1_k2_n0_n1_n2_n3_k3, + a_threadwise_copy.Run(a_grid_desc_k0_k1_k2_m0_m1_m2_m3_k3, + a_grid_buf, + a_thread_desc_k0_k1_k2_m0_m1_m2_m3_k3, make_tuple(I0, I0, I0, I0, I0, I0, I0, I0), - b_thread_buf(Number{})); - b_threadwise_copy.MoveSrcSliceWindow(b_grid_desc_k0_k1_k2_n0_n1_n2_n3_k3, - b_thread_slice_copy_step); + a_thread_buf(Number{})); + a_threadwise_copy.MoveSrcSliceWindow(a_grid_desc_k0_k1_k2_m0_m1_m2_m3_k3, + a_thread_slice_copy_step); }); block_sync_lds(); - a_blockwise_copy.RunWrite(a_block_desc_k0_m_k1, a_block_buf); + b_blockwise_copy.RunWrite(b_block_desc_k0_n_k1, b_block_buf); // move a and b window - a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc_k0_m_k1, - a_block_slice_copy_step); + b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc_k0_n_k1, + b_block_slice_copy_step); i += 1; } while(i < (K0BlockMainLoop - 1)); @@ -520,11 +523,11 @@ struct GridwiseGemm_xdlops_skip_b_lds_multiple_d_cshuffle { block_sync_lds(); - blockwise_gemm.ResetABlockStartWindow(); + blockwise_gemm.ResetBBlockStartWindow(); - static_for<0, BBlockBufferSize, 1>{}([&](auto ii) { - blockwise_gemm.Run(a_block_buf, b_thread_buf(Number{}), c_thread_buf); - blockwise_gemm.MoveABlockSliceWindow(); + static_for<0, ABlockBufferSize, 1>{}([&](auto ii) { + blockwise_gemm.Run(a_thread_buf(Number{}), b_block_buf, c_thread_buf); + blockwise_gemm.MoveBBlockSliceWindow(); }); } }