ayo
7c5d3c9157
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.
80 lines
4.5 KiB
GLSL
80 lines
4.5 KiB
GLSL
// 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);
|
|
}
|