From 0d40254873a51c8ad8e569876ea928e9453aba48 Mon Sep 17 00:00:00 2001
From: Damien Lejeune <damien.lejeune@amd.com>
Date: Tue, 29 Jul 2025 07:01:32 +0000
Subject: [PATCH] Expose skip lds to batch/group gemm kernels api

---
 .../ops/gemm/kernel/batched_gemm_kernel.hpp   | 18 +++++++--
 .../ops/gemm/kernel/grouped_gemm_kernel.hpp   |  4 +-
 .../batched_gemm/test_batched_gemm_util.hpp   |  5 +++
 .../grouped_gemm/test_grouped_gemm_util.hpp   | 39 +++++++++++++------
 4 files changed, 51 insertions(+), 15 deletions(-)

diff --git a/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp
index dccecfc113..7181afa346 100644
--- a/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp
+++ b/include/ck_tile/ops/gemm/kernel/batched_gemm_kernel.hpp
@@ -28,8 +28,18 @@ struct BatchedGemmHostArgs : public ck_tile::GemmHostArgs
                                      ck_tile::index_t batch_stride_B_,
                                      ck_tile::index_t batch_stride_C_,
                                      ck_tile::index_t batch_count_)
-        : GemmHostArgs(
-              a_ptr_, b_ptr_, c_ptr_, k_batch_, M_, N_, K_, stride_A_, stride_B_, stride_C_, skip_a_lds_, skip_b_lds_),
+        : GemmHostArgs(a_ptr_,
+                       b_ptr_,
+                       c_ptr_,
+                       k_batch_,
+                       M_,
+                       N_,
+                       K_,
+                       stride_A_,
+                       stride_B_,
+                       stride_C_,
+                       skip_a_lds_,
+                       skip_b_lds_),
           batch_stride_A(batch_stride_A_),
           batch_stride_B(batch_stride_B_),
           batch_stride_C(batch_stride_C_),
@@ -103,7 +113,9 @@ struct BatchedGemmKernel : public GemmKernel<TilePartitioner_, GemmPipeline_, Ep
                                       hostArgs.stride_A,
                                       hostArgs.stride_B,
                                       hostArgs.stride_C,
-                                      hostArgs.k_batch},
+                                      hostArgs.k_batch,
+                                      hostArgs.skip_a_lds,
+                                      hostArgs.skip_b_lds},
                                      hostArgs.batch_stride_A,
                                      hostArgs.batch_stride_B,
                                      hostArgs.batch_stride_C,
diff --git a/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp b/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp
index 5577cb083a..f24905dc7c 100644
--- a/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp
+++ b/include/ck_tile/ops/gemm/kernel/grouped_gemm_kernel.hpp
@@ -113,7 +113,9 @@ struct GroupedGemmKernel : public GemmKernel<TilePartitioner_, GemmPipeline_, Ep
                                        stride_a,
                                        stride_b,
                                        stride_c,
-                                       gemm_descs[i].k_batch};
+                                       gemm_descs[i].k_batch,
+                                       gemm_descs[i].skip_a_lds,
+                                       gemm_descs[i].skip_b_lds};
 
             gemm_kernel_args_.emplace_back(std::move(karg), block_start, block_end);
         }
diff --git a/test/ck_tile/batched_gemm/test_batched_gemm_util.hpp b/test/ck_tile/batched_gemm/test_batched_gemm_util.hpp
index 4633f23ded..087801b19a 100644
--- a/test/ck_tile/batched_gemm/test_batched_gemm_util.hpp
+++ b/test/ck_tile/batched_gemm/test_batched_gemm_util.hpp
@@ -42,6 +42,9 @@ class TestCkTileBatchedGemm : public ::testing::Test
 
         constexpr bool DoubleSmemBuffer = false;
 
