feat: add d4r inference engines and Windows parity tooling

Add native K/M/RR inference engines, HIP kernels, replay and validation
harnesses, and Windows DLSS setup and runtime support. Include GPU
profiling tools, build and packaging scripts, platform integration
patches, and implementation documentation.
This commit is contained in:
ayo committed 2026-10-16 08:54:47 -04:00
1 parent 25493b3198
commit 7c5d3c9157
430 files changed
+76191

No files matched your search

+79
View File
@@ -0,0 +1,79 @@
// Native part of engine/rr/enc0_m.comp (generated by engine/rr/gen_enc0_m.py): declarations, after the translated
// header. Binding 1 holds NGX's weight image followed by engine/rr/enc0_pack.py's blob (offsets in enc0_m_w.h already
// include its base); binding 2 is the arena.
#extension GL_EXT_shader_8bit_storage : require
#extension GL_KHR_shader_subgroup_vote : require
#extension GL_EXT_control_flow_attributes : require
#extension GL_GOOGLE_include_directive : require
#include "enc0_m_w.h"
#include "rr_exact.glsl"
#ifndef RR_GFX11
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
#endif
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 1, std430) readonly buffer W8U { uint8_t w[]; } w8u;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#ifndef RR_GFX11
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#endif
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CM gl_CooperativeMatrixLayoutColumnMajor
#ifdef RR_GFX11
#include "rr_gfx11_k32.glsl"
#else
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) Q8PAIRS(dst, src)
// One 32-deep step onto an accumulator that holds binary16 values. The reference (and the default here) adds the
// accumulator inside the f32 products and rounds the result to binary16. RR_MIX (build_native_model.sh with
// D4R_RR_MIX=1) keeps the accumulator as binary16, sums the products from zero and takes f16(P + acc), which the
// driver compiles to one v_fma_mix per element in place of two conversions: about 5% faster, but a single rounding
// where the reference has two, so a value can come out one binary16 step away (engine/rr/README.md).
// RAC is the accumulator type, RV(x) a binary16 value as its element; K32H is the step on a binary16 matrix.
#ifdef RR_MIX
#define RAC HAC
#define RV(x) (x)
#define K32(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_] + float((acc)[i_])); }
#define K32H(acc, a0, b0, a1, b1) K32(acc, a0, b0, a1, b1)
#else
#define RAC ACC
#define RV(x) float(x)
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) }
#define K32H(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(acc))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_]); }
#endif
#endif
const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0;
shared uint KVS[2 * 768]; // the current head's K, then V^T; before the attention the key pixels' features
#define XF KVS // the key pixels' 32 FP8 features (plane 0 | plane 1), slot-major
shared uint XT[96 * 8]; // the normalised key tokens
#define RR_KS KVS
#define RR_KS0 0u
#define RR_VS KVS
#define RR_VS0 768u
#define MG XT // the merge rows (16 x 128 FP8) reuse XT
shared float16_t O112[64 * 4]; // surface 112 per token, for surface 128's 2x2 sums
ivec2 keyXY(uint t, int x0, int y0)
{
if (t < 64u) return ivec2(x0 + int(t & 7u), y0 + int(t >> 3u));
const int i = int(t & 7u), side = int((t - 64u) >> 3u);
return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i));
}
#ifndef RR_GFX11
uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d)
{
return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u);
}
#endif
uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); }
// the translated prologue's shared buffer sh: input_noise (10 x 10 pixels, 4 f16) at byte 0, input_tensor (two planes
// of 10 x 10 pixels x 16 FP8 features) at byte 800
float16_t noiseAt(uint gx, uint gy, uint c)
{
const uint b = 80u * gy + 8u * gx + 2u * c;
return hb((sh[b >> 2] >> (8u * (b & 2u))) & 0xffffu);
}