From 354dd5039caa6cb5e82d3d07f425ce8ed7ebdaa9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ville=20Pietil=C3=A4?= <ville.pietila@amd.com>
Date: Fri, 26 Sep 2025 08:39:39 +0000
Subject: [PATCH] Add compile check for assumed row-mjor layout.

---
 include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp
index 1d2208ac39..4f77176be6 100644
--- a/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp
+++ b/include/ck_tile/ops/epilogue/cshuffle_epilogue.hpp
@@ -219,6 +219,13 @@ struct CShuffleEpilogue
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto MakeLdsBlockDescriptor()
     {
+        if constexpr(NumGroupsToMerge > 1)
+        {
+            // We haven't yet tested the ColumnMajor case.
+            static_assert(std::is_same_v<ELayout, tensor_layout::gemm::RowMajor>,
+                        "Currently, the CShuffle Epilogue with NumGroupsToMerge > 1 only supports the Row Major layout");
+        }
+
         // N is contiguous dimension
         if constexpr(std::is_same_v<ELayout, tensor_layout::gemm::RowMajor>)
         {
@@ -366,7 +373,7 @@ struct CShuffleEpilogue
             sequence<0, 3, 0, 3>>{};
         constexpr auto dram_tile_distribution = make_static_tile_distribution(dram_tile_encoding);
 
-        // The LDS data has the following 4D layout
+        // The LDS data has the following 4D layout in the row-major case.
         // linear_index = c + Gs * m + Gs * MPerGroup * n + Gs * MPerGroup * NPerGroup * r
         // for 4D coordinates (r,c,m,n) where (r,c) is the group index and (m,n) is the index within the group.
         // We pick-up only the diagonal blocks where r == c.