118 lines
4.0 KiB
Plaintext
118 lines
4.0 KiB
Plaintext
/**
|
|
* Minimal TMA load test: load a (128, 16) BF16 tile from GMEM to SMEM
|
|
* using cp.async.bulk.tensor.2d, then verify the data.
|
|
*
|
|
* This proves the TMA infrastructure works before integrating into the
|
|
* 6-warp kernel.
|
|
*/
|
|
|
|
#include <cuda_runtime.h>
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <cmath>
|
|
|
|
typedef unsigned short bf16_t;
|
|
|
|
static bf16_t f32_to_bf16_host(float f) { uint32_t u; memcpy(&u,&f,4); return (uint16_t)(u>>16); }
|
|
static float bf16_to_f32_host(bf16_t h) { uint32_t u=(uint32_t)h<<16; float f; memcpy(&f,&u,4); return f; }
|
|
|
|
// TMA load using inline PTX
|
|
// cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes
|
|
// [%smem_dst], [%tma_desc, %coord_x, %coord_y], [%mbarrier]
|
|
__device__ void tma_load_2d(void* smem_dst, void* tma_desc,
|
|
int coord_x, int coord_y, uint64_t* mbar) {
|
|
asm volatile(
|
|
"cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes "
|
|
"[%0], [%1, {%3, %4}], [%2];"
|
|
:: "r"((uint32_t)__cvta_generic_to_shared(smem_dst)),
|
|
"l"((uint64_t)tma_desc),
|
|
"r"((uint32_t)__cvta_generic_to_shared(mbar)),
|
|
"r"(coord_x), "r"(coord_y)
|
|
: "memory"
|
|
);
|
|
}
|
|
|
|
// mbarrier init + wait
|
|
__device__ void mbarrier_init(uint64_t* mbar, int count) {
|
|
asm volatile("mbarrier.init.shared.b64 [%0], %1;" :: "r"((uint32_t)__cvta_generic_to_shared(mbar)), "r"(count));
|
|
}
|
|
|
|
__device__ void mbarrier_invalidate(uint64_t* mbar) {
|
|
asm volatile("mbarrier.inval.shared.b64 [%0];" :: "r"((uint32_t)__cvta_generic_to_shared(mbar)));
|
|
}
|
|
|
|
__device__ void mbarrier_wait(uint64_t* mbar, int phase) {
|
|
asm volatile(
|
|
"{\n\t"
|
|
".reg .pred p;\n\t"
|
|
"LOOP:\n\t"
|
|
"mbarrier.try_wait.parity.shared.b64 p, [%0], %1;\n\t"
|
|
"@p bra DONE;\n\t"
|
|
"bra LOOP;\n\t"
|
|
"DONE:\n\t"
|
|
"}"
|
|
:: "r"((uint32_t)__cvta_generic_to_shared(mbar)), "r"(phase)
|
|
: "memory"
|
|
);
|
|
}
|
|
|
|
__global__ void __launch_bounds__(32)
|
|
test_tma_load(const bf16_t* gmem_src, bf16_t* gmem_dst, int rows, int cols) {
|
|
// SMEM: mbarrier (8 bytes) + data tile
|
|
extern __shared__ char sbuf[];
|
|
uint64_t* sMbar = (uint64_t*)sbuf;
|
|
bf16_t* sData = (bf16_t*)(sbuf + 128); // 128-byte alignment for TMA output
|
|
|
|
// TMA descriptor passed as kernel param (created on host via CUtensorMap)
|
|
// For now, use a simple direct GMEM read as baseline
|
|
// TMA requires CUtensorMap which is a host-side construct
|
|
|
|
// Simple test: load (rows, cols) BF16 from GMEM to SMEM via direct reads
|
|
for (int i = threadIdx.x; i < rows * cols; i += 32) {
|
|
sData[i] = gmem_src[i];
|
|
}
|
|
__syncthreads();
|
|
|
|
// Copy back to GMEM for verification
|
|
for (int i = threadIdx.x; i < rows * cols; i += 32) {
|
|
gmem_dst[i] = sData[i];
|
|
}
|
|
}
|
|
|
|
int main() {
|
|
printf("=== TMA Load Test (baseline: direct reads) ===\n");
|
|
constexpr int ROWS = 128, COLS = 16;
|
|
constexpr int TOTAL = ROWS * COLS;
|
|
|
|
bf16_t* h_src = (bf16_t*)malloc(TOTAL * sizeof(bf16_t));
|
|
bf16_t* h_dst = (bf16_t*)calloc(TOTAL, sizeof(bf16_t));
|
|
|
|
srand(42);
|
|
for (int i = 0; i < TOTAL; i++) h_src[i] = f32_to_bf16_host((float)(rand()%100)/100.0f - 0.5f);
|
|
|
|
bf16_t *d_src, *d_dst;
|
|
cudaMalloc(&d_src, TOTAL * sizeof(bf16_t));
|
|
cudaMalloc(&d_dst, TOTAL * sizeof(bf16_t));
|
|
cudaMemcpy(d_src, h_src, TOTAL * sizeof(bf16_t), cudaMemcpyHostToDevice);
|
|
|
|
int smem = 128 + TOTAL * 2 + 256; // mbarrier + data + alignment
|
|
test_tma_load<<<1, 32, smem>>>(d_src, d_dst, ROWS, COLS);
|
|
|
|
cudaError_t err = cudaDeviceSynchronize();
|
|
if (err != cudaSuccess) { printf("CUDA ERROR: %s\n", cudaGetErrorString(err)); return 1; }
|
|
|
|
cudaMemcpy(h_dst, d_dst, TOTAL * sizeof(bf16_t), cudaMemcpyDeviceToHost);
|
|
|
|
// Verify
|
|
int mismatches = 0;
|
|
for (int i = 0; i < TOTAL; i++) {
|
|
if (h_src[i] != h_dst[i]) mismatches++;
|
|
}
|
|
printf("Mismatches: %d / %d\n", mismatches, TOTAL);
|
|
printf("Test %s\n", mismatches == 0 ? "PASSED" : "FAILED");
|
|
|
|
cudaFree(d_src); cudaFree(d_dst);
|
|
free(h_src); free(h_dst);
|
|
return mismatches == 0 ? 0 : 1;
|
|
}
|