From b7c697172032c14dbccad15eb0c7d95c30043dc6 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Thu, 28 May 2026 13:34:50 +0000
Subject: [PATCH] test: use 32x32b.x8 for A write (avoids 16x256b misalign)

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

diff --git a/tests/unit/test_mma_ts.cu b/tests/unit/test_mma_ts.cu
index 13288e6c..3f3bb926 100644
--- a/tests/unit/test_mma_ts.cu
+++ b/tests/unit/test_mma_ts.cu
@@ -50,34 +50,29 @@ test_mma_ts(float* o_out)
     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)
+    // Write A = all 1.0 into TMEM columns 0-15 using 32x32b.x8
+    // Each 32x32b.x8 writes 32 rows × 8 columns. All lanes write 1.0.
     if (wid == 0) {
-        // First write column 0 and verify
-        float v0 = 1.0f, v1 = 1.0f, v2 = 1.0f, v3 = 1.0f;
-        tmem_store(tb_a, f32_to_u32(v0), f32_to_u32(v1), f32_to_u32(v2), f32_to_u32(v3));
-        tmem_fence_store();
-    }
-    __syncthreads();
-
-    // Read back column 0
-    if (wid == 0) {
-        uint32_t u0, u1, u2, u3;
-        tmem_load(tb_a, u0, u1, u2, u3);
-        tmem_fence_load();
-        if (lane == 0) printf("A[0,0] = %.1f (expect 1.0)\n", u32_to_f32(u0));
-    }
-    __syncthreads();
-
-    // Write remaining columns
-    if (wid == 0) {
-        for (int col = 1; col < 16; col++) {
-            float v0 = 1.0f, v1 = 1.0f, v2 = 1.0f, v3 = 1.0f;
-            tmem_store(tb_a + col, f32_to_u32(v0), f32_to_u32(v1), f32_to_u32(v2), f32_to_u32(v3));
+        for (int n = 0; n < 16 / 8; n++) {
+            float p0=1.0f, p1=1.0f, p2=1.0f, p3=1.0f;
+            float p4=1.0f, p5=1.0f, p6=1.0f, p7=1.0f;
+            asm volatile("tcgen05.st.sync.aligned.32x32b.x8.b32 [%0],{%1,%2,%3,%4,%5,%6,%7,%8};" :: "r"(tb_a+n*8),"f"(p0),"f"(p1),"f"(p2),"f"(p3),"f"(p4),"f"(p5),"f"(p6),"f"(p7));
         }
         tmem_fence_store();
     }
     __syncthreads();
 
+    // Read back A
+    if (wid == 0) {
+        for (int n = 0; n < 16 / 8; 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_a+n*8));
+            asm volatile("tcgen05.wait::ld.sync.aligned;");
+            if (lane == 0) printf("A[0,%d..%d] = %.1f %.1f %.1f %.1f\n", n*8, n*8+7, tmp[0], tmp[1], tmp[2], tmp[3]);
+        }
+    }
+    __syncthreads();
+
     // tcgen05.mma TS: A (TMEM) × B (SMEM) → C (TMEM)
     // A is at tb (columns 0-15)
     // B is at sV (16, 16)