diff --git a/tests/unit/test_p4_tma_bit21_fix.cu b/tests/unit/test_p4_tma_bit21_fix.cu index 73e1c401..219d1261 100644 --- a/tests/unit/test_p4_tma_bit21_fix.cu +++ b/tests/unit/test_p4_tma_bit21_fix.cu @@ -1,16 +1,6 @@ /** - * P4: Test TMA load with the bit-21 workaround from CUTLASS. - * - * Root cause of the TMA hang: driver 13.0 can't read descriptors - * created by toolkit 13.2's cuTensorMapEncodeTiled. CUTLASS clears - * bit 21 of desc[1] as a workaround for driver <= 13.1 with small tensors. - * - * This test: - * 1. Creates a 2D TMA descriptor with NO swizzle - * 2. Dumps the descriptor bytes - * 3. Clears bit 21 of word[1] (the 64-bit word at offset 8) - * 4. Dumps the modified descriptor - * 5. Tests the TMA load with both descriptors + * P4: Test TMA with bit-21 workaround. + * Uses outermost-first dims, 2 strides, 128B-aligned GMEM. */ #include #include @@ -18,12 +8,9 @@ #include #include -__global__ void tma_load_kernel( - const void* tma_desc_ptr, - int* result -) { +__global__ void tma_load_kernel(const void* tma_desc_ptr, int* result) { __shared__ uint64_t mbar; - __shared__ uint16_t smem_out[16 * 16]; + __shared__ uint16_t smem_out[16 * 16] __attribute__((aligned(128))); if (threadIdx.x == 0) { uint32_t mbar_addr = __cvta_generic_to_shared(&mbar); @@ -38,131 +25,90 @@ __global__ void tma_load_kernel( asm volatile( "cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes " "[%0], [%1, {%3, %4}], [%2];" - :: "r"(smem_addr), - "l"(tma_desc_ptr), - "r"(mbar_addr), - "r"(0), - "r"(0) + :: "r"(smem_addr), "l"(tma_desc_ptr), "r"(mbar_addr), + "r"(0), "r"(0) ); } __syncthreads(); if (threadIdx.x == 0) { uint32_t mbar_addr = __cvta_generic_to_shared(&mbar); - int waited = 0; - while (waited < 1000000) { + for (int w = 0; w < 1000000; w++) { uint32_t state; - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" + asm volatile("{\n\t.reg .pred p;\n\t" "mbarrier.try_wait.parity.shared.b64 p, [%0], 0;\n\t" - "selp.b32 %1, 1, 0, p;\n\t" - "}" - : "=r"(state) - : "r"(mbar_addr) - ); + "selp.b32 %1, 1, 0, p;\n\t}" : "=r"(state) : "r"(mbar_addr)); if (state) { *result = 1; return; } - waited++; } *result = -1; } } - -CUtensorMap create_descriptor(void* d_ptr, bool clear_bit21) { - CUtensorMap desc; - // (128, 16) BF16, row-major, NO swizzle - // CUDA 13: globalDim is innermost-first, globalStrides in bytes, rank-1 strides - cuuint64_t globalDim[] = {16, 128}; // (cols, rows) innermost-first - cuuint64_t globalStrides[] = {16 * 2}; // row stride in bytes (rank-1 strides!) - cuuint32_t boxDim[] = {16, 16}; - cuuint32_t elementStrides[] = {1, 1}; - - CUresult res = cuTensorMapEncodeTiled(&desc, - CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, - d_ptr, globalDim, globalStrides, boxDim, elementStrides, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, - CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - - if (res != CUDA_SUCCESS) { - printf(" cuTensorMapEncodeTiled FAILED: %d\n", res); - return desc; - } - - // Apply bit-21 workaround if requested - if (clear_bit21) { - uint64_t* words = reinterpret_cast(&desc); - words[1] &= ~(1ULL << 21); - } - - return desc; -} - - int main() { - const size_t DATA_SIZE = 128 * 16 * 2; // (128, 16) BF16 + const int ROWS = 128, COLS = 16; + const size_t DATA_SIZE = ROWS * COLS * 2; + // Allocate 128B-aligned GMEM void* d_data; - cudaMalloc(&d_data, DATA_SIZE); - cudaMemset(d_data, 1, DATA_SIZE); + cudaMalloc(&d_data, DATA_SIZE + 128); + // Align to 128 bytes + uintptr_t aligned = ((uintptr_t)d_data + 127) & ~127ULL; + void* d_aligned = (void*)aligned; + cudaMemset(d_aligned, 1, DATA_SIZE); int* d_result; cudaMalloc(&d_result, sizeof(int)); - // Test 1: Original descriptor (no bit-21 fix) - printf("=== Test 1: Original descriptor (no fix) ===\n"); - { - CUtensorMap desc = create_descriptor(d_data, false); - // Dump first 16 bytes - auto* b = reinterpret_cast(&desc); - printf(" Bytes [0-7]: "); for (int j=0;j<8;j++) printf("%02x ", b[j]); printf("\n"); - printf(" Bytes [8-15]: "); for (int j=0;j<8;j++) printf("%02x ", b[8+j]); printf("\n"); + // Create descriptor: outermost-first, byte strides + auto make_desc = [&](void* ptr, bool fix_bit21) -> CUtensorMap { + CUtensorMap desc; + cuuint64_t tensorDims[] = {(cuuint64_t)ROWS, (cuuint64_t)COLS}; + cuuint64_t globalStrides[] = {(cuuint64_t)(COLS * 2), (cuuint64_t)2}; + cuuint32_t boxDims[] = {16, 16}; + cuuint32_t elementStrides[] = {1, 1}; + CUresult res = cuTensorMapEncodeTiled(&desc, + CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, + ptr, tensorDims, globalStrides, boxDims, elementStrides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + if (res != CUDA_SUCCESS) printf(" cuTensorMapEncodeTiled FAILED: %d\n", res); + if (fix_bit21) { + uint64_t* w = reinterpret_cast(&desc); + printf(" Before bit21 fix: word[1] = 0x%016lx\n", w[1]); + w[1] &= ~(1ULL << 21); + printf(" After bit21 fix: word[1] = 0x%016lx\n", w[1]); + } + return desc; + }; + + auto test_desc = [&](const char* name, CUtensorMap& desc) { + printf("=== %s ===\n", name); + auto* b = reinterpret_cast(&desc); + printf(" [0-7]: "); for (int j=0;j<8;j++) printf("%02x ", b[j]); printf("\n"); + printf(" [8-15]: "); for (int j=0;j<8;j++) printf("%02x ", b[8+j]); printf("\n"); void* d_desc; cudaMalloc(&d_desc, sizeof(CUtensorMap)); cudaMemcpy(d_desc, &desc, sizeof(CUtensorMap), cudaMemcpyHostToDevice); - cudaMemset(d_result, 0, sizeof(int)); + tma_load_kernel<<<1, 32>>>(d_desc, d_result); - cudaError_t err = cudaDeviceSynchronize(); - int h_result; - cudaMemcpy(&h_result, d_result, sizeof(int), cudaMemcpyDeviceToHost); + int r; cudaMemcpy(&r, d_result, sizeof(int), cudaMemcpyDeviceToHost); - if (err != cudaSuccess) printf(" ERROR: %s (result=%d)\n", cudaGetErrorString(err), h_result); - else if (h_result == 1) printf(" SUCCESS\n"); - else if (h_result == -1) printf(" HANG (mbarrier timeout)\n"); - else printf(" UNKNOWN: result=%d\n", h_result); + if (err != CUDA_SUCCESS) printf(" ERROR: %s (result=%d)\n", cudaGetErrorString(err), r); + else if (r == 1) printf(" SUCCESS\n"); + else if (r == -1) printf(" HANG (mbarrier timeout)\n"); + else printf(" result=%d\n", r); cudaFree(d_desc); - } + }; - // Test 2: Bit-21 cleared (CUTLASS workaround) - printf("\n=== Test 2: Bit-21 cleared (CUTLASS workaround) ===\n"); - { - CUtensorMap desc = create_descriptor(d_data, true); - auto* b = reinterpret_cast(&desc); - printf(" Bytes [0-7]: "); for (int j=0;j<8;j++) printf("%02x ", b[j]); printf("\n"); - printf(" Bytes [8-15]: "); for (int j=0;j<8;j++) printf("%02x ", b[8+j]); printf("\n"); + CUtensorMap desc1 = make_desc(d_aligned, false); + test_desc("Original (no fix)", desc1); - void* d_desc; - cudaMalloc(&d_desc, sizeof(CUtensorMap)); - cudaMemcpy(d_desc, &desc, sizeof(CUtensorMap), cudaMemcpyHostToDevice); - - cudaMemset(d_result, 0, sizeof(int)); - tma_load_kernel<<<1, 32>>>(d_desc, d_result); - - cudaError_t err = cudaDeviceSynchronize(); - int h_result; - cudaMemcpy(&h_result, d_result, sizeof(int), cudaMemcpyDeviceToHost); - - if (err != cudaSuccess) printf(" ERROR: %s (result=%d)\n", cudaGetErrorString(err), h_result); - else if (h_result == 1) printf(" SUCCESS\n"); - else if (h_result == -1) printf(" HANG (mbarrier timeout)\n"); - else printf(" UNKNOWN: result=%d\n", h_result); - - cudaFree(d_desc); - } + CUtensorMap desc2 = make_desc(d_aligned, true); + test_desc("Bit-21 cleared", desc2); cudaFree(d_data); cudaFree(d_result);