Files
d4r/tools/windows/cuda_mma_probe.cpp
T

166 lines
7.3 KiB
C++

#include "hip_api.h"
#include "cuda_api.h"
#include <array>
#include <cmath>
#include <cstring>
// Public synthetic matrix fragments, independent of NVIDIA DLLs and weights.
static constexpr char source[] = R"PTX(
.version 7.0
.target sm_75
.address_size 64
.visible .entry d4r_f16_mma_reference(
.param .u64 A, .param .u64 B, .param .u64 C, .param .u64 D)
{
.reg .b32 %r<24>;
.reg .b64 %rd<12>;
.reg .b16 %h<4>;
ld.param.u64 %rd0, [A];
ld.param.u64 %rd1, [B];
ld.param.u64 %rd2, [C];
ld.param.u64 %rd3, [D];
mov.u32 %r0, %ctaid.x;
mov.u32 %r1, %tid.x;
shr.u32 %r2, %r1, 2;
and.b32 %r3, %r1, 3;
shl.b32 %r3, %r3, 1;
mul.wide.u32 %rd4, %r0, 512;
add.u64 %rd0, %rd0, %rd4;
mul.wide.u32 %rd4, %r0, 256;
add.u64 %rd1, %rd1, %rd4;
add.u64 %rd2, %rd2, %rd4;
add.u64 %rd3, %rd3, %rd4;
mad.lo.u32 %r4, %r2, 16, %r3;
mul.wide.u32 %rd4, %r4, 2;
add.u64 %rd5, %rd0, %rd4;
ld.global.u32 %r8, [%rd5];
ld.global.u32 %r9, [%rd5+256];
ld.global.u32 %r10, [%rd5+16];
ld.global.u32 %r11, [%rd5+272];
mad.lo.u32 %r4, %r3, 8, %r2;
mul.wide.u32 %rd4, %r4, 2;
add.u64 %rd5, %rd1, %rd4;
ld.global.b16 %h0, [%rd5];
ld.global.b16 %h1, [%rd5+16];
mov.b32 %r12, {%h0,%h1};
ld.global.b16 %h2, [%rd5+128];
ld.global.b16 %h3, [%rd5+144];
mov.b32 %r13, {%h2,%h3};
mad.lo.u32 %r4, %r2, 8, %r3;
mul.wide.u32 %rd4, %r4, 2;
add.u64 %rd5, %rd2, %rd4;
ld.global.u32 %r14, [%rd5];
ld.global.u32 %r15, [%rd5+128];
mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16
{%r16,%r17}, {%r8,%r9,%r10,%r11}, {%r12,%r13}, {%r14,%r15};
add.u64 %rd5, %rd3, %rd4;
st.global.u32 [%rd5], %r16;
st.global.u32 [%rd5+128], %r17;
ret;
}
)PTX";
static uint16_t half_bits(double value) {
_Float16 half = static_cast<_Float16>(value);
uint16_t bits;
std::memcpy(&bits, &half, sizeof(bits));
return bits;
}
static double half_value(uint16_t bits) {
_Float16 half;
std::memcpy(&half, &bits, sizeof(bits));
return static_cast<double>(half);
}
int main(int argc, char** argv) {
using namespace d4r::diag;
start();
try {
Args args(argc, argv);
HipApi hip(args.hip_root);
hipDeviceProp_t props{};
hip.select_gfx1201(args.device, props);
CudaApi cuda(args.cuda_dll);
cuda.check(cuda.cuInit(0), "cuInit");
CUdevice device = -1;
int count = 0;
cuda.check(cuda.cuDeviceGetCount(&count), "cuDeviceGetCount");
for (int i = 0; i < count; ++i) {
CUdevice candidate;
int bus, pci, domain;
cuda.check(cuda.cuDeviceGet(&candidate, i), "cuDeviceGet");
cuda.check(cuda.cuDeviceGetAttribute(&bus, 33, candidate), "PCI bus");
cuda.check(cuda.cuDeviceGetAttribute(&pci, 34, candidate), "PCI device");
cuda.check(cuda.cuDeviceGetAttribute(&domain, 50, candidate), "PCI domain");
if (bus == props.pciBusID && pci == props.pciDeviceID && domain == props.pciDomainID) device = candidate;
}
if (device < 0) throw std::runtime_error("CUDA/HIP PCI identity mismatch");
CUcontext context = nullptr;
cuda.check(cuda.cuDevicePrimaryCtxRetain(&context, device), "primary context");
struct Context { CudaApi& api; CUdevice device; ~Context() { api.cuCtxSetCurrent(nullptr); api.cuDevicePrimaryCtxRelease_v2(device); } } context_owner{cuda, device};
cuda.check(cuda.cuCtxSetCurrent(context), "cuCtxSetCurrent");
CUmodule module = nullptr;
cuda.check(cuda.cuModuleLoadData(&module, source), "MMA PTX JIT");
struct Module { CudaApi& api; CUmodule module; ~Module() { api.cuModuleUnload(module); } } module_owner{cuda, module};
CUfunction kernel = nullptr;
cuda.check(cuda.cuModuleGetFunction(&kernel, module, "d4r_f16_mma_reference"), "MMA kernel");
constexpr size_t cases = 512;
std::vector<uint16_t> a(cases * 256), b(cases * 128), c(cases * 128), d(cases * 128);
uint32_t random = 0x192a321u;
auto next = [&]() { random = random * 1664525u + 1013904223u; return random; };
for (size_t n = 0; n < cases; ++n) {
for (size_t i = 0; i < 256; ++i) a[n * 256 + i] = uint16_t(0x2800u + (next() >> 20));
for (size_t i = 0; i < 128; ++i) {
b[n * 128 + i] = uint16_t((next() & 0x8000u) | (0x3000u + (next() >> 20)));
c[n * 128 + i] = half_bits(n % 3 ? 0. : 0.015625);
}
// Sweep large cancellation terms and a tiny surviving product
// across K positions; this exposes loss hidden by integer fixtures.
if (n < 256) {
for (size_t i = 0; i < 256; ++i) a[n * 256 + i] = half_bits(1.);
std::fill(b.begin() + n * 128, b.begin() + (n + 1) * 128, half_bits(0.));
for (size_t col = 0; col < 8; ++col) {
b[n * 128 + (n % 16) * 8 + col] = half_bits(32768.);
b[n * 128 + ((n % 16 + 1) % 16) * 8 + col] = half_bits(-32768.);
b[n * 128 + ((n % 16 + 2 + n / 16 % 14) % 16) * 8 + col] = half_bits(0.0009765625);
}
}
}
struct Allocation {
CudaApi& api; CUdeviceptr pointer = 0;
Allocation(CudaApi& api, size_t bytes) : api(api) { api.check(api.cuMemAlloc_v2(&pointer, bytes), "MMA allocate"); }
~Allocation() { if (pointer) api.cuMemFree_v2(pointer); }
} da(cuda, a.size() * 2), db(cuda, b.size() * 2), dc(cuda, c.size() * 2), dd(cuda, d.size() * 2);
cuda.check(cuda.cuMemcpyHtoD_v2(da.pointer, a.data(), a.size() * 2), "MMA A upload");
cuda.check(cuda.cuMemcpyHtoD_v2(db.pointer, b.data(), b.size() * 2), "MMA B upload");
cuda.check(cuda.cuMemcpyHtoD_v2(dc.pointer, c.data(), c.size() * 2), "MMA C upload");
std::fill(d.begin(), d.end(), uint16_t(0x7e00));
cuda.check(cuda.cuMemcpyHtoD_v2(dd.pointer, d.data(), d.size() * 2), "MMA NaN sentinel");
void* parameters[] = {&da.pointer, &db.pointer, &dc.pointer, &dd.pointer};
cuda.check(cuda.cuLaunchKernel(kernel, cases, 1, 1, 32, 1, 1, 0, nullptr, parameters, nullptr), "MMA launch");
cuda.check(cuda.cuCtxSynchronize(), "MMA synchronize");
cuda.check(cuda.cuMemcpyDtoH_v2(d.data(), dd.pointer, d.size() * 2), "MMA download");
size_t mismatch = 0;
for (size_t n = 0; n < cases; ++n)
for (size_t row = 0; row < 16; ++row)
for (size_t col = 0; col < 8; ++col) {
double sum = half_value(c[n * 128 + row * 8 + col]);
for (size_t k = 0; k < 16; ++k)
sum += half_value(a[n * 256 + row * 16 + k]) * half_value(b[n * 128 + k * 8 + col]);
const uint16_t expected = half_bits(sum);
const uint16_t actual = d[n * 128 + row * 8 + col];
if (expected != actual) {
if (mismatch++ < 8) std::fprintf(stderr, "MMA_DIFFERENT case=%zu row=%zu column=%zu expected=%04x actual=%04x\n", n, row, col, expected, actual);
}
}
std::printf("MMA_REFERENCE cases=%zu elements=%zu differing=%zu\n", cases, d.size(), mismatch);
if (mismatch) return 5;
std::printf("PASS CUDA_F16_MMA architecture=gfx1201 cases=%zu exact_f16=1\n", cases);
return 0;
} catch (const std::exception& error) {
std::fprintf(stderr, "FAIL CUDA_F16_MMA %s\n", error.what());
loaded_modules();
return 4;
}
}