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

test: 32 threads (1 warp), no guards, all participate

Browse Source
This commit is contained in:
biondizzle
2026-05-28 09:56:17 +00:00
parent f0cb71da5c
commit 494149f034
1 changed files with 9 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
tests/unit/test_tmem_cols.cu
View File

@@ -31,15 +31,15 @@ __global__ void test_tmem_2col(float* out) {
int lane = threadIdx.x % WARP;
int wid = threadIdx.x / WARP;
// Alloc 32 TMEM columns
if (wid == 0) {
// Alloc 32 TMEM columns — all 32 threads (1 warp)
{
tmem_alloc(__cvta_generic_to_shared(sBase), 32);
}
__syncthreads();
uint32_t tb = *sBase;
// Store to columns 0 and 1 with fence between
if (wid == 0) {
// Store to columns 0 and 1
{
float v0 = (float)(lane * 4 + 0);
float v1 = (float)(lane * 4 + 1);
float v2 = (float)(lane * 4 + 2);
@@ -47,12 +47,11 @@ __global__ void test_tmem_2col(float* out) {
uint32_t u0, u1, u2, u3;
memcpy(&u0, &v0, 4); memcpy(&u1, &v1, 4);
memcpy(&u2, &v2, 4); memcpy(&u3, &v3, 4);
tmem_store(tb + 0, u0, u1, u2, u3);
}
tmem_fence();
__syncthreads();
if (wid == 0) {
{
float v0 = (float)(lane * 4 + 100);
float v1 = (float)(lane * 4 + 101);
float v2 = (float)(lane * 4 + 102);
@@ -65,20 +64,19 @@ __global__ void test_tmem_2col(float* out) {
tmem_fence();
__syncthreads();
// Read back columns 0 and 1
if (wid == 0) {
// Read back
{
uint32_t r0, r1, r2, r3;
tmem_load(tb + 0, r0, r1, r2, r3);
float f0; memcpy(&f0, &r0, 4);
if (lane == 0) out[0] = f0;
tmem_load(tb + 1, r0, r1, r2, r3);
float f1; memcpy(&f1, &r0, 4);
if (lane == 0) out[1] = f1;
}
__syncthreads();
if (wid == 0) tmem_dealloc(tb, 32);
tmem_dealloc(tb, 32);
}
int main() {
@@ -87,7 +85,7 @@ int main() {
float* d_out; cudaMalloc(&d_out, 2 * sizeof(float));
cudaMemset(d_out, 0, 2 * sizeof(float));
test_tmem_2col<<<1, 64, 1024>>>(d_out);
test_tmem_2col<<<1, 32, 1024>>>(d_out);
cudaError_t err = cudaDeviceSynchronize();
if (err != cudaSuccess) { printf("CUDA ERROR: %s\n", cudaGetErrorString(err)); return 1; }