From 062e16d54ab7053fb0ebba8fda5d61a88162b73b Mon Sep 17 00:00:00 2001 From: aska-0096 Date: Thu, 15 May 2025 02:33:33 +0000 Subject: [PATCH] Improve the pipeline --- example/67_gemm_microscaling/CMakeLists.txt | 4 +- .../67_gemm_microscaling/gemm_mx_common.hpp | 35 -------------- example/67_gemm_microscaling/gemm_mx_fp4.cpp | 4 +- .../blockwise_gemm_pipeline_xdlops_v3_mx.hpp | 46 +++++++++++-------- include/ck/utility/amd_xdlops.hpp | 1 + 5 files changed, 32 insertions(+), 58 deletions(-) diff --git a/example/67_gemm_microscaling/CMakeLists.txt b/example/67_gemm_microscaling/CMakeLists.txt index 643dd1cb6b..6c0c973880 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 -v --save-temps -Wno-gnu-line-marker) +list(APPEND FP4_MXGEMM_OPTIONS "SHELL: -mllvm -greedy-reverse-local-assignment=1") +list(APPEND FP4_MXGEMM_OPTIONS -v --save-temps -Wno-gnu-line-marker -ftemplate-backtrace-limit=0) 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 214218a9dd..bf3c11b96e 100644 --- a/example/67_gemm_microscaling/gemm_mx_common.hpp +++ b/example/67_gemm_microscaling/gemm_mx_common.hpp @@ -315,41 +315,6 @@ bool run_mx_gemm(const ProblemSizeSplitK& problem_size, const ExecutionConfig& c b_k_n_scale.GenerateTensorValue(GeneratorTensor_3{powf(2.0f, -125.0f), 1.0f}); break; - case 3: - - ck::utils::FillConstant{a_data_element(1.0f)}(a_m_k); - b_k_n.GenerateTensorValue(GeneratorTensor_2{-5, 6}); // Z[-5,5] - - if constexpr(ck::is_same_v) - { - 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 - { - ck::utils::FillUniformDistributionIntegerValue{-1.0f, 1.0f}(a_m_k_scale); - ck::utils::FillUniformDistributionIntegerValue{-1.0f, 1.0f}(b_k_n_scale); - } - - break; - case 4: // A random - - a_m_k.GenerateTensorValue(GeneratorTensor_2{-5, 6}); // Z[-5,5] - ck::utils::FillConstant{b_data_element(1.0f)}(b_k_n); - - if constexpr(ck::is_same_v) - { - 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 - { - ck::utils::FillUniformDistributionIntegerValue{-1.0f, 1.0f}(a_m_k_scale); - ck::utils::FillUniformDistributionIntegerValue{-1.0f, 1.0f}(b_k_n_scale); - } - - break; - default: if(config.verbosity > 0) { diff --git a/example/67_gemm_microscaling/gemm_mx_fp4.cpp b/example/67_gemm_microscaling/gemm_mx_fp4.cpp index f47e09443b..0e913f82ec 100644 --- a/example/67_gemm_microscaling/gemm_mx_fp4.cpp +++ b/example/67_gemm_microscaling/gemm_mx_fp4.cpp @@ -53,14 +53,14 @@ using DeviceOpInstance = ck::tensor_operation::device::DeviceGemmMX_Xdl_CShuffle ScaleBlockSize, // ScaleBlockSize: Scaling block size 256, // BlockSize: Thread block size 128, // MPerBlock - 128, // NPerBlock + 256, // NPerBlock KPerBlock, // KPerBlock 32, // AK1 32, // BK1 16, // MPerXDL 16, // NPerXDL 4, // MXdlPerWave - 4, // NXdlPerWave + 8, // NXdlPerWave S<8, 32, 1>, // ABlockTransferThreadClusterLengths_AK0_M_AK1 S<1, 0, 2>, // ABlockTransferThreadClusterArrangeOrder S<1, 0, 2>, // ABlockTransferSrcAccessOrder 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 276b955fb1..345c9de53e 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 @@ -184,11 +184,11 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto i) { - ignore = i; + if constexpr(i < num_buffer_load_a_scale) + { + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + } static_for<0, num_dswrite_per_issue_a, 1>{}([&](auto idswrite) { ignore = idswrite; - __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write }); __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read __builtin_amdgcn_sched_group_barrier( 0x008, num_mfma_per_issue - num_dswrite_per_issue_a, 0); // MFMA }); static_for<0, num_buffer_load_inst_b, 1>{}([&](auto i) { - ignore = i; + if constexpr(i < num_buffer_load_b_scale) + { + __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read + } static_for<0, num_dswrite_per_issue_b, 1>{}([&](auto idswrite) { ignore = idswrite; - __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA + __builtin_amdgcn_sched_group_barrier(0x200, 1, 0); // DS write }); __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read __builtin_amdgcn_sched_group_barrier( @@ -282,6 +292,7 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto m0) { static_for<0, KRepeat / KXdlPack, 1>{}([&](auto k0) { @@ -486,6 +498,9 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx{}([&](auto n0) { static_for<0, KRepeat / KXdlPack, 1>{}([&](auto k0) { @@ -508,14 +523,7 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx::type; - constexpr index_t c_offset = - c_thread_desc_.CalculateOffset(make_tuple(m0, n0, imxdl, inxdl, 0)); + constexpr index_t c_offset = c_thread_desc_.CalculateOffset( + make_tuple(m0, n0, imxdl, inxdl, 0)); // MFMA accumulation xdlops_gemm.template Run::type; - constexpr index_t c_offset = - c_thread_desc_.CalculateOffset(make_tuple(m0, n0, imxdl, inxdl, 0)); + constexpr index_t c_offset = c_thread_desc_.CalculateOffset( + make_tuple(m0, n0, imxdl, inxdl, 0)); // MFMA accumulation xdlops_gemm.template Run::type; - constexpr index_t c_offset = - c_thread_desc_.CalculateOffset(make_tuple(m0, n0, imxdl, inxdl, 0)); + constexpr index_t c_offset = c_thread_desc_.CalculateOffset( + make_tuple(m0, n0, imxdl, inxdl, 0)); // MFMA accumulation xdlops_gemm.template Run #endif } }; +#define BUILTIN_AMDGCN_MFMA_SCALE_F32_16X16X128_F8F6F4_WORKS 1 #ifndef BUILTIN_AMDGCN_MFMA_SCALE_F32_16X16X128_F8F6F4_WORKS #define BUILTIN_AMDGCN_MFMA_SCALE_F32_16X16X128_F8F6F4_WORKS 0