From e809d445dafeeec1f825dc7480a2b20dcf41e3f8 Mon Sep 17 00:00:00 2001
From: OscarXu <huaiguxu@amd.com>
Date: Fri, 27 Jun 2025 03:28:31 -0500
Subject: [PATCH] gemm1 v3 perf optimization WIP. better than buf2vgpr now

---
 .../moe_gemm1_xdl_fp8_blockscale.cpp          |  2 +-
 ...oe_blockscale_b_preshuffle_gufusion_v3.hpp | 29 ++++++++++---------
 2 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8_blockscale.cpp b/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8_blockscale.cpp
index db2359f298..7e8fc665df 100644
--- a/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8_blockscale.cpp
+++ b/example/65_gemm_multiply_multiply/moe_gemm1_xdl_fp8_blockscale.cpp
@@ -180,7 +180,7 @@ static constexpr ck::index_t MPerBlock = 64; using DeviceOpInstance = ck::tensor
                S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 1,
                S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 16, 16, 0,
                4,    2,   S<1, 32, 1, 8>, S<2, 1, 1, 1>,
-               ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v1, ActOP, Nswizzle, true, MulRoutedWeight, int32_t, A0DataType>;
+               ck::BlockGemmPipelineScheduler::Intrawave, ck::BlockGemmPipelineVersion::v3, ActOP, Nswizzle, true, MulRoutedWeight, int32_t, A0DataType>;
 #endif
 // clang-format on
 
diff --git a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v3.hpp b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v3.hpp
index 84896002c6..a9f12a5d92 100644
--- a/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v3.hpp
+++ b/include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_moe_blockscale_b_preshuffle_gufusion_v3.hpp
@@ -261,10 +261,10 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3<
         static_for<0, buffer_load_a_stages, 1>{}([&](auto i) {
             static_for<0, num_mfma_perstage, 1>{}([&](auto imfma) {
                 __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA
-                if constexpr((((i + buffer_load_b_stages) < buffer_load_stages_more) &&
+                if constexpr(((i < buffer_load_stages_more) &&
                               (imfma % buffer_load_issue_point_interval_more ==
                                buffer_load_issue_point_a)) ||
-                             (((i + buffer_load_b_stages) >= buffer_load_stages_more) &&
+                             ((i >= buffer_load_stages_more) &&
                               (imfma % buffer_load_issue_point_interval_less ==
                                buffer_load_issue_point_a)) || 
                               ((num_mfma_perstage < buffer_load_perstage_less) && 
@@ -278,25 +278,18 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3<
                 }
                 // __builtin_amdgcn_sched_group_barrier(0x800, 2, 0); // v_pk_fma
             });
-            // Scale load, 1A
-            __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read
             // __builtin_amdgcn_sched_barrier(0);
         });
 
         // B global read
         static_for<0, buffer_load_b_stages, 1>{}([&](auto i) {
-            // Scale load, 1B
-            // if constexpr (i.value==0){
-            //     __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read
-            // }
-
             static_for<0, num_mfma_perstage, 1>{}([&](auto imfma) {
                 __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA
 
-                if constexpr(((i < buffer_load_stages_more) &&
+                if constexpr((((i+buffer_load_a_stages) < buffer_load_stages_more) &&
                               (imfma % buffer_load_issue_point_interval_more ==
                                buffer_load_issue_point_b)) ||
-                             ((i >= buffer_load_stages_more) &&
+                             (((i+buffer_load_a_stages) >= buffer_load_stages_more) &&
                               (imfma % buffer_load_issue_point_interval_less ==
                                buffer_load_issue_point_b)) || 
                               ((num_mfma_perstage < buffer_load_perstage_less) && 
@@ -311,9 +304,19 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3<
                 }
                 // __builtin_amdgcn_sched_group_barrier(0x800, 2, 0); // v_pk_fma
             });
+            // __builtin_amdgcn_sched_barrier(0);
+        });
+
+        //load scale
+        static_for<0, buffer_load_a_stages, 1>{}([&](auto i) {
+            ignore = i;
+            // Scale load, 1A
+            __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read
+        });
+        static_for<0, buffer_load_b_stages, 1>{}([&](auto i) {
+            ignore = i;
             // Scale load, 1B
             __builtin_amdgcn_sched_group_barrier(0x020, 2, 0); // VMEM read
-            // __builtin_amdgcn_sched_barrier(0);
         });
 
         // lds synchronization, prefetch next loop local A
@@ -788,7 +791,7 @@ struct BlockwiseGemmXdlops_pipeline_moe_blockscale_bpreshuffle_gufusion_v3<
                         c_scale_thread_buf_up(m0) = a_scale_thread_bufs[mfma_reg_buf][m0] * b_scale_thread_bufs_up[mfma_reg_buf][I0];
                     });
 
-                    HotLoopScheduler();
+                    // HotLoopScheduler();
                     __builtin_amdgcn_sched_barrier(0);
                 };