From c05cc1ac93a335af2730f17376287f454fc372ab Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Thu, 28 May 2026 13:32:22 +0000
Subject: [PATCH] test: separate TMEM regions for A and C in TS MMA

---
 tests/unit/test_mma_ts.cu | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/tests/unit/test_mma_ts.cu b/tests/unit/test_mma_ts.cu
index daf1ef4c..9c8e120f 100644
--- a/tests/unit/test_mma_ts.cu
+++ b/tests/unit/test_mma_ts.cu
@@ -42,17 +42,20 @@ test_mma_ts(float* o_out)
     }
     __syncthreads();
 
-    // TMEM alloc — 32 columns (16 for A, 16 for C)
-    if (wid == 0) tmem_alloc(__cvta_generic_to_shared(sTmemBase), 32);
+    // TMEM alloc — 64 columns (16 for A at offset 0, 16 for C at offset 32)
+    // TMEM alloc requires power of 2, minimum 32
+    if (wid == 0) tmem_alloc(__cvta_generic_to_shared(sTmemBase), 64);
     __syncthreads();
     uint32_t tb = *sTmemBase;
+    uint32_t tb_a = tb;       // A starts at column 0
+    uint32_t tb_c = tb + 32;  // C starts at column 32
 
     // Write A = all 1.0 into TMEM columns 0-15 (128 rows × 16 columns)
     if (wid == 0) {
         for (int col = 0; col < 16; col++) {
             // Each column: 128 FP32. Lane i writes positions i*4..i*4+3
             float v0 = 1.0f, v1 = 1.0f, v2 = 1.0f, v3 = 1.0f;
-            tmem_store(tb + col, f32_to_u32(v0), f32_to_u32(v1), f32_to_u32(v2), f32_to_u32(v3));
+            tmem_store(tb_a + col, f32_to_u32(v0), f32_to_u32(v1), f32_to_u32(v2), f32_to_u32(v3));
         }
         tmem_fence_store();
     }
@@ -63,7 +66,7 @@ test_mma_ts(float* o_out)
         float check = 0.0f;
         for (int col = 0; col < 16; col++) {
             uint32_t u0, u1, u2, u3;
-            tmem_load(tb + col, u0, u1, u2, u3);
+            tmem_load(tb_a + col, u0, u1, u2, u3);
             tmem_fence_load();
             check += u32_to_f32(u0);
         }
@@ -82,7 +85,7 @@ test_mma_ts(float* o_out)
     printf("Before MMA: tb=%u, dv=%lu, idesc=%u, tid=%d\n", tb, dv, idesc, tid);
 
     if (tid == 0) {
-        umma_ts_f16(tb, tb, dv, idesc, false);
+        umma_ts_f16(tb_c, tb_a, dv, idesc, false);
     }
     asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory");
     __syncthreads();
@@ -94,7 +97,7 @@ test_mma_ts(float* o_out)
         float c_vals[16];
         for (int n = 0; n < 2; n++) {
             float tmp[8];
-            asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];" : "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]),"=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7]) : "r"(tb + n*8));
+            asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8];" : "=f"(tmp[0]),"=f"(tmp[1]),"=f"(tmp[2]),"=f"(tmp[3]),"=f"(tmp[4]),"=f"(tmp[5]),"=f"(tmp[6]),"=f"(tmp[7]) : "r"(tb_c + n*8));
             asm volatile("tcgen05.wait::ld.sync.aligned;");
             if (lane == 0) for (int c=0;c<8;c++) c_vals[n*8+c] = tmp[c];
         }
@@ -109,7 +112,7 @@ test_mma_ts(float* o_out)
         }
     }
 
-    if (wid == 0) tmem_dealloc(tb, 32);
+    if (wid == 0) tmem_dealloc(tb, 64);
 }
 
 int main() {