biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

test: use 32x32b.x8 for A write (avoids 16x256b misalign)

Browse Source
This commit is contained in:
biondizzle
2026-05-28 13:34:50 +00:00
parent a7c81d66ba
commit b7c6971720
1 changed files with 17 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
tests/unit/test_mma_ts.cu
View File

@@ -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)