88 lines
3.4 KiB
Plaintext
88 lines
3.4 KiB
Plaintext
/**
|
|
* P4: Dump TMA descriptor bytes for comparison.
|
|
* CUDA 13.2 compatible.
|
|
*
|
|
* Signature: cuTensorMapEncodeTiled(
|
|
* CUtensorMap*, CUtensorMapDataType, cuuint32_t tensorRank,
|
|
* void*, cuuint64_t* globalDim, cuuint64_t* globalStrides,
|
|
* cuuint32_t* boxDim, cuuint32_t* elementStrides,
|
|
* CUtensorMapInterleave, CUtensorMapSwizzle,
|
|
* CUtensorMapL2promotion, CUtensorMapFloatOOBfill)
|
|
*/
|
|
#include <cuda.h>
|
|
#include <cuda_runtime.h>
|
|
#include <cstdio>
|
|
#include <cstdint>
|
|
|
|
int main() {
|
|
const int ROWS = 128;
|
|
const int COLS = 16;
|
|
const size_t SIZE = ROWS * COLS * 2;
|
|
|
|
void* d_ptr;
|
|
cudaMalloc(&d_ptr, SIZE);
|
|
cudaMemset(d_ptr, 0, SIZE);
|
|
|
|
// globalDim: tensor dimensions (ROWS, COLS) in elements
|
|
cuuint64_t globalDim[] = {(cuuint64_t)ROWS, (cuuint64_t)COLS};
|
|
// globalStrides: byte strides between rows and between elements
|
|
cuuint64_t globalStrides[] = {(cuuint64_t)(COLS * 2), (cuuint64_t)2};
|
|
// boxDim: TMA tile dimensions (16, 16)
|
|
cuuint32_t boxDim[] = {16, 16};
|
|
// elementStrides: (1, 1) = contiguous
|
|
cuuint32_t elementStrides[] = {1, 1};
|
|
|
|
CUtensorMap tma_desc;
|
|
CUresult res;
|
|
|
|
auto dump = [](const char* label, const CUtensorMap& desc) {
|
|
printf("=== %s ===\n", label);
|
|
auto* b = reinterpret_cast<const uint8_t*>(&desc);
|
|
for (int i = 0; i < 128; i += 16) {
|
|
printf("[%3d-%3d]: ", i, i+15);
|
|
for (int j = 0; j < 16; j++) printf("%02x ", b[i+j]);
|
|
printf("\n");
|
|
}
|
|
};
|
|
|
|
// 1: NO swizzle, OOB_NONE
|
|
res = cuTensorMapEncodeTiled(&tma_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) dump("NO swizzle, OOB_NONE", tma_desc);
|
|
else printf("=== NO swizzle, OOB_NONE: FAILED (%d) ===\n", res);
|
|
|
|
// 2: SWIZZLE_128B, OOB_NONE
|
|
res = cuTensorMapEncodeTiled(&tma_desc,
|
|
CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2,
|
|
d_ptr, globalDim, globalStrides, boxDim, elementStrides,
|
|
CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
|
|
CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
|
|
if (res == CUDA_SUCCESS) dump("SWIZZLE_128B, OOB_NONE", tma_desc);
|
|
else printf("=== SWIZZLE_128B, OOB_NONE: FAILED (%d) ===\n", res);
|
|
|
|
// 3: NO swizzle, OOB_FILL_ZERO
|
|
res = cuTensorMapEncodeTiled(&tma_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_NAN_REQUEST_ZERO_FMA);
|
|
if (res == CUDA_SUCCESS) dump("NO swizzle, OOB_FILL_ZERO", tma_desc);
|
|
else printf("=== NO swizzle, OOB_FILL_ZERO: FAILED (%d) ===\n", res);
|
|
|
|
// 4: SWIZZLE_128B, OOB_FILL_ZERO
|
|
res = cuTensorMapEncodeTiled(&tma_desc,
|
|
CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2,
|
|
d_ptr, globalDim, globalStrides, boxDim, elementStrides,
|
|
CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
|
|
CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA);
|
|
if (res == CUDA_SUCCESS) dump("SWIZZLE_128B, OOB_FILL_ZERO", tma_desc);
|
|
else printf("=== SWIZZLE_128B, OOB_FILL_ZERO: FAILED (%d) ===\n", res);
|
|
|
|
cudaFree(d_ptr);
|
|
printf("\nPASSED\n");
|
|
return 0;
|
|
}
|