From 0959c6582aa6dfa1d29e12c439ab2dfa0c32e8b8 Mon Sep 17 00:00:00 2001
From: Thomas Ning <Thomas.Ning@amd.com>
Date: Thu, 2 Oct 2025 11:15:24 -0700
Subject: [PATCH] add the check of granularity for atomic add (#2959)

[ROCm/composable_kernel commit: cadafde722f838d1fc0b08130cd4fca168acc29c]
---
 .../impl/device_gemm_multiple_d_xdl_cshuffle_v3.hpp       | 4 ++++
 test/ck_tile/gemm/test_gemm_pipeline_ut_cases.inc         | 8 ++++----
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3.hpp b/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3.hpp
index 7e9020d796..02639dbf3e 100644
--- a/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3.hpp
+++ b/include/ck/tensor_operation/gpu/device/impl/device_gemm_multiple_d_xdl_cshuffle_v3.hpp
@@ -682,6 +682,10 @@ struct DeviceGemmMultiD_Xdl_CShuffle_V3 : public DeviceGemmMultipleDSplitK<ALayo
                 return GridwiseGemm64::CheckValidity(arg);
             }
         }
+        if(CDEShuffleBlockTransferScalarPerVectors{}[Number<0>{}] <= 1 && (arg.KBatch > 1))
+        {
+            return false;
+        }
         else
         {
             if constexpr(NXdlPerWave32 > 0)
diff --git a/test/ck_tile/gemm/test_gemm_pipeline_ut_cases.inc b/test/ck_tile/gemm/test_gemm_pipeline_ut_cases.inc
index f793f81cc9..66ef05b0ba 100644
--- a/test/ck_tile/gemm/test_gemm_pipeline_ut_cases.inc
+++ b/test/ck_tile/gemm/test_gemm_pipeline_ut_cases.inc
@@ -11,7 +11,7 @@ TYPED_TEST(TEST_SUITE_NAME, SmallM)
     std::vector<int> Ms{1, 2, 3, 4, 5, 6};
     constexpr int N = 1024;
     std::vector<int> Ks;
-    for (auto K_count: {2, 3, 4, 10, 11}) 
+    for(auto K_count : {2, 3, 4, 10, 11})
     {
         Ks.push_back(K_count * TestFixture::K_Tile);
     }
@@ -36,10 +36,10 @@ TYPED_TEST(TEST_SUITE_NAME, SmallM)
 TYPED_TEST(TEST_SUITE_NAME, MidLargeM)
 {
     std::vector<int> Ms{127, 255, 312, 799, 1573};
-    constexpr int N           = 1024;
+    constexpr int N = 1024;
 
     std::vector<int> Ks;
-    for (auto K_count: {2, 3, 4, 10, 11}) 
+    for(auto K_count : {2, 3, 4, 10, 11})
     {
         Ks.push_back(K_count * TestFixture::K_Tile);
     }
@@ -51,7 +51,7 @@ TYPED_TEST(TEST_SUITE_NAME, MidLargeM)
 
     for(int M : Ms)
     {
-        for (int K: Ks)
+        for(int K : Ks)
         {
             if constexpr(std::is_same_v<typename TestFixture::ALayout,
                                         ck_tile::tensor_layout::gemm::ColumnMajor>)