wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
60da7b9019c712e40723eeab1c856579bd0522ea
d4r/kernels/tools/rrswin_enc0_engine_test.cpp
T
woutermans 50fa65d87b engine/rr: native Ray Reconstruction stages with capture/replay and shim integration 
- engine/rr: d4r's own shaders for the denoiser frame (exposure, six encoders,
  five decoders, hkpn output stage) at NGX layouts, plus PTX->GLSL/HIP packing
  scripts, reference implementations and weight extraction
- runtime.h/runtime_rr.inc, game_rr.{h,cpp}: GameDenoiser wired into the engine
  build (d4r-rr replay binary) and the D3D12 shim's feature-13 path
- shim: answer GetFeatureRequirements from d4r's own state only (the old
  cross-binary CUDA handoff crashed OptiScaler's scan); drop the core's
  CUDA GetFeatureRequirements import
- shim/harness: exact guide capture/replay of a frame's RR inputs
- kernels/rr: expanded enc0/engine prototypes; kernels/tools: ptxsim rewrite
  plus rrsim and an engine test; vulkan: ptx2glsl FP8/cooperative-matrix
  support, kvk inout buffers and KVK_DISABLE_OPTIMIZATION
- bridge: report the surface array in cuSurfObjectCreate traces
- packaging: RR settings and the ICU launch overrides in D4R_README
2026-10-12 12:47:53 +02:00

1014 lines
46 KiB
C++
Raw Blame History

