From 79246e6cb88c69c8380ceccac0f5b0776054df53 Mon Sep 17 00:00:00 2001 From: aska-0096 Date: Mon, 12 May 2025 16:54:44 +0000 Subject: [PATCH] function pass with inline asm hacky --- example/67_gemm_microscaling/CMakeLists.txt | 2 +- .../67_gemm_microscaling/gemm_mx_common.hpp | 92 +++---- example/67_gemm_microscaling/gemm_mx_fp4.cpp | 2 +- ...blockwise_gemm_mx_pipeline_xdlops_base.hpp | 120 ++++++--- .../blockwise_gemm_pipeline_xdlops_v3_mx.hpp | 229 ++++++++---------- .../impl/device_gemm_xdl_cshuffle_v3_mx.hpp | 10 +- .../grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp | 198 ++++++++------- .../tensor_operation/gpu/warp/xdlops_gemm.hpp | 43 ++++ include/ck/utility/amd_xdlops.hpp | 44 +++- 9 files changed, 449 insertions(+), 291 deletions(-) diff --git a/example/67_gemm_microscaling/CMakeLists.txt b/example/67_gemm_microscaling/CMakeLists.txt index 5baf87fc85..643dd1cb6b 100644 --- a/example/67_gemm_microscaling/CMakeLists.txt +++ b/example/67_gemm_microscaling/CMakeLists.txt @@ -13,6 +13,6 @@ add_example_executable(example_gemm_mx_fp4 gemm_mx_fp4.cpp) add_example_dependencies(example_gemm_mx example_gemm_mx_fp4) set(FP4_MXGEMM_OPTIONS) -list(APPEND FP4_MXGEMM_OPTIONS "SHELL: -mllvm -greedy-reverse-local-assignment=1 -mllvm --slp-threshold=-32") +# list(APPEND FP4_MXGEMM_OPTIONS "SHELL: -mllvm -greedy-reverse-local-assignment=1 -mllvm --slp-threshold=-32") list(APPEND FP4_MXGEMM_OPTIONS -v --save-temps -Wno-gnu-line-marker) target_compile_options(example_gemm_mx_fp4 PRIVATE ${FP4_MXGEMM_OPTIONS}) \ No newline at end of file diff --git a/example/67_gemm_microscaling/gemm_mx_common.hpp b/example/67_gemm_microscaling/gemm_mx_common.hpp index d4b7e562f0..fed60db508 100644 --- a/example/67_gemm_microscaling/gemm_mx_common.hpp +++ b/example/67_gemm_microscaling/gemm_mx_common.hpp @@ -104,7 +104,7 @@ bool parse_cmd_args(int argc, } #if 1 -void preShuffleScaleBuffer(const ck::e8m0_bexp_t* src, ck::e8m0_bexp_t* dst, int MN, int K) +void preShuffleScaleBuffer(ck::e8m0_bexp_t* src, ck::e8m0_bexp_t* dst, int MN, int K) { int MNXdlPack = 2; int KXdlPack = 2; @@ -128,8 +128,8 @@ void preShuffleScaleBuffer(const ck::e8m0_bexp_t* src, ck::e8m0_bexp_t* dst, int { int n0 = n / (XdlMNThread * MNXdlPack); // i MNRepeat int tempn = n % (XdlMNThread * MNXdlPack); - int n1 = tempn / MNXdlPack; // i XdlMNThread - int n2 = tempn % MNXdlPack; // i MNXdlPack + int n1 = tempn % XdlMNThread; // i XdlMNThread + int n2 = tempn / XdlMNThread; // i MNXdlPack int k0 = k / (XdlKThread * KXdlPack); // i KRepeat int tempk = k % (XdlKThread * KXdlPack); @@ -140,8 +140,10 @@ void preShuffleScaleBuffer(const ck::e8m0_bexp_t* src, ck::e8m0_bexp_t* dst, int k0 * MNXdlPack * KXdlPack * XdlMNThread * XdlKThread + k1 * MNXdlPack * KXdlPack * XdlMNThread + n1 * MNXdlPack * KXdlPack + k2 * MNXdlPack + n2; - + // src[n * K + k] = ck::type_convert(static_cast(powf(2.0f, n2 + + // k2 * MNXdlPack))); dst[outputIndex] = src[n * K + k]; + // printf("Src: %0d, Dst: %d\n", n * K + k, outputIndex);; } } } @@ -280,13 +282,17 @@ bool run_mx_gemm(const ProblemSizeSplitK& problem_size, const ExecutionConfig& c a_m_k.GenerateTensorValue(GeneratorTensor_2{-5, 6}); // Z[-5,5] b_k_n.GenerateTensorValue(GeneratorTensor_2{-5, 6}); // Z[-5,5] + // ck::utils::FillConstant{a_data_element(1.0f)}(a_m_k); + // ck::utils::FillConstant{b_data_element(1.0f)}(b_k_n); if constexpr(ck::is_same_v) { a_m_k_scale.GenerateTensorValue( - GeneratorTensor_2{125, 129}); // scales: {0.25, 0.5, 1, 2} + GeneratorTensor_2{120, 135}); // scales: {0.25, 0.5, 1, 2} b_k_n_scale.GenerateTensorValue( GeneratorTensor_2{125, 129}); // scales: {0.25, 0.5, 1, 2} + // ck::utils::FillConstant{ck::type_convert(1.0f)}(a_m_k_scale); + // ck::utils::FillConstant{ck::type_convert(1.0f)}(b_k_n_scale); } else { @@ -347,28 +353,6 @@ bool run_mx_gemm(const ProblemSizeSplitK& problem_size, const ExecutionConfig& c std::cout << "NOTE: No input data initialization." << std::endl; } } - printf("a_scale:\n"); - for(ck::index_t i = 0; i < M; i++) - { - for(ck::index_t j = 0; j < K / ScaleBlockSize; j++) - { - // a_m_k_scale(i, j) = - // ck::type_convert(static_cast(powf(2.0f, (j / 4) % 4))); - printf("%02x ", *reinterpret_cast(&a_m_k_scale(i, j))); - } - printf("\n"); - } - printf("b_scale:\n"); - for(ck::index_t i = 0; i < N; i++) - { - for(ck::index_t j = 0; j < K / ScaleBlockSize; j++) - { - // b_k_n_scale(j, i) = - // ck::type_convert(static_cast(powf(2.0f, (j / 4) % 4))); - printf("%02x ", *reinterpret_cast(&b_k_n_scale(j, i))); - } - printf("\n"); - } #if 1 preShuffleScaleBuffer( @@ -376,21 +360,47 @@ bool run_mx_gemm(const ProblemSizeSplitK& problem_size, const ExecutionConfig& c preShuffleScaleBuffer( b_k_n_scale.mData.data(), b_shuffled_scale.mData.data(), N, K / ScaleBlockSize); #endif + // printf("a_scale:\n"); + // for(ck::index_t i = 0; i < M; i++) + // { + // for(ck::index_t j = 0; j < K / ScaleBlockSize; j++) + // { + // // a_m_k_scale(i, j) = + // // ck::type_convert(static_cast(powf(2.0f, (j / 4) % 4))); + // a_m_k_scale(i, j) =ck::type_convert(static_cast(1.0f)); + // a_shuffled_scale(i, j) =ck::type_convert(static_cast(1.0f)); + // printf("%02x ", *reinterpret_cast(&a_m_k_scale(i, j))); + // } + // printf("\n"); + // } + // printf("b_scale:\n"); + // for(ck::index_t i = 0; i < N; i++) + // { + // for(ck::index_t j = 0; j < K / ScaleBlockSize; j++) + // { + // // // b_k_n_scale(j, i) = + // // // ck::type_convert(static_cast(powf(2.0f, (j / 4) % 4))); + // // b_k_n_scale(j, i) =ck::type_convert(static_cast(1.0f)); + // // b_shuffled_scale(j, i) =ck::type_convert(static_cast(1.0f)); + // printf("%02x ", *reinterpret_cast(&b_k_n_scale(j, i))); + // } + // printf("\n"); + // } - printf("a_shuffled_scale:\n"); - for(ck::index_t i = 0; i < M * K / ScaleBlockSize; i++) - { - printf("%02x ", *reinterpret_cast(&(a_shuffled_scale.mData.data()[i]))); - if(i % 64 == 63) - printf("\n"); - } - printf("b_shuffled_scale:\n"); - for(ck::index_t i = 0; i < N * K / ScaleBlockSize; i++) - { - printf("%02x ", *reinterpret_cast(&(b_shuffled_scale.mData.data()[i]))); - if(i % 64 == 63) - printf("\n"); - } + // printf("a_shuffled_scale:\n"); + // for(ck::index_t i = 0; i < M * K / ScaleBlockSize; i++) + // { + // printf("%02x ", *reinterpret_cast(&(a_shuffled_scale.mData.data()[i]))); + // if(i % 64 == 63) + // printf("\n"); + // } + // printf("b_shuffled_scale:\n"); + // for(ck::index_t i = 0; i < N * K / ScaleBlockSize; i++) + // { + // printf("%02x ", *reinterpret_cast(&(b_shuffled_scale.mData.data()[i]))); + // if(i % 64 == 63) + // printf("\n"); + // } if(config.verbosity > 0) std::cout << "Device memory allocation..." << std::endl; diff --git a/example/67_gemm_microscaling/gemm_mx_fp4.cpp b/example/67_gemm_microscaling/gemm_mx_fp4.cpp index 01d57d0471..f47e09443b 100644 --- a/example/67_gemm_microscaling/gemm_mx_fp4.cpp +++ b/example/67_gemm_microscaling/gemm_mx_fp4.cpp @@ -75,7 +75,7 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMX_Xdl_CShuffle 32, // BBlockTransferSrcScalarPerVector 32, // BBlockTransferDstScalarPerVector_BK1 false, // BBlockLdsExtraN - 1, // CShuffleMXdlPerWavePerShuffle + 2, // CShuffleMXdlPerWavePerShuffle 2, // CShuffleNXdlPerWavePerShuffle S<1, 32, 1, 8>, // CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock 8, // CShuffleBlockTransferScalarPerVector_NPerBlock diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_mx_pipeline_xdlops_base.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_mx_pipeline_xdlops_base.hpp index 0c3fc8ffdd..3f7740a290 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_mx_pipeline_xdlops_base.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_mx_pipeline_xdlops_base.hpp @@ -129,7 +129,7 @@ struct BlockwiseGemmXdlops_mx_pipeline_base const auto xdlops_a_idx = xdlops_gemm.CalculateAThreadOriginDataIndex(); - return make_tuple(0, waveId_m, xdlops_a_idx[I1], KThreadChunk * xdlops_a_idx[I0]); + return make_tuple(0, waveId_m, 0, xdlops_a_idx[I1], KThreadChunk * xdlops_a_idx[I0]); } __device__ static auto CalculateBThreadOriginDataIndex() @@ -140,7 +140,7 @@ struct BlockwiseGemmXdlops_mx_pipeline_base const auto xdlops_b_idx = xdlops_gemm.CalculateBThreadOriginDataIndex(); - return make_tuple(0, waveId_n, xdlops_b_idx[I1], KThreadChunk * xdlops_b_idx[I0]); + return make_tuple(0, waveId_n, 0, xdlops_b_idx[I1], KThreadChunk * xdlops_b_idx[I0]); } template @@ -155,24 +155,27 @@ struct BlockwiseGemmXdlops_mx_pipeline_base const auto blk_idx = xdlops_gemm.GetBeginOfThreadBlk(xdlops_i, blk_i); constexpr auto mrepeat_mwave_mperxdl_to_m_adaptor = make_single_stage_tensor_adaptor( - make_tuple(make_unmerge_transform(make_tuple(MRepeat, MWaves, MPerXDL))), + make_tuple( + make_unmerge_transform(make_tuple(MRepeat / MXdlPack, MWaves, MXdlPack, MPerXDL))), make_tuple(Sequence<0>{}), - make_tuple(Sequence<0, 1, 2>{})); + make_tuple(Sequence<0, 1, 2, 3>{})); constexpr auto nrepeat_nwave_nperxdl_to_n_adaptor = make_single_stage_tensor_adaptor( - make_tuple(make_unmerge_transform(make_tuple(NRepeat, NWaves, NPerXDL))), + make_tuple( + make_unmerge_transform(make_tuple(NRepeat / NXdlPack, NWaves, NXdlPack, NPerXDL))), make_tuple(Sequence<0>{}), - make_tuple(Sequence<0, 1, 2>{})); + make_tuple(Sequence<0, 1, 2, 3>{})); + // We pack 2 mfma in M/N direction, so we need to divide by 2 const index_t c_thread_m = mrepeat_mwave_mperxdl_to_m_adaptor.CalculateBottomIndex( - make_tuple(m0, waveId_m, blk_idx[I0]))[I0]; + make_tuple(m0 / MXdlPack, waveId_m, m0 % MXdlPack, blk_idx[I0]))[I0]; const index_t c_thread_n = nrepeat_nwave_nperxdl_to_n_adaptor.CalculateBottomIndex( - make_tuple(n0, waveId_n, blk_idx[I1]))[I0]; + make_tuple(n0 / NXdlPack, waveId_n, n0 % NXdlPack, blk_idx[I1]))[I0]; return make_tuple(c_thread_m, c_thread_n); } - using Tuple4 = decltype(CalculateAThreadOriginDataIndex()); + using Tuple5 = decltype(CalculateAThreadOriginDataIndex()); /** * @brief Constructor for BlockwiseGemmXdlops_mx_pipeline_base. @@ -192,8 +195,8 @@ struct BlockwiseGemmXdlops_mx_pipeline_base * repeat dimensions. */ __host__ __device__ - BlockwiseGemmXdlops_mx_pipeline_base(Tuple4 a_origin = CalculateAThreadOriginDataIndex(), - Tuple4 b_origin = CalculateBThreadOriginDataIndex()) + BlockwiseGemmXdlops_mx_pipeline_base(Tuple5 a_origin = CalculateAThreadOriginDataIndex(), + Tuple5 b_origin = CalculateBThreadOriginDataIndex()) : a_thread_copy_(a_origin), b_thread_copy_(b_origin) { static_assert(AMmaTileDesc::IsKnownAtCompileTime() && BMmaTileDesc::IsKnownAtCompileTime(), @@ -234,6 +237,28 @@ struct BlockwiseGemmXdlops_mx_pipeline_base make_tuple(Number{}, Number{}, I1, I1, M0, M1, M2, N)); } + // XDL output supporting C_xdl = A_xdl * B_xdl, packed mfma + __host__ __device__ static constexpr auto GetCThreadDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3() + { + constexpr auto c_m0_m1_m2_n_tblk_lens = xdlops_gemm.GetCM0M1M2NThreadBlkLengths(); + + constexpr auto M0 = c_m0_m1_m2_n_tblk_lens[I0]; + constexpr auto M1 = c_m0_m1_m2_n_tblk_lens[I1]; + constexpr auto M2 = c_m0_m1_m2_n_tblk_lens[I2]; + constexpr auto N = c_m0_m1_m2_n_tblk_lens[I3]; + + return make_naive_tensor_descriptor_packed(make_tuple(Number{}, + Number{}, + I1, + I1, + Number{}, + Number{}, + M0, + M1, + M2, + N)); + } + __host__ __device__ static constexpr auto GetCThreadDescriptor_G_M0_N0_M1_N1_M2_M3_M4_N2() { constexpr auto c_m0_m1_m2_n_tblk_lens = xdlops_gemm.GetCM0M1M2NThreadBlkLengths(); @@ -275,6 +300,23 @@ struct BlockwiseGemmXdlops_mx_pipeline_base return xdlops_gemm.MakeCDescriptor_M0_N0_M1_N1_M2_M3_M4_N2(c_block_desc_m0_n0_m1_n1_m2_n2); } + // XDL output supporting C_xdl = A_xdl * B_xdl_packed mfma + __host__ __device__ static constexpr auto GetCBlockDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3() + { + constexpr auto c_block_desc_m0_n0_m1_n1_m2_n2 = + make_naive_tensor_descriptor_packed(make_tuple(Number{}, + Number{}, + Number{}, + Number{}, + Number{}, + Number{}, + Number{}, + Number{})); + + return xdlops_gemm.MakeCDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3( + c_block_desc_m0_n0_m1_n1_m2_n2); + } + __host__ __device__ static constexpr auto GetCBlockDescriptor_G_M0_N0_M1_N1_M2_M3_M4_N2() { constexpr auto c_block_desc_g_m0_n0_m1_n1_m2_n2 = @@ -327,49 +369,59 @@ struct BlockwiseGemmXdlops_mx_pipeline_base c_grid_desc_g_m0_n0_m1_n1_m2_n2); } - static constexpr AMmaTileDesc a_block_desc_m0_m1_m2_k; - static constexpr BMmaTileDesc b_block_desc_n0_n1_n2_k; + static constexpr AMmaTileDesc a_block_desc_m0_m1_m2_m3_k; + static constexpr BMmaTileDesc b_block_desc_n0_n1_n2_n3_k; protected: // M1, N1 as double buffer index // Read buffer + Compute buffer // A[M0, M1, M2, KPack] - static constexpr auto a_thread_desc_ = make_naive_tensor_descriptor( - make_tuple(Number{}, I1, Number{}, Number{}), - make_tuple(Number{}, - Number{}, - Number{}, - I1)); + static constexpr auto a_thread_desc_ = + make_naive_tensor_descriptor(make_tuple(Number{}, + I1, + Number{}, + Number{}, + Number{}), + make_tuple(Number{}, + Number{}, + Number{}, + Number{}, + I1)); // B[N0, N1, N2, KPack] - static constexpr auto b_thread_desc_ = make_naive_tensor_descriptor( - make_tuple(Number{}, I1, Number{}, Number{}), - make_tuple(Number{}, - Number{}, - Number{}, - I1)); + static constexpr auto b_thread_desc_ = + make_naive_tensor_descriptor(make_tuple(Number{}, + I1, + Number{}, + Number{}, + Number{}), + make_tuple(Number{}, + Number{}, + Number{}, + Number{}, + I1)); // C[M, N, NumRegXdlops] static constexpr auto c_thread_desc_ = make_naive_tensor_descriptor_packed( - make_tuple(Number{}, Number{}, xdlops_gemm.GetRegSizePerXdlops())); + make_tuple(Number{}, Number{}, Number{}, Number{}, xdlops_gemm.GetRegSizePerXdlops())); using AThreadCopy = ThreadwiseTensorSliceTransfer_v4, - Sequence<0, 1, 2, 3>, - 3, + Sequence<1, 1, 1, 1, KThreadChunk>, + Sequence<0, 1, 2, 3, 4>, + 4, A_K1, A_K1>; using BThreadCopy = ThreadwiseTensorSliceTransfer_v4, - Sequence<0, 1, 2, 3>, - 3, + Sequence<1, 1, 1, 1, KThreadChunk>, + Sequence<0, 1, 2, 3, 4>, + 4, B_K1, B_K1>; diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_mx.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_mx.hpp index 7bb0643f46..276b955fb1 100644 --- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_mx.hpp +++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3_mx.hpp @@ -137,8 +137,8 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx(&a_scale_thread_bufs(I0)[Number<0>{}]), - *reinterpret_cast(&a_scale_thread_bufs(I0)[Number<1>{}]), - *reinterpret_cast(&a_scale_thread_bufs(I0)[Number<2>{}]), - *reinterpret_cast(&a_scale_thread_bufs(I0)[Number<3>{}])); - } -#endif + // restore row id and advance to the next set of scales a_scale_thread_copy.MoveSrcSliceWindow( - a_scale_grid_desc, make_multi_index(-MWaves * MRepeat / MXdlPack, KRepeat / KXdlPack, 0)); + a_scale_grid_desc, + make_multi_index(-MWaves * MRepeat / MXdlPack, KRepeat / KXdlPack, 0)); // Prefetch b_scales static_for<0, NRepeat / NXdlPack, 1>{}([&](auto n0) { @@ -397,21 +388,12 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx(&b_scale_thread_bufs(I0)[Number<0>{}]), - *reinterpret_cast(&b_scale_thread_bufs(I0)[Number<1>{}]), - *reinterpret_cast(&b_scale_thread_bufs(I0)[Number<2>{}]), - *reinterpret_cast(&b_scale_thread_bufs(I0)[Number<3>{}])); - } -#endif + // restore col id and advance to the next set of scales // NWaves * NPerXDL * NRepeat == NPerBlock b_scale_thread_copy.MoveSrcSliceWindow( - b_scale_grid_desc, make_multi_index(-NWaves * NRepeat / NXdlPack, KRepeat / KXdlPack, 0)); + b_scale_grid_desc, + make_multi_index(-NWaves * NRepeat / NXdlPack, KRepeat / KXdlPack, 0)); // Local prefill 1 a_blockwise_copy.RunWrite(a_block_desc, a_block_buf); @@ -432,11 +414,19 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto chunk) { constexpr auto a_k_step_chunk = k_step + chunk * KThreadChunk * xdlops_gemm.mfma_instr.num_input_blks; - a_thread_copy_.Run(a_block_desc_m0_m1_m2_k, - make_tuple(m0, I0, I0, Number{}), + a_thread_copy_.Run(a_block_desc_m0_m1_m2_m3_k, + make_tuple(Number{}, + I0, + Number{}, + I0, + Number{}), a_block_buf, a_thread_desc_, - make_tuple(m0, I0, k, Number{}), + make_tuple(Number{}, + I0, + Number{}, + k, + Number{}), a_thread_buf); }); }); @@ -445,11 +435,19 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto chunk) { constexpr auto b_k_step_chunk = k_step + chunk * KThreadChunk * xdlops_gemm.mfma_instr.num_input_blks; - b_thread_copy_.Run(b_block_desc_n0_n1_n2_k, - make_tuple(n0, I0, I0, Number{}), + b_thread_copy_.Run(b_block_desc_n0_n1_n2_n3_k, + make_tuple(Number{}, + I0, + Number{}, + I0, + Number{}), b_block_buf, b_thread_desc_, - make_tuple(n0, I0, k, Number{}), + make_tuple(Number{}, + I0, + Number{}, + k, + Number{}), b_thread_buf); }); }); @@ -485,7 +483,8 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto n0) { @@ -506,7 +505,8 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto imxdl) { static_for<0, NXdlPack, 1>{}([&](auto inxdl) { constexpr auto kxdl = ikxdl + k0 * KXdlPack; - constexpr auto mxdl = imxdl + m0 * MXdlPack; - constexpr auto nxdl = inxdl + n0 * NXdlPack; vector_type a_thread_vec; @@ -562,14 +560,14 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto ik) { - a_thread_vec.template AsType()(ik) = - a_thread_buf - [Number{}]; - b_thread_vec.template AsType()(ik) = - b_thread_buf - [Number{}]; + a_thread_vec.template AsType()( + ik) = a_thread_buf + [Number{}]; + b_thread_vec.template AsType()( + ik) = b_thread_buf + [Number{}]; }); using mfma_input_type_a = @@ -591,7 +589,7 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto k) { constexpr auto k_step = k * xdlops_gemm.KPerXdlops * (KPack / xdlops_gemm.K1PerXdlops); - static_for<0, MRepeat, 1>{}([&](auto m0) { static_for<0, xdlops_gemm.K1PerXdlops / KThreadChunk, 1>{}( [&](auto chunk) { constexpr auto a_k_step_chunk = k_step + chunk * KThreadChunk * xdlops_gemm.mfma_instr.num_input_blks; - a_thread_copy_.Run( - a_block_desc_m0_m1_m2_k, - make_tuple(m0, I0, I0, Number{}), - a_block_buf, - a_thread_desc_, - make_tuple(m0, I0, k, Number{}), - a_thread_buf); + a_thread_copy_.Run(a_block_desc_m0_m1_m2_m3_k, + make_tuple(Number{}, + I0, + Number{}, + I0, + Number{}), + a_block_buf, + a_thread_desc_, + make_tuple(Number{}, + I0, + Number{}, + k, + Number{}), + a_thread_buf); }); }); static_for<0, NRepeat, 1>{}([&](auto n0) { @@ -649,13 +653,20 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}), - b_block_buf, - b_thread_desc_, - make_tuple(n0, I0, k, Number{}), - b_thread_buf); + b_thread_copy_.Run(b_block_desc_n0_n1_n2_n3_k, + make_tuple(Number{}, + I0, + Number{}, + I0, + Number{}), + b_block_buf, + b_thread_desc_, + make_tuple(Number{}, + I0, + Number{}, + k, + Number{}), + b_thread_buf); }); }); }); @@ -705,27 +716,7 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx(&a_scale_thread_bufs(I1)[Number<0>{}]), - *reinterpret_cast(&a_scale_thread_bufs(I1)[Number<1>{}]), - *reinterpret_cast(&a_scale_thread_bufs(I1)[Number<2>{}]), - *reinterpret_cast(&a_scale_thread_bufs(I1)[Number<3>{}])); - } - if(get_thread_local_1d_id()) - { - printf("2stGMEM Tid: %03d, Scale B: %02x %02x %02x %02x\n", - get_thread_local_1d_id(), - *reinterpret_cast(&b_scale_thread_bufs(I1)[Number<0>{}]), - *reinterpret_cast(&b_scale_thread_bufs(I1)[Number<1>{}]), - *reinterpret_cast(&b_scale_thread_bufs(I1)[Number<2>{}]), - *reinterpret_cast(&b_scale_thread_bufs(I1)[Number<3>{}])); - } -#endif block_sync_lds(); a_blockwise_copy.RunWrite(a_block_desc, a_block_buf); b_blockwise_copy.RunWrite(b_block_desc, b_block_buf); @@ -760,8 +751,6 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto imxdl) { static_for<0, NXdlPack, 1>{}([&](auto inxdl) { constexpr auto kxdl = ikxdl + k0 * KXdlPack; - constexpr auto mxdl = imxdl + m0 * MXdlPack; - constexpr auto nxdl = inxdl + n0 * NXdlPack; vector_type a_thread_vec; vector_type b_thread_vec; @@ -769,10 +758,10 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto ik) { a_thread_vec.template AsType()(ik) = a_thread_buf[Number{}]; + make_tuple(m0, I0, imxdl, kxdl, ik))>{}]; b_thread_vec.template AsType()(ik) = b_thread_buf[Number{}]; + make_tuple(n0, I0, inxdl, kxdl, ik))>{}]; }); using mfma_input_type_a = @@ -793,7 +782,7 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx::type; constexpr index_t c_offset = - c_thread_desc_.CalculateOffset(make_tuple(mxdl, nxdl, 0)); + c_thread_desc_.CalculateOffset(make_tuple(m0, n0, imxdl, inxdl, 0)); // MFMA accumulation xdlops_gemm.template Run{}], - c_thread_buf[Number<1>{}], - c_thread_buf[Number<2>{}], - c_thread_buf[Number<3>{}]); - } -#endif + block_sync_lds(); static_for<0, KRepeat, 1>{}([&](auto k) { constexpr auto k_step = k * xdlops_gemm.KPerXdlops * (KPack / xdlops_gemm.K1PerXdlops); - static_for<0, MRepeat, 1>{}([&](auto m0) { static_for<0, xdlops_gemm.K1PerXdlops / KThreadChunk, 1>{}([&](auto chunk) { constexpr auto a_k_step_chunk = k_step + chunk * KThreadChunk * xdlops_gemm.mfma_instr.num_input_blks; - a_thread_copy_.Run(a_block_desc_m0_m1_m2_k, - make_tuple(m0, I0, I0, Number{}), + a_thread_copy_.Run(a_block_desc_m0_m1_m2_m3_k, + make_tuple(Number{}, + I0, + Number{}, + I0, + Number{}), a_block_buf, a_thread_desc_, - make_tuple(m0, I0, k, Number{}), + make_tuple(Number{}, + I0, + Number{}, + k, + Number{}), a_thread_buf); }); }); @@ -845,11 +831,19 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto chunk) { constexpr auto b_k_step_chunk = k_step + chunk * KThreadChunk * xdlops_gemm.mfma_instr.num_input_blks; - b_thread_copy_.Run(b_block_desc_n0_n1_n2_k, - make_tuple(n0, I0, I0, Number{}), + b_thread_copy_.Run(b_block_desc_n0_n1_n2_n3_k, + make_tuple(Number{}, + I0, + Number{}, + I0, + Number{}), b_block_buf, b_thread_desc_, - make_tuple(n0, I0, k, Number{}), + make_tuple(Number{}, + I0, + Number{}, + k, + Number{}), b_thread_buf); }); }); @@ -885,8 +879,6 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto imxdl) { static_for<0, NXdlPack, 1>{}([&](auto inxdl) { constexpr auto kxdl = ikxdl + k0 * KXdlPack; - constexpr auto mxdl = imxdl + m0 * MXdlPack; - constexpr auto nxdl = inxdl + n0 * NXdlPack; vector_type a_thread_vec; vector_type b_thread_vec; @@ -894,10 +886,10 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto ik) { a_thread_vec.template AsType()(ik) = a_thread_buf[Number{}]; + make_tuple(m0, I0, imxdl, kxdl, ik))>{}]; b_thread_vec.template AsType()(ik) = b_thread_buf[Number{}]; + make_tuple(n0, I0, inxdl, kxdl, ik))>{}]; }); using mfma_input_type_a = @@ -918,7 +910,7 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx::type; constexpr index_t c_offset = - c_thread_desc_.CalculateOffset(make_tuple(mxdl, nxdl, 0)); + c_thread_desc_.CalculateOffset(make_tuple(m0, n0, imxdl, inxdl, 0)); // MFMA accumulation xdlops_gemm.template Run{}], - c_thread_buf[Number<1>{}], - c_thread_buf[Number<2>{}], - c_thread_buf[Number<3>{}]); - } -#endif } else if constexpr(TailNum == TailNumber::Odd) { @@ -980,8 +961,6 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto imxdl) { static_for<0, NXdlPack, 1>{}([&](auto inxdl) { constexpr auto kxdl = ikxdl + k0 * KXdlPack; - constexpr auto mxdl = imxdl + m0 * MXdlPack; - constexpr auto nxdl = inxdl + n0 * NXdlPack; vector_type a_thread_vec; vector_type b_thread_vec; @@ -989,10 +968,10 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto ik) { a_thread_vec.template AsType()(ik) = a_thread_buf[Number{}]; + make_tuple(m0, I0, imxdl, kxdl, ik))>{}]; b_thread_vec.template AsType()(ik) = b_thread_buf[Number{}]; + make_tuple(n0, I0, inxdl, kxdl, ik))>{}]; }); using mfma_input_type_a = @@ -1013,7 +992,7 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx::type; constexpr index_t c_offset = - c_thread_desc_.CalculateOffset(make_tuple(mxdl, nxdl, 0)); + c_thread_desc_.CalculateOffset(make_tuple(m0, n0, imxdl, inxdl, 0)); // MFMA accumulation xdlops_gemm.template Run 1) { if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd) @@ -388,12 +388,6 @@ struct DeviceGemmMX_Xdl_CShuffleV3 : public DeviceGemmMX; - Run(kernel); } else { @@ -426,6 +420,7 @@ struct DeviceGemmMX_Xdl_CShuffleV3 : public DeviceGemmMX; Run(kernel); } +#endif } } diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp index 77d554d18e..063783b6e2 100644 --- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp +++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp @@ -152,6 +152,8 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 static constexpr auto I5 = Number<5>{}; static constexpr auto I6 = Number<6>{}; static constexpr auto I7 = Number<7>{}; + static constexpr auto I8 = Number<8>{}; + static constexpr auto I9 = Number<9>{}; // K1 should be Number<...> static constexpr auto AK0Number = Number{}; @@ -254,7 +256,11 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 return math::integer_divide_ceil(N, NPerBlock); } - template + template __host__ __device__ static constexpr auto MakeGemmMmaTileDescriptor(const TileDesc_K0_MN_K1&) { constexpr index_t K0 = TileDesc_K0_MN_K1{}.GetLength(Number<0>{}); @@ -263,10 +269,12 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 return transform_tensor_descriptor( TileDesc_K0_MN_K1{}, make_tuple(make_merge_transform_v3_division_mod(make_tuple(Number{}, Number{})), - make_unmerge_transform(make_tuple( - Number{}, Number{}, Number{}))), + make_unmerge_transform(make_tuple(Number{}, + Number{}, + Number{}, + Number{}))), make_tuple(Sequence<0, 2>{}, Sequence<1>{}), - make_tuple(Sequence<3>{}, Sequence<0, 1, 2>{})); + make_tuple(Sequence<4>{}, Sequence<0, 1, 2, 3>{})); } __host__ __device__ static auto MakeAGridDescriptor_AK0_M_AK1( @@ -466,20 +474,22 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 template __host__ __device__ static constexpr auto - MakeAMmaTileDescriptor_M0_M1_M2_K(const ABlockDesc_AK0_M_AK1&) + MakeAMmaTileDescriptor_M0_M1_M2_M3_K(const ABlockDesc_AK0_M_AK1&) { constexpr index_t MWaves = MPerBlock / (MXdlPerWave * MPerXdl); - return MakeGemmMmaTileDescriptor(ABlockDesc_AK0_M_AK1{}); + return MakeGemmMmaTileDescriptor( + ABlockDesc_AK0_M_AK1{}); } template __host__ __device__ static constexpr auto - MakeBMmaTileDescriptor_N0_N1_N2_K(const BBlockDesc_BK0_N_BK1&) + MakeBMmaTileDescriptor_N0_N1_N2_N3_K(const BBlockDesc_BK0_N_BK1&) { constexpr index_t NWaves = NPerBlock / (NXdlPerWave * NPerXdl); - return MakeGemmMmaTileDescriptor(BBlockDesc_BK0_N_BK1{}); + return MakeGemmMmaTileDescriptor( + BBlockDesc_BK0_N_BK1{}); } __host__ __device__ static auto @@ -768,15 +778,14 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 // in some cases. else if constexpr(is_same::value) { - constexpr auto a_lds_block_desc = make_naive_tensor_descriptor( - make_tuple( - AK0Number, Number{}, AK1Number), - make_tuple(AK1Number, Number{}, I1)); + constexpr auto a_lds_block_desc = + make_naive_tensor_descriptor(make_tuple(AK0Number, Number{}, AK1Number), + make_tuple(AK1Number, Number{}, I1)); constexpr auto a_lds_block_desc_permuted = transform_tensor_descriptor( a_lds_block_desc, - make_tuple(make_xor_with_modulo_transform(make_tuple( - Number{}, Number{})), + make_tuple(make_xor_with_modulo_transform( + make_tuple(Number{}, Number{})), make_pass_through_transform(AK1Number)), make_tuple(Sequence<1, 0>{}, Sequence<2>{}), make_tuple(Sequence<1, 0>{}, Sequence<2>{})); @@ -885,15 +894,14 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 else if constexpr(is_same::value) { // NLdsLayer * K0 as logical Bank - constexpr auto b_lds_block_desc = make_naive_tensor_descriptor( - make_tuple( - BK0Number, Number{}, BK1Number), - make_tuple(BK1Number, Number{}, I1)); + constexpr auto b_lds_block_desc = + make_naive_tensor_descriptor(make_tuple(BK0Number, Number{}, BK1Number), + make_tuple(BK1Number, Number{}, I1)); constexpr auto b_lds_block_desc_permuted = transform_tensor_descriptor( b_lds_block_desc, - make_tuple(make_xor_with_modulo_transform(make_tuple( - Number{}, Number{})), + make_tuple(make_xor_with_modulo_transform( + make_tuple(Number{}, Number{})), make_pass_through_transform(BK1Number)), make_tuple(Sequence<1, 0>{}, Sequence<2>{}), make_tuple(Sequence<1, 0>{}, Sequence<2>{})); @@ -1016,9 +1024,9 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 AccDataType, decltype(GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1()), decltype(GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1()), - decltype(MakeAMmaTileDescriptor_M0_M1_M2_K( + decltype(MakeAMmaTileDescriptor_M0_M1_M2_M3_K( GetABlockDescriptor_AK0PerBlock_MPerBlock_AK1())), - decltype(MakeBMmaTileDescriptor_N0_N1_N2_K( + decltype(MakeBMmaTileDescriptor_N0_N1_N2_N3_K( GetBBlockDescriptor_BK0PerBlock_NPerBlock_BK1())), ABlockTransferSrcScalarPerVector, BBlockTransferSrcScalarPerVector, @@ -1480,7 +1488,8 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 // mfma.selected_mfma.num_threads_per_blk; // A wave access continuous memory - auto thread_offset_shuffled = get_thread_local_1d_id() % BlockwiseGemmPipe::WaveSize * KXdlPack * MXdlPack; + auto thread_offset_shuffled = + get_thread_local_1d_id() % BlockwiseGemmPipe::WaveSize * KXdlPack * MXdlPack; auto a_thread_offset_m = waveId_m; @@ -1496,8 +1505,9 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 1, // SrcScalarStrideInVector true>( a_scale_grid_desc_am_ak, - make_multi_index( - block_m_id * MPerBlock + a_thread_offset_m, 0, thread_offset_shuffled)); + make_multi_index(block_m_id * MPerBlock / MPerXdl / MXdlPack + a_thread_offset_m, + 0, + thread_offset_shuffled)); auto b_thread_offset_n = waveId_n; @@ -1513,8 +1523,9 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 1, true>( b_scale_grid_desc_bn_ak, - make_multi_index( - block_n_id * NPerBlock + b_thread_offset_n, 0, thread_offset_shuffled)); + make_multi_index(block_n_id * NPerBlock / NPerXdl / NXdlPack + b_thread_offset_n, + 0, + thread_offset_shuffled)); blockwise_gemm_pipeline.template Run(a_grid_desc_ak0_m_ak1, a_block_desc_ak0_m_ak1, @@ -1543,27 +1554,32 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 static_assert(MXdlPerWave % CShuffleMXdlPerWavePerShuffle == 0 && NXdlPerWave % CShuffleNXdlPerWavePerShuffle == 0, "wrong!"); + static_assert(CShuffleMXdlPerWavePerShuffle % MXdlPack == 0 && + CShuffleNXdlPerWavePerShuffle % NXdlPack == 0, + "wrong!"); constexpr index_t MWave = MPerBlock / (MXdlPerWave * MPerXdl); constexpr index_t NWave = NPerBlock / (NXdlPerWave * NPerXdl); // TODO: hacky, fix it! constexpr auto c_thread_desc_m0_n0_m1_n1_m2_m3_m4_n2 = - blockwise_gemm_pipeline.GetCThreadDescriptor_M0_N0_M1_N1_M2_M3_M4_N2(); + blockwise_gemm_pipeline.GetCThreadDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3(); // TODO: hacky, fix it! // c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp is only used to get lengths constexpr auto c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp = - blockwise_gemm_pipeline.GetCBlockDescriptor_M0_N0_M1_N1_M2_M3_M4_N2(); + blockwise_gemm_pipeline.GetCBlockDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3(); constexpr auto M0 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I0); constexpr auto N0 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I1); constexpr auto M1 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I2); constexpr auto N1 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I3); constexpr auto M2 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I4); - constexpr auto M3 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I5); - constexpr auto M4 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I6); - constexpr auto N2 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I7); + constexpr auto N2 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I5); + constexpr auto M3 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I6); + constexpr auto M4 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I7); + constexpr auto M5 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I8); + constexpr auto N3 = c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2_tmp.GetLength(I9); constexpr auto c_shuffle_block_desc_mblock_mperblock_nblock_nperblock = GetCShuffleBlockDescriptor_MBlock_MPerBlock_NBlock_NPerBlock(); @@ -1577,19 +1593,25 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 make_tuple( make_freeze_transform(I0), make_unmerge_transform(make_tuple( - Number{}, // M0 (MXdlPerWave) per shuffle - M1, // M1 = MWave - M2, // M2 * M3 * M4 = MPerXdl - M3, - M4)), + Number{}, // M0 (MXdlPerWave) per + // shuffle + M1, // M1 = MWave + M2, // M2 = MXdlPack + M3, // M3 * M4 * M5 = MPerXdl + M4, + M5)), make_freeze_transform(I0), make_unmerge_transform(make_tuple( - Number{}, // N0 (NXdlPerWave) per shuffle - N1, // N1 = NWave - N2))), // N2 = NPerXdl + Number{}, // N0 (NXdlPerWave) per + // shuffle + N1, // N1 = NWave + N2, // N2 = NXdlPack + N3))), // N3 = NPerXdl make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}), - make_tuple( - Sequence<>{}, Sequence<0, 2, 4, 5, 6>{}, Sequence<>{}, Sequence<1, 3, 7>{})); + make_tuple(Sequence<>{}, + Sequence<0, 2, 4, 6, 7, 8>{}, + Sequence<>{}, + Sequence<1, 3, 5, 9>{})); // calculate origin of thread output tensor on global memory // blockwise GEMM c matrix starting index @@ -1601,8 +1623,8 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 const auto m_thread_data_on_block_to_m0_m1_m2_m3_m4_adaptor = make_single_stage_tensor_adaptor( - make_tuple(make_merge_transform(make_tuple(M0, M1, M2, M3, M4))), - make_tuple(Sequence<0, 1, 2, 3, 4>{}), + make_tuple(make_merge_transform(make_tuple(M0, M1, M2, M3, M4, M5))), + make_tuple(Sequence<0, 1, 2, 3, 4, 5>{}), make_tuple(Sequence<0>{})); const auto m_thread_data_on_block_idx = @@ -1611,8 +1633,8 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 const auto n_thread_data_on_block_to_n0_n1_n2_adaptor = make_single_stage_tensor_adaptor( - make_tuple(make_merge_transform(make_tuple(N0, N1, N2))), - make_tuple(Sequence<0, 1, 2>{}), + make_tuple(make_merge_transform(make_tuple(N0, N1, N2, N3))), + make_tuple(Sequence<0, 1, 2, 3>{}), make_tuple(Sequence<0>{})); const auto n_thread_data_on_block_idx = @@ -1620,36 +1642,39 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 make_multi_index(n_thread_data_on_block)); // shuffle: threadwise copy C from VGPR to LDS - auto c_thread_copy_vgpr_to_lds = - ThreadwiseTensorSliceTransfer_v1r3, - Sequence<0, 1, 2, 3, 4, 5, 6, 7>, - 7, - 1, - InMemoryDataOperationEnum::Set, - 1, - true>{ - c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2, - make_multi_index(0, - 0, - m_thread_data_on_block_idx[I1], - n_thread_data_on_block_idx[I1], - m_thread_data_on_block_idx[I2], - m_thread_data_on_block_idx[I3], - m_thread_data_on_block_idx[I4], - n_thread_data_on_block_idx[I2]), - ck::tensor_operation::element_wise::PassThrough{}}; + auto c_thread_copy_vgpr_to_lds = ThreadwiseTensorSliceTransfer_v1r3< + AccDataType, + CShuffleDataType, + decltype(c_thread_desc_m0_n0_m1_n1_m2_m3_m4_n2), + decltype(c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2), + ck::tensor_operation::element_wise::PassThrough, + Sequence, + Sequence<0, 1, 2, 3, 4, 5, 6, 7, 8, 9>, + 9, + 1, + InMemoryDataOperationEnum::Set, + 1, + true>{c_block_desc_m0_n0_m1_n1_m2_m3_m4_n2, + make_multi_index(0, + 0, + m_thread_data_on_block_idx[I1], + n_thread_data_on_block_idx[I1], + m_thread_data_on_block_idx[I2], + n_thread_data_on_block_idx[I2], + m_thread_data_on_block_idx[I3], + m_thread_data_on_block_idx[I4], + m_thread_data_on_block_idx[I5], + n_thread_data_on_block_idx[I3]), + ck::tensor_operation::element_wise::PassThrough{}}; // shuffle: blockwise copy C from LDS to global auto c_shuffle_block_copy_lds_to_global = ThreadGroupTensorSliceTransfer_v6r1< @@ -1679,12 +1704,23 @@ struct GridwiseGemmMX_xdl_cshuffle_v3 // space filling curve for threadwise C in VGPR constexpr auto sfc_c_vgpr = - SpaceFillingCurve, - Sequence<0, 1, 2, 3, 4, 5, 6, 7>, - Sequence, + Sequence<0, 1, 2, 3, 4, 5, 6, 7, 8, 9>, + Sequence{})); } + // XDL output supporting C = A * B + // M3_N3 -> M3_M4_M5_N3 + template + __host__ __device__ static constexpr auto MakeCDescriptor_M0_N0_M1_N1_M2_N2_M3_M4_M5_N3( + const CDesc_M0_N0_M1_N1_M2_N2& c_desc_m0_n0_m1_n1_m2_n2) + { + const auto M0 = c_desc_m0_n0_m1_n1_m2_n2.GetLength(I0); + const auto N0 = c_desc_m0_n0_m1_n1_m2_n2.GetLength(I1); + const auto M1 = c_desc_m0_n0_m1_n1_m2_n2.GetLength(I2); + const auto N1 = c_desc_m0_n0_m1_n1_m2_n2.GetLength(I3); + const auto M2 = c_desc_m0_n0_m1_n1_m2_n2.GetLength(I4); + const auto N2 = c_desc_m0_n0_m1_n1_m2_n2.GetLength(I5); + + return transform_tensor_descriptor( + c_desc_m0_n0_m1_n1_m2_n2, + make_tuple(make_pass_through_transform(M0), + make_pass_through_transform(N0), + make_pass_through_transform(M1), + make_pass_through_transform(N1), + make_pass_through_transform(M2), + make_pass_through_transform(N2), + make_unmerge_transform(make_tuple(Number{}, + Number{}, + Number{})), + make_pass_through_transform(Number{})), + make_tuple(Sequence<0>{}, + Sequence<1>{}, + Sequence<2>{}, + Sequence<3>{}, + Sequence<4>{}, + Sequence<5>{}, + Sequence<6>{}, + Sequence<7>{}), + make_tuple(Sequence<0>{}, + Sequence<1>{}, + Sequence<2>{}, + Sequence<3>{}, + Sequence<4>{}, + Sequence<5>{}, + Sequence<6, 7, 8>{}, + Sequence<9>{})); + } + // transposed XDL output supporting C' = B' * A' // M2_N2 -> M2_N2_N3_N4 template diff --git a/include/ck/utility/amd_xdlops.hpp b/include/ck/utility/amd_xdlops.hpp index 85edf6fc7c..d944068f74 100644 --- a/include/ck/utility/amd_xdlops.hpp +++ b/include/ck/utility/amd_xdlops.hpp @@ -769,8 +769,8 @@ struct intrin_mfma_scale_f32_16x16x128f8f6f4<16, 16, OpselA, OpselB> int32x4_t arg_a = bit_cast(reg_a); int32x4_t arg_b = bit_cast(reg_b); +#if 0 using arg_type = int32x8_t; - reg_c.template AsType()(Number<0>{}) = __builtin_amdgcn_mfma_scale_f32_16x16x128_f8f6f4( arg_type{arg_a[0], arg_a[1], arg_a[2], arg_a[3], 0, 0, 0, 0}, @@ -782,6 +782,48 @@ struct intrin_mfma_scale_f32_16x16x128f8f6f4<16, 16, OpselA, OpselB> scale_a, OpselB, // OPSEL scale_b); +#else + using arg_type = int32x4_t; +#define v_mfma_scale(OPSEL_A_L, OPSEL_A_H, OPSEL_B_L, OPSEL_B_H) \ + else if constexpr((OpselA == 1 * OPSEL_A_L + 2 * OPSEL_A_H) && \ + (OpselB == 1 * OPSEL_B_L + 2 * OPSEL_B_H)) \ + { \ + asm volatile("v_mfma_scale_f32_16x16x128_f8f6f4 %0, %1, %2, %3, %4, %5 " \ + "op_sel:[" #OPSEL_A_L "," #OPSEL_A_H "] " \ + "op_sel_hi:[" #OPSEL_B_L "," #OPSEL_B_H "] " \ + "cbsz:4 blgp:4" \ + : "+v"(reg_c.template AsType()(Number<0>{})) \ + : "v"(arg_type{arg_a[0], arg_a[1], arg_a[2], arg_a[3]}), \ + "v"(arg_type{arg_b[0], arg_b[1], arg_b[2], arg_b[3]}), \ + "v"(reg_c.template AsType()[Number<0>{}]), \ + "v"(scale_a), \ + "v"(scale_b)); \ + } + using arg_type = int32x4_t; + if constexpr(false) {} + v_mfma_scale(0, 0, 0, 0) // + v_mfma_scale(0, 0, 0, 1) // + v_mfma_scale(0, 0, 1, 0) // + v_mfma_scale(0, 0, 1, 1) // + v_mfma_scale(0, 1, 0, 0) // + v_mfma_scale(0, 1, 0, 1) // + v_mfma_scale(0, 1, 1, 0) // + v_mfma_scale(0, 1, 1, 1) // + v_mfma_scale(1, 0, 0, 0) // + v_mfma_scale(1, 0, 0, 1) // + v_mfma_scale(1, 0, 1, 0) // + v_mfma_scale(1, 0, 1, 1) // + v_mfma_scale(1, 1, 0, 0) // + v_mfma_scale(1, 1, 0, 1) // + v_mfma_scale(1, 1, 1, 0) // + v_mfma_scale(1, 1, 1, 1) // + else + { + static_assert(0, "Unsupported op_sel"); + } +#undef v_mfma_scale +#endif + #else ignore = reg_a; ignore = scale_a;