diff --git a/example/65_gemm_multiply_multiply/gemm_multiply_multiply_xdl_fp8_bpreshuffle_padding.cpp b/example/65_gemm_multiply_multiply/gemm_multiply_multiply_xdl_fp8_bpreshuffle_padding.cpp index 4dc08fe1f5..178255d2b1 100644 --- a/example/65_gemm_multiply_multiply/gemm_multiply_multiply_xdl_fp8_bpreshuffle_padding.cpp +++ b/example/65_gemm_multiply_multiply/gemm_multiply_multiply_xdl_fp8_bpreshuffle_padding.cpp @@ -129,7 +129,7 @@ void preShuffleBuffer(const DataType* src, DataType* dst, int N, int K, int NXdl } } } -int GetPreShufflePadded(int K) { return (K + 64 - 1) / 64 * 64; } +int GetPreShufflePadded(int K) { return (K + ShufflePadded - 1) / ShufflePadded * ShufflePadded; } using PassThrough = ck::tensor_operation::element_wise::PassThrough; using AElementOp = PassThrough; @@ -140,38 +140,38 @@ static constexpr auto GemmSpec = ck::tensor_operation::device::GemmSpecializatio using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMultiD_Xdl_CShuffle_V3_BPreshuffle // clang-format off - // < Row, Col, DsLayout, ELayout, A0DataType, B0DataType, DsDataType, EDataType, AccDataType, CShuffleDataType, - // AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, - // 256, 256, 128, - // 16, 16, - // 16, 16, - // 8, 8, - // S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, - // S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, - // 1, 2, S<1, 32, 1, 8>, S<8, 8, 1>, - // ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v3, B0DataType>; - - < Row, Col, DsLayout, ELayout, A0DataType, B0DataType, DsDataType, EDataType, AccDataType, CShuffleDataType, + < Row, Col, DsLayout, ELayout, A0DataType, B0DataType, DsDataType, EDataType, AccDataType, CShuffleDataType, AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, - 128, 128, 128, - 16, 16, - 32, 32, - 2, 2, - S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, - S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, - 1, 1, S<1, 32, 1, 8>, S<8, 8, 1>, - ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v2, B0DataType>; + 256, 256, 128, + 16, 16, + 16, 16, + 8, 8, + S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, + S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, + 1, 2, S<1, 32, 1, 8>, S<8, 8, 1>, + ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v3, B0DataType>; + + // < Row, Col, DsLayout, ELayout, A0DataType, B0DataType, DsDataType, EDataType, AccDataType, CShuffleDataType, + // AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, + // 128, 128, 128, + // 16, 16, + // 32, 32, + // 2, 2, + // S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, + // S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, + // 1, 1, S<1, 32, 1, 8>, S<8, 8, 1>, + // ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v2, B0DataType>; // < Row, Col, DsLayout, ELayout, A0DataType, B0DataType, DsDataType, EDataType, AccDataType, CShuffleDataType, // AElementOp, BElementOp, CDEElementOp, GemmSpec, 256, -// 256, 128, 128, -// 16, 16, -// 32, 32, -// 4, 2, -// S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, -// S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, -// 1, 1, S<1, 32, 1, 8>, S<8, 8, 1>, -// ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v2, B0DataType>; +// 128, 256, 128, +// 16, 16, +// 32, 32, +// 4, 2, +// S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, +// S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0, +// 1, 1, S<1, 32, 1, 8>, S<8, 8, 1>, +// ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v3, B0DataType>; // clang-format on @@ -268,7 +268,7 @@ int main(int argc, char* argv[]) } }; - auto Knew = GetPreShufflePadded(K); + auto Knew = ck::tensor_operation::device::GetPreShufflePadded(K); auto StrideBnew = Knew; Tensor a0_m_k(f_host_tensor_descriptor(M, K, StrideA, A0Layout{})); Tensor b0_k_n(f_host_tensor_descriptor(K, N, StrideB, B0Layout{})); diff --git a/include/ck/tensor_operation/gpu/device/device_gemm_multiple_d.hpp b/include/ck/tensor_operation/gpu/device/device_gemm_multiple_d.hpp index 3c79b92ec8..a6b18a9d06 100644 --- a/include/ck/tensor_operation/gpu/device/device_gemm_multiple_d.hpp +++ b/include/ck/tensor_operation/gpu/device/device_gemm_multiple_d.hpp @@ -147,6 +147,7 @@ struct DeviceGemmMultipleDSplitKBPreShuffle : public BaseOperator virtual int GetPreShuffleParameters() = 0; }; +#define ShufflePadded 256 } // namespace device } // namespace tensor_operation } // namespace ck diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_b_preshuffle.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_b_preshuffle.hpp index 767fd2382e..8a0aaa7b0c 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_b_preshuffle.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_multi_d_b_preshuffle.hpp @@ -214,7 +214,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle __host__ __device__ static auto CalculateBKShufflePadded(index_t K) { - return (K + 64 - 1) / 64 * 64; + return (K + ShufflePadded - 1) / ShufflePadded * ShufflePadded; } __host__ __device__ static auto CalculateKPadded(index_t K) @@ -605,7 +605,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle MBlock{CalculateMBlock(M_)}, NBlock{CalculateNBlock(N_)}, BN0Shuffled{CalculateBN0Shuffled(N_)}, - BK0Shuffled{CalculateBK0Shuffled((K_ + 64 - 1) / 64 * 64)} + BK0Shuffled{CalculateBK0Shuffled(CalculateBKShufflePadded(K_))} { } @@ -936,7 +936,7 @@ struct GridwiseGemmMultiD_xdl_cshuffle_v3_b_preshuffle return false; } - if(karg.N % NPerBlock != 0 && NXdlPerWave == 1) + if(karg.N % NPerBlock != 0 && karg.N % (NPerXdl * 2) != 0) { return false; } diff --git a/profiler/include/profiler/profile_gemm_multiply_multiply_weight_preshuffle_impl.hpp b/profiler/include/profiler/profile_gemm_multiply_multiply_weight_preshuffle_impl.hpp index ddada8c4c6..c7209e3c8d 100644 --- a/profiler/include/profiler/profile_gemm_multiply_multiply_weight_preshuffle_impl.hpp +++ b/profiler/include/profiler/profile_gemm_multiply_multiply_weight_preshuffle_impl.hpp @@ -56,7 +56,8 @@ void preShuffleBuffer( } } } -int GetPreShufflePadded(int K) { return (K + 64 - 1) / 64 * 64; } + +int GetPreShufflePadded(int K) { return (K + ShufflePadded - 1) / ShufflePadded * ShufflePadded; } template