1875 lines
75 KiB
C++
1875 lines
75 KiB
C++
// Checks ZLUDA's emulation of the FP8/matrix PTX instructions used by the
|
|
// RRLite DLSS network (rrlite_enc*/dec*) against CPU references built from
|
|
// the PTX ISA fragment layouts:
|
|
// cvt.rn.f16x2.e4m3x2 (exhaustive over all 65536 inputs)
|
|
// cvt.rn.satfinite.e4m3x2.f16x2 (exhaustive over all f16 lanes)
|
|
// cvt.rn.satfinite.relu.e4m3x2.{f16x2,f32} (exhaustive over all f16 values)
|
|
// mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16
|
|
// mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16
|
|
// movmatrix.sync.aligned.m8n8.trans.b16
|
|
// Matrix inputs are small values whose products and sums are exact in f16,
|
|
// so results must match bit for bit regardless of accumulation order.
|
|
#include <dlfcn.h>
|
|
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <random>
|
|
#include <vector>
|
|
|
|
using CUresult = int;
|
|
using CUdevice = int;
|
|
using CUcontext = void*;
|
|
using CUmodule = void*;
|
|
using CUfunction = void*;
|
|
using CUstream = void*;
|
|
using CUdeviceptr = uint64_t;
|
|
|
|
template <typename Function> static Function load_function(void* library, const char* name)
|
|
{
|
|
void* raw = dlsym(library, name);
|
|
Function function{};
|
|
static_assert(sizeof(function) == sizeof(raw));
|
|
std::memcpy(&function, &raw, sizeof(function));
|
|
return function;
|
|
}
|
|
|
|
static const char* ptx = R"PTX(
|
|
.version 8.7
|
|
.target sm_89
|
|
.address_size 64
|
|
|
|
.visible .entry d4r_cvt_f16x2_from_e4m3x2(.param .u64 input, .param .u64 output, .param .u32 count)
|
|
{
|
|
.reg .pred %p<2>;
|
|
.reg .b16 %rs<2>;
|
|
.reg .b32 %r<8>;
|
|
.reg .b64 %rd<8>;
|
|
mov.u32 %r0, %ctaid.x;
|
|
mov.u32 %r1, %ntid.x;
|
|
mov.u32 %r2, %tid.x;
|
|
mad.lo.s32 %r3, %r0, %r1, %r2;
|
|
ld.param.u32 %r4, [count];
|
|
setp.ge.u32 %p0, %r3, %r4;
|
|
@%p0 bra $done;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mul.wide.u32 %rd2, %r3, 2;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
ld.global.u16 %rs0, [%rd3];
|
|
cvt.rn.f16x2.e4m3x2 %r5, %rs0;
|
|
mul.wide.u32 %rd4, %r3, 4;
|
|
add.s64 %rd5, %rd1, %rd4;
|
|
st.global.u32 [%rd5], %r5;
|
|
$done:
|
|
ret;
|
|
}
|
|
|
|
.visible .entry d4r_cvt_e4m3x2_from_f16x2(.param .u64 input, .param .u64 output, .param .u32 count)
|
|
{
|
|
.reg .pred %p<2>;
|
|
.reg .b16 %rs<2>;
|
|
.reg .b32 %r<8>;
|
|
.reg .b64 %rd<8>;
|
|
mov.u32 %r0, %ctaid.x;
|
|
mov.u32 %r1, %ntid.x;
|
|
mov.u32 %r2, %tid.x;
|
|
mad.lo.s32 %r3, %r0, %r1, %r2;
|
|
ld.param.u32 %r4, [count];
|
|
setp.ge.u32 %p0, %r3, %r4;
|
|
@%p0 bra $done;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mul.wide.u32 %rd2, %r3, 4;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
ld.global.u32 %r5, [%rd3];
|
|
cvt.rn.satfinite.e4m3x2.f16x2 %rs0, %r5;
|
|
mul.wide.u32 %rd4, %r3, 2;
|
|
add.s64 %rd5, %rd1, %rd4;
|
|
st.global.u16 [%rd5], %rs0;
|
|
$done:
|
|
ret;
|
|
}
|
|
|
|
.visible .entry d4r_cvt_e4m3x2_relu_from_f16x2(.param .u64 input, .param .u64 output, .param .u32 count)
|
|
{
|
|
.reg .pred %p<2>;
|
|
.reg .b16 %rs<2>;
|
|
.reg .b32 %r<8>;
|
|
.reg .b64 %rd<8>;
|
|
mov.u32 %r0, %ctaid.x;
|
|
mov.u32 %r1, %ntid.x;
|
|
mov.u32 %r2, %tid.x;
|
|
mad.lo.s32 %r3, %r0, %r1, %r2;
|
|
ld.param.u32 %r4, [count];
|
|
setp.ge.u32 %p0, %r3, %r4;
|
|
@%p0 bra $done;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mul.wide.u32 %rd2, %r3, 4;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
ld.global.u32 %r5, [%rd3];
|
|
cvt.rn.satfinite.relu.e4m3x2.f16x2 %rs0, %r5;
|
|
mul.wide.u32 %rd4, %r3, 2;
|
|
add.s64 %rd5, %rd1, %rd4;
|
|
st.global.u16 [%rd5], %rs0;
|
|
$done:
|
|
ret;
|
|
}
|
|
|
|
// The same f16x2 inputs widened to f32; the high half is the first operand.
|
|
.visible .entry d4r_cvt_e4m3x2_relu_from_f32(.param .u64 input, .param .u64 output, .param .u32 count)
|
|
{
|
|
.reg .pred %p<2>;
|
|
.reg .b16 %rs<4>;
|
|
.reg .b32 %r<8>;
|
|
.reg .f32 %f<2>;
|
|
.reg .b64 %rd<8>;
|
|
mov.u32 %r0, %ctaid.x;
|
|
mov.u32 %r1, %ntid.x;
|
|
mov.u32 %r2, %tid.x;
|
|
mad.lo.s32 %r3, %r0, %r1, %r2;
|
|
ld.param.u32 %r4, [count];
|
|
setp.ge.u32 %p0, %r3, %r4;
|
|
@%p0 bra $done;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mul.wide.u32 %rd2, %r3, 4;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
ld.global.u32 %r5, [%rd3];
|
|
mov.b32 {%rs1, %rs2}, %r5;
|
|
cvt.f32.f16 %f0, %rs1;
|
|
cvt.f32.f16 %f1, %rs2;
|
|
cvt.rn.satfinite.relu.e4m3x2.f32 %rs0, %f1, %f0;
|
|
mul.wide.u32 %rd4, %r3, 2;
|
|
add.s64 %rd5, %rd1, %rd4;
|
|
st.global.u16 [%rd5], %rs0;
|
|
$done:
|
|
ret;
|
|
}
|
|
|
|
// Per lane: 4 A words, 2 B words, 2 C words in; 2 D words out.
|
|
.visible .entry d4r_mma_m16n8k16_f16(.param .u64 input, .param .u64 output)
|
|
{
|
|
.reg .b32 %r<16>;
|
|
.reg .b64 %rd<8>;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mov.u32 %r0, %laneid;
|
|
mul.wide.u32 %rd2, %r0, 32;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
ld.global.v4.u32 {%r1, %r2, %r3, %r4}, [%rd3];
|
|
ld.global.v4.u32 {%r5, %r6, %r7, %r8}, [%rd3+16];
|
|
mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%r9, %r10}, {%r1, %r2, %r3, %r4}, {%r5, %r6}, {%r7, %r8};
|
|
mul.wide.u32 %rd4, %r0, 8;
|
|
add.s64 %rd5, %rd1, %rd4;
|
|
st.global.v2.u32 [%rd5], {%r9, %r10};
|
|
ret;
|
|
}
|
|
|
|
// Chained MMAs in a loop, two per step sharing A (the pattern of DLSS's
|
|
// convolutions): per lane 8 steps x (4 A, 2 B0, 2 B1 words), then C0 and C1
|
|
// (2 words each); out D0 and D1. Exercises accumulators carried through a
|
|
// loop and pairs of MMAs that share A.
|
|
.visible .entry d4r_mma_chain_f16(.param .u64 input, .param .u64 output)
|
|
{
|
|
.reg .b32 %r<16>;
|
|
.reg .b64 %rd<8>;
|
|
.reg .pred %p<2>;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mov.u32 %r0, %laneid;
|
|
mul.wide.u32 %rd2, %r0, 272;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
ld.global.v4.u32 {%r9, %r10, %r11, %r12}, [%rd3+256];
|
|
mov.u32 %r13, 0;
|
|
$L_chain:
|
|
ld.global.v4.u32 {%r1, %r2, %r3, %r4}, [%rd3];
|
|
ld.global.v4.u32 {%r5, %r6, %r7, %r8}, [%rd3+16];
|
|
mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%r9, %r10}, {%r1, %r2, %r3, %r4}, {%r5, %r6}, {%r9, %r10};
|
|
mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%r11, %r12}, {%r1, %r2, %r3, %r4}, {%r7, %r8}, {%r11, %r12};
|
|
add.s64 %rd3, %rd3, 32;
|
|
add.u32 %r13, %r13, 1;
|
|
setp.lt.u32 %p0, %r13, 8;
|
|
@%p0 bra $L_chain;
|
|
mul.wide.u32 %rd4, %r0, 16;
|
|
add.s64 %rd5, %rd1, %rd4;
|
|
st.global.v4.u32 [%rd5], {%r9, %r10, %r11, %r12};
|
|
ret;
|
|
}
|
|
|
|
// The same chain with m16n8k32 e4m3 MMAs (A 4 words, B 2 words: same sizes).
|
|
.visible .entry d4r_mma_chain_e4m3(.param .u64 input, .param .u64 output)
|
|
{
|
|
.reg .b32 %r<16>;
|
|
.reg .b64 %rd<8>;
|
|
.reg .pred %p<2>;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mov.u32 %r0, %laneid;
|
|
mul.wide.u32 %rd2, %r0, 272;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
ld.global.v4.u32 {%r9, %r10, %r11, %r12}, [%rd3+256];
|
|
mov.u32 %r13, 0;
|
|
$L_chain8:
|
|
ld.global.v4.u32 {%r1, %r2, %r3, %r4}, [%rd3];
|
|
ld.global.v4.u32 {%r5, %r6, %r7, %r8}, [%rd3+16];
|
|
mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 {%r9, %r10}, {%r1, %r2, %r3, %r4}, {%r5, %r6}, {%r9, %r10};
|
|
mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 {%r11, %r12}, {%r1, %r2, %r3, %r4}, {%r7, %r8}, {%r11, %r12};
|
|
add.s64 %rd3, %rd3, 32;
|
|
add.u32 %r13, %r13, 1;
|
|
setp.lt.u32 %p0, %r13, 8;
|
|
@%p0 bra $L_chain8;
|
|
mul.wide.u32 %rd4, %r0, 16;
|
|
add.s64 %rd5, %rd1, %rd4;
|
|
st.global.v4.u32 [%rd5], {%r9, %r10, %r11, %r12};
|
|
ret;
|
|
}
|
|
|
|
// Two e4m3 MMAs sharing A whose second accumulator is one zeroed register
|
|
// used for both halves ({%r20, %r20}), as DLSS's first layers do. In the
|
|
// "late" kernel that register is rewritten between the two MMAs, which must
|
|
// keep them apart. Per lane: 4 A, 2 B0, 2 B1, 2 C0 words and a value for
|
|
// %r20; out D0, D1.
|
|
.visible .entry d4r_mma_pair_alias_e4m3(.param .u64 input, .param .u64 output, .param .u32 late)
|
|
{
|
|
.reg .b32 %r<24>;
|
|
.reg .b64 %rd<8>;
|
|
.reg .pred %p<2>;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
ld.param.u32 %r22, [late];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mov.u32 %r0, %laneid;
|
|
mul.wide.u32 %rd2, %r0, 48;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
ld.global.v4.u32 {%r1, %r2, %r3, %r4}, [%rd3];
|
|
ld.global.v4.u32 {%r5, %r6, %r7, %r8}, [%rd3+16];
|
|
ld.global.v4.u32 {%r9, %r10, %r11, %r12}, [%rd3+32];
|
|
setp.ne.u32 %p0, %r22, 0;
|
|
mov.u32 %r20, 0;
|
|
mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 {%r13, %r14}, {%r1, %r2, %r3, %r4}, {%r5, %r6}, {%r9, %r10};
|
|
mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 {%r15, %r16}, {%r1, %r2, %r3, %r4}, {%r7, %r8}, {%r20, %r20};
|
|
mul.wide.u32 %rd4, %r0, 16;
|
|
add.s64 %rd5, %rd1, %rd4;
|
|
st.global.v4.u32 [%rd5], {%r13, %r14, %r15, %r16};
|
|
ret;
|
|
}
|
|
|
|
.visible .entry d4r_mma_pair_alias_late_e4m3(.param .u64 input, .param .u64 output, .param .u32 late)
|
|
{
|
|
.reg .b32 %r<24>;
|
|
.reg .b64 %rd<8>;
|
|
.reg .pred %p<2>;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
ld.param.u32 %r22, [late];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mov.u32 %r0, %laneid;
|
|
mul.wide.u32 %rd2, %r0, 48;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
ld.global.v4.u32 {%r1, %r2, %r3, %r4}, [%rd3];
|
|
ld.global.v4.u32 {%r5, %r6, %r7, %r8}, [%rd3+16];
|
|
ld.global.v4.u32 {%r9, %r10, %r11, %r12}, [%rd3+32];
|
|
setp.ne.u32 %p0, %r22, 0;
|
|
mov.u32 %r20, 0;
|
|
mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 {%r13, %r14}, {%r1, %r2, %r3, %r4}, {%r5, %r6}, {%r9, %r10};
|
|
selp.b32 %r20, %r11, 0, %p0;
|
|
mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 {%r15, %r16}, {%r1, %r2, %r3, %r4}, {%r7, %r8}, {%r20, %r20};
|
|
mul.wide.u32 %rd4, %r0, 16;
|
|
add.s64 %rd5, %rd1, %rd4;
|
|
st.global.v4.u32 [%rd5], {%r13, %r14, %r15, %r16};
|
|
ret;
|
|
}
|
|
|
|
.visible .entry d4r_mma_m16n8k8_f16(.param .u64 input, .param .u64 output)
|
|
{
|
|
.reg .b32 %r<16>;
|
|
.reg .b64 %rd<8>;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mov.u32 %r0, %laneid;
|
|
mul.wide.u32 %rd2, %r0, 32;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
ld.global.v4.u32 {%r1, %r2, %r3, %r4}, [%rd3];
|
|
ld.global.v4.u32 {%r5, %r6, %r7, %r8}, [%rd3+16];
|
|
mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16
|
|
{%r9, %r10},
|
|
{%r1, %r2},
|
|
{%r5},
|
|
{%r7, %r8};
|
|
mul.wide.u32 %rd4, %r0, 8;
|
|
add.s64 %rd5, %rd1, %rd4;
|
|
st.global.v2.u32 [%rd5], {%r9, %r10};
|
|
ret;
|
|
}
|
|
|
|
// fma.rn.sat.f16x2 and fma.rn.relu.f16x2 as used by the DLTSS model kernels.
|
|
// Per thread: 3 input words, 3 output words (DLTSS writes .relu after the type).
|
|
.visible .entry d4r_f16x2_fma_modes(.param .u64 input, .param .u64 output, .param .u32 count)
|
|
{
|
|
.reg .pred %p<2>;
|
|
.reg .b32 %r<16>;
|
|
.reg .b64 %rd<8>;
|
|
mov.u32 %r0, %ctaid.x;
|
|
mov.u32 %r1, %ntid.x;
|
|
mov.u32 %r2, %tid.x;
|
|
mad.lo.s32 %r3, %r0, %r1, %r2;
|
|
ld.param.u32 %r4, [count];
|
|
setp.ge.u32 %p0, %r3, %r4;
|
|
@%p0 bra $fma_done;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mul.wide.u32 %rd2, %r3, 12;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
ld.global.u32 %r5, [%rd3];
|
|
ld.global.u32 %r6, [%rd3+4];
|
|
ld.global.u32 %r7, [%rd3+8];
|
|
fma.rn.sat.f16x2 %r8, %r5, %r6, %r7;
|
|
fma.rn.relu.f16x2 %r9, %r5, %r6, %r7;
|
|
fma.rn.f16x2.relu %r10, %r5, %r6, %r7;
|
|
mul.wide.u32 %rd4, %r3, 12;
|
|
add.s64 %rd5, %rd1, %rd4;
|
|
st.global.v2.u32 [%rd5], {%r8, %r9};
|
|
st.global.u32 [%rd5+8], %r10;
|
|
$fma_done:
|
|
ret;
|
|
}
|
|
|
|
// Forms used only by the DLTSS model kernels: ldmatrix x4/x2 on 32-bit shared
|
|
// addresses, prmt selectors, set.equ.u32.f16x2, dp2a.lo and cp.async.cg with a
|
|
// source size. One warp; per lane 16 input words and 20 output words.
|
|
.visible .entry d4r_dltss_misc(.param .u64 input, .param .u64 output)
|
|
{
|
|
.reg .b32 %r<64>;
|
|
.reg .b64 %rd<16>;
|
|
.shared .align 16 .b8 smem_[1024];
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mov.u32 %r0, %laneid;
|
|
mul.wide.u32 %rd2, %r0, 64;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
mul.wide.u32 %rd4, %r0, 80;
|
|
add.s64 %rd5, %rd1, %rd4;
|
|
ld.global.v4.u32 {%r1, %r2, %r3, %r4}, [%rd3];
|
|
mov.u32 %r5, smem_;
|
|
shl.b32 %r6, %r0, 4;
|
|
add.s32 %r7, %r5, %r6;
|
|
st.shared.v4.b32 [%r7], {%r1, %r2, %r3, %r4};
|
|
bar.sync 0;
|
|
// Row address: lane l points at 16-byte row (5 * l) mod 32.
|
|
mul.lo.u32 %r8, %r0, 5;
|
|
and.b32 %r8, %r8, 31;
|
|
shl.b32 %r8, %r8, 4;
|
|
add.s32 %r9, %r5, %r8;
|
|
ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%r10, %r11, %r12, %r13}, [%r9];
|
|
ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%r14, %r15}, [%r9];
|
|
st.global.v4.u32 [%rd5], {%r10, %r11, %r12, %r13};
|
|
st.global.v2.u32 [%rd5+16], {%r14, %r15};
|
|
ld.global.v4.u32 {%r16, %r17, %r18, %r19}, [%rd3+16];
|
|
prmt.b32 %r20, %r16, 0, 0x3424;
|
|
prmt.b32 %r21, %r16, 0, 0x1404;
|
|
prmt.b32 %r22, %r16, 0, 0x4431;
|
|
prmt.b32 %r23, %r16, %r17, 0x7632;
|
|
prmt.b32 %r24, %r16, %r17, 0x5410;
|
|
prmt.b32 %r25, %r16, %r17, 0x7531;
|
|
set.equ.u32.f16x2 %r26, %r16, %r19;
|
|
dp2a.lo.u32.u32 %r27, %r16, %r17, %r18;
|
|
st.global.v4.u32 [%rd5+24], {%r20, %r21, %r22, %r23};
|
|
st.global.v4.u32 [%rd5+40], {%r24, %r25, %r26, %r27};
|
|
// cp.async of 16 bytes with src-size = w12 (zero fill past it).
|
|
ld.global.u32 %r28, [%rd3+48];
|
|
add.s64 %rd6, %rd3, 32;
|
|
add.s32 %r29, %r7, 512;
|
|
cp.async.cg.shared.global [%r29], [%rd6], 16, %r28;
|
|
cp.async.commit_group;
|
|
cp.async.wait_group 0;
|
|
bar.sync 0;
|
|
ld.shared.v4.b32 {%r30, %r31, %r32, %r33}, [%r29];
|
|
st.global.v4.u32 [%rd5+56], {%r30, %r31, %r32, %r33};
|
|
ret;
|
|
}
|
|
|
|
// Rounding forms from the DLTSS output kernel's source-row mapping
|
|
// (round half away from zero via add.rz + cvt.rzi). Per thread: 2 input
|
|
// floats, 6 output words.
|
|
.visible .entry d4r_round_ops(.param .u64 input, .param .u64 output, .param .u32 count)
|
|
{
|
|
.reg .pred %p<4>;
|
|
.reg .b32 %r<16>;
|
|
.reg .f32 %f<16>;
|
|
.reg .b64 %rd<8>;
|
|
mov.u32 %r0, %ctaid.x;
|
|
mov.u32 %r1, %ntid.x;
|
|
mov.u32 %r2, %tid.x;
|
|
mad.lo.s32 %r3, %r0, %r1, %r2;
|
|
ld.param.u32 %r4, [count];
|
|
setp.ge.u32 %p0, %r3, %r4;
|
|
@%p0 bra $round_done;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mul.wide.u32 %rd2, %r3, 8;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
ld.global.v2.f32 {%f0, %f1}, [%rd3];
|
|
add.rz.ftz.f32 %f2, %f0, %f1;
|
|
mov.b32 %r5, %f0;
|
|
and.b32 %r6, %r5, -2147483648;
|
|
or.b32 %r7, %r6, 1056964608;
|
|
mov.b32 %f3, %r7;
|
|
add.rz.ftz.f32 %f4, %f0, %f3;
|
|
cvt.rzi.f32.f32 %f5, %f4;
|
|
cvt.rzi.ftz.s32.f32 %r8, %f5;
|
|
cvt.rzi.f32.f32 %f6, %f0;
|
|
add.ftz.f32 %f7, %f0, %f1;
|
|
mov.u32 %r9, 0;
|
|
mov.u32 %r10, 31;
|
|
mov.u32 %r11, -1;
|
|
shfl.sync.idx.b32 %r12|%p1, %r3, %r9, %r10, %r11;
|
|
selp.b32 %r13, 1, 0, %p1;
|
|
// bfi as in the DLTSS output kernel's ldmatrix row addressing.
|
|
and.b32 %r5, %r3, 8;
|
|
and.b32 %r6, %r3, 7;
|
|
bfi.b32 %r7, %r5, %r6, 3, 29;
|
|
bfi.b32 %r9, %r3, %r6, 4, 8;
|
|
mul.wide.u32 %rd4, %r3, 36;
|
|
add.s64 %rd5, %rd1, %rd4;
|
|
mov.b32 %r14, %f2;
|
|
mov.b32 %r15, %f5;
|
|
st.global.v4.u32 [%rd5], {%r14, %r15, %r8, %r12};
|
|
mov.b32 %r14, %f6;
|
|
mov.b32 %r15, %f7;
|
|
st.global.v4.u32 [%rd5+16], {%r14, %r15, %r13, %r7};
|
|
st.global.u32 [%rd5+32], %r9;
|
|
$round_done:
|
|
ret;
|
|
}
|
|
|
|
.visible .entry d4r_mma_m16n8k32_e4m3(.param .u64 input, .param .u64 output)
|
|
{
|
|
.reg .b32 %r<16>;
|
|
.reg .b64 %rd<8>;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mov.u32 %r0, %laneid;
|
|
mul.wide.u32 %rd2, %r0, 32;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
ld.global.v4.u32 {%r1, %r2, %r3, %r4}, [%rd3];
|
|
ld.global.v4.u32 {%r5, %r6, %r7, %r8}, [%rd3+16];
|
|
mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 {%r9, %r10}, {%r1, %r2, %r3, %r4}, {%r5, %r6}, {%r7, %r8};
|
|
mul.wide.u32 %rd4, %r0, 8;
|
|
add.s64 %rd5, %rd1, %rd4;
|
|
st.global.v2.u32 [%rd5], {%r9, %r10};
|
|
ret;
|
|
}
|
|
|
|
.visible .entry d4r_movmatrix_trans(.param .u64 input, .param .u64 output)
|
|
{
|
|
.reg .b32 %r<4>;
|
|
.reg .b64 %rd<8>;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mov.u32 %r0, %laneid;
|
|
mul.wide.u32 %rd2, %r0, 4;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
ld.global.u32 %r1, [%rd3];
|
|
movmatrix.sync.aligned.m8n8.trans.b16 %r2, %r1;
|
|
add.s64 %rd5, %rd1, %rd2;
|
|
st.global.u32 [%rd5], %r2;
|
|
ret;
|
|
}
|
|
|
|
.visible .entry d4r_global_v2_u16(.param .u64 input, .param .u64 output)
|
|
{
|
|
.reg .b16 %rs<2>;
|
|
.reg .b32 %r<3>;
|
|
.reg .b64 %rd<6>;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mov.u32 %r0, %laneid;
|
|
mul.wide.u32 %rd2, %r0, 4;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
ld.global.v2.u16 {%rs0, %rs1}, [%rd3];
|
|
mov.b32 %r1, {%rs0, %rs1};
|
|
mul.wide.u32 %rd4, %r0, 8;
|
|
add.s64 %rd5, %rd1, %rd4;
|
|
st.global.u32 [%rd5], %r1;
|
|
st.global.v2.u16 [%rd5+4], {%rs0, %rs1};
|
|
ret;
|
|
}
|
|
|
|
.visible .entry d4r_ex2_f16x2(.param .u64 input, .param .u64 output, .param .u32 count)
|
|
{
|
|
.reg .pred %p<2>;
|
|
.reg .b32 %r<8>;
|
|
.reg .b64 %rd<6>;
|
|
mov.u32 %r0, %ctaid.x;
|
|
mov.u32 %r1, %ntid.x;
|
|
mov.u32 %r2, %tid.x;
|
|
mad.lo.u32 %r3, %r0, %r1, %r2;
|
|
ld.param.u32 %r4, [count];
|
|
setp.ge.u32 %p0, %r3, %r4;
|
|
@%p0 bra $done;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
mul.wide.u32 %rd2, %r3, 4;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
add.s64 %rd4, %rd1, %rd2;
|
|
ld.global.u32 %r5, [%rd3];
|
|
ex2.approx.f16x2 %r6, %r5;
|
|
st.global.u32 [%rd4], %r6;
|
|
$done:
|
|
ret;
|
|
}
|
|
|
|
.visible .entry d4r_multidim_lane(.param .u64 output)
|
|
.maxntid 32, 1, 8
|
|
{
|
|
.reg .b32 %r<8>;
|
|
.reg .b64 %rd<5>;
|
|
ld.param.u64 %rd0, [output];
|
|
mov.u32 %r0, %tid.x;
|
|
mov.u32 %r1, %tid.y;
|
|
mov.u32 %r2, %tid.z;
|
|
mov.u32 %r3, %laneid;
|
|
mad.lo.u32 %r4, %r2, 32, %r0;
|
|
mul.wide.u32 %rd1, %r4, 16;
|
|
add.s64 %rd2, %rd0, %rd1;
|
|
st.global.v4.u32 [%rd2], {%r0, %r1, %r2, %r3};
|
|
ret;
|
|
}
|
|
// Packed and scalar f16 arithmetic as emitted by cuda_fp16.h inline asm in
|
|
// the RRLite kernels, including nested register scopes that redeclare the
|
|
// same names. Per thread: 3 input words, 13 output words.
|
|
.visible .entry d4r_f16x2_ops(.param .u64 input, .param .u64 output, .param .u32 count)
|
|
{
|
|
.reg .pred %p<2>;
|
|
.reg .b16 %rs<8>;
|
|
.reg .b32 %r<32>;
|
|
.reg .b64 %rd<8>;
|
|
.reg .f32 %f<4>;
|
|
.reg .f64 %fd<2>;
|
|
mov.u32 %r0, %ctaid.x;
|
|
mov.u32 %r1, %ntid.x;
|
|
mov.u32 %r2, %tid.x;
|
|
mad.lo.s32 %r3, %r0, %r1, %r2;
|
|
ld.param.u32 %r4, [count];
|
|
setp.ge.u32 %p0, %r3, %r4;
|
|
@%p0 bra $done;
|
|
ld.param.u64 %rd0, [input];
|
|
ld.param.u64 %rd1, [output];
|
|
cvta.to.global.u64 %rd0, %rd0;
|
|
cvta.to.global.u64 %rd1, %rd1;
|
|
mul.wide.u32 %rd2, %r3, 12;
|
|
add.s64 %rd3, %rd0, %rd2;
|
|
ld.global.u32 %r5, [%rd3];
|
|
ld.global.u32 %r6, [%rd3+4];
|
|
ld.global.u32 %r7, [%rd3+8];
|
|
{abs.f16x2 %r10,%r5;
|
|
}
|
|
{neg.f16x2 %r11,%r5;
|
|
}
|
|
{min.f16x2 %r12,%r5,%r6;
|
|
}
|
|
{max.f16x2 %r13,%r5,%r6;
|
|
}
|
|
{add.f16x2 %r14,%r5,%r6;
|
|
}
|
|
{sub.f16x2 %r15,%r5,%r6;
|
|
}
|
|
{mul.f16x2 %r16,%r5,%r6;
|
|
}
|
|
{fma.rn.f16x2 %r17,%r5,%r6,%r7;
|
|
}
|
|
mov.b32 {%rs0, %rs1}, %r5;
|
|
mov.b32 {%rs2, %rs3}, %r6;
|
|
{add.f16 %rs4,%rs0,%rs2;
|
|
}
|
|
{mul.f16 %rs5,%rs1,%rs3;
|
|
}
|
|
mov.b32 %r18, {%rs4, %rs5};
|
|
mov.f32 %f0, 0f3F000000;
|
|
{.reg .f16 low;
|
|
cvt.rn.f16.f32 low, %f0;
|
|
mov.b32 %r20, {low,low};}
|
|
mov.f32 %f1, 0f40000000;
|
|
{.reg .f16 low;
|
|
cvt.rn.f16.f32 low, %f1;
|
|
mov.b32 %r21, {low,low};}
|
|
mov.f64 %fd0, 0dBFE1CC0000000000;
|
|
{ cvt.rn.f16.f64 %rs6, %fd0;}
|
|
mov.f64 %fd1, 0d3FF0000000000000;
|
|
{ cvt.rn.f16.f64 %rs7, %fd1;}
|
|
mov.b32 %r22, {%rs6, %rs7};
|
|
{.reg.b16 hl, hu;
|
|
.reg.b32 fl, fu;
|
|
mov.b32 {hl, hu}, %r10;
|
|
cvt.f32.f16 fl, hl;
|
|
cvt.f32.f16 fu, hu;
|
|
rsqrt.approx.ftz.f32 fl, fl;
|
|
rsqrt.approx.ftz.f32 fu, fu;
|
|
cvt.rn.f16.f32 hl, fl;
|
|
cvt.rn.f16.f32 hu, fu;
|
|
mov.b32 %r23, {hl, hu};
|
|
}
|
|
mul.wide.u32 %rd4, %r3, 52;
|
|
add.s64 %rd5, %rd1, %rd4;
|
|
st.global.v4.u32 [%rd5], {%r10, %r11, %r12, %r13};
|
|
st.global.v4.u32 [%rd5+16], {%r14, %r15, %r16, %r17};
|
|
st.global.v4.u32 [%rd5+32], {%r18, %r20, %r21, %r22};
|
|
st.global.u32 [%rd5+48], %r23;
|
|
$done:
|
|
ret;
|
|
}
|
|
)PTX";
|
|
|
|
// --- CPU references ---------------------------------------------------------
|
|
|
|
static float half_to_float(uint16_t bits)
|
|
{
|
|
const uint32_t sign = (bits & 0x8000u) << 16;
|
|
const uint32_t exponent = (bits >> 10) & 0x1fu;
|
|
const uint32_t mantissa = bits & 0x3ffu;
|
|
if (exponent == 0)
|
|
return std::ldexp(static_cast<float>(mantissa), -24) * (sign ? -1.0f : 1.0f);
|
|
uint32_t result = exponent == 0x1f ? sign | 0x7f800000u | (mantissa << 13)
|
|
: sign | ((exponent + 112u) << 23) | (mantissa << 13);
|
|
float value;
|
|
std::memcpy(&value, &result, sizeof(value));
|
|
return value;
|
|
}
|
|
|
|
// Exact for every value representable in f16 (all E4M3 values and all the
|
|
// matrix results used below).
|
|
static uint16_t float_to_half_exact(float value)
|
|
{
|
|
for (uint32_t bits = 0; bits < 0x10000u; ++bits)
|
|
if (half_to_float(static_cast<uint16_t>(bits)) == value &&
|
|
std::signbit(half_to_float(static_cast<uint16_t>(bits))) == std::signbit(value))
|
|
return static_cast<uint16_t>(bits);
|
|
return 0x7e00u;
|
|
}
|
|
|
|
// OCP FP8 E4M3 (NVIDIA e4m3): bias 7, no infinities, S.1111.111 is NaN.
|
|
static float e4m3_to_float(uint8_t bits)
|
|
{
|
|
const float sign = (bits & 0x80u) ? -1.0f : 1.0f;
|
|
const int exponent = (bits >> 3) & 0xf;
|
|
const int mantissa = bits & 0x7;
|
|
if (exponent == 0xf && mantissa == 0x7)
|
|
return std::nanf("");
|
|
if (exponent == 0)
|
|
return sign * std::ldexp(static_cast<float>(mantissa), -9);
|
|
return sign * std::ldexp(1.0f + mantissa / 8.0f, exponent - 7);
|
|
}
|
|
|
|
// cvt.rn.satfinite: round to nearest even, finite overflow and infinities
|
|
// saturate to +-448, NaN stays NaN.
|
|
static uint8_t float_to_e4m3_satfinite(float value)
|
|
{
|
|
if (std::isnan(value))
|
|
return 0x7f;
|
|
const uint8_t sign = std::signbit(value) ? 0x80u : 0u;
|
|
const float magnitude = std::fabs(value);
|
|
if (magnitude >= 448.0f)
|
|
return sign | 0x7e;
|
|
uint8_t best = 0;
|
|
float best_error = INFINITY;
|
|
for (uint8_t code = 0; code <= 0x7e; ++code)
|
|
{
|
|
const float error = std::fabs(e4m3_to_float(code) - magnitude);
|
|
if (error < best_error || (error == best_error && (code & 1u) == 0))
|
|
{
|
|
best = code;
|
|
best_error = error;
|
|
}
|
|
}
|
|
return sign | best;
|
|
}
|
|
|
|
// PTX ISA fragment layouts. groupID = lane / 4, threadID_in_group = lane % 4.
|
|
// mma.m16n8k16 .f16 A (16x16): a[i] for i in 0..7, packed two per register.
|
|
static void layout_a_f16_k16(int lane, int i, int& row, int& col)
|
|
{
|
|
const int group = lane >> 2, thread = lane & 3;
|
|
row = (i < 2 || (i >= 4 && i < 6)) ? group : group + 8;
|
|
col = thread * 2 + (i & 1) + (i >= 4 ? 8 : 0);
|
|
}
|
|
// mma.m16n8k16 .f16 B (16x8): b[i] for i in 0..3.
|
|
static void layout_b_f16_k16(int lane, int i, int& row, int& col)
|
|
{
|
|
const int group = lane >> 2, thread = lane & 3;
|
|
row = thread * 2 + (i & 1) + (i >= 2 ? 8 : 0);
|
|
col = group;
|
|
}
|
|
// f16 accumulator C/D (16x8): c[i] for i in 0..3.
|
|
static void layout_c_f16(int lane, int i, int& row, int& col)
|
|
{
|
|
const int group = lane >> 2, thread = lane & 3;
|
|
row = i < 2 ? group : group + 8;
|
|
col = thread * 2 + (i & 1);
|
|
}
|
|
// mma.m16n8k8 .f16 A (16x8): a[i] for i in 0..3; B (8x8): b[i] for i in 0..1.
|
|
static void layout_a_f16_k8(int lane, int i, int& row, int& col)
|
|
{
|
|
const int group = lane >> 2, thread = lane & 3;
|
|
row = i < 2 ? group : group + 8;
|
|
col = thread * 2 + (i & 1);
|
|
}
|
|
static void layout_b_f16_k8(int lane, int i, int& row, int& col)
|
|
{
|
|
const int group = lane >> 2, thread = lane & 3;
|
|
row = thread * 2 + (i & 1);
|
|
col = group;
|
|
}
|
|
// mma.m16n8k32 8-bit A (16x32): a[i] for i in 0..15, four per register.
|
|
static void layout_a_8bit_k32(int lane, int i, int& row, int& col)
|
|
{
|
|
const int group = lane >> 2, thread = lane & 3;
|
|
row = (i < 4 || (i >= 8 && i < 12)) ? group : group + 8;
|
|
col = thread * 4 + (i & 3) + (i >= 8 ? 16 : 0);
|
|
}
|
|
// mma.m16n8k32 8-bit B (32x8): b[i] for i in 0..7.
|
|
static void layout_b_8bit_k32(int lane, int i, int& row, int& col)
|
|
{
|
|
const int group = lane >> 2, thread = lane & 3;
|
|
row = thread * 4 + (i & 3) + (i >= 4 ? 16 : 0);
|
|
col = group;
|
|
}
|
|
|
|
// --- driver plumbing --------------------------------------------------------
|
|
|
|
struct Driver
|
|
{
|
|
using InitFn = CUresult (*)(unsigned int);
|
|
using DeviceGetFn = CUresult (*)(CUdevice*, int);
|
|
using ContextCreateFn = CUresult (*)(CUcontext*, unsigned int, CUdevice);
|
|
using ModuleLoadFn = CUresult (*)(CUmodule*, const void*);
|
|
using ModuleGetFunctionFn = CUresult (*)(CUfunction*, CUmodule, const char*);
|
|
using AllocFn = CUresult (*)(CUdeviceptr*, size_t);
|
|
using FreeFn = CUresult (*)(CUdeviceptr);
|
|
using CopyHtoDFn = CUresult (*)(CUdeviceptr, const void*, size_t);
|
|
using CopyDtoHFn = CUresult (*)(void*, CUdeviceptr, size_t);
|
|
using LaunchKernelFn = CUresult (*)(CUfunction, unsigned int, unsigned int, unsigned int,
|
|
unsigned int, unsigned int, unsigned int, unsigned int,
|
|
CUstream, void**, void**);
|
|
using SynchronizeFn = CUresult (*)(void);
|
|
|
|
InitFn init;
|
|
DeviceGetFn deviceGet;
|
|
ContextCreateFn contextCreate;
|
|
ModuleLoadFn moduleLoad;
|
|
ModuleGetFunctionFn moduleGetFunction;
|
|
AllocFn alloc;
|
|
FreeFn free;
|
|
CopyHtoDFn copyHtoD;
|
|
CopyDtoHFn copyDtoH;
|
|
LaunchKernelFn launch;
|
|
SynchronizeFn synchronize;
|
|
CUmodule module = nullptr;
|
|
|
|
bool load(void* library)
|
|
{
|
|
init = load_function<InitFn>(library, "cuInit");
|
|
deviceGet = load_function<DeviceGetFn>(library, "cuDeviceGet");
|
|
contextCreate = load_function<ContextCreateFn>(library, "cuCtxCreate_v2");
|
|
moduleLoad = load_function<ModuleLoadFn>(library, "cuModuleLoadData");
|
|
moduleGetFunction = load_function<ModuleGetFunctionFn>(library, "cuModuleGetFunction");
|
|
alloc = load_function<AllocFn>(library, "cuMemAlloc_v2");
|
|
free = load_function<FreeFn>(library, "cuMemFree_v2");
|
|
copyHtoD = load_function<CopyHtoDFn>(library, "cuMemcpyHtoD_v2");
|
|
copyDtoH = load_function<CopyDtoHFn>(library, "cuMemcpyDtoH_v2");
|
|
launch = load_function<LaunchKernelFn>(library, "cuLaunchKernel");
|
|
synchronize = load_function<SynchronizeFn>(library, "cuCtxSynchronize");
|
|
return init && deviceGet && contextCreate && moduleLoad && moduleGetFunction && alloc &&
|
|
free && copyHtoD && copyDtoH && launch && synchronize;
|
|
}
|
|
|
|
// Runs kernel(input, output[, count]) over one buffer each way.
|
|
bool run(const char* name, const void* input, size_t inputBytes, void* output,
|
|
size_t outputBytes, unsigned int blocks, unsigned int threads, int count)
|
|
{
|
|
CUfunction function = nullptr;
|
|
CUdeviceptr deviceInput = 0, deviceOutput = 0;
|
|
CUresult result = moduleGetFunction(&function, module, name);
|
|
if (result == 0)
|
|
result = alloc(&deviceInput, inputBytes);
|
|
if (result == 0)
|
|
result = alloc(&deviceOutput, outputBytes);
|
|
if (result == 0)
|
|
result = copyHtoD(deviceInput, input, inputBytes);
|
|
if (result == 0)
|
|
{
|
|
void* withCount[] = {&deviceInput, &deviceOutput, &count};
|
|
void* withoutCount[] = {&deviceInput, &deviceOutput};
|
|
result = launch(function, blocks, 1, 1, threads, 1, 1, 0, nullptr,
|
|
count >= 0 ? withCount : withoutCount, nullptr);
|
|
}
|
|
if (result == 0)
|
|
result = synchronize();
|
|
if (result == 0)
|
|
result = copyDtoH(output, deviceOutput, outputBytes);
|
|
if (deviceInput != 0)
|
|
free(deviceInput);
|
|
if (deviceOutput != 0)
|
|
free(deviceOutput);
|
|
if (result != 0)
|
|
std::printf("%s: CUDA result=%d\n", name, result);
|
|
return result == 0;
|
|
}
|
|
};
|
|
|
|
// --- tests ------------------------------------------------------------------
|
|
|
|
static bool test_cvt_from_e4m3(Driver& driver)
|
|
{
|
|
std::vector<uint16_t> input(65536);
|
|
for (uint32_t index = 0; index < input.size(); ++index)
|
|
input[index] = static_cast<uint16_t>(index);
|
|
std::vector<uint32_t> output(input.size());
|
|
if (!driver.run("d4r_cvt_f16x2_from_e4m3x2", input.data(), input.size() * 2, output.data(),
|
|
output.size() * 4, 256, 256, static_cast<int>(input.size())))
|
|
return false;
|
|
size_t mismatches = 0;
|
|
for (uint32_t index = 0; index < input.size(); ++index)
|
|
{
|
|
for (int lane = 0; lane < 2; ++lane)
|
|
{
|
|
const float expected = e4m3_to_float(static_cast<uint8_t>(index >> (8 * lane)));
|
|
const uint16_t got_bits = static_cast<uint16_t>(output[index] >> (16 * lane));
|
|
const float got = half_to_float(got_bits);
|
|
const bool ok = std::isnan(expected) ? std::isnan(got)
|
|
: got == expected && std::signbit(got) == std::signbit(expected);
|
|
if (!ok && mismatches++ < 8)
|
|
std::printf(" e4m3x2 0x%04x lane %d: got f16 0x%04x (%g) expected %g\n", index,
|
|
lane, got_bits, got, expected);
|
|
}
|
|
}
|
|
std::printf("cvt.rn.f16x2.e4m3x2 (65536 inputs): %s (%zu mismatches)\n",
|
|
mismatches == 0 ? "PASS" : "FAIL", mismatches);
|
|
return mismatches == 0;
|
|
}
|
|
|
|
static bool test_cvt_to_e4m3(Driver& driver, const char* kernel = "d4r_cvt_e4m3x2_from_f16x2",
|
|
const char* label = "cvt.rn.satfinite.e4m3x2.f16x2", bool relu = false)
|
|
{
|
|
// Lane 0 sweeps every f16 value while lane 1 sweeps them in reverse.
|
|
std::vector<uint32_t> input(65536);
|
|
for (uint32_t index = 0; index < input.size(); ++index)
|
|
input[index] = index | ((0xffffu - index) << 16);
|
|
std::vector<uint16_t> output(input.size());
|
|
if (!driver.run(kernel, input.data(), input.size() * 4, output.data(),
|
|
output.size() * 2, 256, 256, static_cast<int>(input.size())))
|
|
return false;
|
|
size_t mismatches = 0;
|
|
for (uint32_t index = 0; index < input.size(); ++index)
|
|
{
|
|
for (int lane = 0; lane < 2; ++lane)
|
|
{
|
|
const uint16_t half = static_cast<uint16_t>(input[index] >> (16 * lane));
|
|
// .relu: negative results (and -0) become +0, NaN stays NaN.
|
|
const float value = half_to_float(half);
|
|
const uint8_t expected = relu && !std::isnan(value) && std::signbit(value)
|
|
? 0 : float_to_e4m3_satfinite(value);
|
|
const uint8_t got = static_cast<uint8_t>(output[index] >> (8 * lane));
|
|
const bool ok = (expected & 0x7f) == 0x7f ? (got & 0x7f) == 0x7f : got == expected;
|
|
if (!ok && mismatches++ < 8)
|
|
std::printf(" f16 0x%04x (%g) lane %d: got e4m3 0x%02x expected 0x%02x\n", half,
|
|
half_to_float(half), lane, got, expected);
|
|
}
|
|
}
|
|
std::printf("%s (131072 lanes): %s (%zu mismatches)\n", label,
|
|
mismatches == 0 ? "PASS" : "FAIL", mismatches);
|
|
return mismatches == 0;
|
|
}
|
|
|
|
// Runs one MMA through the fragment layouts and compares D with A*B+C.
|
|
template <typename ALayout, typename BLayout>
|
|
static bool test_mma(Driver& driver, const char* kernel, const char* label, int k, int elementBits,
|
|
ALayout aLayout, BLayout bLayout, const std::vector<float>& a,
|
|
const std::vector<float>& b, const std::vector<float>& c,
|
|
const std::vector<uint8_t>& aCodes, const std::vector<uint8_t>& bCodes)
|
|
{
|
|
uint32_t input[32][8] = {};
|
|
// 16x k A and k x 8 B fragments spread over 32 lanes.
|
|
const int aElements = k / 2;
|
|
const int bElements = k / 4;
|
|
const int perWord = 32 / elementBits;
|
|
for (int lane = 0; lane < 32; ++lane)
|
|
{
|
|
for (int i = 0; i < aElements; ++i)
|
|
{
|
|
int row, col;
|
|
aLayout(lane, i, row, col);
|
|
const uint32_t bits = elementBits == 16 ? float_to_half_exact(a[row * k + col])
|
|
: aCodes[row * k + col];
|
|
input[lane][i / perWord] |= bits << (elementBits * (i % perWord));
|
|
}
|
|
for (int i = 0; i < bElements; ++i)
|
|
{
|
|
int row, col;
|
|
bLayout(lane, i, row, col);
|
|
const uint32_t bits = elementBits == 16 ? float_to_half_exact(b[row * 8 + col])
|
|
: bCodes[row * 8 + col];
|
|
input[lane][4 + i / perWord] |= bits << (elementBits * (i % perWord));
|
|
}
|
|
for (int i = 0; i < 4; ++i)
|
|
{
|
|
int row, col;
|
|
layout_c_f16(lane, i, row, col);
|
|
input[lane][6 + i / 2] |= static_cast<uint32_t>(float_to_half_exact(c[row * 8 + col]))
|
|
<< (16 * (i % 2));
|
|
}
|
|
}
|
|
uint32_t output[32][2] = {};
|
|
if (!driver.run(kernel, input, sizeof(input), output, sizeof(output), 1, 32, -1))
|
|
return false;
|
|
size_t mismatches = 0;
|
|
for (int lane = 0; lane < 32; ++lane)
|
|
{
|
|
for (int i = 0; i < 4; ++i)
|
|
{
|
|
int row, col;
|
|
layout_c_f16(lane, i, row, col);
|
|
float expected = c[row * 8 + col];
|
|
for (int kk = 0; kk < k; ++kk)
|
|
expected += a[row * k + kk] * b[kk * 8 + col];
|
|
const float got = half_to_float(static_cast<uint16_t>(output[lane][i / 2] >> (16 * (i % 2))));
|
|
if (got != expected && mismatches++ < 6)
|
|
std::printf(" %s D[%d][%d] (lane %d, d%d): got %g expected %g\n", label, row, col,
|
|
lane, i, got, expected);
|
|
}
|
|
}
|
|
std::printf("%s: %s (%zu of 128 outputs mismatched)\n", label, mismatches == 0 ? "PASS" : "FAIL",
|
|
mismatches);
|
|
return mismatches == 0;
|
|
}
|
|
|
|
// Order-independent f16 distance in units in the last place.
|
|
static int half_ulp_distance(uint16_t a, uint16_t b)
|
|
{
|
|
auto ordered = [](uint16_t bits) { return (bits & 0x8000u) ? -static_cast<int>(bits & 0x7fffu) : static_cast<int>(bits); };
|
|
return std::abs(ordered(a) - ordered(b));
|
|
}
|
|
|
|
// Random f16 inputs through d4r_mma_chain_f16, against a reference that sums
|
|
// each MMA exactly (in double) and rounds the accumulator to f16 once per MMA.
|
|
// Tensor cores and the translations differ only in summation order inside one
|
|
// MMA, so a result may be off by a few ulp; layout errors are off by far more.
|
|
static bool test_mma_chain(Driver& driver, int trial)
|
|
{
|
|
std::mt19937 random(1234 + trial);
|
|
std::uniform_real_distribution<float> uniform(-2.0f, 2.0f);
|
|
auto pick_half = [&]() { return static_cast<float>(static_cast<_Float16>(uniform(random))); };
|
|
constexpr int kSteps = 8;
|
|
std::vector<float> a(kSteps * 16 * 16), b0(kSteps * 16 * 8), b1(kSteps * 16 * 8), c0(16 * 8), c1(16 * 8);
|
|
for (float& value : a) value = pick_half();
|
|
for (float& value : b0) value = pick_half();
|
|
for (float& value : b1) value = pick_half();
|
|
for (float& value : c0) value = pick_half();
|
|
for (float& value : c1) value = pick_half();
|
|
static uint32_t input[32][68];
|
|
std::memset(input, 0, sizeof(input));
|
|
for (int lane = 0; lane < 32; ++lane)
|
|
{
|
|
for (int step = 0; step < kSteps; ++step)
|
|
{
|
|
uint32_t* words = input[lane] + step * 8;
|
|
for (int i = 0; i < 8; ++i)
|
|
{
|
|
int row, col;
|
|
layout_a_f16_k16(lane, i, row, col);
|
|
words[i / 2] |= static_cast<uint32_t>(float_to_half_exact(a[step * 256 + row * 16 + col])) << (16 * (i % 2));
|
|
}
|
|
for (int i = 0; i < 4; ++i)
|
|
{
|
|
int row, col;
|
|
layout_b_f16_k16(lane, i, row, col);
|
|
words[4 + i / 2] |= static_cast<uint32_t>(float_to_half_exact(b0[step * 128 + row * 8 + col])) << (16 * (i % 2));
|
|
words[6 + i / 2] |= static_cast<uint32_t>(float_to_half_exact(b1[step * 128 + row * 8 + col])) << (16 * (i % 2));
|
|
}
|
|
}
|
|
for (int i = 0; i < 4; ++i)
|
|
{
|
|
int row, col;
|
|
layout_c_f16(lane, i, row, col);
|
|
input[lane][64 + i / 2] |= static_cast<uint32_t>(float_to_half_exact(c0[row * 8 + col])) << (16 * (i % 2));
|
|
input[lane][66 + i / 2] |= static_cast<uint32_t>(float_to_half_exact(c1[row * 8 + col])) << (16 * (i % 2));
|
|
}
|
|
}
|
|
uint32_t output[32][4] = {};
|
|
if (!driver.run("d4r_mma_chain_f16", input, sizeof(input), output, sizeof(output), 1, 32, -1))
|
|
return false;
|
|
std::vector<uint16_t> reference0(128), reference1(128);
|
|
std::vector<double> acc0(c0.begin(), c0.end()), acc1(c1.begin(), c1.end());
|
|
for (int step = 0; step < kSteps; ++step)
|
|
for (int row = 0; row < 16; ++row)
|
|
for (int col = 0; col < 8; ++col)
|
|
{
|
|
double sum0 = acc0[row * 8 + col], sum1 = acc1[row * 8 + col];
|
|
for (int k = 0; k < 16; ++k)
|
|
{
|
|
sum0 += static_cast<double>(a[step * 256 + row * 16 + k]) * b0[step * 128 + k * 8 + col];
|
|
sum1 += static_cast<double>(a[step * 256 + row * 16 + k]) * b1[step * 128 + k * 8 + col];
|
|
}
|
|
acc0[row * 8 + col] = static_cast<double>(static_cast<_Float16>(sum0));
|
|
acc1[row * 8 + col] = static_cast<double>(static_cast<_Float16>(sum1));
|
|
}
|
|
int worst = 0, exact = 0;
|
|
for (int lane = 0; lane < 32; ++lane)
|
|
for (int i = 0; i < 4; ++i)
|
|
{
|
|
int row, col;
|
|
layout_c_f16(lane, i, row, col);
|
|
const uint16_t got0 = static_cast<uint16_t>(output[lane][i / 2] >> (16 * (i % 2)));
|
|
const uint16_t got1 = static_cast<uint16_t>(output[lane][2 + i / 2] >> (16 * (i % 2)));
|
|
const _Float16 want0 = static_cast<_Float16>(acc0[row * 8 + col]);
|
|
const _Float16 want1 = static_cast<_Float16>(acc1[row * 8 + col]);
|
|
uint16_t bits0, bits1;
|
|
std::memcpy(&bits0, &want0, 2);
|
|
std::memcpy(&bits1, &want1, 2);
|
|
const int d0 = half_ulp_distance(got0, bits0), d1 = half_ulp_distance(got1, bits1);
|
|
worst = std::max(worst, std::max(d0, d1));
|
|
exact += (d0 == 0) + (d1 == 0);
|
|
}
|
|
const bool ok = worst <= 8;
|
|
std::printf("mma.m16n8k16.f16 chained x%d, paired, trial %d: %s (worst %d ulp, %d of 256 exact)\n", kSteps,
|
|
trial, ok ? "PASS" : "FAIL", worst, exact);
|
|
return ok;
|
|
}
|
|
|
|
// Full-range e4m3 codes (normals, subnormals, zeros; no NaN) through one
|
|
// m16n8k32 MMA against a double reference rounded to f16 once. The f32
|
|
// summation order may differ from the reference, so allow 2 ulp.
|
|
static bool test_mma_e4m3_codes(Driver& driver, const char* label, const std::vector<uint8_t>& aCodes,
|
|
const std::vector<uint8_t>& bCodes, const std::vector<float>& c)
|
|
{
|
|
uint32_t input[32][8] = {};
|
|
for (int lane = 0; lane < 32; ++lane)
|
|
{
|
|
for (int i = 0; i < 16; ++i)
|
|
{
|
|
int row, col;
|
|
layout_a_8bit_k32(lane, i, row, col);
|
|
input[lane][i / 4] |= static_cast<uint32_t>(aCodes[row * 32 + col]) << (8 * (i % 4));
|
|
}
|
|
for (int i = 0; i < 8; ++i)
|
|
{
|
|
int row, col;
|
|
layout_b_8bit_k32(lane, i, row, col);
|
|
input[lane][4 + i / 4] |= static_cast<uint32_t>(bCodes[row * 8 + col]) << (8 * (i % 4));
|
|
}
|
|
for (int i = 0; i < 4; ++i)
|
|
{
|
|
int row, col;
|
|
layout_c_f16(lane, i, row, col);
|
|
input[lane][6 + i / 2] |= static_cast<uint32_t>(float_to_half_exact(c[row * 8 + col])) << (16 * (i % 2));
|
|
}
|
|
}
|
|
uint32_t output[32][2] = {};
|
|
if (!driver.run("d4r_mma_m16n8k32_e4m3", input, sizeof(input), output, sizeof(output), 1, 32, -1))
|
|
return false;
|
|
int worst = 0, exact = 0;
|
|
for (int lane = 0; lane < 32; ++lane)
|
|
for (int i = 0; i < 4; ++i)
|
|
{
|
|
int row, col;
|
|
layout_c_f16(lane, i, row, col);
|
|
double sum = c[row * 8 + col];
|
|
for (int k = 0; k < 32; ++k)
|
|
sum += static_cast<double>(e4m3_to_float(aCodes[row * 32 + k])) * e4m3_to_float(bCodes[k * 8 + col]);
|
|
const _Float16 want = static_cast<_Float16>(sum);
|
|
uint16_t bits;
|
|
std::memcpy(&bits, &want, 2);
|
|
const uint16_t got = static_cast<uint16_t>(output[lane][i / 2] >> (16 * (i % 2)));
|
|
const int distance = half_ulp_distance(got, bits);
|
|
if (distance > 2 && worst <= 2)
|
|
std::printf(" %s D[%d][%d]: got 0x%04x expected 0x%04x (%g)\n", label, row, col, got, bits, sum);
|
|
worst = std::max(worst, distance);
|
|
exact += distance == 0;
|
|
}
|
|
const bool ok = worst <= 2;
|
|
std::printf("%s: %s (worst %d ulp, %d of 128 exact)\n", label, ok ? "PASS" : "FAIL", worst, exact);
|
|
return ok;
|
|
}
|
|
|
|
// Random full-range e4m3 codes through d4r_mma_chain_e4m3 (8 steps, pairs that
|
|
// share A), against a reference that rounds the accumulator to f16 per MMA.
|
|
static bool test_mma_chain_e4m3(Driver& driver, int trial)
|
|
{
|
|
std::mt19937 random(4321 + trial);
|
|
auto pick_code = [&]() {
|
|
uint8_t code;
|
|
do
|
|
code = static_cast<uint8_t>(random());
|
|
while ((code & 0x7f) == 0x7f || (code & 0x78) >= 0x50); // no NaN, |x| < 32
|
|
return code;
|
|
};
|
|
constexpr int kSteps = 8;
|
|
std::vector<uint8_t> a(kSteps * 16 * 32), b0(kSteps * 32 * 8), b1(kSteps * 32 * 8);
|
|
for (uint8_t& code : a) code = pick_code();
|
|
for (uint8_t& code : b0) code = pick_code();
|
|
for (uint8_t& code : b1) code = pick_code();
|
|
std::uniform_real_distribution<float> uniform(-4.0f, 4.0f);
|
|
std::vector<float> c0(16 * 8), c1(16 * 8);
|
|
for (float& value : c0) value = static_cast<float>(static_cast<_Float16>(uniform(random)));
|
|
for (float& value : c1) value = static_cast<float>(static_cast<_Float16>(uniform(random)));
|
|
static uint32_t input[32][68];
|
|
std::memset(input, 0, sizeof(input));
|
|
for (int lane = 0; lane < 32; ++lane)
|
|
{
|
|
for (int step = 0; step < kSteps; ++step)
|
|
{
|
|
uint32_t* words = input[lane] + step * 8;
|
|
for (int i = 0; i < 16; ++i)
|
|
{
|
|
int row, col;
|
|
layout_a_8bit_k32(lane, i, row, col);
|
|
words[i / 4] |= static_cast<uint32_t>(a[step * 512 + row * 32 + col]) << (8 * (i % 4));
|
|
}
|
|
for (int i = 0; i < 8; ++i)
|
|
{
|
|
int row, col;
|
|
layout_b_8bit_k32(lane, i, row, col);
|
|
words[4 + i / 4] |= static_cast<uint32_t>(b0[step * 256 + row * 8 + col]) << (8 * (i % 4));
|
|
words[6 + i / 4] |= static_cast<uint32_t>(b1[step * 256 + row * 8 + col]) << (8 * (i % 4));
|
|
}
|
|
}
|
|
for (int i = 0; i < 4; ++i)
|
|
{
|
|
int row, col;
|
|
layout_c_f16(lane, i, row, col);
|
|
input[lane][64 + i / 2] |= static_cast<uint32_t>(float_to_half_exact(c0[row * 8 + col])) << (16 * (i % 2));
|
|
input[lane][66 + i / 2] |= static_cast<uint32_t>(float_to_half_exact(c1[row * 8 + col])) << (16 * (i % 2));
|
|
}
|
|
}
|
|
uint32_t output[32][4] = {};
|
|
if (!driver.run("d4r_mma_chain_e4m3", input, sizeof(input), output, sizeof(output), 1, 32, -1))
|
|
return false;
|
|
std::vector<double> acc0(c0.begin(), c0.end()), acc1(c1.begin(), c1.end());
|
|
for (int step = 0; step < kSteps; ++step)
|
|
for (int row = 0; row < 16; ++row)
|
|
for (int col = 0; col < 8; ++col)
|
|
{
|
|
double sum0 = acc0[row * 8 + col], sum1 = acc1[row * 8 + col];
|
|
for (int k = 0; k < 32; ++k)
|
|
{
|
|
const double av = e4m3_to_float(a[step * 512 + row * 32 + k]);
|
|
sum0 += av * e4m3_to_float(b0[step * 256 + k * 8 + col]);
|
|
sum1 += av * e4m3_to_float(b1[step * 256 + k * 8 + col]);
|
|
}
|
|
acc0[row * 8 + col] = static_cast<double>(static_cast<_Float16>(sum0));
|
|
acc1[row * 8 + col] = static_cast<double>(static_cast<_Float16>(sum1));
|
|
}
|
|
int worst = 0, exact = 0;
|
|
for (int lane = 0; lane < 32; ++lane)
|
|
for (int i = 0; i < 4; ++i)
|
|
{
|
|
int row, col;
|
|
layout_c_f16(lane, i, row, col);
|
|
const uint16_t got0 = static_cast<uint16_t>(output[lane][i / 2] >> (16 * (i % 2)));
|
|
const uint16_t got1 = static_cast<uint16_t>(output[lane][2 + i / 2] >> (16 * (i % 2)));
|
|
const _Float16 want0 = static_cast<_Float16>(acc0[row * 8 + col]);
|
|
const _Float16 want1 = static_cast<_Float16>(acc1[row * 8 + col]);
|
|
uint16_t bits0, bits1;
|
|
std::memcpy(&bits0, &want0, 2);
|
|
std::memcpy(&bits1, &want1, 2);
|
|
const int d0 = half_ulp_distance(got0, bits0), d1 = half_ulp_distance(got1, bits1);
|
|
worst = std::max(worst, std::max(d0, d1));
|
|
exact += (d0 == 0) + (d1 == 0);
|
|
}
|
|
const bool ok = worst <= 8;
|
|
std::printf("mma.m16n8k32.e4m3 chained x%d, paired, trial %d: %s (worst %d ulp, %d of 256 exact)\n", kSteps,
|
|
trial, ok ? "PASS" : "FAIL", worst, exact);
|
|
return ok;
|
|
}
|
|
|
|
// Checks d4r_mma_pair_alias_e4m3 against a double reference (2 ulp).
|
|
static bool test_mma_pair_alias(Driver& driver, bool late)
|
|
{
|
|
std::mt19937 random(77 + late);
|
|
auto code = [&]() {
|
|
uint8_t value;
|
|
do
|
|
value = static_cast<uint8_t>(random());
|
|
while ((value & 0x7f) == 0x7f || (value & 0x78) >= 0x50);
|
|
return value;
|
|
};
|
|
std::vector<uint8_t> a(16 * 32), b0(32 * 8), b1(32 * 8);
|
|
for (uint8_t& value : a) value = code();
|
|
for (uint8_t& value : b0) value = code();
|
|
for (uint8_t& value : b1) value = code();
|
|
std::uniform_real_distribution<float> uniform(-4.0f, 4.0f);
|
|
std::vector<float> c0(16 * 8);
|
|
for (float& value : c0) value = static_cast<float>(static_cast<_Float16>(uniform(random)));
|
|
// The late value of %r20: one f16 pair per lane, so C1 = {v, v} there.
|
|
const uint16_t lateHalf0 = 0x3c00, lateHalf1 = 0xc000; // 1.0, -2.0
|
|
uint32_t input[32][12] = {};
|
|
for (int lane = 0; lane < 32; ++lane)
|
|
{
|
|
for (int i = 0; i < 16; ++i)
|
|
{
|
|
int row, col;
|
|
layout_a_8bit_k32(lane, i, row, col);
|
|
input[lane][i / 4] |= static_cast<uint32_t>(a[row * 32 + col]) << (8 * (i % 4));
|
|
}
|
|
for (int i = 0; i < 8; ++i)
|
|
{
|
|
int row, col;
|
|
layout_b_8bit_k32(lane, i, row, col);
|
|
input[lane][4 + i / 4] |= static_cast<uint32_t>(b0[row * 8 + col]) << (8 * (i % 4));
|
|
input[lane][6 + i / 4] |= static_cast<uint32_t>(b1[row * 8 + col]) << (8 * (i % 4));
|
|
}
|
|
for (int i = 0; i < 4; ++i)
|
|
{
|
|
int row, col;
|
|
layout_c_f16(lane, i, row, col);
|
|
input[lane][8 + i / 2] |= static_cast<uint32_t>(float_to_half_exact(c0[row * 8 + col])) << (16 * (i % 2));
|
|
}
|
|
input[lane][10] = lateHalf0 | (static_cast<uint32_t>(lateHalf1) << 16);
|
|
}
|
|
uint32_t output[32][4] = {};
|
|
if (!driver.run(late ? "d4r_mma_pair_alias_late_e4m3" : "d4r_mma_pair_alias_e4m3", input, sizeof(input), output,
|
|
sizeof(output), 1, 32, late ? 1 : 0))
|
|
return false;
|
|
int worst = 0;
|
|
for (int lane = 0; lane < 32; ++lane)
|
|
for (int i = 0; i < 4; ++i)
|
|
{
|
|
int row, col;
|
|
layout_c_f16(lane, i, row, col);
|
|
double sum0 = c0[row * 8 + col];
|
|
// C1 = {%r20, %r20}: element i reads half (i % 2) of %r20.
|
|
double sum1 = late ? half_to_float(i % 2 == 0 ? lateHalf0 : lateHalf1) : 0.0;
|
|
for (int k = 0; k < 32; ++k)
|
|
{
|
|
sum0 += static_cast<double>(e4m3_to_float(a[row * 32 + k])) * e4m3_to_float(b0[k * 8 + col]);
|
|
sum1 += static_cast<double>(e4m3_to_float(a[row * 32 + k])) * e4m3_to_float(b1[k * 8 + col]);
|
|
}
|
|
const _Float16 want0 = static_cast<_Float16>(sum0), want1 = static_cast<_Float16>(sum1);
|
|
uint16_t bits0, bits1;
|
|
std::memcpy(&bits0, &want0, 2);
|
|
std::memcpy(&bits1, &want1, 2);
|
|
const uint16_t got0 = static_cast<uint16_t>(output[lane][i / 2] >> (16 * (i % 2)));
|
|
const uint16_t got1 = static_cast<uint16_t>(output[lane][2 + i / 2] >> (16 * (i % 2)));
|
|
worst = std::max(worst, std::max(half_ulp_distance(got0, bits0), half_ulp_distance(got1, bits1)));
|
|
}
|
|
const bool ok = worst <= 2;
|
|
std::printf("mma.m16n8k32.e4m3 pair, aliased zero accumulator%s: %s (worst %d ulp)\n",
|
|
late ? " rewritten between" : "", ok ? "PASS" : "FAIL", worst);
|
|
return ok;
|
|
}
|
|
|
|
static bool test_movmatrix(Driver& driver)
|
|
{
|
|
// Source fragment: lane holds row lane/4, columns 2*(lane%4) and +1.
|
|
uint32_t input[32], output[32];
|
|
for (int lane = 0; lane < 32; ++lane)
|
|
{
|
|
const int row = lane >> 2, col = (lane & 3) * 2;
|
|
input[lane] = static_cast<uint32_t>(row * 8 + col) | (static_cast<uint32_t>(row * 8 + col + 1) << 16);
|
|
}
|
|
if (!driver.run("d4r_movmatrix_trans", input, sizeof(input), output, sizeof(output), 1, 32, -1))
|
|
return false;
|
|
size_t mismatches = 0;
|
|
for (int lane = 0; lane < 32; ++lane)
|
|
{
|
|
const int row = lane >> 2, col = (lane & 3) * 2;
|
|
// Destination holds the transpose in the same fragment layout.
|
|
const uint32_t expected = static_cast<uint32_t>(col * 8 + row) |
|
|
(static_cast<uint32_t>((col + 1) * 8 + row) << 16);
|
|
if (output[lane] != expected && mismatches++ < 6)
|
|
std::printf(" movmatrix lane %d: got 0x%08x expected 0x%08x\n", lane, output[lane], expected);
|
|
}
|
|
std::printf("movmatrix.sync.aligned.m8n8.trans.b16: %s (%zu of 32 lanes mismatched)\n",
|
|
mismatches == 0 ? "PASS" : "FAIL", mismatches);
|
|
return mismatches == 0;
|
|
}
|
|
|
|
static bool test_global_v2_u16(Driver& driver)
|
|
{
|
|
uint32_t input[32], output[32][2] = {};
|
|
for (uint32_t lane = 0; lane < 32; ++lane)
|
|
input[lane] = (0x9137u + lane * 37u) | ((0xe5c3u - lane * 29u) << 16);
|
|
if (!driver.run("d4r_global_v2_u16", input, sizeof(input), output, sizeof(output), 1, 32, -1))
|
|
return false;
|
|
size_t mismatches = 0;
|
|
for (int lane = 0; lane < 32; ++lane)
|
|
for (int copy = 0; copy < 2; ++copy)
|
|
if (output[lane][copy] != input[lane] && mismatches++ < 6)
|
|
std::printf(" global.v2.u16 lane %d copy %d: got 0x%08x expected 0x%08x\n",
|
|
lane, copy, output[lane][copy], input[lane]);
|
|
std::printf("ld/st.global.v2.u16: %s (%zu of 64 words mismatched)\n",
|
|
mismatches == 0 ? "PASS" : "FAIL", mismatches);
|
|
return mismatches == 0;
|
|
}
|
|
|
|
static bool test_ex2_f16x2(Driver& driver)
|
|
{
|
|
constexpr int count = 4096;
|
|
std::vector<uint32_t> input(count), output(count);
|
|
for (int index = 0; index < count; ++index)
|
|
{
|
|
const _Float16 lo = static_cast<_Float16>((index % 97 - 48) * 0.125f);
|
|
const _Float16 hi = static_cast<_Float16>((index % 113 - 56) * 0.125f);
|
|
uint16_t loBits, hiBits;
|
|
std::memcpy(&loBits, &lo, sizeof(loBits));
|
|
std::memcpy(&hiBits, &hi, sizeof(hiBits));
|
|
input[index] = static_cast<uint32_t>(loBits) | (static_cast<uint32_t>(hiBits) << 16);
|
|
}
|
|
if (!driver.run("d4r_ex2_f16x2", input.data(), input.size() * 4, output.data(), output.size() * 4,
|
|
16, 256, count))
|
|
return false;
|
|
size_t mismatches = 0;
|
|
double worstRelative = 0.0;
|
|
for (int index = 0; index < count; ++index)
|
|
for (int half = 0; half < 2; ++half)
|
|
{
|
|
const auto bits = static_cast<uint16_t>(input[index] >> (16 * half));
|
|
const auto gotBits = static_cast<uint16_t>(output[index] >> (16 * half));
|
|
const float expected = std::exp2(half_to_float(bits));
|
|
const float got = half_to_float(gotBits);
|
|
const double relative = std::fabs(got - expected) / std::max(1.0f, expected);
|
|
worstRelative = std::max(worstRelative, relative);
|
|
if ((!std::isfinite(got) || relative > 0.01) && mismatches++ < 6)
|
|
std::printf(" ex2.f16x2 input %g: got %g expected %g\n", half_to_float(bits), got, expected);
|
|
}
|
|
std::printf("ex2.approx.f16x2: %s (%zu of %d lanes mismatched, worst relative %.6g)\n",
|
|
mismatches == 0 ? "PASS" : "FAIL", mismatches, count * 2, worstRelative);
|
|
return mismatches == 0;
|
|
}
|
|
|
|
static bool test_multidim_lane(Driver& driver)
|
|
{
|
|
CUfunction function = nullptr;
|
|
CUdeviceptr deviceOutput = 0;
|
|
uint32_t output[256][4] = {};
|
|
CUresult result = driver.moduleGetFunction(&function, driver.module, "d4r_multidim_lane");
|
|
if (result == 0)
|
|
result = driver.alloc(&deviceOutput, sizeof(output));
|
|
if (result == 0)
|
|
{
|
|
void* arguments[] = {&deviceOutput};
|
|
result = driver.launch(function, 1, 1, 1, 32, 1, 8, 0, nullptr, arguments, nullptr);
|
|
}
|
|
if (result == 0)
|
|
result = driver.synchronize();
|
|
if (result == 0)
|
|
result = driver.copyDtoH(output, deviceOutput, sizeof(output));
|
|
if (deviceOutput != 0)
|
|
driver.free(deviceOutput);
|
|
size_t mismatches = 0;
|
|
for (int z = 0; z < 8; ++z)
|
|
for (int x = 0; x < 32; ++x)
|
|
{
|
|
const uint32_t* got = output[z * 32 + x];
|
|
if ((got[0] != static_cast<uint32_t>(x) || got[1] != 0 ||
|
|
got[2] != static_cast<uint32_t>(z) || got[3] != static_cast<uint32_t>(x)) &&
|
|
mismatches++ < 6)
|
|
std::printf(" lane block x=%d z=%d: got %u,%u,%u,%u\n", x, z, got[0], got[1], got[2], got[3]);
|
|
}
|
|
std::printf("block 32x1x8 tid/laneid: %s (CUDA result=%d, %zu of 256 lanes mismatched)\n",
|
|
result == 0 && mismatches == 0 ? "PASS" : "FAIL", result, mismatches);
|
|
return result == 0 && mismatches == 0;
|
|
}
|
|
|
|
// Rounds a double to the nearest f16 (ties to even), with overflow to
|
|
// infinity; exact for sums and products of two f16 values.
|
|
static uint16_t double_to_half_rne(double value)
|
|
{
|
|
if (std::isnan(value))
|
|
return 0x7e00u;
|
|
const uint16_t sign = std::signbit(value) ? 0x8000u : 0u;
|
|
double magnitude = std::fabs(value);
|
|
if (magnitude >= 65520.0)
|
|
return sign | 0x7c00u;
|
|
if (magnitude < std::ldexp(1.0, -14))
|
|
{
|
|
// Subnormal: units of 2^-24.
|
|
const double units = magnitude * 16777216.0;
|
|
double rounded = std::nearbyint(units); // default mode is ties-to-even
|
|
return sign | static_cast<uint16_t>(rounded);
|
|
}
|
|
int exponent;
|
|
double fraction = std::frexp(magnitude, &exponent); // magnitude = fraction * 2^exponent, fraction in [0.5,1)
|
|
double mantissa = std::nearbyint((fraction * 2.0 - 1.0) * 1024.0);
|
|
int biased = exponent - 1 + 15;
|
|
if (mantissa >= 1024.0)
|
|
{
|
|
mantissa = 0.0;
|
|
++biased;
|
|
}
|
|
if (biased >= 31)
|
|
return sign | 0x7c00u;
|
|
return sign | static_cast<uint16_t>((biased << 10) | static_cast<int>(mantissa));
|
|
}
|
|
|
|
static bool half_bits_equal(uint16_t got, uint16_t expected)
|
|
{
|
|
const bool gotNan = (got & 0x7c00u) == 0x7c00u && (got & 0x3ffu) != 0;
|
|
const bool expectedNan = (expected & 0x7c00u) == 0x7c00u && (expected & 0x3ffu) != 0;
|
|
if (gotNan || expectedNan)
|
|
return gotNan && expectedNan;
|
|
// min/max of +0 and -0 may return either zero.
|
|
if ((got & 0x7fffu) == 0 && (expected & 0x7fffu) == 0)
|
|
return true;
|
|
return got == expected;
|
|
}
|
|
|
|
static bool test_f16x2_fma_modes(Driver& driver)
|
|
{
|
|
std::mt19937 random(7);
|
|
constexpr int count = 8192;
|
|
std::vector<uint32_t> input(count * 3);
|
|
for (uint32_t& word : input)
|
|
{
|
|
auto random_half = [&]() -> uint16_t {
|
|
for (;;)
|
|
{
|
|
const uint16_t bits = static_cast<uint16_t>(random());
|
|
if ((bits & 0x7c00u) != 0x7c00u && std::fabs(half_to_float(bits)) < 4.0f)
|
|
return bits;
|
|
}
|
|
};
|
|
word = random_half() | (static_cast<uint32_t>(random_half()) << 16);
|
|
}
|
|
// NaN and infinity operands.
|
|
input[0] = 0x7e00u | (0xfc00u << 16);
|
|
input[1] = 0x3c00u | (0x3c00u << 16);
|
|
input[2] = 0x0000u | (0x0000u << 16);
|
|
std::vector<uint32_t> output(count * 3);
|
|
if (!driver.run("d4r_f16x2_fma_modes", input.data(), input.size() * 4, output.data(),
|
|
output.size() * 4, count / 256, 256, count))
|
|
return false;
|
|
const char* names[] = {"fma.rn.sat.f16x2", "fma.rn.relu.f16x2", "fma.rn.f16x2.relu"};
|
|
size_t mismatches[3] = {};
|
|
for (int thread = 0; thread < count; ++thread)
|
|
{
|
|
for (int half = 0; half < 2; ++half)
|
|
{
|
|
const uint16_t ah = static_cast<uint16_t>(input[thread * 3] >> (16 * half));
|
|
const uint16_t bh = static_cast<uint16_t>(input[thread * 3 + 1] >> (16 * half));
|
|
const uint16_t ch = static_cast<uint16_t>(input[thread * 3 + 2] >> (16 * half));
|
|
const double result = static_cast<double>(half_to_float(ah)) * half_to_float(bh) + half_to_float(ch);
|
|
const uint16_t rounded = double_to_half_rne(result);
|
|
const bool nan = std::isnan(result);
|
|
// .sat: clamp to [0, 1], NaN to +0. .relu: negative to +0, NaN stays NaN.
|
|
const uint16_t sat = nan ? 0 : (rounded & 0x8000u) ? 0 : half_to_float(rounded) > 1.0f ? 0x3c00u : rounded;
|
|
const uint16_t relu = nan ? 0x7fffu : (rounded & 0x8000u) ? 0 : rounded;
|
|
const uint16_t expected[3] = {sat, relu, relu};
|
|
for (int op = 0; op < 3; ++op)
|
|
{
|
|
const uint16_t got = static_cast<uint16_t>(output[thread * 3 + op] >> (16 * half));
|
|
if (!half_bits_equal(got, expected[op]) && mismatches[op]++ < 3)
|
|
std::printf(" %s half %d: a=0x%04x b=0x%04x c=0x%04x got 0x%04x expected 0x%04x\n",
|
|
names[op], half, ah, bh, ch, got, expected[op]);
|
|
}
|
|
}
|
|
}
|
|
for (int op = 0; op < 3; ++op)
|
|
std::printf("%s (%d lanes): %s (%zu mismatches)\n", names[op], count * 2,
|
|
mismatches[op] == 0 ? "PASS" : "FAIL", mismatches[op]);
|
|
return mismatches[0] == 0 && mismatches[1] == 0 && mismatches[2] == 0;
|
|
}
|
|
|
|
static uint32_t reference_prmt(uint32_t a, uint32_t b, uint32_t selector)
|
|
{
|
|
const uint64_t bytes = (static_cast<uint64_t>(b) << 32) | a;
|
|
uint32_t result = 0;
|
|
for (int index = 0; index < 4; ++index)
|
|
{
|
|
const uint32_t nibble = (selector >> (4 * index)) & 0xfu;
|
|
uint8_t value = static_cast<uint8_t>(bytes >> (8 * (nibble & 7u)));
|
|
if (nibble & 8u)
|
|
value = (value & 0x80u) ? 0xffu : 0x00u;
|
|
result |= static_cast<uint32_t>(value) << (8 * index);
|
|
}
|
|
return result;
|
|
}
|
|
|
|
static bool test_dltss_misc(Driver& driver)
|
|
{
|
|
std::mt19937 random(31);
|
|
uint32_t input[32][16] = {};
|
|
for (int lane = 0; lane < 32; ++lane)
|
|
{
|
|
for (int word = 0; word < 12; ++word)
|
|
input[lane][word] = static_cast<uint32_t>(random());
|
|
// Make some f16 halves compare equal, including +0 against -0.
|
|
if (lane % 3 == 0)
|
|
input[lane][7] = input[lane][4];
|
|
if (lane % 5 == 0)
|
|
input[lane][7] = (input[lane][7] & 0xffff0000u) | (input[lane][4] & 0xffffu);
|
|
if (lane == 1)
|
|
{
|
|
input[lane][4] = 0x80000000u;
|
|
input[lane][7] = 0x00000000u;
|
|
}
|
|
input[lane][12] = static_cast<uint32_t>(lane % 17);
|
|
}
|
|
uint32_t output[32][20] = {};
|
|
if (!driver.run("d4r_dltss_misc", input, sizeof(input), output, sizeof(output), 1, 32, -1))
|
|
return false;
|
|
// Shared tile: row r (16 bytes) comes from lane r's words 0..3. Matrix m row q
|
|
// is addressed by lane 8m+q, which points at tile row (5 * (8m+q)) mod 32.
|
|
auto tile_half = [&](int row, int column) -> uint32_t {
|
|
const uint32_t word = input[row][column / 2];
|
|
return (word >> (16 * (column % 2))) & 0xffffu;
|
|
};
|
|
const char* names[] = {"ldmatrix.x4", "ldmatrix.x2", "prmt 0x3424", "prmt 0x1404", "prmt 0x4431",
|
|
"prmt 0x7632", "prmt 0x5410", "prmt 0x7531", "set.equ.u32.f16x2",
|
|
"dp2a.lo.u32.u32", "cp.async.cg src-size"};
|
|
size_t mismatches[11] = {};
|
|
auto check = [&](int op, int lane, uint32_t got, uint32_t expected) {
|
|
if (got != expected && mismatches[op]++ < 3)
|
|
std::printf(" %s lane %d: got 0x%08x expected 0x%08x\n", names[op], lane, got, expected);
|
|
};
|
|
for (int lane = 0; lane < 32; ++lane)
|
|
{
|
|
const int group = lane / 4, pair = lane % 4;
|
|
for (int matrix = 0; matrix < 4; ++matrix)
|
|
{
|
|
const int row = (5 * (8 * matrix + group)) % 32;
|
|
const uint32_t expected = tile_half(row, 2 * pair) | (tile_half(row, 2 * pair + 1) << 16);
|
|
check(0, lane, output[lane][matrix], expected);
|
|
if (matrix < 2)
|
|
check(1, lane, output[lane][4 + matrix], expected);
|
|
}
|
|
const uint32_t a = input[lane][4], b = input[lane][5], c = input[lane][6], e = input[lane][7];
|
|
check(2, lane, output[lane][6], reference_prmt(a, 0, 0x3424));
|
|
check(3, lane, output[lane][7], reference_prmt(a, 0, 0x1404));
|
|
check(4, lane, output[lane][8], reference_prmt(a, 0, 0x4431));
|
|
check(5, lane, output[lane][9], reference_prmt(a, b, 0x7632));
|
|
check(6, lane, output[lane][10], reference_prmt(a, b, 0x5410));
|
|
check(7, lane, output[lane][11], reference_prmt(a, b, 0x7531));
|
|
uint32_t equal = 0;
|
|
for (int half = 0; half < 2; ++half)
|
|
{
|
|
const float x = half_to_float(static_cast<uint16_t>(a >> (16 * half)));
|
|
const float y = half_to_float(static_cast<uint16_t>(e >> (16 * half)));
|
|
// equ is the unordered comparison: NaN operands compare equal.
|
|
if (x == y || std::isnan(x) || std::isnan(y))
|
|
equal |= 0xffffu << (16 * half);
|
|
}
|
|
check(8, lane, output[lane][12], equal);
|
|
const uint32_t dot = c + (a & 0xffffu) * (b & 0xffu) + (a >> 16) * ((b >> 8) & 0xffu);
|
|
check(9, lane, output[lane][13], dot);
|
|
const uint32_t size = input[lane][12];
|
|
for (int word = 0; word < 4; ++word)
|
|
{
|
|
uint32_t expected = 0;
|
|
for (int byte = 0; byte < 4; ++byte)
|
|
if (static_cast<uint32_t>(word * 4 + byte) < size)
|
|
expected |= input[lane][8 + word] & (0xffu << (8 * byte));
|
|
check(10, lane, output[lane][14 + word], expected);
|
|
}
|
|
}
|
|
bool passed = true;
|
|
for (int op = 0; op < 11; ++op)
|
|
{
|
|
std::printf("%s: %s (%zu mismatches)\n", names[op], mismatches[op] == 0 ? "PASS" : "FAIL",
|
|
mismatches[op]);
|
|
passed &= mismatches[op] == 0;
|
|
}
|
|
return passed;
|
|
}
|
|
|
|
static float round_toward_zero_sum(float a, float b)
|
|
{
|
|
// Exact sum in double, then truncate toward zero to float.
|
|
const double exact = static_cast<double>(a) + static_cast<double>(b);
|
|
float nearest = static_cast<float>(exact);
|
|
if (std::fabs(static_cast<double>(nearest)) > std::fabs(exact))
|
|
nearest = std::nextafter(nearest, 0.0f);
|
|
return nearest;
|
|
}
|
|
|
|
static bool test_round_ops(Driver& driver)
|
|
{
|
|
std::mt19937 random(5);
|
|
std::uniform_real_distribution<float> coordinate(-40.0f, 40.0f);
|
|
constexpr int count = 4096;
|
|
std::vector<float> input(count * 2);
|
|
for (int index = 0; index < count; ++index)
|
|
{
|
|
// Mix exact quarter/half values (row mappings) with arbitrary ones.
|
|
const float quarter = static_cast<float>(static_cast<int>(random() % 321) - 160) * 0.25f;
|
|
input[index * 2] = index % 2 ? quarter : coordinate(random);
|
|
input[index * 2 + 1] = index % 3 ? 1e-8f * static_cast<float>(random() % 7) : coordinate(random);
|
|
}
|
|
std::vector<uint32_t> output(count * 9);
|
|
if (!driver.run("d4r_round_ops", input.data(), input.size() * 4, output.data(), output.size() * 4,
|
|
count / 256, 256, count))
|
|
return false;
|
|
const char* names[] = {"add.rz.ftz.f32", "round half away (add.rz+cvt.rzi.f32)", "cvt.rzi.ftz.s32.f32",
|
|
"shfl.sync.idx lane 0", "cvt.rzi.f32.f32", "add.ftz.f32", "shfl.sync pred",
|
|
"bfi.b32 pos 3 len 29", "bfi.b32 pos 4 len 8"};
|
|
size_t mismatches[9] = {};
|
|
auto bits = [](float value) { uint32_t word; std::memcpy(&word, &value, 4); return word; };
|
|
for (int index = 0; index < count; ++index)
|
|
{
|
|
const float a = input[index * 2], b = input[index * 2 + 1];
|
|
const float away = std::trunc(round_toward_zero_sum(a, std::copysign(0.5f, a)));
|
|
const uint32_t expected[9] = {bits(round_toward_zero_sum(a, b)), bits(away),
|
|
static_cast<uint32_t>(static_cast<int32_t>(away)),
|
|
static_cast<uint32_t>(index & ~31), bits(std::trunc(a)),
|
|
bits(a + b), 1u,
|
|
(static_cast<uint32_t>(index) & 7u) | ((static_cast<uint32_t>(index) & 8u) << 3),
|
|
((static_cast<uint32_t>(index) & 7u) & ~0xff0u) | ((static_cast<uint32_t>(index) & 0xffu) << 4)};
|
|
const uint32_t* row = &output[index * 9];
|
|
const uint32_t got[9] = {row[0], row[1], row[2], row[3], row[4], row[5], row[6], row[7], row[8]};
|
|
for (int op = 0; op < 9; ++op)
|
|
{
|
|
// Signed zeros compare equal.
|
|
const bool zeroes = op != 2 && op != 3 && op < 6 && (got[op] & 0x7fffffffu) == 0 &&
|
|
(expected[op] & 0x7fffffffu) == 0;
|
|
if (got[op] != expected[op] && !zeroes && mismatches[op]++ < 3)
|
|
std::printf(" %s: a=%.9g b=%.9g got 0x%08x expected 0x%08x\n", names[op], a, b, got[op],
|
|
expected[op]);
|
|
}
|
|
}
|
|
bool passed = true;
|
|
for (int op = 0; op < 9; ++op)
|
|
{
|
|
std::printf("%s: %s (%zu mismatches)\n", names[op], mismatches[op] == 0 ? "PASS" : "FAIL", mismatches[op]);
|
|
passed &= mismatches[op] == 0;
|
|
}
|
|
return passed;
|
|
}
|
|
|
|
static bool test_f16x2_ops(Driver& driver)
|
|
{
|
|
std::mt19937 random(99);
|
|
constexpr int count = 8192;
|
|
std::vector<uint32_t> input(count * 3);
|
|
auto random_half = [&]() -> uint16_t {
|
|
// Finite values across the normal and subnormal range, magnitudes below 64.
|
|
for (;;)
|
|
{
|
|
const uint16_t bits = static_cast<uint16_t>(random());
|
|
if ((bits & 0x7c00u) != 0x7c00u && std::fabs(half_to_float(bits)) < 64.0f)
|
|
return bits;
|
|
}
|
|
};
|
|
for (uint32_t& word : input)
|
|
word = random_half() | (static_cast<uint32_t>(random_half()) << 16);
|
|
std::vector<uint32_t> output(count * 13);
|
|
if (!driver.run("d4r_f16x2_ops", input.data(), input.size() * 4, output.data(), output.size() * 4,
|
|
count / 256, 256, count))
|
|
return false;
|
|
const char* names[] = {"abs.f16x2", "neg.f16x2", "min.f16x2", "max.f16x2", "add.f16x2",
|
|
"sub.f16x2", "mul.f16x2", "fma.rn.f16x2", "add.f16/mul.f16",
|
|
"nested .reg cvt 0.5", "nested .reg cvt 2.0", "cvt.rn.f16.f64",
|
|
"h2rsqrt (nested .reg)"};
|
|
size_t mismatches[13] = {};
|
|
for (int thread = 0; thread < count; ++thread)
|
|
{
|
|
const uint32_t a = input[thread * 3], b = input[thread * 3 + 1], c = input[thread * 3 + 2];
|
|
for (int half = 0; half < 2; ++half)
|
|
{
|
|
const uint16_t ah = static_cast<uint16_t>(a >> (16 * half));
|
|
const uint16_t bh = static_cast<uint16_t>(b >> (16 * half));
|
|
const uint16_t ch = static_cast<uint16_t>(c >> (16 * half));
|
|
const double av = half_to_float(ah), bv = half_to_float(bh), cv = half_to_float(ch);
|
|
const uint16_t expected[13] = {
|
|
static_cast<uint16_t>(ah & 0x7fffu),
|
|
static_cast<uint16_t>(ah ^ 0x8000u),
|
|
double_to_half_rne(av < bv ? av : bv),
|
|
double_to_half_rne(av > bv ? av : bv),
|
|
double_to_half_rne(av + bv),
|
|
double_to_half_rne(av - bv),
|
|
double_to_half_rne(av * bv),
|
|
double_to_half_rne(av * bv + cv),
|
|
half == 0 ? double_to_half_rne(av + bv) : double_to_half_rne(av * bv),
|
|
0x3800u,
|
|
0x4000u,
|
|
half == 0 ? double_to_half_rne(-0.55615234375) : 0x3c00u,
|
|
double_to_half_rne(1.0 / std::sqrt(std::fabs(av))),
|
|
};
|
|
for (int op = 0; op < 13; ++op)
|
|
{
|
|
const uint16_t got = static_cast<uint16_t>(output[thread * 13 + op] >> (16 * half));
|
|
bool ok = half_bits_equal(got, expected[op]);
|
|
if (op == 12 && !ok)
|
|
{
|
|
// rsqrt.approx: allow one f16 ulp; zero/subnormal inputs flush (.ftz) to +inf.
|
|
ok = std::abs(static_cast<int>(got) - static_cast<int>(expected[op])) <= 1 ||
|
|
((ah & 0x7c00u) == 0 && got == 0x7c00u);
|
|
}
|
|
if (!ok && mismatches[op]++ < 3)
|
|
std::printf(" %s half %d: a=0x%04x b=0x%04x c=0x%04x got 0x%04x expected 0x%04x\n",
|
|
names[op], half, ah, bh, ch, got, expected[op]);
|
|
}
|
|
}
|
|
}
|
|
bool passed = true;
|
|
for (int op = 0; op < 13; ++op)
|
|
{
|
|
std::printf("%s (%d lanes): %s (%zu mismatches)\n", names[op], count * 2,
|
|
mismatches[op] == 0 ? "PASS" : "FAIL", mismatches[op]);
|
|
passed &= mismatches[op] == 0;
|
|
}
|
|
return passed;
|
|
}
|
|
|
|
int main()
|
|
{
|
|
void* library = dlopen("libcuda.so", RTLD_NOW | RTLD_LOCAL);
|
|
if (library == nullptr)
|
|
{
|
|
std::fprintf(stderr, "dlopen(libcuda.so) failed: %s\n", dlerror());
|
|
return 1;
|
|
}
|
|
Driver driver{};
|
|
if (!driver.load(library))
|
|
{
|
|
std::fprintf(stderr, "ZLUDA libcuda.so is missing a required export\n");
|
|
return 1;
|
|
}
|
|
CUdevice device = 0;
|
|
CUcontext context = nullptr;
|
|
if (driver.init(0) != 0 || driver.deviceGet(&device, 0) != 0 ||
|
|
driver.contextCreate(&context, 0, device) != 0)
|
|
{
|
|
std::fprintf(stderr, "CUDA context setup failed\n");
|
|
return 1;
|
|
}
|
|
const CUresult loadResult = driver.moduleLoad(&driver.module, ptx);
|
|
if (loadResult != 0)
|
|
{
|
|
std::fprintf(stderr, "cuModuleLoadData failed: %d\n", loadResult);
|
|
return 1;
|
|
}
|
|
|
|
std::mt19937 random(1234);
|
|
const float smallValues[] = {-2.0f, -1.0f, -0.5f, 0.0f, 0.5f, 1.0f, 2.0f};
|
|
auto pick = [&]() { return smallValues[random() % 7]; };
|
|
|
|
int failures = 0;
|
|
failures += !test_cvt_from_e4m3(driver);
|
|
failures += !test_cvt_to_e4m3(driver);
|
|
failures += !test_cvt_to_e4m3(driver, "d4r_cvt_e4m3x2_relu_from_f16x2", "cvt.rn.satfinite.relu.e4m3x2.f16x2", true);
|
|
failures += !test_cvt_to_e4m3(driver, "d4r_cvt_e4m3x2_relu_from_f32", "cvt.rn.satfinite.relu.e4m3x2.f32", true);
|
|
failures += !test_movmatrix(driver);
|
|
failures += !test_global_v2_u16(driver);
|
|
failures += !test_ex2_f16x2(driver);
|
|
failures += !test_multidim_lane(driver);
|
|
failures += !test_f16x2_ops(driver);
|
|
failures += !test_f16x2_fma_modes(driver);
|
|
failures += !test_dltss_misc(driver);
|
|
failures += !test_round_ops(driver);
|
|
|
|
for (int trial = 0; trial < 3; ++trial)
|
|
{
|
|
std::vector<float> a(16 * 8), b(8 * 8), c(16 * 8);
|
|
for (float& value : a) value = pick();
|
|
for (float& value : b) value = pick();
|
|
for (float& value : c) value = pick();
|
|
char label[96];
|
|
std::snprintf(label, sizeof(label), "mma.m16n8k8.f16.f16 trial %d", trial);
|
|
failures += !test_mma(driver, "d4r_mma_m16n8k8_f16", label, 8, 16, layout_a_f16_k8,
|
|
layout_b_f16_k8, a, b, c, {}, {});
|
|
}
|
|
|
|
for (int trial = 0; trial < 3; ++trial)
|
|
{
|
|
std::vector<float> a(16 * 16), b(16 * 8), c(16 * 8);
|
|
for (float& value : a) value = pick();
|
|
for (float& value : b) value = pick();
|
|
for (float& value : c) value = pick();
|
|
char label[96];
|
|
std::snprintf(label, sizeof(label), "mma.m16n8k16.f16.f16 trial %d", trial);
|
|
failures += !test_mma(driver, "d4r_mma_m16n8k16_f16", label, 16, 16, layout_a_f16_k16,
|
|
layout_b_f16_k16, a, b, c, {}, {});
|
|
}
|
|
for (int trial = 0; trial < 4; ++trial)
|
|
failures += !test_mma_chain(driver, trial);
|
|
for (int trial = 0; trial < 3; ++trial)
|
|
{
|
|
std::vector<float> a(16 * 32), b(32 * 8), c(16 * 8);
|
|
std::vector<uint8_t> aCodes(a.size()), bCodes(b.size());
|
|
for (size_t index = 0; index < a.size(); ++index)
|
|
{
|
|
a[index] = pick();
|
|
aCodes[index] = float_to_e4m3_satfinite(a[index]);
|
|
}
|
|
for (size_t index = 0; index < b.size(); ++index)
|
|
{
|
|
b[index] = pick();
|
|
bCodes[index] = float_to_e4m3_satfinite(b[index]);
|
|
}
|
|
for (float& value : c) value = pick();
|
|
char label[96];
|
|
std::snprintf(label, sizeof(label), "mma.m16n8k32.f16.e4m3.e4m3 trial %d", trial);
|
|
failures += !test_mma(driver, "d4r_mma_m16n8k32_e4m3", label, 32, 8, layout_a_8bit_k32,
|
|
layout_b_8bit_k32, a, b, c, aCodes, bCodes);
|
|
}
|
|
for (int trial = 0; trial < 3; ++trial)
|
|
{
|
|
// Any finite code with |x| < 32 (so sums stay well inside f16).
|
|
std::mt19937 codes(99 + trial);
|
|
auto code = [&]() {
|
|
uint8_t value;
|
|
do
|
|
value = static_cast<uint8_t>(codes());
|
|
while ((value & 0x7f) == 0x7f || (value & 0x78) >= 0x50);
|
|
return value;
|
|
};
|
|
std::vector<uint8_t> aCodes(16 * 32), bCodes(32 * 8);
|
|
for (uint8_t& value : aCodes) value = code();
|
|
for (uint8_t& value : bCodes) value = code();
|
|
std::vector<float> c(16 * 8);
|
|
for (float& value : c) value = pick();
|
|
char label[96];
|
|
std::snprintf(label, sizeof(label), "mma.m16n8k32.e4m3 full-range codes trial %d", trial);
|
|
failures += !test_mma_e4m3_codes(driver, label, aCodes, bCodes, c);
|
|
}
|
|
{
|
|
// Subnormal A (m * 2^-9) times large B: a flushed denormal reads as 0.
|
|
std::mt19937 codes(7);
|
|
std::vector<uint8_t> aCodes(16 * 32), bCodes(32 * 8);
|
|
for (uint8_t& value : aCodes) value = static_cast<uint8_t>((codes() & 0x80) | (1 + codes() % 7));
|
|
for (uint8_t& value : bCodes) value = static_cast<uint8_t>((codes() & 0x80) | 0x70 | (codes() % 8)); // 256..480
|
|
for (uint8_t& value : bCodes) if ((value & 0x7f) == 0x7f) value &= 0xfe;
|
|
std::vector<float> c(16 * 8, 0.0f);
|
|
failures += !test_mma_e4m3_codes(driver, "mma.m16n8k32.e4m3 subnormal A x large B", aCodes, bCodes, c);
|
|
}
|
|
for (int trial = 0; trial < 4; ++trial)
|
|
failures += !test_mma_chain_e4m3(driver, trial);
|
|
failures += !test_mma_pair_alias(driver, false);
|
|
failures += !test_mma_pair_alias(driver, true);
|
|
std::printf("FP8/matrix instruction tests: %s (%d failing)\n", failures == 0 ? "PASS" : "FAIL",
|
|
failures);
|
|
return failures == 0 ? 0 : 1;
|
|
}
|