diff --git a/projects/composablekernel/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_data_multiple_d_wmma_cshuffle_v3.hpp b/projects/composablekernel/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_data_multiple_d_wmma_cshuffle_v3.hpp index 13b892857c12..27cc0ffef1fa 100644 --- a/projects/composablekernel/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_data_multiple_d_wmma_cshuffle_v3.hpp +++ b/projects/composablekernel/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_data_multiple_d_wmma_cshuffle_v3.hpp @@ -60,7 +60,9 @@ template + ck::index_t max_transpose_transfer_dst_scalar_per_vector, + bool UseLdsTranspose, + bool TransposeC> struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3; } // namespace ck::tensor_operation::device @@ -120,7 +122,9 @@ template + ck::index_t max_transpose_transfer_dst_scalar_per_vector, + bool UseLdsTranspose, + bool TransposeC> struct InstanceTraits< ck::tensor_operation::device::DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3< NDimSpatial, @@ -173,7 +177,9 @@ struct InstanceTraits< ComputeTypeA_, ComputeTypeB_, max_transpose_transfer_src_scalar_per_vector, - max_transpose_transfer_dst_scalar_per_vector>> + max_transpose_transfer_dst_scalar_per_vector, + UseLdsTranspose, + TransposeC>> { static constexpr auto kTensorOpName = "DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3"; diff --git a/projects/composablekernel/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp b/projects/composablekernel/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp index cfc8b4e05af7..be7f0919142a 100644 --- a/projects/composablekernel/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp +++ b/projects/composablekernel/experimental/builder/include/ck_tile/builder/reflect/instance_traits_device_grouped_conv_bwd_weight_wmma_cshuffle_v3.hpp @@ -54,7 +54,9 @@ template + ck::index_t MaxTransposeTransferDstScalarPerVector, + bool UseLdsTranspose, + bool TransposeC> struct DeviceGroupedConvBwdWeight_Wmma_CShuffleV3; } // namespace ck::tensor_operation::device @@ -112,7 +114,9 @@ template + ck::index_t MaxTransposeTransferDstScalarPerVector, + bool UseLdsTranspose, + bool TransposeC> struct InstanceTraits> + MaxTransposeTransferDstScalarPerVector, + UseLdsTranspose, + TransposeC>> { static constexpr auto kTensorOpName = "DeviceGroupedConvBwdWeight_Wmma_CShuffleV3"; using device_kernel_tag = DeviceGroupedConvBwdWeight_Wmma_CShuffle_V3_Tag; diff --git a/projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmma_selector.hpp b/projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmma_selector.hpp index fbd794f4770a..57681441af44 100644 --- a/projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmma_selector.hpp +++ b/projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmma_selector.hpp @@ -29,8 +29,10 @@ template + bool TransposeC = false, + bool BSkipLDS = false, + bool UseLdsTransposeA = false, + bool UseLdsTransposeB = false> constexpr auto BlockGemmPipeline_Selector() { if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1) @@ -56,7 +58,9 @@ constexpr auto BlockGemmPipeline_Selector() KPack, KInner, TransposeC, - BSkipLDS>{}; + BSkipLDS, + UseLdsTransposeA, + UseLdsTransposeB>{}; } else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v3) { @@ -81,7 +85,9 @@ constexpr auto BlockGemmPipeline_Selector() KPack, KInner, TransposeC, - BSkipLDS>{}; + BSkipLDS, + UseLdsTransposeA, + UseLdsTransposeB>{}; } else { diff --git a/projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_base.hpp b/projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_base.hpp index 878bb9b3b10f..d3b3f718eccf 100644 --- a/projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_base.hpp +++ b/projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_base.hpp @@ -36,7 +36,9 @@ template + bool TransposeC = false, + bool UseLdsTransposeA = false, + bool UseLdsTransposeB = false> struct BlockwiseGemmWmmaops_pipeline_base { static constexpr auto I0 = Number<0>{}; @@ -48,6 +50,8 @@ struct BlockwiseGemmWmmaops_pipeline_base using ThisThreadBlock = ThisThreadBlock; + static constexpr bool TransposeC_ = TransposeC; + static constexpr index_t WaveSize = 32; static constexpr index_t MWaves = MPerBlock / (MRepeat * MPerWmma); @@ -70,8 +74,12 @@ struct BlockwiseGemmWmmaops_pipeline_base TransposeC>{}; static constexpr index_t KPerThread = wmma_gemm.wmma_instr.k_per_blk * KInner; - static constexpr index_t A_K1 = ck::math::min(AWmmaTileDesc{}.GetLength(I6), KPerThread); - static constexpr index_t B_K1 = ck::math::min(BWmmaTileDesc{}.GetLength(I6), KPerThread); + static constexpr index_t A_K1 = UseLdsTransposeA + ? ck::math::min(AWmmaTileDesc{}.GetLength(I5), KPerThread) + : ck::math::min(AWmmaTileDesc{}.GetLength(I6), KPerThread); + static constexpr index_t B_K1 = UseLdsTransposeB + ? ck::math::min(BWmmaTileDesc{}.GetLength(I5), KPerThread) + : ck::math::min(BWmmaTileDesc{}.GetLength(I6), KPerThread); static_assert(KPack % (A_K1 * A_KRow) == 0, "wrong!"); static_assert(KPack % (B_K1 * B_KRow) == 0, "wrong!"); @@ -282,8 +290,20 @@ struct BlockwiseGemmWmmaops_pipeline_base #else const auto wmma_krow = 0; #endif - - return make_tuple(0, 0, 0, waveId_m, wmma_krow, wmma_a_idx, 0); + if constexpr(UseLdsTransposeA) + { + return make_tuple(0, + 0, + 0, + waveId_m, + wmma_krow, + wmma_a_idx % A_K1, + (wmma_a_idx / A_K1) * (MPerWmma / 2)); + } + else + { + return make_tuple(0, 0, 0, waveId_m, wmma_krow, wmma_a_idx, 0); + } } __device__ static auto CalculateBThreadOriginDataIndex() @@ -300,7 +320,20 @@ struct BlockwiseGemmWmmaops_pipeline_base const auto wmma_krow = 0; #endif - return make_tuple(0, 0, 0, waveId_n, wmma_krow, wmma_b_idx, 0); + if constexpr(UseLdsTransposeB) + { + return make_tuple(0, + 0, + 0, + waveId_n, + wmma_krow, + wmma_b_idx % B_K1, + (wmma_b_idx / B_K1) * (NPerWmma / 2)); + } + else + { + return make_tuple(0, 0, 0, waveId_n, wmma_krow, wmma_b_idx, 0); + } } template @@ -405,6 +438,7 @@ struct BlockwiseGemmWmmaops_pipeline_base AccStride)); } + // C = A * B __host__ __device__ static constexpr auto GetCBlockDescriptor_MRepeat_MWave_MSubGroup_NRepeat_NWave_NThreadPerSubGroup_MAccVgprs() { @@ -421,6 +455,23 @@ struct BlockwiseGemmWmmaops_pipeline_base c_block_desc_mrepeat_mwave_mperwmma_nrepeat_nwave_nperwmma); } + // C' = B' * A' + __host__ __device__ static constexpr auto + GetCBlockDescriptor_MRepeat_MWave_MThreadPerSubGroup_NRepeat_NWave_NSubGroup_MAccVgprs() + { + constexpr auto c_block_desc_mrepeat_mwave_mperwmma_nrepeat_nwave_nperwmma = + make_naive_tensor_descriptor_packed(make_tuple(Number{}, + Number{}, + Number{}, + Number{}, + Number{}, + Number{})); + + return wmma_gemm + .MakeCDesc_MBlockxRepeat_MWave_MThreadPerSubGroup_NBlockxRepeat_NWave_NSubGroup_NAccVgprs( + c_block_desc_mrepeat_mwave_mperwmma_nrepeat_nwave_nperwmma); + } + // Describe how data allocated in thread copy src buffer // M0_M1_M2 = MRepeat_MWave_MPerWmma, N0_N1_N2 = NRepeat_NWave_NPerWmma static constexpr AWmmaTileDesc a_block_desc_k0_m0_m1_m2_k1; @@ -472,7 +523,8 @@ struct BlockwiseGemmWmmaops_pipeline_base Sequence<0, 1, 2, 3, 4, 5, 6>, 6, A_K1, - A_K1>; + A_K1, + UseLdsTransposeA>; using BThreadCopy = ThreadwiseTensorSliceTransfer_v4, 6, B_K1, - B_K1>; + B_K1, + UseLdsTransposeB>; AThreadCopy a_thread_copy_; BThreadCopy b_thread_copy_; diff --git a/projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v1.hpp b/projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v1.hpp index 26053b57ed21..37be93a5ec37 100644 --- a/projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v1.hpp +++ b/projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v1.hpp @@ -29,8 +29,10 @@ template + bool TransposeC = false, + bool BSkipLDS = false, + bool UseLdsTransposeA = false, + bool UseLdsTransposeB = false> struct BlockwiseGemmWmmaops_pipeline_v1 { }; @@ -54,7 +56,9 @@ template + bool TransposeC, + bool UseLdsTransposeA, + bool UseLdsTransposeB> struct BlockwiseGemmWmmaops_pipeline_v1 + false, + UseLdsTransposeA, + UseLdsTransposeB> : BlockwiseGemmWmmaops_pipeline_base + TransposeC, + UseLdsTransposeA, + UseLdsTransposeB> { // GlobalPrefetchStages: 1 // LocalPreFillStages: 1 @@ -121,7 +129,9 @@ struct BlockwiseGemmWmmaops_pipeline_v1; + TransposeC, + UseLdsTransposeA, + UseLdsTransposeB>; using Base::I0; using Base::I1; using typename Base::HotLoopInstList; @@ -585,7 +595,8 @@ struct BlockwiseGemmWmmaops_pipeline_v1, 6, A_K1, - A_K1>; + A_K1, + UseLdsTransposeA>; using BThreadCopy = ThreadwiseTensorSliceTransfer_v4, 6, B_K1, - B_K1>; + B_K1, + UseLdsTransposeB>; AThreadCopy a_thread_copy_{Base::CalculateAThreadOriginDataIndex()}; BThreadCopy b_thread_copy_{Base::CalculateBThreadOriginDataIndex()}; @@ -622,7 +634,9 @@ template + bool TransposeC, + bool UseLdsTransposeA, + bool UseLdsTransposeB> struct BlockwiseGemmWmmaops_pipeline_v1 + false, + UseLdsTransposeA, + UseLdsTransposeB> : BlockwiseGemmWmmaops_pipeline_base + TransposeC, + UseLdsTransposeA, + UseLdsTransposeB> { // GlobalPrefetchStages: 1 // LocalPreFillStages: 1 @@ -689,7 +707,9 @@ struct BlockwiseGemmWmmaops_pipeline_v1; + TransposeC, + UseLdsTransposeA, + UseLdsTransposeB>; using Base::I0; using Base::I1; @@ -1020,7 +1040,8 @@ struct BlockwiseGemmWmmaops_pipeline_v1, 6, A_K1, - A_K1>; + A_K1, + UseLdsTransposeA>; using BThreadCopy = ThreadwiseTensorSliceTransfer_v4, 6, B_K1, - B_K1>; + B_K1, + UseLdsTransposeB>; AThreadCopy a_thread_copy_{Base::CalculateAThreadOriginDataIndex()}; BThreadCopy b_thread_copy_{Base::CalculateBThreadOriginDataIndex()}; @@ -1057,7 +1079,9 @@ template + bool TransposeC, + bool UseLdsTransposeA, + bool UseLdsTransposeB> struct BlockwiseGemmWmmaops_pipeline_v1 + true, + UseLdsTransposeA, + UseLdsTransposeB> : BlockwiseGemmWmmaops_pipeline_base + TransposeC, + UseLdsTransposeA, + UseLdsTransposeB> { // GlobalPrefetchStages: 2 // LocalPreFillStages: 1 @@ -1124,7 +1152,12 @@ struct BlockwiseGemmWmmaops_pipeline_v1; + TransposeC, + UseLdsTransposeA, + UseLdsTransposeB>; + + static_assert(!UseLdsTransposeB, + "Lds Transpose not possible: preshuffleB doesn't use LDS for B matrix"); using Base::I0; using Base::I1; using Base::MWaves; diff --git a/projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v3.hpp b/projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v3.hpp index ff024c1a1c98..bcc0c7652efc 100644 --- a/projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v3.hpp +++ b/projects/composablekernel/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_wmmaops_v3.hpp @@ -33,8 +33,10 @@ template + bool TransposeC = false, + bool BSkipLDS = false, + bool UseLdsTransposeA = false, + bool UseLdsTransposeB = false> struct BlockwiseGemmWmmaops_pipeline_v3 { }; @@ -58,7 +60,9 @@ template + bool TransposeC, + bool UseLdsTransposeA, + bool UseLdsTransposeB> struct BlockwiseGemmWmmaops_pipeline_v3 + false, + UseLdsTransposeA, + UseLdsTransposeB> : BlockwiseGemmWmmaops_pipeline_base + TransposeC, + UseLdsTransposeA, + UseLdsTransposeB> { using Base = BlockwiseGemmWmmaops_pipeline_base; + TransposeC, + UseLdsTransposeA, + UseLdsTransposeB>; using Base::I0; using Base::I1; using Base::I2; diff --git a/projects/composablekernel/include/ck/tensor_operation/gpu/device/impl/device_gemm_wmma_cshuffle_v3.hpp b/projects/composablekernel/include/ck/tensor_operation/gpu/device/impl/device_gemm_wmma_cshuffle_v3.hpp index 47b4da88e64d..3967d33565f5 100644 --- a/projects/composablekernel/include/ck/tensor_operation/gpu/device/impl/device_gemm_wmma_cshuffle_v3.hpp +++ b/projects/composablekernel/include/ck/tensor_operation/gpu/device/impl/device_gemm_wmma_cshuffle_v3.hpp @@ -166,7 +166,8 @@ template + bool PermuteB = false, + bool UseLdsTranspose = false> struct DeviceGemm_Wmma_CShuffleV3 : public DeviceGemmV2; + PermuteB, + false, // IsBPreShuffled + false, // ForceThreadTileTransfer + false, // IsFusedKernel + UseLdsTranspose>; // UseLdsTranspose using Argument = typename GridwiseGemm::Argument; diff --git a/projects/composablekernel/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_wmma_cshuffle_v3.hpp b/projects/composablekernel/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_wmma_cshuffle_v3.hpp index 8dfd81bfb2b9..4ad83284712a 100644 --- a/projects/composablekernel/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_wmma_cshuffle_v3.hpp +++ b/projects/composablekernel/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_bwd_data_multiple_d_wmma_cshuffle_v3.hpp @@ -267,7 +267,9 @@ template + index_t MaxTransposeTransferOutScalarPerVector = 1, + bool UseLdsTranspose = false, + bool TransposeC = false> struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3 : public DeviceGroupedConvBwdDataMultipleD::value); + // GridwiseGemm using GridwiseGemm = GridwiseGemm_wmma_cshuffle_v3< - ALayout, - BLayout, + ck::tensor_layout::gemm::RowMajor, // ALayout + ck::tensor_layout::gemm::RowMajor, // BLayout DsLayout, ELayout, Tuple, @@ -475,23 +479,27 @@ struct DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3 false, // PermuteA false, // PermuteB false, // IsBPreShuffled - true>; // ForceThreadTileTransfer - -#define GridwiseGemmCTransposeTemplateParameters \ - ALayout, BLayout, DsLayout, ELayout, Tuple, Tuple, AccDataType, \ - CShuffleDataType, DsDataType, EDataType, BElementwiseOp, AElementwiseOp, CDEElementwiseOp, \ - GemmSpec, BlockSize, NPerBlock, MPerBlock, KPerBlock, BK1, AK1, NPerWmma, MPerWmma, \ - NRepeat, MRepeat, BBlockTransferThreadClusterLengths_BK0_N_BK1, \ - BBlockTransferThreadClusterArrangeOrder, BBlockTransferSrcAccessOrder, \ - BBlockTransferSrcVectorDim, BBlockTransferSrcScalarPerVector, \ - BBlockTransferDstScalarPerVector_BK1, false, BBlockLdsExtraN, \ - ABlockTransferThreadClusterLengths_AK0_M_AK1, ABlockTransferThreadClusterArrangeOrder, \ - ABlockTransferSrcAccessOrder, ABlockTransferSrcVectorDim, \ - ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, false, \ - ABlockLdsExtraM, CShuffleMRepeatPerShuffle, CShuffleNRepeatPerShuffle, \ - CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, \ - CShuffleBlockTransferScalarPerVector, BlkGemmPipeSched, BlkGemmPipelineVer, BComputeType, \ - AComputeType, false, false, false, true + true, // ForceThreadTileTransfer + false, // IsFusedKernel + UseLdsTranspose, + TransposeC>; + +#define GridwiseGemmCTransposeTemplateParameters \ + ck::tensor_layout::gemm::RowMajor, ck::tensor_layout::gemm::RowMajor, DsLayout, ELayout, \ + Tuple, Tuple, AccDataType, CShuffleDataType, DsDataType, EDataType, \ + BElementwiseOp, AElementwiseOp, CDEElementwiseOp, GemmSpec, BlockSize, NPerBlock, \ + MPerBlock, KPerBlock, BK1, AK1, NPerWmma, MPerWmma, NRepeat, MRepeat, \ + BBlockTransferThreadClusterLengths_BK0_N_BK1, BBlockTransferThreadClusterArrangeOrder, \ + BBlockTransferSrcAccessOrder, BBlockTransferSrcVectorDim, \ + BBlockTransferSrcScalarPerVector, BBlockTransferDstScalarPerVector_BK1, false, \ + BBlockLdsExtraN, ABlockTransferThreadClusterLengths_AK0_M_AK1, \ + ABlockTransferThreadClusterArrangeOrder, ABlockTransferSrcAccessOrder, \ + ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, \ + ABlockTransferDstScalarPerVector_AK1, false, ABlockLdsExtraM, CShuffleMRepeatPerShuffle, \ + CShuffleNRepeatPerShuffle, \ + CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock, \ + CShuffleBlockTransferScalarPerVector, BlkGemmPipeSched, BlkGemmPipelineVer, BComputeType, \ + AComputeType, false, false, false, true, false, UseLdsTranspose, TransposeC using GridwiseGemmCTranspose = std::conditional_t + index_t MaxTransposeTransferDstScalarPerVector = 1, + bool UseLdsTranspose = false, + bool TransposeC = false> struct DeviceGroupedConvBwdWeight_Wmma_CShuffleV3 : public DeviceGroupedConvBwdWeight; // ForceThreadTileTransfer + true, // ForceThreadTileTransfer + false, // IsFusedKernel + UseLdsTranspose, + TransposeC>; // Argument using CGridDesc_MBlock_MPerBlock_NBlock_NPerBlock = diff --git a/projects/composablekernel/include/ck/tensor_operation/gpu/grid/epilogue_cshuffle_v3_wmma_base.hpp b/projects/composablekernel/include/ck/tensor_operation/gpu/grid/epilogue_cshuffle_v3_wmma_base.hpp index dd12cdca8c23..3a20950f7eec 100644 --- a/projects/composablekernel/include/ck/tensor_operation/gpu/grid/epilogue_cshuffle_v3_wmma_base.hpp +++ b/projects/composablekernel/include/ck/tensor_operation/gpu/grid/epilogue_cshuffle_v3_wmma_base.hpp @@ -36,6 +36,8 @@ struct EpilogueCShuffleBase static constexpr auto I5 = Number<5>{}; static constexpr auto I6 = Number<6>{}; + static constexpr bool TransposeC = BlockwiseGemmPipe::TransposeC_; + static constexpr index_t NumDTensor = DsDataType::Size(); static constexpr auto EShuffleBlockTransferScalarPerVector = CDEShuffleBlockTransferScalarPerVectors{}[I0]; @@ -80,123 +82,254 @@ struct EpilogueCShuffleBase __device__ static constexpr auto GetCShuffleLDSDescriptor() { - // C mapping in single block - constexpr auto c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp = - BlockwiseGemmPipe:: - GetCBlockDescriptor_MRepeat_MWave_MSubGroup_NRepeat_NWave_NThreadPerSubGroup_MAccVgprs(); - - constexpr auto MWave = - c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp - .GetLength(I1); - constexpr auto MSubGroup = - c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp - .GetLength(I2); - constexpr auto NWave = - c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp - .GetLength(I4); - constexpr auto NThreadPerSubGroup = - c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp - .GetLength(I5); - constexpr auto MAccVgprs = - c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp - .GetLength(I6); - - return transform_tensor_descriptor( - GetCShuffleBlockDescriptor_MShRepeat_MPerShRepeat_NShRepeat_NPerShRepeat(), - make_tuple(make_freeze_transform(I0), - make_unmerge_transform(make_tuple( - Number{}, // MRepeat per shuffle repeat - MWave, // MWave - MSubGroup, // MSubGroup * MAccVgprs = MPerWmma - MAccVgprs)), - make_freeze_transform(I0), - make_unmerge_transform(make_tuple( - Number{}, // NRepeat per shuffle repeat - NWave, // NWave - NThreadPerSubGroup))), // NThreadPerSubGroup = NPerWmma - make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}), - make_tuple(Sequence<>{}, Sequence<0, 1, 2, 6>{}, Sequence<>{}, Sequence<3, 4, 5>{})); + if constexpr(TransposeC) + { + // C mapping in single block + constexpr auto + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp = + BlockwiseGemmPipe:: + GetCBlockDescriptor_MRepeat_MWave_MThreadPerSubGroup_NRepeat_NWave_NSubGroup_MAccVgprs(); + + constexpr auto MWave = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I1); + constexpr auto NSubGroup = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I5); + constexpr auto NWave = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I4); + constexpr auto MThreadPerSubGroup = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I2); + constexpr auto NAccVprs = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I6); + + return transform_tensor_descriptor( + GetCShuffleBlockDescriptor_MShRepeat_MPerShRepeat_NShRepeat_NPerShRepeat(), + make_tuple(make_freeze_transform(I0), + make_unmerge_transform(make_tuple( + Number{}, // MRepeat per shuffle repeat + MWave, // MWave + MThreadPerSubGroup)), + make_freeze_transform(I0), + make_unmerge_transform(make_tuple( + Number{}, // NRepeat per shuffle repeat + NWave, // NWave + NSubGroup, + NAccVprs))), + make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}), + make_tuple( + Sequence<>{}, Sequence<0, 1, 2>{}, Sequence<>{}, Sequence<3, 4, 5, 6>{})); + } + else + { + // C mapping in single block + constexpr auto + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp = + BlockwiseGemmPipe:: + GetCBlockDescriptor_MRepeat_MWave_MSubGroup_NRepeat_NWave_NThreadPerSubGroup_MAccVgprs(); + + constexpr auto MWave = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I1); + constexpr auto MSubGroup = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I2); + constexpr auto NWave = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I4); + constexpr auto NThreadPerSubGroup = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I5); + constexpr auto MAccVgprs = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I6); + + return transform_tensor_descriptor( + GetCShuffleBlockDescriptor_MShRepeat_MPerShRepeat_NShRepeat_NPerShRepeat(), + make_tuple(make_freeze_transform(I0), + make_unmerge_transform(make_tuple( + Number{}, // MRepeat per shuffle repeat + MWave, // MWave + MSubGroup, // MSubGroup * MAccVgprs = MPerWmma + MAccVgprs)), + make_freeze_transform(I0), + make_unmerge_transform(make_tuple( + Number{}, // NRepeat per shuffle repeat + NWave, // NWave + NThreadPerSubGroup))), // NThreadPerSubGroup = NPerWmma + make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}), + make_tuple( + Sequence<>{}, Sequence<0, 1, 2, 6>{}, Sequence<>{}, Sequence<3, 4, 5>{})); + } } __device__ static auto GetVgprToLDSEpilogueDescriptor() { - // C mapping in single block - constexpr auto c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp = - BlockwiseGemmPipe:: - GetCBlockDescriptor_MRepeat_MWave_MSubGroup_NRepeat_NWave_NThreadPerSubGroup_MAccVgprs(); - - constexpr auto MWave = - c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp - .GetLength(I1); - constexpr auto MSubGroup = - c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp - .GetLength(I2); - constexpr auto NWave = - c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp - .GetLength(I4); - constexpr auto NThreadPerSubGroup = - c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp - .GetLength(I5); - constexpr auto MAccVgprs = - c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp - .GetLength(I6); - - // calculate origin of thread output tensor on global memory - // blockwise GEMM c matrix starting index - const auto c_thread_mtx_on_block = - BlockwiseGemmPipe::CalculateCThreadOriginDataIndex(I0, I0); - - const index_t m_thread_data_on_block = c_thread_mtx_on_block[I0]; - const index_t n_thread_data_on_block = c_thread_mtx_on_block[I1]; - - const auto m_thread_data_on_block_to_mrepeat_mwave_msubgroup_maccvgprs_adaptor = - make_single_stage_tensor_adaptor( - make_tuple(make_merge_transform(make_tuple(MRepeat, MWave, MSubGroup, MAccVgprs))), - make_tuple(Sequence<0, 1, 2, 3>{}), - make_tuple(Sequence<0>{})); - - const auto m_thread_data_on_block_idx = - m_thread_data_on_block_to_mrepeat_mwave_msubgroup_maccvgprs_adaptor - .CalculateBottomIndex(make_multi_index(m_thread_data_on_block)); - - const auto n_thread_data_on_block_to_nrepeat_nwave_nthreadpersubgroup_adaptor = - make_single_stage_tensor_adaptor( - make_tuple(make_merge_transform(make_tuple(NRepeat, NWave, NThreadPerSubGroup))), - make_tuple(Sequence<0, 1, 2>{}), - make_tuple(Sequence<0>{})); - - const auto n_thread_data_on_block_idx = - n_thread_data_on_block_to_nrepeat_nwave_nthreadpersubgroup_adaptor.CalculateBottomIndex( - make_multi_index(n_thread_data_on_block)); - - return ThreadwiseTensorSliceTransfer_v1r3< - AccDataType, - CShuffleDataType, - decltype(BlockwiseGemmPipe:: - GetCThreadDescriptor_MRepeat_MWave_MSubGroup_NRepeat_NWave_NThreadPerSubGroup_MAccVgprs()), - decltype(GetCShuffleLDSDescriptor()), - ck::tensor_operation::element_wise::PassThrough, - Sequence, - Sequence<0, 1, 2, 3, 4, 5, 6>, - 6, - 1, - InMemoryDataOperationEnum::Set, - 1, - true>{GetCShuffleLDSDescriptor(), - make_multi_index(0, - m_thread_data_on_block_idx[I1], - m_thread_data_on_block_idx[I2], - 0, - n_thread_data_on_block_idx[I1], - n_thread_data_on_block_idx[I2], - m_thread_data_on_block_idx[I3]), - ck::tensor_operation::element_wise::PassThrough{}}; + if constexpr(TransposeC) + { + // C mapping in single block + constexpr auto + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp = + BlockwiseGemmPipe:: + GetCBlockDescriptor_MRepeat_MWave_MThreadPerSubGroup_NRepeat_NWave_NSubGroup_MAccVgprs(); + + constexpr auto MWave = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I1); + constexpr auto NSubGroup = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I5); + constexpr auto NWave = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I4); + constexpr auto MThreadPerSubGroup = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I2); + constexpr auto NAccVgprs = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I6); + + // calculate origin of thread output tensor on global memory + // blockwise GEMM c matrix starting index + const auto c_thread_mtx_on_block = + BlockwiseGemmPipe::CalculateCThreadOriginDataIndex(I0, I0); + + const index_t m_thread_data_on_block = c_thread_mtx_on_block[I0]; + const index_t n_thread_data_on_block = c_thread_mtx_on_block[I1]; + + const auto m_thread_data_on_block_to_mrepeat_mwave_msubgroup_maccvgprs_adaptor = + make_single_stage_tensor_adaptor(make_tuple(make_merge_transform(make_tuple( + MRepeat, MWave, MThreadPerSubGroup))), + make_tuple(Sequence<0, 1, 2>{}), + make_tuple(Sequence<0>{})); + + const auto m_thread_data_on_block_idx = + m_thread_data_on_block_to_mrepeat_mwave_msubgroup_maccvgprs_adaptor + .CalculateBottomIndex(make_multi_index(m_thread_data_on_block)); + + const auto n_thread_data_on_block_to_nrepeat_nwave_nthreadpersubgroup_adaptor = + make_single_stage_tensor_adaptor(make_tuple(make_merge_transform(make_tuple( + NRepeat, NWave, NSubGroup, NAccVgprs))), + make_tuple(Sequence<0, 1, 2, 3>{}), + make_tuple(Sequence<0>{})); + + const auto n_thread_data_on_block_idx = + n_thread_data_on_block_to_nrepeat_nwave_nthreadpersubgroup_adaptor + .CalculateBottomIndex(make_multi_index(n_thread_data_on_block)); + + return ThreadwiseTensorSliceTransfer_v1r3< + AccDataType, + CShuffleDataType, + decltype(BlockwiseGemmPipe:: + GetCThreadDescriptor_MRepeat_MWave_MSubGroup_NRepeat_NWave_NThreadPerSubGroup_MAccVgprs()), + decltype(GetCShuffleLDSDescriptor()), + ck::tensor_operation::element_wise::PassThrough, + Sequence, + Sequence<0, 1, 2, 3, 4, 5, 6>, + 6, + 1, + InMemoryDataOperationEnum::Set, + 1, + true>{GetCShuffleLDSDescriptor(), + make_multi_index(0, + m_thread_data_on_block_idx[I1], + m_thread_data_on_block_idx[I2], + 0, + n_thread_data_on_block_idx[I1], + n_thread_data_on_block_idx[I2], + n_thread_data_on_block_idx[I3]), + ck::tensor_operation::element_wise::PassThrough{}}; + } + else + { + // C mapping in single block + constexpr auto + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp = + BlockwiseGemmPipe:: + GetCBlockDescriptor_MRepeat_MWave_MSubGroup_NRepeat_NWave_NThreadPerSubGroup_MAccVgprs(); + + constexpr auto MWave = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I1); + constexpr auto MSubGroup = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I2); + constexpr auto NWave = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I4); + constexpr auto NThreadPerSubGroup = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I5); + constexpr auto MAccVgprs = + c_block_desc_mrepeat_mwave_msubgroup_nrepeat_nwave_nthreadpersubgroup_maccvgprs_tmp + .GetLength(I6); + + // calculate origin of thread output tensor on global memory + // blockwise GEMM c matrix starting index + const auto c_thread_mtx_on_block = + BlockwiseGemmPipe::CalculateCThreadOriginDataIndex(I0, I0); + + const index_t m_thread_data_on_block = c_thread_mtx_on_block[I0]; + const index_t n_thread_data_on_block = c_thread_mtx_on_block[I1]; + + const auto m_thread_data_on_block_to_mrepeat_mwave_msubgroup_maccvgprs_adaptor = + make_single_stage_tensor_adaptor(make_tuple(make_merge_transform(make_tuple( + MRepeat, MWave, MSubGroup, MAccVgprs))), + make_tuple(Sequence<0, 1, 2, 3>{}), + make_tuple(Sequence<0>{})); + + const auto m_thread_data_on_block_idx = + m_thread_data_on_block_to_mrepeat_mwave_msubgroup_maccvgprs_adaptor + .CalculateBottomIndex(make_multi_index(m_thread_data_on_block)); + + const auto n_thread_data_on_block_to_nrepeat_nwave_nthreadpersubgroup_adaptor = + make_single_stage_tensor_adaptor(make_tuple(make_merge_transform(make_tuple( + NRepeat, NWave, NThreadPerSubGroup))), + make_tuple(Sequence<0, 1, 2>{}), + make_tuple(Sequence<0>{})); + + const auto n_thread_data_on_block_idx = + n_thread_data_on_block_to_nrepeat_nwave_nthreadpersubgroup_adaptor + .CalculateBottomIndex(make_multi_index(n_thread_data_on_block)); + + return ThreadwiseTensorSliceTransfer_v1r3< + AccDataType, + CShuffleDataType, + decltype(BlockwiseGemmPipe:: + GetCThreadDescriptor_MRepeat_MWave_MSubGroup_NRepeat_NWave_NThreadPerSubGroup_MAccVgprs()), + decltype(GetCShuffleLDSDescriptor()), + ck::tensor_operation::element_wise::PassThrough, + Sequence, + Sequence<0, 1, 2, 3, 4, 5, 6>, + 6, + 1, + InMemoryDataOperationEnum::Set, + 1, + true>{GetCShuffleLDSDescriptor(), + make_multi_index(0, + m_thread_data_on_block_idx[I1], + m_thread_data_on_block_idx[I2], + 0, + n_thread_data_on_block_idx[I1], + n_thread_data_on_block_idx[I2], + m_thread_data_on_block_idx[I3]), + ck::tensor_operation::element_wise::PassThrough{}}; + } } template + bool ABThreadTransferSrcResetCoordinateAfterRun, + bool UseLdsTranspose> struct ABTransferThreadTiles { __device__ static constexpr bool IsLDSNeeded() { return true; } @@ -153,9 +154,27 @@ struct ABTransferThreadTiles { // bank conflict when writting the data into LDS, but don't worry, we have whole entire // loop to hide it in v4. it may give you some benefit from less valu in compute address - return make_naive_tensor_descriptor( - make_tuple(ABK0Number, Number{}, ABK1Number), - make_tuple(Number{} * ABK1Number, ABK1Number, I1)); + if constexpr(!UseLdsTranspose) + { + return make_naive_tensor_descriptor( + make_tuple(ABK0Number, Number{}, ABK1Number), + make_tuple(Number{} * ABK1Number, ABK1Number, I1)); + } + else + { + constexpr index_t MN1 = MNPerWmma / 2; + constexpr auto base_desc = make_naive_tensor_descriptor( + make_tuple(Number{}, Number{}, Number{}), + make_tuple(Number{} * Number{}, Number{}, I1)); + + return transform_tensor_descriptor( + base_desc, + make_tuple( + make_merge_transform(make_tuple(Number{}, Number{})), + make_unmerge_transform(make_tuple(ABK0Number, ABK1Number))), + make_tuple(Sequence<0, 2>{}, Sequence<1>{}), + make_tuple(Sequence<1>{}, Sequence<0, 2>{})); + } } // xor tensor transformation request more unnecessary vgpr usage, would cause register spill // in some cases. @@ -199,6 +218,7 @@ struct ABTransferThreadTiles } else { + static_assert(!UseLdsTranspose, "UseLdsTranspose is not supported for swizzled layout"); // kfold and mpair dimension is not always required. // more dimension in merge_transform increase the difficulty of generating immarg offset // for compiler. @@ -300,6 +320,7 @@ struct ABTransferThreadTiles constexpr index_t NumABTensor = ABsDataType::Size(); const index_t mn_block_data_idx_on_grid = __builtin_amdgcn_readfirstlane(block_mn_id * MNPerBlock); + // workaround because v7r2 is not as general as v4r1 if constexpr(NumABTensor > 1) { @@ -334,6 +355,9 @@ struct ABTransferThreadTiles } else { + using LdsDimAccessOrder = + std::conditional_t, Sequence<0, 1, 2>>; + constexpr index_t VectorDim = UseLdsTranspose ? 1 : 2; return ThreadGroupTensorSliceTransfer_v4r1< ThisThreadBlock, ABElementwiseOperation, @@ -347,9 +371,9 @@ struct ABTransferThreadTiles decltype(grid_descriptor[I0]), decltype(block_descriptor), ABBlockTransferSrcAccessOrder, - Sequence<0, 1, 2>, + LdsDimAccessOrder, ABBlockTransferSrcVectorDim, - 2, + VectorDim, ABBlockTransferSrcScalarPerVector, ABBlockTransferDstScalarPerVector_ABK1, 1, @@ -382,6 +406,7 @@ struct ABTransferThreadTiles #endif if constexpr(KInner > 1) { + static_assert(!UseLdsTranspose, "KInner > 1 doesn't have a use case for LDS transpose"); // KPack = KInner * KPerWmma // K1 = KInner * KPerWmmaBlk // Each thread loads multiple tiles with one instruction @@ -407,6 +432,8 @@ struct ABTransferThreadTiles // (rest of the single WMMA tile for single thread) and then over KRow // (rest of the single WMMA tile for single wave) // KPack / KRow / K1 - MNRepeat - K0 / KRow - MNWaves - KRow - MNPerWmma - K1 + constexpr index_t MNPerWmma_Dim = UseLdsTranspose ? 6 : 5; + constexpr index_t ABK1_Dim = UseLdsTranspose ? 5 : 6; return transform_tensor_descriptor( BlockDesc{}, make_tuple( @@ -416,10 +443,14 @@ struct ABTransferThreadTiles make_tuple(Number{}, Number{}, Number{})), make_pass_through_transform(Number{})), make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}), - make_tuple(Sequence<2, 4, 0>{}, Sequence<1, 3, 5>{}, Sequence<6>{})); + make_tuple(Sequence<2, 4, 0>{}, + Sequence<1, 3, MNPerWmma_Dim>{}, + Sequence{})); } else { + static_assert(!UseLdsTranspose, + "ABK1 > KPerWmmaBlk doesn't have a use case for LDS transpose"); // K1 > single tile (KPerWmmaBlk) // Each thread will load KPerWmmaBlk for the WMMA instruction // Since K1 > single tile, each thread loads KPerWmmaBlk and the next diff --git a/projects/composablekernel/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3.hpp b/projects/composablekernel/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3.hpp index 897891ea4cbf..5ea9032154eb 100644 --- a/projects/composablekernel/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3.hpp +++ b/projects/composablekernel/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3.hpp @@ -185,7 +185,9 @@ template + bool IsFusedKernel = false, + bool UseLdsTranspose = false, + bool TransposeC = false> struct GridwiseGemm_wmma_cshuffle_v3 : GridwiseGemm_wmma_cshuffle_v3_base< ALayout, @@ -240,7 +242,9 @@ struct GridwiseGemm_wmma_cshuffle_v3 PermuteB, IsBPreShuffled, ForceThreadTileTransfer, - IsFusedKernel> + IsFusedKernel, + UseLdsTranspose, + TransposeC> { using Base = GridwiseGemm_wmma_cshuffle_v3_base< ALayout, @@ -295,7 +299,9 @@ struct GridwiseGemm_wmma_cshuffle_v3 PermuteB, IsBPreShuffled, ForceThreadTileTransfer, - IsFusedKernel>; + IsFusedKernel, + UseLdsTranspose, + TransposeC>; using Base::I0; using Base::I1; diff --git a/projects/composablekernel/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_common.hpp b/projects/composablekernel/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_common.hpp index d1edf5fc2b1c..e571ca31239f 100644 --- a/projects/composablekernel/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_common.hpp +++ b/projects/composablekernel/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3_common.hpp @@ -81,7 +81,9 @@ template + bool IsFusedKernel = false, + bool UseLdsTranspose = false, + bool TransposeC = false> struct GridwiseGemm_wmma_cshuffle_v3_base { @@ -215,8 +217,20 @@ struct GridwiseGemm_wmma_cshuffle_v3_base static constexpr bool UseDirectStore = false; #endif +#if defined(__gfx1250__) + static constexpr bool UseLdsTransposeA = + UseLdsTranspose && !is_same_v; + static constexpr bool UseLdsTransposeB = + UseLdsTranspose && !is_same_v; +#else + static constexpr bool UseLdsTransposeA = false; + static constexpr bool UseLdsTransposeB = false; +#endif + static constexpr bool UseBlockPaddingA = ABlockLdsExtraM || BlkGemmPipelineVer == BlockGemmPipelineVersion::v4; + static_assert(!UseLdsTransposeA || + (UseBlockPaddingA && AK1Value == 8 && NumATensor == 1 && sizeof(LDSTypeA) == 2)); using ATransfer = typename std::conditional< IsAWaveTransferApplicable, ATransferWaveTiles, @@ -239,10 +253,13 @@ struct GridwiseGemm_wmma_cshuffle_v3_base ABlockTransferSrcVectorDim, ABlockTransferSrcScalarPerVector, ABlockTransferDstScalarPerVector_AK1, - AThreadTransferSrcResetCoordinateAfterRun>>::type; + AThreadTransferSrcResetCoordinateAfterRun, + UseLdsTransposeA>>::type; static constexpr bool UseBlockPaddingB = BBlockLdsExtraN || BlkGemmPipelineVer == BlockGemmPipelineVersion::v4; + static_assert(!UseLdsTransposeB || + (UseBlockPaddingB && BK1Value == 8 && NumBTensor == 1 && sizeof(LDSTypeB) == 2)); using BTransfer = typename std::conditional< IsBPreShuffled, ABTransferThreadTilesPreShuffle>::type>::type; + BThreadTransferSrcResetCoordinateAfterRun, + UseLdsTransposeB>>::type>::type; static_assert(!(is_same_v, pk_i4_t> && GemmSpec != tensor_operation::device::GemmSpecialization::Default), @@ -690,8 +708,10 @@ struct GridwiseGemm_wmma_cshuffle_v3_base NRepeat, KPack, KInner, - false, - IsBPreShuffled>())>; + TransposeC, + IsBPreShuffled, + UseLdsTransposeA, + UseLdsTransposeB>())>; struct Traits { diff --git a/projects/composablekernel/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp b/projects/composablekernel/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp index 765b395cfe3c..55b5bbf5b27f 100644 --- a/projects/composablekernel/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp +++ b/projects/composablekernel/include/ck/tensor_operation/gpu/thread/threadwise_tensor_slice_transfer.hpp @@ -1260,6 +1260,7 @@ template ::type = false> struct ThreadwiseTensorSliceTransfer_v4 @@ -1395,8 +1396,8 @@ struct ThreadwiseTensorSliceTransfer_v4 if constexpr(SrcBuffer::IsDynamicBuffer()) { src_tmp_vector.template AsType()(Number<0>{}) = - src_buf.template Get(src_data_coord.GetOffset() / PackedSize, - is_src_valid); + src_buf.template Get( + src_data_coord.GetOffset() / PackedSize, is_src_valid); } else if constexpr(SrcBuffer::IsStaticBuffer()) { diff --git a/projects/composablekernel/include/ck/utility/amd_transpose_load.hpp b/projects/composablekernel/include/ck/utility/amd_transpose_load.hpp index 09e1782d27ea..aa99fd3223e4 100644 --- a/projects/composablekernel/include/ck/utility/amd_transpose_load.hpp +++ b/projects/composablekernel/include/ck/utility/amd_transpose_load.hpp @@ -34,4 +34,39 @@ __device__ auto amd_global_load_transpose_to_vgpr(const T* in_ptr) } #endif +#if defined(__gfx1250__) +template +__device__ auto amd_lds_load_transpose_to_vgpr(const T* __restrict__ in_ptr) +{ +#define CK_LDS_ADDR __attribute__((address_space(3))) +#ifdef __clang__ +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wold-style-cast" +#endif + // Use C-style cast to change address space without dropping llvm noalias attribute + const auto in_ptr_ = (CK_LDS_ADDR T*)(const_cast(in_ptr)); +#ifdef __clang__ +#pragma clang diagnostic pop +#endif + + using vector_t = typename vector_type::type; + if constexpr(is_same::value) + { + using llvm_fp16x8_t = __fp16 __attribute__((ext_vector_type(8))); + auto lds_ptr = reinterpret_cast(in_ptr_); + return bit_cast(__builtin_amdgcn_ds_load_tr16_b128_v8f16(lds_ptr)); + } + else if constexpr(is_same::value) + { + auto lds_ptr = reinterpret_cast(in_ptr_); + return bit_cast(__builtin_amdgcn_ds_load_tr16_b128_v8bf16(lds_ptr)); + } + else + { + static_assert(false, "not implemented"); + } +#undef CK_LDS_ADDR +} +#endif + } // namespace ck diff --git a/projects/composablekernel/include/ck/utility/dynamic_buffer.hpp b/projects/composablekernel/include/ck/utility/dynamic_buffer.hpp index 31d241fba579..7b5f0c0406bf 100644 --- a/projects/composablekernel/include/ck/utility/dynamic_buffer.hpp +++ b/projects/composablekernel/include/ck/utility/dynamic_buffer.hpp @@ -127,6 +127,14 @@ struct DynamicBuffer return amd_global_load_transpose_to_vgpr(p_data_ + i); #else static_assert(!DoTranspose, "load-with-transpose only supported on gfx12+"); +#endif + } + else if constexpr(GetAddressSpace() == AddressSpaceEnum::Lds && DoTranspose) + { +#ifdef __gfx1250__ + return amd_lds_load_transpose_to_vgpr(p_data_ + i); +#else + static_assert(!DoTranspose, "load-with-transpose only supported on gfx12+"); #endif } else diff --git a/projects/composablekernel/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_bwd_data/device_grouped_conv_bwd_data_wmma_v3_instances.hpp b/projects/composablekernel/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_bwd_data/device_grouped_conv_bwd_data_wmma_v3_instances.hpp index 2a159bf19fb2..c64d1b0162ab 100644 --- a/projects/composablekernel/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_bwd_data/device_grouped_conv_bwd_data_wmma_v3_instances.hpp +++ b/projects/composablekernel/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_bwd_data/device_grouped_conv_bwd_data_wmma_v3_instances.hpp @@ -134,8 +134,13 @@ using device_grouped_conv_bwd_data_wmma_v3_bf16_large_tiles_instances = std::tup // ##############################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | // generic instance DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<16, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 8, 1, 1, 1, S<1, 32, 1, 16>, S<8,8,8>>, + DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S< 2, 32, 8>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 8, 1, 1, 1, S<1, 32, 1, 16>, S<8,8,8>, ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, BF16, BF16, 1, 1, true>, + DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S< 2, 16, 8>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 8, 1, 1, 1, S<1, 32, 1, 8>, S<8,8,8>, ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, BF16, BF16, 1, 1, true>, DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 8, 1, 1, 1, S<1, 64, 1, 4>, S<8,8,8>>, DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 2, 8, 1, 1, 1, S<1, 32, 1, 8>, S<8,8,8>>, + DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S< 1, 32, 8>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 8, 1, 1, 1, S<1, 32, 1, 8>, S<8,8,8>, ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, BF16, BF16, 1, 1, true, true>, + DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S< 4, 8, 8>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 8, 1, 1, 1, S<1, 64, 1, 4>, S<8,8,8>, ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, BF16, BF16, 1, 1, true, true>, + DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S< 4, 16, 8>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 8, 1, 1, 1, S<1,128, 1, 4>, S<8,8,8>, ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, BF16, BF16, 1, 1, true, true>, DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 8, 1, 1, 1, S<1, 64, 1, 4>, S<8,8,8>>, DeviceGroupedConvBwdDataMultipleD_Wmma_CShuffleV3, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, 1, S<4, 64, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 1, 8, 1, 1, 1, S<1, 16, 1, 16>, S<8,8,8>>, diff --git a/projects/composablekernel/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_bwd_weight/device_grouped_conv_bwd_weight_v3_wmma_instance.hpp b/projects/composablekernel/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_bwd_weight/device_grouped_conv_bwd_weight_v3_wmma_instance.hpp index d06eceaa094f..93c1ad32c049 100644 --- a/projects/composablekernel/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_bwd_weight/device_grouped_conv_bwd_weight_v3_wmma_instance.hpp +++ b/projects/composablekernel/library/include/ck/library/tensor_operation_instance/gpu/grouped_conv_bwd_weight/device_grouped_conv_bwd_weight_v3_wmma_instance.hpp @@ -99,8 +99,11 @@ using device_grouped_conv_bwd_weight_v3_wmma_c_shuffle_bf16_instances = std::tup DeviceGroupedConvBwdWeight_Wmma_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, BF16, BF16, BF16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 256, 64, 256, 256, 8, 16, 16, 2, 4, S<32, 8, 1>, S<2, 0, 1>, S<1, 0, 2>, 1, 8, 8, 1, S<32, 8, 1>, S<2, 0, 1>, S<1, 0, 2>, 1, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, 8, Scheduler, PipelineVersion>, DeviceGroupedConvBwdWeight_Wmma_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, BF16, BF16, BF16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 256, 128, 256, 128, 8, 16, 16, 4, 4, S<16, 16, 1>, S<2, 0, 1>, S<1, 0, 2>, 1, 8, 8, 1, S<16, 16, 1>, S<2, 0, 1>, S<1, 0, 2>, 1, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, 8, Scheduler, PipelineVersion>, DeviceGroupedConvBwdWeight_Wmma_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, BF16, BF16, BF16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 256, 256, 256, 64, 8, 16, 16, 8, 4, S<8, 32, 1>, S<2, 0, 1>, S<1, 0, 2>, 1, 8, 8, 1, S<8, 32, 1>, S<2, 0, 1>, S<1, 0, 2>, 1, 8, 8, 0, 1, 1, S<1, 32, 1, 8>, 8, Scheduler, PipelineVersion>, - DeviceGroupedConvBwdWeight_Wmma_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, BF16, BF16, BF16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 512, 128, 256, 256, 8, 16, 16, 4, 2, S<32, 16, 1>, S<2, 0, 1>, S<1, 0, 2>, 1, 8, 8, 1, S<32, 16, 1>, S<2, 0, 1>, S<1, 0, 2>, 1, 8, 8, 0, 1, 1, S<1, 32, 1, 16>, 8, Scheduler, PipelineVersion> -#endif + DeviceGroupedConvBwdWeight_Wmma_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, BF16, BF16, BF16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 512, 128, 256, 256, 8, 16, 16, 4, 2, S<32, 16, 1>, S<2, 0, 1>, S<1, 0, 2>, 1, 8, 8, 1, S<32, 16, 1>, S<2, 0, 1>, S<1, 0, 2>, 1, 8, 8, 0, 1, 1, S<1, 32, 1,16>, 8, Scheduler, PipelineVersion>, + DeviceGroupedConvBwdWeight_Wmma_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, BF16, BF16, BF16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 256, 256, 256, 64, 8, 16, 16, 8, 4, S< 1, 32, 8>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 8, 1, S< 1, 32, 8>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 8, 1, 1, 1, S<1, 32, 1, 8>, 8, Scheduler, PipelineVersion, BF16, BF16, 1, 1, true, true>, + DeviceGroupedConvBwdWeight_Wmma_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, BF16, BF16, BF16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 512, 128, 256, 128, 8, 16, 16, 4, 2, S< 4, 16, 8>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 8, 1, S< 2, 32, 8>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 8, 1, 1, 1, S<1, 32, 1,16>, 8, Scheduler, PipelineVersion, BF16, BF16, 1, 1, true, true>, + DeviceGroupedConvBwdWeight_Wmma_CShuffleV3< NDimSpatial, ALayout, BLayout, ELayout, BF16, BF16, BF16, F32, PassThrough, PassThrough, PassThrough, ConvSpec, 512, 128, 128, 64, 8, 16, 16, 4, 1, S< 4, 16, 8>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 8, 1, S< 4, 16, 8>, S<0, 2, 1>, S<0, 2, 1>, 1, 8, 8, 1, 1, 1, S<1, 32, 1,16>, 8, Scheduler, PipelineVersion, BF16, BF16, 1, 1, true, true> +#endif //clang-format on >;