// Steps 1-5 of the staged-image engine's enc0-port validation: each new primitive in
// kernels/rr/rrswin_engine.h run on the GPU with checked HIP calls and compared against
// an independently written host model.
//
// step 1 swin_slots_from_wimg -- the slot bytes directly, then through a m16n16
// swin_gemm_tile against a random staged image (the "from wimg" part only
// reshuffles layout, the GEMM proves the (row, k) map as the k16 test's step 4
// does); 32 rows, two m-tiles, so slot_step = n_rows is exercised
// step 2 the dcell exchange (writer, wave barrier, reader): bit-exact against the D
// cells the host model derives from the engine-tile values; a twelve-tile
// phase (n_mt = 2, n_nt = 6), also driven in the body's (2,2)-tile chunks, and
// an eight-tile phase (n_mt = 2, n_nt = 4)
// step 3 swin_stage_operand_p modes 0/1/2: the staged image bytes compared bit for
// byte against the host's own requantisation (an exhaustive u16->u8 e4m3 table
// built at startup from the documented satfinite semantics), the activation
// function at a one-ulp f16 bar per half (the k16 test's step 6 style)
// step 4 an integrated engine phase: per-column f16 C seeds for a 32-row phase in
// the column-major f16 blob form (§2), two m-tiles x six n-tiles, k32, pairs
// rounded once to f16 like the engine loop does
// step 5 the phase1->2 chain: D1 (per-column C), relu-staged (mode 1; the staged
// image is dumped and compared bit for byte), phase 2 on the staged image with
// its own per-column C, pairs rounded once to f16
//
// The host models are written from the contract (kernels/rr/rrswin_enc0_engine.md) and
// the translated body's documented layouts, not from the device code: the enc0 weight
// waddr, the fragment->(row, k) map, the dcell exchange's lds addresses and the
// satfinite-e4m3 rounding. They share no code with the device side.
//
// Build and run:
// hipcc --offload-arch=gfx1201 -DD4R_TEX_FP8 -DD4R_WMMA_LAYOUT=12 -O2 \
// -o /tmp/enc0_engine_test kernels/tools/rrswin_enc0_engine_test.cpp \
// && /tmp/enc0_engine_test
#include <hip/hip_runtime.h>
#include <cstdio>
#include <cstdint>
#include <cmath>
#include <cstring>
#include <cstdlib>
#include <vector>
#include "../rr/rrswin_engine.h"
#define CHECK(expr) \
do \
{ \
hipError_t e_ = (expr); \
if (e_ != hipSuccess) \
{ \
std::printf("HIP error at %s:%d: %s -> %s\n", __FILE__, __LINE__, #expr, \
hipGetErrorString(e_)); \
std::exit(2); \
} \
} while (0)
// ---------------------------------------------------------------------------
// independent host models -----------------------------------------------------------
// the enc0 weight-image layout, from the contract (rrswin_enc0.hip's "byte j of a
// lane's 16-byte weight tile holds (k, n) at addr(n, k) = ...")
static int waddr_host(int k, int n)
{
return 512 * (n >> 4) + 64 * (n & 7) + 16 * ((k >> 1) & 3) + 8 * ((n >> 3) & 1)
+ 4 * ((k >> 3) & 1) + 2 * ((k >> 4) & 1) + (k & 1);
}
// exact e4m3 decode: full f32 (subnormals and the NaN included), from the format's
// definition (3 mantissa bits, bias 7, subnormals 2^-9 apart)
static double dec(uint8_t c)
{
const uint32_t sign = c & 0x80u, e = (c >> 3) & 0xfu, m = c & 7u;
const double mag = e == 0 ? (m / 8.0) * std::pow(2.0, -6.0)
: (1.0 + m / 8.0) * std::pow(2.0, (int)e - 7.0);
return sign ? -mag : mag;
}
// f16 bits <-> value, host-side
__attribute__((device, host)) static inline _Float16 h_of(uint16_t b)
{
_Float16 h;
std::memcpy(&h, &b, sizeof(h));
return h;
}
__attribute__((device, host)) static inline uint16_t b_of(_Float16 h)
{
uint16_t b;
std::memcpy(&b, &h, sizeof(b));
return b;
}
// ---------------------------------------------------------------------------
// the RNE-satfinite e4m3 rounding table ---------------------------------------------
// For f16 bits: the exact f64 value, then, at every binary step up 2^-24..2^25, bits
// 0..7: bit 7 into code (the merged exponent covers normal and subnormal), the bit
// below the cut -> a half, anything below -> sticky. At the top of the walk the code
// saturates at 0x7e; at the bottom subnormals emit 2^-9-spaced codes. Overflow and
// NaN are the doc's two cases (0x7e and the whole-tab 0x7f).
static uint8_t tab[0x10000];
static void tab_build()
{
for (uint32_t b = 0; b < 0x10000u; ++b)
{
const uint16_t bits = (uint16_t)b;
const uint32_t mag = bits & 0x7fffu, sign = (bits >> 8) & 0x80u;
if (mag > 0x7c00u)
{
tab[b] = (uint8_t)(sign | 0x7fu); // NaN keeps its sign
continue;
}
if (mag >= 0x7c00u)
{
tab[b] = (uint8_t)(sign | 0x7eu); // inf: out of range -> max finite
continue;
}
if (mag == 0)
{
tab[b] = (uint8_t)sign; // +/-0 preserved
continue;
}
const int e = (int)(bits >> 10) & 0x1f, m = (int)(bits & 0x3ffu);
// the exact f64 of the f16 magnitude
double v;
if (e > 0)
v = (1.0 + m / 1024.0) * std::pow(2.0, e - 15.0);
else
v = (m / 1024.0) * std::pow(2.0, -14.0);
// the walk: the leading power of two (2^9 .. 2^-6) fixes the code's merged
// exponent, the three steps below it pick the mantissa bits, and the first
// dropped power of two is the rounding bit, the rest sticky. Under the whole
// subnormal grid (v < 2^-9) the same walk on the 2^-9-spaced subnormal codes
// decides, rounding on 2^-10.
uint32_t code = 0;
uint32_t round = 0, sticky = 0;
int base = 0;
int hit = 0;
for (int q = 9; q >= -6; --q) // the merged exponent, q >= -6
{
const double step = std::pow(2.0, q);
if (v >= step)
{
code = (uint32_t)((q + 7) << 3);
base = q;
v -= step;
hit = 1;
break;
}
}
if (hit)
{
for (int q = base - 1; q >= base - 3; --q)
{
const double step = std::pow(2.0, q);
if (v >= step)
{
v -= step;
code |= 1u << (q - (base - 3));
}
}
// the first dropped power of two is the rounding bit, the rest sticky
const double halfstep = std::pow(2.0, (double)(base - 4));
if (v >= halfstep)
round = 1;
if (v > halfstep)
sticky = 1;
}
else
{
// subnormal: codes m * 2^-9; walk the three steps 2^-7, 2^-8, 2^-9 (m bits
// 2,1,0), then the half (2^-10) and sticky
for (int q = -7; q >= -9; --q)
{
const double step = std::pow(2.0, q);
if (v >= step)
{
v -= step;
code |= 1u << (q + 9);
}
}
const double halfstep = std::pow(2.0, -10.0);
if (v >= halfstep)
round = 1;
if (v > halfstep)
sticky = 1;
}
if (round && (sticky || (code & 1u)))
++code; // RNE: tie -> even
if (code > 0x7eu)
code = 0x7eu;
tab[b] = (uint8_t)(sign | code);
}
}
// RNE-satfinite e4m3 from f16 bits, through the table (the device's
// cvt.rn.satfinite.e4m3x2.f16x2, low half -> low byte)
static uint8_t requant_code(uint16_t bits) { return tab[bits]; }
static uint8_t requant_code_bytes(double v) { return requant_code(b_of((_Float16)v)); }
// The requantisation the device performs (dec->f16->satfinite->dec, the k16 test's
// requant with the table behind it)
static double requant(double v) { return dec(requant_code_bytes(v)); }
// relu_keep_nan_halves per half, host model: negative or -0 becomes +0, NaN kept
static uint16_t relu_half(uint16_t bits)
{
return (bits & 0x8000u) && (bits & 0x7fffu) <= 0x7c00u ? 0u : bits;
}
// The clamped cubic activation per half, host model from the spec formula in
// rrswin_enc0.hip's comment, f16 throughout (each intermediate cast back), one-ulp bar
// for the packed-f16x2 device ops
static uint16_t act_model(uint16_t bits)
{
const _Float16 x = h_of(bits);
const _Float16 two = h_of(0x4000), slope = h_of(0x2d30), mid = h_of(0x3698), half0 = h_of(0x3800);
_Float16 y = x > -two ? x : -two;
y = y < two ? y : two;
y = y > -y ? y : -y;
const _Float16 g = (_Float16)(half0 + y * (mid - slope * y));
return b_of((_Float16)(x * g));
}
// ---------------------------------------------------------------------------
// step 1 kernels: the slots (dumped for the byte check), and one m16n16 GEMM per
// m-tile against a random staged image (tile (mt, nt) of a 32-row x 32-token block)
__global__ void step1_kernel(const uint8_t* wimg, uint8_t* slots_dump, const uint8_t* ximg,
float* out)
{
__shared__ uint8_t slots[2048];
__shared__ uint8_t img[32 * 32];
for (int i = threadIdx.x; i < 32 * 32; i += blockDim.x)
img[i] = ximg[i];
__syncthreads();
rrswin::swin_slots_from_wimg((lds_u8*)slots, wimg, 0, 32);
__syncthreads();
for (int i = threadIdx.x; i < 2048; i += blockDim.x)
slots_dump[i] = slots[i];
for (int mt = 0; mt < 2; ++mt)
{
f8v c = {0, 0, 0, 0, 0, 0, 0, 0};
c = rrswin::swin_gemm_tile(c, (const lds_u8*)slots, 16 * mt, 32, (const lds_u8*)img, 32, 0);
for (int i = 0; i < 8; ++i)
out[256 * mt + i * 32 + threadIdx.x] = c[i];
}
}
// ---------------------------------------------------------------------------
// step 2 kernel: the dcell exchange. The value source is random f32 D values, shaped
// as the lane's own [tiles][8] VGPRs (pair j covers VGPRs 2j, 2j + 1). The kernel packs
// pairs per tile, writes, barriers, reads, and dumps d flat as u32 words per lane.
// `chunked`: 0 = the whole twelve-tile buffer in one call, 1 = the body's (2,2)-tile
// chunk driver over the same [12][8] source (d's [6][2][4] accumulated through three
// chunks), 2 = the eight-tile phase (its own [8][8] source buffer).
__global__ void step2_kernel(const float* vin, uint32_t* out, int chunked)
{
__shared__ __attribute__((aligned(1024))) unsigned char sraw[12 * 128 * 4];
rrswin::lds_u32* const lds = (rrswin::lds_u32*)(uintptr_t)sraw;
const uint32_t l = threadIdx.x;
const uint32_t ntiles = chunked ? 12u : (uint32_t)(chunked == 0 ? 12u : 8u);
(void)ntiles;
// pack pairs per tile
if (chunked == 0 || chunked == 1)
{
const int tiles = chunked == 0 ? 12 : 12;
uint32_t pairs[12 * 4];
for (int tl = 0; tl < tiles; ++tl)
for (int j = 0; j < 4; ++j)
{
const _Float16 lo = (_Float16)vin[((size_t)l * tiles + tl) * 8 + 2 * j];
const _Float16 hi = (_Float16)vin[((size_t)l * tiles + tl) * 8 + 2 * j + 1];
uint16_t bl, bh;
std::memcpy(&bl, &lo, sizeof(bl));
std::memcpy(&bh, &hi, sizeof(bh));
pairs[tl * 4 + j] = (uint32_t)bl | ((uint32_t)bh << 16);
}
for (int i = 0; i < 12 * 128; ++i)
lds[i] = 0u;
__builtin_amdgcn_wave_barrier();
if (chunked == 0)
{
rrswin::swin_dcell_exchange_write(lds, pairs, 2, 6);
__builtin_amdgcn_wave_barrier();
uint32_t d[6 * 2 * 4];
rrswin::swin_dcell_exchange_read(lds, d, 6, 4, 6);
for (int i = 0; i < 48; ++i)
out[i * 32 + threadIdx.x] = d[i];
}
else
{
// the body's chunk shape: value source [2][6][4], d + 16c per chunk
uint32_t xp[16], d[6 * 2 * 4];
for (int c = 0; c < 3; ++c)
{
for (int mt = 0; mt < 2; ++mt)
for (int l2 = 0; l2 < 2; ++l2)
for (int p = 0; p < 4; ++p)
xp[(mt * 2 + l2) * 4 + p] = pairs[(mt * 6 + 2 * c + l2) * 4 + p];
rrswin::swin_dcell_exchange_write(lds, xp, 2, 2);
__builtin_amdgcn_wave_barrier();
rrswin::swin_dcell_exchange_read(lds, d + 16 * c, 2, 4, 2);
__builtin_amdgcn_wave_barrier();
}
for (int i = 0; i < 48; ++i)
out[i * 32 + threadIdx.x] = d[i];
}
}
else
{
// eight-tile phase: [2][4][4] pairs -> d[4][2][4]; the value source is laid out
// as [8][8] per lane, at the same (lane, tile, vgpr) offsets as the 12-tile one
uint32_t pairs[8 * 4];
for (int tl = 0; tl < 8; ++tl)
for (int j = 0; j < 4; ++j)
{
const _Float16 lo = (_Float16)vin[((size_t)l * 8 + tl) * 8 + 2 * j];
const _Float16 hi = (_Float16)vin[((size_t)l * 8 + tl) * 8 + 2 * j + 1];
uint16_t bl, bh;
std::memcpy(&bl, &lo, sizeof(bl));
std::memcpy(&bh, &hi, sizeof(bh));
pairs[tl * 4 + j] = (uint32_t)bl | ((uint32_t)bh << 16);
}
for (int i = 0; i < 8 * 128; ++i)
lds[i] = 0u;
__builtin_amdgcn_wave_barrier();
rrswin::swin_dcell_exchange_write(lds, pairs, 2, 4);
__builtin_amdgcn_wave_barrier();
uint32_t d[4 * 2 * 4];
rrswin::swin_dcell_exchange_read(lds, d, 4, 4, 4);
for (int i = 0; i < 32; ++i)
out[i * 32 + threadIdx.x] = d[i];
}
}
// ---------------------------------------------------------------------------
// step 3 kernel: swin_stage_operand_p's three modes into three images.
__global__ void step3_kernel(const uint32_t* pairs, uint8_t* img0, uint8_t* img1, uint8_t* img2)
{
__shared__ uint8_t s0[32 * 32], s1[32 * 32], s2[32 * 32];
for (int i = threadIdx.x; i < 32 * 32; i += blockDim.x)
{
s0[i] = 0;
s1[i] = 0;
s2[i] = 0;
}
__syncthreads();
const uint32_t* mine = pairs + threadIdx.x * 32;
for (int tile = 0; tile < 2; ++tile)
{
const uint32_t* p = mine + tile * 4;
rrswin::swin_stage_operand_p((lds_u8*)s0, 32, 16 * tile, 0, p, 0);
rrswin::swin_stage_operand_p((lds_u8*)s1, 32, 16 * tile, 0, p, 1);
rrswin::swin_stage_operand_p((lds_u8*)s2, 32, 16 * tile, 0, p, 2);
}
__syncthreads();
for (int i = threadIdx.x; i < 32 * 32; i += blockDim.x)
{
img0[i] = s0[i];
img1[i] = s1[i];
img2[i] = s2[i];
}
}
// ---------------------------------------------------------------------------
// step 4 kernel: the integrated engine phase. Per-column f16 C seeds at w+5120, two
// m-tiles x six n-tiles, k32, each tile's pairs rounded once to f16 (enc0_pack).
__global__ void step4_kernel(const uint8_t* wimg, uint8_t* slots_dump, const uint8_t* ximg,
const uint32_t* cseed, uint32_t* out)
{
__shared__ uint8_t slots[2048];
__shared__ uint8_t img[96 * 32];
for (int i = threadIdx.x; i < 96 * 32; i += blockDim.x)
img[i] = ximg[i];
__syncthreads();
rrswin::swin_slots_from_wimg((lds_u8*)slots, wimg, 0, 32);
__syncthreads();
for (int i = threadIdx.x; i < 2048; i += blockDim.x)
slots_dump[i] = slots[i];
const uint32_t h = threadIdx.x >> 4;
for (int mt = 0; mt < 2; ++mt)
for (int nt = 0; nt < 6; ++nt)
{
// the per-column C seed: c[i] = f32(f16 at w + 2*(i + 8h + 16mt)), §2
f8v c;
for (int i = 0; i < 8; ++i)
{
// the per-column blob's column 2*(i + 8h) + 32mt sits at byte 2*(...):
// two adjacent f16, in one u32 when the byte offset is even
const uint32_t w32 = cseed[(2 * (i + 8 * h + 16 * mt)) >> 2];
const uint16_t bits = (uint16_t)(((2 * (i + 8 * h + 16 * mt)) & 2) ? w32 >> 16
: w32 & 0xffffu);
c[i] = (float)h_of(bits);
}
c = rrswin::swin_gemm_tile(c, (const lds_u8*)slots, 16 * mt, 32, (const lds_u8*)img, 32, 16 * nt);
uint32_t pairs[4];
for (int p = 0; p < 4; ++p)
{
const _Float16 lo = (_Float16)c[2 * p], hi = (_Float16)c[2 * p + 1];
uint16_t bl, bh;
std::memcpy(&bl, &lo, sizeof(bl));
std::memcpy(&bh, &hi, sizeof(bh));
pairs[p] = (uint32_t)bl | ((uint32_t)bh << 16);
}
for (int p = 0; p < 4; ++p)
out[((mt * 6 + nt) * 4 + p) * 32 + threadIdx.x] = pairs[p];
}
}
// ---------------------------------------------------------------------------
// step 5 kernel: the phase1->2 chain. Phase 1: per-column C at w+2048, weights at
// w+4096; its pairs relu-staged (mode 1) into a fresh image; phase 2: per-column C at
// w+5120, weights at w+6144, on the staged image. The C seeds are blobs this test
// generates at w+2048 (phase 1) and w+5120 (phase 2), read column-major (§2).
__global__ void step5_kernel(const uint8_t* wimg, const uint8_t* ximg, const uint32_t* cseed1,
const uint32_t* cseed2, uint8_t* staged, uint32_t* out)
{
__shared__ uint8_t slots[2048];
__shared__ uint8_t img[96 * 32];
__shared__ uint8_t img2[96 * 32];
for (int i = threadIdx.x; i < 96 * 32; i += blockDim.x)
img[i] = ximg[i];
for (int i = threadIdx.x; i < 96 * 32; i += blockDim.x)
img2[i] = 0;
__syncthreads();
rrswin::swin_slots_from_wimg((lds_u8*)slots, wimg, 4096, 32);
__syncthreads();
uint32_t pairs1[2][6][4];
const uint32_t h = threadIdx.x >> 4;
for (int mt = 0; mt < 2; ++mt)
for (int nt = 0; nt < 6; ++nt)
{
f8v c;
for (int i = 0; i < 8; ++i)
{
const uint32_t w32 = cseed1[(2 * (i + 8 * h + 16 * mt)) >> 2];
const uint16_t bits = (uint16_t)(((2 * (i + 8 * h + 16 * mt)) & 2) ? w32 >> 16
: w32 & 0xffffu);
c[i] = (float)h_of(bits);
}
c = rrswin::swin_gemm_tile(c, (const lds_u8*)slots, 16 * mt, 32, (const lds_u8*)img, 32, 16 * nt);
for (int p = 0; p < 4; ++p)
{
const _Float16 lo = (_Float16)c[2 * p], hi = (_Float16)c[2 * p + 1];
uint16_t bl, bh;
std::memcpy(&bl, &lo, sizeof(bl));
std::memcpy(&bh, &hi, sizeof(bh));
pairs1[mt][nt][p] = (uint32_t)bl | ((uint32_t)bh << 16);
}
}
__syncthreads();
for (int mt = 0; mt < 2; ++mt)
for (int nt = 0; nt < 6; ++nt)
rrswin::swin_stage_operand_p((lds_u8*)img2, 32, 16 * nt, mt, pairs1[mt][nt], 1);
__syncthreads();
for (int i = threadIdx.x; i < 96 * 32; i += blockDim.x)
staged[i] = img2[i];
__syncthreads();
rrswin::swin_slots_from_wimg((lds_u8*)slots, wimg, 6144, 32);
__syncthreads();
for (int mt = 0; mt < 2; ++mt)
for (int nt = 0; nt < 6; ++nt)
{
f8v c;
for (int i = 0; i < 8; ++i)
{
const uint32_t w32 = cseed2[(2 * (i + 8 * h + 16 * mt)) >> 2];
const uint16_t bits = (uint16_t)(((2 * (i + 8 * h + 16 * mt)) & 2) ? w32 >> 16
: w32 & 0xffffu);
c[i] = (float)h_of(bits);
}
c = rrswin::swin_gemm_tile(c, (const lds_u8*)slots, 16 * mt, 32, (const lds_u8*)img2, 32, 16 * nt);
uint32_t pairs[4];
for (int p = 0; p < 4; ++p)
{
const _Float16 lo = (_Float16)c[2 * p], hi = (_Float16)c[2 * p + 1];
uint16_t bl, bh;
std::memcpy(&bl, &lo, sizeof(bl));
std::memcpy(&bh, &hi, sizeof(bh));
pairs[p] = (uint32_t)bl | ((uint32_t)bh << 16);
}
for (int p = 0; p < 4; ++p)
out[((mt * 6 + nt) * 4 + p) * 32 + threadIdx.x] = pairs[p];
}
}
// pair q's halves are VGPRs 2q, 2q + 1; VGPR v of a lane's half h is row v + 8h.
static int i_vgpr(int q, int half, int h) { return 2 * q + half + 8 * h; }
// ---------------------------------------------------------------------------
// main ---------------------------------------------------------------------------
// well-spread e4m3 codes (the k16 test's next_code: exponents 0..7 and mantissas, no
// NaN codes), so a wrong (row, column) mapping shows up as a real mismatch
static uint32_t lcg_main = 0x12345678u;
static uint8_t next_code()
{
lcg_main = lcg_main * 1664525u + 1013904223u;
uint32_t mag = (lcg_main >> 13) & 0x3fu; // exponents 0..7 and mantissas
if (mag == 0)
mag = 0x21;
const uint32_t sign = (lcg_main >> 7) & 1u;
return (uint8_t)((sign << 7) | mag);
}
int main()
{
tab_build();
// ---------------------------------------------------------------- step 1
{
std::vector<uint8_t> wimg(4096, 0);
for (auto& v : wimg)
v = next_code();
std::vector<uint8_t> ximg(32 * 32, 0);
for (auto& v : ximg)
v = next_code();
uint8_t *dwimg, *dslots, *dximg;
float* dout;
CHECK(hipMalloc(&dwimg, wimg.size()));
CHECK(hipMalloc(&dslots, 2048));
CHECK(hipMalloc(&dximg, ximg.size()));
CHECK(hipMalloc(&dout, 512 * sizeof(float)));
CHECK(hipMemset(dslots, 0, 2048));
CHECK(hipMemcpy(dwimg, wimg.data(), wimg.size(), hipMemcpyHostToDevice));
CHECK(hipMemcpy(dximg, ximg.data(), ximg.size(), hipMemcpyHostToDevice));
step1_kernel<<<1, 32>>>(dwimg, dslots, dximg, dout);
CHECK(hipGetLastError());
CHECK(hipDeviceSynchronize());
std::vector<uint8_t> slots(2048, 0);
CHECK(hipMemcpy(slots.data(), dslots, slots.size(), hipMemcpyDeviceToHost));
std::vector<float> got(512);
CHECK(hipMemcpy(got.data(), dout, got.size() * sizeof(float), hipMemcpyDeviceToHost));
// 1a: every slot byte = img[waddr(k, n)]
int bad = 0, checked = 0;
for (int bank = 0; bank < 64; ++bank)
{
// the slot layout is n-indexed: bank n holds k 0..15 (byte j = element
// (j, n), one full waddr each — the image's records do NOT store
// consecutive k contiguously) and bank 32 + n holds k 16..31
const int n = bank & 31;
const int k0 = bank < 32 ? 0 : 16;
for (int j = 0; j < 16; ++j)
{
const uint8_t have = slots[(size_t)bank * 16 + j];
const uint8_t want = wimg[waddr_host(k0 + j, n)];
++checked;
if (have != want)
{
if (bad < 4)
std::printf(" slot[h%d n%2d j%2d] = 0x%02x, expected 0x%02x\n",
bank >= 32, n, j, have, want);
++bad;
}
}
}
std::printf("step 1a slot bytes vs waddr: checked %d, mismatched %d\n", checked, bad);
// 1b: the GEMM through the slots, against the double-precision m16n16 model
// the device mma's accumulator layout: lane l, VGPR i holds D row i + 8h,
// column l & 15; the model's ref is the contract's own double sum for D[wr][tk]
int gbad = 0, gchecked = 0;
for (int mt = 0; mt < 2; ++mt)
for (int l = 0; l < 32; ++l)
{
const int hh = l >> 4, col = l & 15;
for (int i = 0; i < 8; ++i)
{
const int wr = 16 * mt + i + 8 * hh;
double ref = 0.0;
for (int k = 0; k < 32; ++k)
ref += dec(wimg[waddr_host(k, wr)]) * dec(ximg[(size_t)col * 32 + k]);
const float g = got[256 * mt + i * 32 + l];
++gchecked;
if (std::fabs((double)g - ref) > 2e-3 * (1.0 + std::fabs(ref)))
{
if (gbad < 4)
std::printf(" step1 GEMM tile %d lane %2d vgpr %d: got %+.6f ref %+.6f\n",
mt, l, i, g, ref);
++gbad;
}
}
}
std::printf("step 1b GEMM through the slots: checked %d, mismatched %d\n", gchecked, gbad);
CHECK(hipFree(dwimg));
CHECK(hipFree(dslots));
CHECK(hipFree(dximg));
CHECK(hipFree(dout));
if (bad || gbad)
std::exit(1);
}
// ---------------------------------------------------------------- step 2
{
// 12-tile phase: per-lane [12][8] random f32 D values; 8-tile: [8][8]
std::vector<float> v12(32 * 12 * 8), v8(32 * 8 * 8);
uint32_t lcg = 0x0f0f0f0fu;
auto nextr = [&lcg]() {
lcg = lcg * 1664525u + 1013904223u;
return (float)((int)(lcg >> 8) - 8388608) * (3.0f / 8388608.0f);
};
for (auto& v : v12)
v = nextr() * 60.0f;
for (auto& v : v8)
v = nextr() * 60.0f;
float *dv, *dv8;
uint32_t* dout;
CHECK(hipMalloc(&dv, v12.size() * 4));
CHECK(hipMalloc(&dv8, v8.size() * 4));
CHECK(hipMalloc(&dout, 48 * 32 * sizeof(uint32_t)));
CHECK(hipMemcpy(dv, v12.data(), v12.size() * 4, hipMemcpyHostToDevice));
CHECK(hipMemcpy(dv8, v8.data(), v8.size() * 4, hipMemcpyHostToDevice));
CHECK(hipMemset(dout, 0, 48 * 32 * sizeof(uint32_t)));
for (int pass = 0; pass < 3; ++pass)
{
const int chunked = pass == 0 ? 0 : (pass == 1 ? 1 : 2);
step2_kernel<<<1, 32>>>(chunked == 2 ? dv8 : dv, dout, chunked);
CHECK(hipGetLastError());
CHECK(hipDeviceSynchronize());
std::vector<uint32_t> got(48 * 32);
CHECK(hipMemcpy(got.data(), dout, got.size() * 4, hipMemcpyDeviceToHost));
const int tiles = chunked == 2 ? 8 : 12;
const float* src = chunked == 2 ? v8.data() : v12.data();
const int n_nt = chunked == 2 ? 4 : 6;
int bad = 0, checked = 0;
for (int lane = 0; lane < 32; ++lane)
{
const int g = (lane >> 2) & 7, t = lane & 3;
for (int mi = 0; mi < tiles / 2; ++mi) // n_mi
for (int rh = 0; rh < 2; ++rh)
for (int nt = 0; nt < 4; ++nt) // n_nt4
{
// the wanted cell is D[wr = 8nt + 2t + cb][tk = 16mi + 8rh + g]:
// tile (nt>>1, mi), lane 8rh + g + 16*(nt&1), VGPRs 2t, 2t+1
const int src_lane = 8 * rh + g + 16 * (nt & 1);
const int src_tile = (nt >> 1) * n_nt + mi;
const float vlo = src[((size_t)src_lane * tiles + src_tile) * 8 + 2 * t];
const float vhi = src[((size_t)src_lane * tiles + src_tile) * 8 + 2 * t + 1];
const uint32_t want = (uint32_t)b_of((_Float16)vlo)
| ((uint32_t)b_of((_Float16)vhi) << 16);
const uint32_t have = got[((mi * 2 + rh) * 4 + nt) * 32 + lane];
++checked;
if (have != want)
{
if (bad < 4)
std::printf(" dcell lane %2d mi %d rh %d nt %d: got %08x want %08x\n",
lane, mi, rh, nt, have, want);
++bad;
}
}
}
std::printf("step 2 %s exchange: checked %d, mismatched %d\n",
chunked == 0 ? "twelve-tile whole-buffer"
: (chunked == 1 ? "twelve-tile (2,2)-chunk" : "eight-tile"),
checked, bad);
if (bad)
std::exit(1);
}
CHECK(hipFree(dv));
CHECK(hipFree(dv8));
CHECK(hipFree(dout));
}
// ---------------------------------------------------------------- step 3
{
// every f16 of the whole 16-bit space (all 65536, incl. -0.0, negatives, NaN
// halves, values beyond +-2 and subnormals), shuffled into the 64 pairs per lane
// x 4 tiles the kernel stages (mode 0/1's bar is bit exact, so the whole space)
uint16_t allbits[65536];
for (uint32_t b = 0; b < 0x10000u; ++b)
allbits[b] = (uint16_t)b;
uint32_t lcg = 0x00c0ffeeu;
auto nextr = [&lcg]() {
lcg = lcg * 1664525u + 1013904223u;
return lcg >> 16;
};
for (uint32_t b = 0xffffu; b >= 1; --b)
{
const uint32_t j = nextr() % (b + 1);
const uint16_t t = allbits[b];
allbits[b] = allbits[j];
allbits[j] = t;
}
std::vector<uint32_t> pairs(32 * 32);
for (size_t i = 0; i < pairs.size(); ++i)
{
const uint16_t lo = allbits[2 * i];
const uint16_t hi = allbits[2 * i + 1];
pairs[i] = (uint32_t)lo | ((uint32_t)hi << 16);
}
uint32_t* dp;
uint8_t *di0, *di1, *di2;
CHECK(hipMalloc(&dp, pairs.size() * 4));
CHECK(hipMalloc(&di0, 32 * 32));
CHECK(hipMalloc(&di1, 32 * 32));
CHECK(hipMalloc(&di2, 32 * 32));
CHECK(hipMemcpy(dp, pairs.data(), pairs.size() * 4, hipMemcpyHostToDevice));
CHECK(hipMemset(di0, 0, 32 * 32));
CHECK(hipMemset(di1, 0, 32 * 32));
CHECK(hipMemset(di2, 0, 32 * 32));
step3_kernel<<<1, 32>>>(dp, di0, di1, di2);
CHECK(hipGetLastError());
CHECK(hipDeviceSynchronize());
std::vector<uint8_t> i0(32 * 32), i1(32 * 32), i2(32 * 32);
CHECK(hipMemcpy(i0.data(), di0, i0.size(), hipMemcpyDeviceToHost));
CHECK(hipMemcpy(i1.data(), di1, i1.size(), hipMemcpyDeviceToHost));
CHECK(hipMemcpy(i2.data(), di2, i2.size(), hipMemcpyDeviceToHost));
auto within_ulp = [](uint16_t a, uint16_t b) {
// compare as f16 by ordering the sign-magnitude representation
const auto key = [](uint16_t v) -> int {
return (v & 0x8000) ? (int)(0x8000 - (v & 0x7fff)) : (int)(v & 0x7fff);
};
const int d = key(a) - key(b);
return d >= -1 && d <= 1;
};
// mode 0 and mode 1: bit exact through the table
int bad0 = 0, bad1 = 0, checked = 0;
for (int tile = 0; tile < 2; ++tile)
for (int lane = 0; lane < 32; ++lane)
{
const int h = lane >> 4, col = lane & 15;
const uint32_t* p = pairs.data() + (size_t)(lane * 32) + tile * 4;
for (int q = 0; q < 4; ++q)
for (int half = 0; half < 2; ++half)
{
const uint16_t bits = (uint16_t)(half ? p[q] >> 16 : p[q] & 0xffffu);
const uint16_t re = relu_half(bits);
const uint8_t want0 = requant_code(bits);
const uint8_t want1 = requant_code(re);
// the staged byte sits at image row = row_base + (lane & 15),
// byte = 16*chunk + 8h + (VGPR = 2q + half)
const int rowi = 16 * tile + col;
const size_t off = (size_t)rowi * 32 + i_vgpr(q, half, h);
++checked;
if (i0[off] != want0)
{
if (bad0 < 4)
std::printf(" mode0 img row %2d byte %2d (tile %d lane %2d q %d h%d):"
" got 0x%02x want 0x%02x\n",
rowi, (int)off % 32, tile, lane, q, half,
i0[off], want0);
++bad0;
}
++checked;
if (i1[off] != want1)
{
if (bad1 < 4)
std::printf(" mode1 img row %2d byte %2d (tile %d lane %2d q %d h%d):"
" got 0x%02x want 0x%02x\n",
rowi, (int)off % 32, tile, lane, q, half,
i1[off], want1);
++bad1;
}
}
}
std::printf("step 3a stage modes 0/1 vs the table: checked %d, mode0 %d, mode1 %d mismatched\n",
checked, bad0, bad1);
// mode 2: the clamped cubic then encode, one ulp per half (the device applies
// the formula with packed f16x2 ops, the scalar _Float16 model can land one ulp
// away -- the k16 test's step 6 bar; a logic error moves whole clamps)
int bad2 = 0, checked2 = 0;
for (int tile = 0; tile < 2; ++tile)
for (int lane = 0; lane < 32; ++lane)
{
const int h = lane >> 4, col = lane & 15;
const uint32_t* p = pairs.data() + (size_t)(lane * 32) + tile * 4;
for (int q = 0; q < 4; ++q)
for (int half = 0; half < 2; ++half)
{
const uint16_t bits = (uint16_t)(half ? p[q] >> 16 : p[q] & 0xffffu);
const uint16_t acted = act_model(bits);
const uint8_t want2 = requant_code(acted);
const int rowi = 16 * tile + col;
const size_t off = (size_t)rowi * 32 + i_vgpr(q, half, h);
++checked2;
if (!within_ulp(i2[off], want2))
{
if (bad2 < 4)
std::printf(" mode2 img row %2d byte %2d: got 0x%02x want 0x%02x (act %04x)\n",
rowi, (int)off % 32, i2[off], want2, acted);
++bad2;
}
}
}
std::printf("step 3b stage mode 2 (1 ulp): checked %d, mismatched %d\n", checked2, bad2);
CHECK(hipFree(dp));
CHECK(hipFree(di0));
CHECK(hipFree(di1));
CHECK(hipFree(di2));
if (bad0 || bad1 || bad2)
std::exit(1);
}
// ---------------------------------------------------------------- step 4
{
// random weights in the waddr layout (32 rows x 32 k), random staged image
// (96 rows x 32 bytes), per-column f16 C seeds (32 columns, column c at byte 2c)
std::vector<uint8_t> wimg(4096, 0), ximg(96 * 32, 0);
for (int k = 0; k < 32; ++k)
for (int n = 0; n < 32; ++n)
wimg[waddr_host(k, n)] = next_code();
for (auto& v : ximg)
v = next_code();
std::vector<uint8_t> cseed(64, 0);
for (int c = 0; c < 32; ++c)
{
const uint16_t b = (uint16_t)(0x3c00u - (next_code() & 0xffu)); // +-small f16
cseed[2 * c] = (uint8_t)(b & 0xffu);
cseed[2 * c + 1] = (uint8_t)(b >> 8);
}
uint8_t *dw, *dx, *dsl;
uint32_t* dseed;
uint32_t* dout;
CHECK(hipMalloc(&dw, wimg.size()));
CHECK(hipMalloc(&dx, ximg.size()));
CHECK(hipMalloc(&dseed, cseed.size()));
CHECK(hipMalloc(&dsl, 2048));
CHECK(hipMalloc(&dout, 48 * 4 * 32 * 4));
CHECK(hipMemcpy(dw, wimg.data(), wimg.size(), hipMemcpyHostToDevice));
CHECK(hipMemcpy(dx, ximg.data(), ximg.size(), hipMemcpyHostToDevice));
CHECK(hipMemcpy(dseed, cseed.data(), cseed.size(), hipMemcpyHostToDevice));
CHECK(hipMemset(dout, 0, 48 * 4 * 32 * 4));
step4_kernel<<<1, 32>>>(dw, dsl, dx, dseed, dout);
CHECK(hipGetLastError());
CHECK(hipDeviceSynchronize());
std::vector<uint32_t> got(48 * 4 * 32);
CHECK(hipMemcpy(got.data(), dout, got.size() * 4, hipMemcpyDeviceToHost));
int bad = 0, checked = 0;
for (int mt = 0; mt < 2; ++mt)
for (int nt = 0; nt < 6; ++nt)
for (int lane = 0; lane < 32; ++lane)
{
const int h = lane >> 4, col = lane & 15;
for (int p = 0; p < 4; ++p)
for (int half = 0; half < 2; ++half)
{
// the engine cell D[wr = 16mt + 2p + half + 8h][tk = 16nt + col]:
// double sum over k of dec(w[waddr(k, wr)]) * dec(img[tk*32 + k])
// + f32(C seed f16), rounded once to f16
const int wr = 16 * mt + 2 * p + half + 8 * h;
const int tk = 16 * nt + col;
double ref = 0.0;
for (int k = 0; k < 32; ++k)
ref += dec(wimg[waddr_host(k, wr)]) * dec(ximg[(size_t)tk * 32 + k]);
const uint16_t seedbits =
(uint16_t)(cseed[2 * wr] | (cseed[2 * wr + 1] << 8));
ref += (double)h_of(seedbits);
const uint16_t want = b_of((_Float16)ref);
const uint16_t bits = (uint16_t)(half ? got[((mt * 6 + nt) * 4 + p) * 32 + lane] >> 16
: got[((mt * 6 + nt) * 4 + p) * 32 + lane] & 0xffffu);
++checked;
if (bits != want)
{
if (bad < 6)
std::printf(" D[wr %2d][tk %2d] (mt %d nt %d lane %2d p %d h%d):"
" got %04x want %04x\n",
wr, tk, mt, nt, lane, p, half, bits, want);
++bad;
}
}
}
std::printf("step 4 integrated phase with C seeds: checked %d, mismatched %d\n",
checked, bad);
CHECK(hipFree(dw));
CHECK(hipFree(dx));
CHECK(hipFree(dseed));
CHECK(hipFree(dsl));
CHECK(hipFree(dout));
if (bad)
std::exit(1);
}
// ---------------------------------------------------------------- step 5
{
// phase 1: per-column C at w+2048, weights at w+4096; phase 2: per-column C at
// w+5120, weights at w+6144, off the relu-staged image
std::vector<uint8_t> wimg(10240, 0), ximg(96 * 32, 0);
for (auto& v : ximg)
v = next_code();
for (int k = 0; k < 32; ++k)
for (int n = 0; n < 32; ++n)
wimg[4096 + waddr_host(k, n)] = next_code();
for (int k = 0; k < 32; ++k)
for (int n = 0; n < 32; ++n)
wimg[6144 + waddr_host(k, n)] = next_code();
std::vector<uint8_t> cs1(64, 0), cs2(64, 0);
for (int c = 0; c < 32; ++c)
{
const uint16_t b = (uint16_t)(0x3c00u - (next_code() & 0xffu));
cs1[2 * c] = (uint8_t)(b & 0xffu);
cs1[2 * c + 1] = (uint8_t)(b >> 8);
const uint16_t b2 = (uint16_t)(0x3c00u - (next_code() & 0xffu));
cs2[2 * c] = (uint8_t)(b2 & 0xffu);
cs2[2 * c + 1] = (uint8_t)(b2 >> 8);
}
uint8_t *dw, *dx, *dst;
uint32_t *ds1, *ds2;
uint32_t* dout;
CHECK(hipMalloc(&dw, wimg.size()));
CHECK(hipMalloc(&dx, ximg.size()));
CHECK(hipMalloc(&ds1, cs1.size()));
CHECK(hipMalloc(&ds2, cs2.size()));
CHECK(hipMalloc(&dst, 96 * 32));
CHECK(hipMalloc(&dout, 48 * 4 * 32 * 4));
CHECK(hipMemcpy(dw, wimg.data(), wimg.size(), hipMemcpyHostToDevice));
CHECK(hipMemcpy(dx, ximg.data(), ximg.size(), hipMemcpyHostToDevice));
CHECK(hipMemcpy(ds1, cs1.data(), cs1.size(), hipMemcpyHostToDevice));
CHECK(hipMemcpy(ds2, cs2.data(), cs2.size(), hipMemcpyHostToDevice));
CHECK(hipMemset(dst, 0, 96 * 32));
CHECK(hipMemset(dout, 0, 48 * 4 * 32 * 4));
step5_kernel<<<1, 32>>>(dw, dx, ds1, ds2, dst, dout);
CHECK(hipGetLastError());
CHECK(hipDeviceSynchronize());
std::vector<uint8_t> st(96 * 32, 0);
CHECK(hipMemcpy(st.data(), dst, st.size(), hipMemcpyDeviceToHost));
std::vector<uint32_t> got(48 * 4 * 32);
CHECK(hipMemcpy(got.data(), dout, got.size() * 4, hipMemcpyDeviceToHost));
// 5a: the staged image BIT-exactly: D1 cells -> relu per element -> RNE-satfinite
// e4m3 code per element
int sbad = 0, schecked = 0;
for (int mt = 0; mt < 2; ++mt)
for (int nt = 0; nt < 6; ++nt)
for (int lane = 0; lane < 32; ++lane)
{
const int h = lane >> 4, col = lane & 15;
for (int p = 0; p < 4; ++p)
for (int half = 0; half < 2; ++half)
{
const int wr = 16 * mt + 2 * p + half + 8 * h;
const int tk = 16 * nt + col;
double ref = 0.0;
for (int k = 0; k < 32; ++k)
ref += dec(wimg[4096 + waddr_host(k, wr)]) * dec(ximg[(size_t)tk * 32 + k]);
const uint16_t seedbits = (uint16_t)(cs1[2 * wr] | (cs1[2 * wr + 1] << 8));
ref += (double)h_of(seedbits);
uint16_t d1bits = b_of((_Float16)ref);
const uint16_t re = relu_half(d1bits);
const uint8_t want = requant_code(re);
const size_t off = (size_t)tk * 32 + (16 * mt + 2 * p + half + 8 * h);
++schecked;
if (st[off] != want)
{
if (sbad < 4)
std::printf(" staged[col %2d][k %2d] = 0x%02x, expected 0x%02x\n",
tk, wr, st[off], want);
++sbad;
}
}
}
std::printf("step 5a staged image (relu requant) bit exact: checked %d, mismatched %d\n",
schecked, sbad);
// 5b: phase 2's pairs off the staged image, with phase 2's own per-column C
int bad = 0, checked = 0;
for (int mt = 0; mt < 2; ++mt)
for (int nt = 0; nt < 6; ++nt)
for (int lane = 0; lane < 32; ++lane)
{
const int h = lane >> 4, col = lane & 15;
for (int p = 0; p < 4; ++p)
for (int half = 0; half < 2; ++half)
{
const int wr = 16 * mt + 2 * p + half + 8 * h;
const int tk = 16 * nt + col;
// phase 2's GEMM: sum over k of dec(w2[k][wr]) * the staged
// byte's decode (relu(D1 cell (k, tk)) requantised to e4m3,
// the operand the device's GEMM reads) + f32(C seed f16),
// rounded once to f16
double ref = 0.0;
for (int k = 0; k < 32; ++k)
{
double d1 = 0.0;
for (int j = 0; j < 32; ++j)
d1 += dec(wimg[4096 + waddr_host(j, k)]) * dec(ximg[(size_t)tk * 32 + j]);
const uint16_t seedbits = (uint16_t)(cs1[2 * k] | (cs1[2 * k + 1] << 8));
d1 += (double)h_of(seedbits);
const uint16_t re = relu_half(b_of((_Float16)d1));
ref += dec(wimg[6144 + waddr_host(k, wr)]) * dec(requant_code(re));
}
const uint16_t seed2bits = (uint16_t)(cs2[2 * wr] | (cs2[2 * wr + 1] << 8));
ref += (double)h_of(seed2bits);
const uint16_t want = b_of((_Float16)ref);
const uint16_t bits = (uint16_t)(half ? got[((mt * 6 + nt) * 4 + p) * 32 + lane] >> 16
: got[((mt * 6 + nt) * 4 + p) * 32 + lane] & 0xffffu);
++checked;
if (bits != want)
{
if (bad < 4)
std::printf(" D2[wr %2d][tk %2d] (mt %d nt %d lane %2d p %d h%d):"
" got %04x want %04x\n",
wr, tk, mt, nt, lane, p, half, bits, want);
++bad;
}
}
}
std::printf("step 5b phase 2 pairs off the staged image: checked %d, mismatched %d\n",
checked, bad);
CHECK(hipFree(dw));
CHECK(hipFree(dx));
CHECK(hipFree(ds1));
CHECK(hipFree(ds2));
CHECK(hipFree(dst));
CHECK(hipFree(dout));
if (sbad || bad)
std::exit(1);
}
return 0;
}
Reference in new issue View Git Blame Copy Permalink