+        constexpr bool kSkipALds = false;
+        constexpr bool kSkipBLds = false;
+
         constexpr bool kPadM = false;
         constexpr bool kPadN = false;
         constexpr bool kPadK = false;
@@ -64,6 +67,8 @@ class TestCkTileBatchedGemm : public ::testing::Test
                                                                      kPadN,
                                                                      kPadK,
                                                                      DoubleSmemBuffer,
+                                                                     kSkipALds,
+                                                                     kSkipBLds,
                                                                      ALayout,
                                                                      BLayout,
                                                                      CLayout,
diff --git a/test/ck_tile/grouped_gemm/test_grouped_gemm_util.hpp b/test/ck_tile/grouped_gemm/test_grouped_gemm_util.hpp
index 3dec229643..86dd5d78a6 100644
--- a/test/ck_tile/grouped_gemm/test_grouped_gemm_util.hpp
+++ b/test/ck_tile/grouped_gemm/test_grouped_gemm_util.hpp
@@ -30,6 +30,9 @@ class TestCkTileGroupedGemm : public ::testing::Test
         static const bool kPadN = false;
         static const bool kPadK = false;
 
+        static const bool SkipALds = false;
+        static const bool SkipBLds = false;
+
         static const int kBlockPerCu         = 1;
         static const ck_tile::index_t M_Tile = 128;
         static const ck_tile::index_t N_Tile = 128;
@@ -74,20 +77,23 @@ class TestCkTileGroupedGemm : public ::testing::Test
         using TilePartitioner = ck_tile::
             GemmSpatiallyLocalTilePartitioner<GemmShape, TileParitionerGroupNum, TileParitionerM01>;
 
-        using Traits              = ck_tile::TileGemmTraits<GroupedGemKernelParam::kPadM,
+        using Traits = ck_tile::TileGemmTraits<GroupedGemKernelParam::kPadM,
                                                GroupedGemKernelParam::kPadN,
                                                GroupedGemKernelParam::kPadK,
                                                ALayout,
                                                BLayout,
                                                CLayout>;
-        using GemmUniversalTraits = ck_tile::TileGemmUniversalTraits<GroupedGemKernelParam::kPadM,
-                                                                     GroupedGemKernelParam::kPadN,
-                                                                     GroupedGemKernelParam::kPadK,
-                                                                     DoubleSmemBuffer,
-                                                                     ALayout,
-                                                                     BLayout,
-                                                                     CLayout,
-                                                                     TransposeC>;
+        using GemmUniversalTraits =
+            ck_tile::TileGemmUniversalTraits<GroupedGemKernelParam::kPadM,
+                                             GroupedGemKernelParam::kPadN,
+                                             GroupedGemKernelParam::kPadK,
+                                             DoubleSmemBuffer,
+                                             GroupedGemKernelParam::SkipALds,
+                                             GroupedGemKernelParam::SkipBLds,
+                                             ALayout,
+                                             BLayout,
+                                             CLayout,
+                                             TransposeC>;
         using GemmPipelineProblem =
             ck_tile::GemmPipelineProblem<ADataType, BDataType, AccDataType, GemmShape, Traits>;
 
@@ -317,8 +323,19 @@ class TestCkTileGroupedGemm : public ::testing::Test
 
             // TODO add support for kbatch > 1
             static constexpr ck_tile::index_t k_batch = 1;
-            gemm_descs.push_back(
-                {p_a, p_b, p_c, k_batch, M, N, K, stride_As[i], stride_Bs[i], stride_Cs[i]});
+
+            gemm_descs.push_back({p_a,
+                                  p_b,
+                                  p_c,
+                                  k_batch,
+                                  M,
+                                  N,
+                                  K,
+                                  stride_As[i],
+                                  stride_Bs[i],
+                                  stride_Cs[i],
+                                  GroupedGemKernelParam::SkipALds,
+                                  GroupedGemKernelParam::SkipBLds});
         }
 
         ck_tile::DeviceMem gemm_workspace;