From 7151ef86f8632858ef1ed4bc560bf8c26b69c555 Mon Sep 17 00:00:00 2001
From: aska-0096 <haocwang@amd.com>
Date: Fri, 23 May 2025 07:51:26 +0000
Subject: [PATCH] terrible agpr spill observed

---
 ...blockwise_gemm_mx_pipeline_xdlops_base.hpp | 16 ++++++++------
 .../blockwise_gemm_pipeline_xdlops_v3_mx.hpp  |  7 +++---
 ...roup_tensor_slice_transfer_direct_load.hpp |  9 +++-----
 .../impl/device_gemm_xdl_cshuffle_v3_mx.hpp   | 22 +++++++++++++++++++
 .../grid/gridwise_gemm_xdl_cshuffle_v3_mx.hpp |  6 +++++
 5 files changed, 44 insertions(+), 16 deletions(-)

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 cee0b3f7d5..b7ecafce7e 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
@@ -404,17 +404,19 @@ struct BlockwiseGemmXdlops_mx_pipeline_base
     protected:
     // M1, N1 as double buffer index
     // Read buffer + Compute buffer
-    // A[M0, M1, M2, KPack]
+    // A[MRepeat/MXdlPack, MWave, MXdlPack, KRepeat, KPack]
+    //   1                 0      3         2           4
+    // 
     static constexpr auto a_thread_desc_ =
         make_naive_tensor_descriptor(make_tuple(Number<MRepeat / MXdlPack>{},
                                                 I1,
                                                 Number<MXdlPack>{},
                                                 Number<KRepeat>{},
                                                 Number<KPack>{}),
-                                     make_tuple(Number<KPack * MXdlPack>{},
-                                                Number<KRepeat * MRepeat * KPack>{},
-                                                Number<MRepeat * KPack>{},
+                                     make_tuple(Number<KRepeat * MXdlPack* KPack>{},
+                                                Number<MRepeat * KRepeat * KPack>{},
                                                 Number<KPack>{},
+                                                Number<MXdlPack*KPack>{},
                                                 I1));
 
     // B[N0, N1, N2, KPack]
@@ -424,10 +426,10 @@ struct BlockwiseGemmXdlops_mx_pipeline_base
                                                 Number<KRepeat>{},
                                                 Number<NXdlPack>{},
                                                 Number<KPack>{}),
-                                     make_tuple(Number<KPack * NXdlPack>{},
-                                                Number<KRepeat * NRepeat * KPack>{},
-                                                Number<NRepeat * KPack>{},
+                                     make_tuple(Number<KRepeat * NXdlPack* KPack>{},
+                                                Number<NRepeat * KRepeat * KPack>{},
                                                 Number<KPack>{},
+                                                Number<NXdlPack*KPack>{},
                                                 I1));
 
     // C[M, N, NumRegXdlops]
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 0459240fda..c5ce2bb416 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
@@ -246,8 +246,9 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx<BlockGemmPipelineScheduler::Intrawave,
         constexpr auto num_ds_read_a_scale = MRepeat / MXdlPack * KRepeat / KXdlPack;
         constexpr auto num_ds_read_b_scale = NRepeat / NXdlPack * KRepeat / KXdlPack;
 
-        constexpr auto num_dsread_a_scale_mfma = num_ds_read_a_scale;
-        constexpr auto num_dsread_b_scale_mfma = num_ds_read_b_scale;
+        // seems like compiler will use ds_read2_b32 for it.
+        constexpr auto num_dsread_a_scale_mfma = num_ds_read_a_scale/2;
+        constexpr auto num_dsread_b_scale_mfma = num_ds_read_b_scale/2;
         // To relieve TA backport pressure:
         // give small interval for buffer_load_dword
         // give large interval for buffer_load_dwordx4
@@ -513,7 +514,7 @@ struct BlockwiseGemmXdlops_pipeline_v3_mx<BlockGemmPipelineScheduler::Intrawave,
 
         a_blockwise_copy.MoveSrcSliceWindow(a_grid_desc, a_block_copy_step);
         b_blockwise_copy.MoveSrcSliceWindow(b_grid_desc, b_block_copy_step);
