From 51a4ae44ef1b2e675291d6ae65401dcf9504ba5b Mon Sep 17 00:00:00 2001
From: kiefer <kiefer.van.teutem@streamhpc.com>
Date: Mon, 10 Nov 2025 13:13:04 +0000
Subject: [PATCH] Post-merge fixes. Make sure the new gridwise gemm wmma v3
 common Run function can be used. Remove splitK, and forceThreadTileTransfer
 for now. Also add CShuffle epilogue argument.

---
 ..._conv_fwd_multiple_abd_wmma_cshuffle_v3.hpp | 18 ++++++++++++------
 .../grid/gridwise_gemm_wmma_cshuffle_v3.hpp    | 16 +++++++++-------
 2 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_wmma_cshuffle_v3.hpp b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_wmma_cshuffle_v3.hpp
index 7e076c15b3..a0a84fa0a0 100644
--- a/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_wmma_cshuffle_v3.hpp
+++ b/include/ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_wmma_cshuffle_v3.hpp
@@ -82,7 +82,10 @@ __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy)
                     std::is_same_v<e_data_type, ck::bhalf_t>)))
     {
 #endif
-        __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
+        __shared__ char p_shared[GridwiseGemm::template GetSharedMemoryNumberOfByte<
+            typename GridwiseGemm::EpilogueCShuffle>()];
+
+        auto epilogue_args = typename GridwiseGemm::EpilogueCShuffle{};
 
         GridwiseGemm::template Run<AGridDesc_AK0_M_AK1,
                                    BGridDesc_BK0_N_BK1,
@@ -100,7 +103,8 @@ __launch_bounds__(CK_MAX_THREAD_PER_BLOCK, MinimumOccupancy)
                                             compute_ptr_offset_of_batch,
                                             compute_ptr_offset_of_n,
                                             num_k_per_block,
-                                            karg);
+                                            karg,
+                                            epilogue_args);
 #if defined(__gfx11__)
     }
 #endif
@@ -448,8 +452,9 @@ struct DeviceGroupedConvFwdMultipleABD_Wmma_CShuffle_V3
         BlkGemmPipelineVer,
         AComputeDataType,
         BComputeDataType,
-        false,  // PermuteA
-        false>; // PermuteB
+        false, // PermuteA
+        false, // PermuteB
+        true>; // ForceThreadTileTransfer
 
     // TODO: Previously available template param DoElementwiseBeforeCShuffle!
 
@@ -520,8 +525,9 @@ struct DeviceGroupedConvFwdMultipleABD_Wmma_CShuffle_V3
         AComputeDataType, // TODO: Swapped these but will probably never get verified because the
                           // only mixed precision instances are not NCHW.
 
-        false,  // PermuteB
-        false>; // PermuteA
+        false, // PermuteB
+        false, // PermuteA
+        true>; // ForceThreadTileTransfer
 
     using GridwiseGemmCTranspose =
         std::conditional_t<CTranspose, GridwiseGemmSwappedParams, GridwiseGemm>;
diff --git a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3.hpp b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3.hpp
index d3ea19a112..b92fbe6601 100644
--- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3.hpp
+++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_wmma_cshuffle_v3.hpp
@@ -682,7 +682,8 @@ struct GridwiseGemm_wmma_cshuffle_v3
               typename ComputePtrOffsetOfN,
               bool HasMainKBlockLoop,
               InMemoryDataOperationEnum EGlobalMemoryDataOperation,
-              TailNumber TailNum>
+              TailNumber TailNum,
+              typename EpilogueArgument>
     __device__ static void Run(void* p_shared,
                                const AGridDesc_AK0_M_K1& a_grid_desc_ak0_m_ak1,
                                const BGridDesc_BK0_N_K1& b_grid_desc_bk0_n_bk1,
@@ -692,13 +693,14 @@ struct GridwiseGemm_wmma_cshuffle_v3
                                    e_grid_desc_mblock_mperblock_nblock_nperblock,
                                const ComputePtrOffsetOfBatch& compute_ptr_offset_of_batch,
                                const ComputePtrOffsetOfN& compute_ptr_offset_of_n,
-                               const index_t num_k_per_block,
-                               Argument& karg)
+                               [[maybe_unused]] const index_t num_k_per_block,
+                               Argument& karg,
+                               EpilogueArgument& epilogue_args)
     {
         const index_t g_idx = __builtin_amdgcn_readfirstlane(blockIdx.y);
         const index_t n_idx = __builtin_amdgcn_readfirstlane(blockIdx.z / karg.KBatch);
-        const index_t k_idx =
-            __builtin_amdgcn_readfirstlane((blockIdx.z - n_idx * karg.KBatch) * num_k_per_block);
+        // const index_t k_idx =
+        //     __builtin_amdgcn_readfirstlane((blockIdx.z - n_idx * karg.KBatch) * num_k_per_block);
 
         // offset base pointer for each work-group
         const long_index_t a_batch_offset =
@@ -783,6 +785,7 @@ struct GridwiseGemm_wmma_cshuffle_v3
                            decltype(ds_grid_desc_mblock_mperblock_nblock_nperblock),
                            decltype(e_grid_desc_mblock_mperblock_nblock_nperblock),
                            decltype(b_scale_struct),
+                           decltype(epilogue_args),
                            HasMainKBlockLoop,
                            EGlobalMemoryDataOperation,
                            TailNum>(p_as_grid_,
@@ -801,8 +804,7 @@ struct GridwiseGemm_wmma_cshuffle_v3
                                     block_n_id,
                                     num_k_block_per_scale,
                                     b_scale_struct,
-                                    karg.KBatch,
-                                    k_idx);
+                                    epilogue_args);
     }
 };