-
+        
         a_scale_blockwise_copy.Run(
             a_scale_grid_desc, a_scale_grid_buf, a_scale_block_desc, a_scale_block_bufs(I0));
         b_scale_blockwise_copy.Run(
diff --git a/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_direct_load.hpp b/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_direct_load.hpp
index 38040c1974..fc8dbcfdd4 100644
--- a/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_direct_load.hpp
+++ b/include/ck/tensor_operation/gpu/block/thread_group_tensor_slice_transfer_direct_load.hpp
@@ -41,6 +41,7 @@ namespace ck {
  */
 template <typename ThreadGroup,
           typename BlockSliceLengths,
+          typename WaveClusterLengths, // Ugly, FIX me.
           typename ThreadClusterLengths,
           typename ThreadClusterArrangeOrder,
           typename SrcData,
@@ -68,12 +69,8 @@ struct ThreadGroupTensorSliceTransfer_DirectLoad
 
     static constexpr auto block_slice_lengths    = BlockSliceLengths{};
     static constexpr auto thread_cluster_lengths = ThreadClusterLengths{};
-    static constexpr auto wave_thread_cluster_lengths =
-        Sequence<ThreadClusterLengths{}.At(I0),
-                 ThreadClusterLengths{}.At(I1) * 64 / ThreadGroup::GetNumOfThread(),
-                 1>{};
-    static constexpr auto wave_cluster_lengths =
-        Sequence<1, ThreadGroup::GetNumOfThread() / 64, 1>{};
+    static constexpr auto wave_cluster_lengths = WaveClusterLengths{};
+    static constexpr auto wave_thread_cluster_lengths = thread_cluster_lengths/wave_cluster_lengths;
 
     static constexpr auto thread_single_load_size = generate_sequence(
         detail::lambda_scalar_per_access<DstVectorDim, ScalarPerVector>{}, Number<nDim>{});
diff --git a/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3_mx.hpp b/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3_mx.hpp
index 05ab739785..528a4d6ae5 100644
--- a/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3_mx.hpp
+++ b/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle_v3_mx.hpp
@@ -416,6 +416,28 @@ struct DeviceGemmMX_Xdl_CShuffleV3 : public DeviceGemmMX<ALayout,
                         Run(kernel);
                     }
                 }
+                else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v3){
+                    if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd)
+                        {
+                            const auto kernel =
+                                kernel_gemm_xdl_cshuffle_v3_2lds<GridwiseGemm,
+                                                            false,
+                                                            InMemoryDataOperationEnum::Set,
+                                                            minimum_occupancy,
+                                                            TailNumber::Odd>;
+                            Run(kernel);
+                        }
+                        else
+                        {
+                            const auto kernel =
+                                kernel_gemm_xdl_cshuffle_v3_2lds<GridwiseGemm,
+                                                            false,
+                                                            InMemoryDataOperationEnum::Set,
+                                                            minimum_occupancy,
+                                                            TailNumber::Even>;
+                            Run(kernel);
+                        }
+                }
             }
 
             return ave_time;
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 fb623d8bb9..b5a50feebd 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
@@ -1253,6 +1253,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3
 
         auto a_blockwise_copy =
             ThreadGroupTensorSliceTransfer_DirectLoad<ThisThreadBlock,
+                                                      Sequence<1, BlockSize/64, 1>,
                                                       Sequence<AK0Number, MPerBlock, AK1Number>,
                                                       ABlockTransferThreadClusterLengths_AK0_M_AK1,
                                                       ABlockTransferThreadClusterArrangeOrder,
@@ -1272,6 +1273,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3
         // B matrix blockwise copy
         auto b_blockwise_copy =
             ThreadGroupTensorSliceTransfer_DirectLoad<ThisThreadBlock,
+                                                      Sequence<1, BlockSize/64, 1>,
                                                       Sequence<BK0Number, NPerBlock, BK1Number>,
                                                       BBlockTransferThreadClusterLengths_BK0_N_BK1,
                                                       BBlockTransferThreadClusterArrangeOrder,
@@ -1755,6 +1757,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3
         auto a_blockwise_copy =
             ThreadGroupTensorSliceTransfer_DirectLoad<ThisThreadBlock,
                                                       Sequence<AK0Number, MPerBlock, AK1Number>,
+                                                      Sequence<1, BlockSize/64, 1>,
                                                       ABlockTransferThreadClusterLengths_AK0_M_AK1,
                                                       ABlockTransferThreadClusterArrangeOrder,
                                                       ADataType,
@@ -1774,6 +1777,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3
         auto b_blockwise_copy =
             ThreadGroupTensorSliceTransfer_DirectLoad<ThisThreadBlock,
                                                       Sequence<BK0Number, NPerBlock, BK1Number>,
+                                                      Sequence<1, BlockSize/64, 1>,
                                                       BBlockTransferThreadClusterLengths_BK0_N_BK1,
                                                       BBlockTransferThreadClusterArrangeOrder,
                                                       BDataType,
@@ -1828,6 +1832,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3
             Sequence<MPerBlock/MPerXdl / MXdlPack,
                      KRepeat / KXdlPack,
                      64* KXdlPack * MXdlPack / scale_pack_size_b>,
+            Sequence<BlockSize / 64, 1, 1>,
             Sequence<BlockSize / 64, 1, 64>,
             Sequence<0, 1, 2>,
             AScaleDataType,
@@ -1847,6 +1852,7 @@ struct GridwiseGemmMX_xdl_cshuffle_v3
             Sequence<NPerBlock/NPerXdl / NXdlPack,
                      KRepeat / KXdlPack,
                      64* KXdlPack * NXdlPack / scale_pack_size_b>,
+            Sequence<BlockSize / 64, 1, 1>,
             Sequence<BlockSize / 64, 1, 64>,
             Sequence<0, 1, 2>,
             BScaleDataType,