wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add d4r inference engines and Windows parity tooling

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

engine/output_k.comp
+397
View File
@@ -0,0 +1,397 @@
#version 460
// d4r engine, preset K: the output stage (temporal accumulation and upsampling), written from the algorithm
// described in OUTPUT_K_NOTES.md. One 16x16 workgroup shades a 16x16 block of output pixels.
// 1. The block's footprint of the render-resolution colour (up to 24x24 texels) is range-compressed, converted to
// YCoCg and staged in shared memory (with depth, unless the motion is display-resolution).
// 2. Motion comes from the nearest of the diagonal depth neighbours (display-resolution: the pixel's own vector); it locates the history.
// 3. The 2x-resolution luma history is resampled with a sharpened Catmull-Rom filter at the four sub-pixels.
// 4. The current luma is filtered at the four sub-pixels with twelve taps of an anisotropic Gaussian whose shape
// comes from a table indexed by two network outputs.
// 5. History and current are blended per sub-pixel with a confidence from the tap geometry; chroma follows luma.
// Not handled (the reference kernel has them, the captured frames do not use them): blocks whose footprint exceeds
// 24x24 texels, and the alternative output modes selected by three flag bytes.
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_GOOGLE_include_directive : require
#include "output_store.glsl"
#include "round_half.glsl"
layout(local_size_x = 16, local_size_y = 16) in;
layout(set = 0, binding = 0, std430) readonly buffer Params
{
vec2 out2rnd, rndOffs; // render position = (pixel + .5) * out2rnd + rndOffs
vec2 near2rnd, sub2rnd; // the same scale for the nearest-texel lookup; for 2x sub-pixels
vec2 out2tok; // token column = int(pixel * out2tok) / 2
vec2 mvScale, mvOffs; // previous position = pixel + .5 + mvScale * (mv + mvOffs)
vec2 invOut, invOutValid; // 1 / output size, for the history uv and for the on-screen test
ivec2 colOrigin, colMax, depOrigin, depMax, mvOrigin, mvMax, outOrigin, outSize;
int tokW, tokH;
uint flags; // bit 0: reset; bit 1: regular depth (larger is farther), otherwise inverted;
// bit 2: display-resolution motion (one vector per output pixel, no depth);
// bits 8..11: which of the model's reconstruction tables this frame uses
float expo, expoRatio; // exposure scale of this frame; previous / current
float sharp, minAlpha, alphaThr;
float lutScale, kA, kB, c184, c186;
} p;
layout(set = 0, binding = 1) uniform sampler2D colorTex; // render resolution, linear RGB + alpha
layout(set = 0, binding = 2) uniform sampler2D histTex; // output resolution: compressed Y, Co, Cg, alpha (bilinear)
layout(set = 0, binding = 3) uniform sampler2D hiresTex; // 2x output resolution: compressed luma
layout(set = 0, binding = 4) uniform sampler2D motionTex;
layout(set = 0, binding = 5) uniform sampler2D depthTex;
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D histOut;
layout(set = 0, binding = 7, r16ui) uniform writeonly uimage2D hiresOut;
layout(set = 0, binding = 8, rgba16f) uniform writeonly image2D tokOut;
layout(set = 0, binding = 9, FINAL_OUTPUT_FORMAT) uniform writeonly image2D finalOut;
layout(set = 0, binding = 10, std430) readonly buffer Head { float16_t h[]; } head; // 40 values per token: 4 sub-blocks x 10
layout(set = 0, binding = 11, std430) readonly buffer Lut { uvec4 e[]; } lut;
shared f16vec4 tileC[576]; // compressed Y, Co, Cg and alpha of the footprint, 24 texels per row
shared float tileD[576];
// LDR (NGX's *_ldr_* kernels, the game did not set the HDR flag): display-referred colour, no range compression
#ifndef LDR
#define LDR 0
#endif
#if LDR
float compressK(float L) { return 1.0; }
float expandK(float v) { return 1.0; }
#else
// The reference's constants and association (output_color_pack): rcp.approx(y + 1), or above 1000
// rcp.approx(y) * (lg2.approx(fma(y, 9, 1)) * ln2 / 9 + 0xBC4FAE29). Its offset differs by 1.7e-5
// relative from 0.012676, and folding ln2/9 into one constant moves the result a further ulp.
float compressK(float L) // L * compressK(L) maps [0, inf) to [0, 1): Reinhard below 1000, logarithmic above
{
if (L < 1000.0)
return 1.0 / (L + 1.0);
precise const float t = (log2(fma(L, 9.0, 1.0)) * uintBitsToFloat(0x3F317218u)) / 9.0
+ uintBitsToFloat(0xBC4FAE29u);
return t / L;
}
float expandK(float v) // v * expandK(v) is the inverse
{
if (v < uintBitsToFloat(0x3F7FBE88u))
return 1.0 / (1.0 - v);
precise const float y = (v + uintBitsToFloat(0x3C4FAE29u)) * 9.0;
return (exp2(y * uintBitsToFloat(0x3FB8AA3Bu)) - 1.0) / (9.0 * v);
}
#endif
float roundHalfAway(float x) { const float t = trunc(x); return abs(x - t) >= 0.5 ? t + sign(x) : t; }
ivec2 clampOrigin(ivec2 v, ivec2 origin, ivec2 hi) // the reference's clamp: origin for v <= 0, else min(v + origin, hi)
{
return ivec2(v.x > 0 ? min(v.x + origin.x, hi.x) : origin.x, v.y > 0 ? min(v.y + origin.y, hi.y) : origin.y);
}
vec4 softIn(vec4 v) { return mix(0.99 + (v - 0.99) * 0.001, v, lessThan(v, vec4(0.99))); } // stored luma -> working luma
// Catmull-Rom weights of the four taps around a fraction t
void crWeights(float t, out float w0, out float w1, out float w2, out float w3)
{
const float t2 = t * t, t3 = t2 * t;
w0 = t2 - 0.5 * (t + t3);
w1 = 1.5 * t3 - 2.5 * t2 + 1.0;
w3 = 0.5 * (t3 - t2);
w2 = 1.0 - w0 - w1 - w3;
}
// The reference's six coefficients for one filter axis, computed once per pixel: the bilinear split
// mid = w1 / (w1 + w2) and the four sharpened weights, kept separate from the bilinear pair (folding
// sharp*(w1-mid) into mid changes the row result by up to a few ulp, which the half-encoded highlight
// range then amplifies).
struct CrAxis { float mid, split, s0, s1, s2, s3; };
CrAxis crAxis(float w0, float w1, float w2, float w3, float mid)
{
return CrAxis(mid, 1.0 - mid, p.sharp * w0, p.sharp * (w1 - mid), p.sharp * (w2 - (1.0 - mid)), p.sharp * w3);
}
// One pass along one axis: the two middle taps are bilinear, the four sharpened taps are saturated, and
// every term is the reference's own rounded product/fma.
float crPass(float t0, float t1, float t2, float t3, CrAxis c)
{
float t = c.split * t2;
t = fma(c.mid, t1, t);
t = fma(c.s0, clamp(t0, 0.0, 1.0), t);
t = fma(c.s1, clamp(t1, 0.0, 1.0), t);
t = fma(c.s2, clamp(t2, 0.0, 1.0), t);
t = fma(c.s3, clamp(t3, 0.0, 1.0), t);
return t;
}
void main()
{
const ivec2 pix = ivec2(gl_GlobalInvocationID.xy);
const vec2 pc = vec2(pix) + 0.5;
const bool displayMv = (p.flags & 4u) != 0u;
// ---------------------------------------------------------------- 1. footprint of the block
const vec2 b0 = (vec2(gl_WorkGroupID.xy * 16u) + 0.5) * p.out2rnd + p.rndOffs;
const vec2 b1 = (vec2(gl_WorkGroupID.xy * 16u + 15u) + 0.5) * p.out2rnd + p.rndOffs;
const ivec2 t0 = ivec2(vec2(roundHalfAway(b0.x), roundHalfAway(b0.y)) - 4.0);
const ivec2 t1 = ivec2(vec2(roundHalfAway(b1.x), roundHalfAway(b1.y)) + 4.0);
const int tw = t1.x - t0.x + 1, th = t1.y - t0.y + 1;
for (int i = int(gl_LocalInvocationIndex); i < tw * th; i += 256)
{
const ivec2 t = ivec2(i % tw, i / tw), g = t + t0;
const vec4 c = texelFetch(colorTex, clampOrigin(g, p.colOrigin, p.colMax), 0);
const vec3 rgb = max(c.rgb, 0.0);
const float Y = 0.25 * rgb.r + 0.5 * rgb.g + 0.25 * rgb.b;
const float L = p.expo * Y;
const float k = compressK(L);
tileC[t.y * 24 + t.x] = f16vec4(L * k, p.expo * (0.5 * rgb.r - 0.5 * rgb.b) * k, p.expo * (0.5 * rgb.g - 0.25 * rgb.r - 0.25 * rgb.b) * k, c.a);
// larger = nearer: regular depth is negated, which turns its "nearest" comparisons into the inverted ones
if (!displayMv)
{
const float depth = texelFetch(depthTex, clampOrigin(g, p.depOrigin, p.depMax), 0).r;
tileD[t.y * 24 + t.x] = (p.flags & 2u) != 0u ? -depth : depth;
}
}
memoryBarrierShared(); barrier();
// ---------------------------------------------------------------- network outputs of this pixel's 2x2 sub-block
const ivec2 tk = ivec2(vec2(pix) * p.out2tok);
const uint hbase = uint((tk.y / 2) * p.tokW + tk.x / 2) * 40u + uint((tk.y & 1) * 2 + (tk.x & 1)) * 10u;
// sigmoid of a head word (the reference's fma chain: y = x * (x^2 + 27) / fma(x^2, 9, 27), sat(0.5*y + 0.5))
const float x = 0.5 * float(head.h[hbase + 9u]);
precise const float x2 = x * x;
const float sig = clamp(x < -3.0 ? 0.0 : x > 3.0 ? 1.0
: fma((x * (x2 + 27.0)) / fma(x2, 9.0, 27.0), 0.5, 0.5), 0.0, 1.0);
// ---------------------------------------------------------------- 2. motion
// Render-resolution vectors come from the nearest diagonal depth neighbour; display-resolution ones belong to the pixel.
ivec2 mvAt = pix;
if (!displayMv)
{
const vec2 nr = pc * p.near2rnd + p.rndOffs;
const ivec2 ni = ivec2(roundHalfAway(nr.x), roundHalfAway(nr.y));
const int di = (ni.y - t0.y) * 24 + (ni.x - t0.x);
const float d0 = tileD[di], dmm = tileD[di - 25], dpp = tileD[di + 25], dpm = tileD[di - 23], dmp = tileD[di + 23];
ivec2 dsel = ivec2(0);
if (d0 < 0.5 * (dmm + dpp) - 0.001 * abs(d0) || d0 < 0.5 * (dpm + dmp) - 0.001 * abs(d0))
{
float m = d0;
if (dmm > m) { m = dmm; dsel = ivec2(-1, -1); }
if (dpm > m) { m = dpm; dsel = ivec2(1, -1); }
if (dmp > m) { m = dmp; dsel = ivec2(-1, 1); }
if (dpp > m) { dsel = ivec2(1, 1); }
}
mvAt = ni + dsel;
}
const vec2 mv = texelFetch(motionTex, clampOrigin(mvAt, p.mvOrigin, p.mvMax), 0).xy;
const vec2 q = pc + p.mvScale * (mv + p.mvOffs); // where this pixel was, in output pixels
const vec2 qv = q * p.invOutValid;
bool noHist = (p.flags & 1u) != 0u || qv.x < 0.0 || qv.x > 1.0 || qv.y < 0.0 || qv.y > 1.0;
const vec4 hist = textureLod(histTex, q * p.invOut, 0.0);
// ---------------------------------------------------------------- 3. luma history at the four sub-pixels
vec4 Hq;
{
const vec2 h = (q - 0.5) * 2.0 + 0.5;
const vec2 base = floor(h - 0.5) + 0.5;
const vec2 fr = clamp(h - base, 0.0, 1.0);
float x0, x1, x2, x3, y0, y1, y2, y3;
crWeights(fr.x, x0, x1, x2, x3); crWeights(fr.y, y0, y1, y2, y3);
const float mx = x1 / (x1 + x2), my = y1 / (y1 + y2);
const vec2 ts = p.invOut * 0.5;
// 5x5 texels T[col][row] from nine gathers (gather: x = (0,1), y = (1,1), z = (1,0), w = (0,0))
vec4 g00 = textureGather(hiresTex, ts * (base + vec2(-0.5, -0.5)), 0);
vec4 g10 = textureGather(hiresTex, ts * (base + vec2(1.5, -0.5)), 0);
vec4 g20 = textureGather(hiresTex, ts * (base + vec2(3.5, -0.5)), 0);
vec4 g01 = textureGather(hiresTex, ts * (base + vec2(-0.5, 1.5)), 0);
vec4 g11 = textureGather(hiresTex, ts * (base + vec2(1.5, 1.5)), 0);
vec4 g21 = textureGather(hiresTex, ts * (base + vec2(3.5, 1.5)), 0);
vec4 g02 = textureGather(hiresTex, ts * (base + vec2(-0.5, 3.5)), 0);
vec4 g12 = textureGather(hiresTex, ts * (base + vec2(1.5, 3.5)), 0);
vec4 g22 = textureGather(hiresTex, ts * (base + vec2(3.5, 3.5)), 0);
// stored luma at 0.99 and above is an encoding of [0.99, 1): rare enough to branch on
const vec4 gm = max(max(max(g00, g10), max(g20, g01)), max(max(g11, g21), max(max(g02, g12), g22)));
if (!(max(max(gm.x, gm.y), max(gm.z, gm.w)) < 0.99))
{
g00 = softIn(g00); g10 = softIn(g10); g20 = softIn(g20); g01 = softIn(g01); g11 = softIn(g11);
g21 = softIn(g21); g02 = softIn(g02); g12 = softIn(g12); g22 = softIn(g22);
}
float T[5][5]; // [row][col]
T[0][0] = g00.w; T[0][1] = g00.z; T[1][0] = g00.x; T[1][1] = g00.y;
T[0][2] = g10.w; T[0][3] = g10.z; T[1][2] = g10.x; T[1][3] = g10.y;
T[0][4] = g20.w; T[1][4] = g20.x;
T[2][0] = g01.w; T[2][1] = g01.z; T[3][0] = g01.x; T[3][1] = g01.y;
T[2][2] = g11.w; T[2][3] = g11.z; T[3][2] = g11.x; T[3][3] = g11.y;
T[2][4] = g21.w; T[3][4] = g21.x;
T[4][0] = g02.w; T[4][1] = g02.z; T[4][2] = g12.w; T[4][3] = g12.z; T[4][4] = g22.w;
// Each row filtered for the two sub-pixel columns. Stored luma lies in [0, 1), so the clamps of the
// sharpened taps do nothing here; the column pass below keeps them (rows can overshoot).
const CrAxis cx = crAxis(x0, x1, x2, x3, mx), cy = crAxis(y0, y1, y2, y3, my);
float R[5][2];
for (int r = 0; r < 5; ++r)
for (int c = 0; c < 2; ++c)
R[r][c] = crPass(T[r][c], T[r][c + 1], T[r][c + 2], T[r][c + 3], cx);
for (int sy = 0; sy < 2; ++sy)
for (int sx = 0; sx < 2; ++sx)
{
const float v = crPass(R[sy][sx], R[sy + 1][sx], R[sy + 2][sx], R[sy + 3][sx], cy);
const float lo = min(min(T[sy + 1][sx + 1], T[sy + 1][sx + 2]), min(T[sy + 2][sx + 1], T[sy + 2][sx + 2]));
const float hi = max(max(T[sy + 1][sx + 1], T[sy + 1][sx + 2]), max(T[sy + 2][sx + 1], T[sy + 2][sx + 2]));
Hq[sy * 2 + sx] = max(min(v, hi), lo);
}
// the history was compressed with last frame's exposure. The reference always re-exposes
// (v' = L * k(L) with L = ratio * (v * g(v))); skipping it when the ratio is exactly 1 leaves the
// pre-rounding value, which is not what k(g(v)) costs on a real frame.
for (int s = 0; s < 4; ++s)
{
const float L = p.expoRatio * (Hq[s] * expandK(Hq[s]));
Hq[s] = L * compressK(L);
}
}
{
precise const float chk = fma(sig, hist.b, fma(hist.r, hist.g, fma(Hq.x, Hq.y, Hq.z * Hq.w)));
noHist = noHist || !(abs(chk) < uintBitsToFloat(0x7F800000u));
}
const float s1 = noHist ? 1.0 : sig, ha = noHist ? 1.0 : hist.a;
precise const float f343 = fma(s1, 1.0 - ha, ha);
precise const float f344 = fma(ha, 1.0 - s1 * s1, s1 * s1);
if (noHist) Hq = vec4(0.0);
const vec3 hc = noHist ? vec3(0.0) : hist.rgb;
// ---------------------------------------------------------------- current colour, bilinear from the tile
// Positions, weights and sums of this part are half floats, as in the reference: the filter below is sensitive
// to where the samples sit, and half precision moves them by up to 1/64 texel.
const f16vec2 u = f16vec2(pc * p.out2rnd + p.rndOffs - vec2(t0));
const ivec2 ui = ivec2(floor(vec2(u)));
const f16vec2 uf = clamp(u - f16vec2(ui), float16_t(0.0), float16_t(1.0)), ug = float16_t(1.0) - uf;
const int ci = ui.y * 24 + ui.x;
f16vec4 curh = tileC[ci] * (ug.x * ug.y);
curh = fma(tileC[ci + 1], f16vec4(uf.x * ug.y), curh);
curh = fma(tileC[ci + 24], f16vec4(ug.x * uf.y), curh);
curh = fma(tileC[ci + 25], f16vec4(uf.x * uf.y), curh);
vec4 cur = vec4(curh);
// "bright" over the 2x2 pixel quad: the lanes of a subgroup are the pixels of two rows, x + 16 * (y & 1)
uint bright = cur.x > 0.8 ? 1u : 0u;
bright |= subgroupShuffleXor(bright, 16u);
bright |= subgroupShuffleXor(bright, 1u);
// ---------------------------------------------------------------- 4. kernel-predicted luma filter
const f16vec2 hm = f16vec2(head.h[hbase], head.h[hbase + 1u]);
const f16vec2 m = f16vec2(vec2(hm) / vec2(float16_t(1.0) + abs(hm)));
const ivec2 mi = clamp(ivec2(roundEven(vec2(m * float16_t(p.lutScale)))) + 8, 0, 15);
const uint fl = (bright << 2) | ((sig < 0.25 && f343 < 0.25) ? 2u : 0u) | (p.alphaThr > f343 ? 1u : 0u);
const uvec4 le = lut.e[(((p.flags >> 8) & 15u) << 11) | (fl << 8) | uint(mi.x << 4) | uint(mi.y)];
const float16_t e0 = uint16BitsToFloat16(uint16_t(le.x & 0xffffu)), e1 = uint16BitsToFloat16(uint16_t(le.x >> 16));
const float16_t e2 = uint16BitsToFloat16(uint16_t(le.y & 0xffffu)), slope = uint16BitsToFloat16(uint16_t(le.y >> 16));
const bool alongX = (le.z & 0xffffu) == 0u; // rows of taps step in y, taps within a row step in x
const float16_t e5 = uint16BitsToFloat16(uint16_t(le.z >> 16)), e6 = uint16BitsToFloat16(uint16_t(le.w & 0xffffu));
// sub-pixel positions in tile coordinates: s = sy * 2 + sx
// The reference fuses the scale and the offset and subtracts the tile origin afterwards. The order matters
// where a position falls on a half-float rounding tie (a whole output column at some size ratios): any other
// order moves those samples by one half-float step, which the recurrent history then spreads over the frame.
precise const vec2 sa0 = fma(vec2(pix * 2) + 0.5, p.sub2rnd, p.rndOffs), sa1 = fma(vec2(pix * 2 + 1) + 0.5, p.sub2rnd, p.rndOffs);
precise const vec2 sq0 = sa0 - vec2(t0), sq1 = sa1 - vec2(t0);
const f16vec2 sp0 = f16vec2(sq0), sp1 = f16vec2(sq1);
const f16vec4 subX = f16vec4(sp0.x, sp1.x, sp0.x, sp1.x), subY = f16vec4(sp0.y, sp0.y, sp1.y, sp1.y);
const f16vec4 subA = alongX ? subY : subX, subB = alongX ? subX : subY;
const float16_t cA = alongX ? u.y : u.x, cB = alongX ? u.x : u.y;
const int stepB = alongX ? 1 : 24, stepA = alongX ? 24 : 1;
const int a0 = int(floor(float(cA))) - 1;
// Every product and sum of the filter is rounded to half precision, as in the reference: with single precision
// the result differs visibly more (49 dB against 116 dB on the luma history of the test frames).
f16vec4 sumY = f16vec4(0.0), sumW = f16vec4(0.0);
const float16_t line0 = cB + (float16_t(a0) - cA) * slope;
float16_t line = line0;
for (int r = 0; r < 4; ++r)
{
const int bs = int(roundEven(float(line))) - 1;
const f16vec4 dA = float16_t(a0 + r) - subA;
f16vec4 db = float16_t(bs) - subB;
const f16vec4 cross = dA * e2;
f16vec4 lin = fma(db, cross, (dA * dA) * e0);
const int ti = (a0 + r) * stepA + bs * stepB;
for (int j = 0; j < 3; ++j)
{
// Keep the rounded exponent in f16 through exp2; avoids f16/f32 conversions.
const f16vec4 w = exp2(fma(db * db, f16vec4(e1), lin));
sumY = fma(w, f16vec4(tileC[ti + j * stepB].x), sumY); sumW += w;
db += float16_t(1.0); lin += cross;
}
line += slope;
}
vec4 F = vec4(sumY) / vec4(sumW);
// Where every exponent is below -12 the result is the tap with the largest exponent instead. That needs all
// twelve weights below 2^-12, so the sum of weights rules it out for almost every pixel.
if (any(lessThan(sumW, f16vec4(12.5 / 4096.0))))
{
f16vec4 eMax = f16vec4(-60000.0), yMax = f16vec4(0.0);
line = line0;
for (int r = 0; r < 4; ++r)
{
const int bs = int(roundEven(float(line))) - 1;
const f16vec4 dA = float16_t(a0 + r) - subA;
f16vec4 db = float16_t(bs) - subB;
const f16vec4 cross = dA * e2;
f16vec4 lin = fma(db, cross, (dA * dA) * e0);
const int ti = (a0 + r) * stepA + bs * stepB;
for (int j = 0; j < 3; ++j)
{
const f16vec4 E = fma(db * db, f16vec4(e1), lin);
const bvec4 gt = not(greaterThan(eMax, E)); // a later tap wins ties
eMax = mix(eMax, E, gt); yMax = mix(yMax, f16vec4(tileC[ti + j * stepB].x), gt);
db += float16_t(1.0); lin += cross;
}
line += slope;
}
F = mix(F, vec4(yMax), lessThan(eMax, f16vec4(-12.0)));
}
// The reference holds the filtered current luma in half precision: the blend at the end of the stage reads
// it as an f16 value widened to f32 (dlss-0046 converts every lane with cvt.rn.f16.f32 before the
// `fma bl, F - Hq, Hq` that forms N), and the normalisation and the low-exponent fallback also produce
// halves. Keeping it f32 pulls N a fraction of a half code towards the current luma wherever the history
// and the current frame differ, which the recurrent history then carries into the next frames.
F = roundHalfValue(F);
// ---------------------------------------------------------------- 5. blend
f16vec4 conf;
{
const f16vec4 dA = roundEven(subA) - subA;
const f16vec4 ln = fma(dA, f16vec4(slope), subB);
const f16vec4 dB = abs(roundEven(ln) - ln);
const float16_t kA = float16_t(alongX ? p.kB : p.kA), kB = float16_t(alongX ? p.kA : p.kB);
const f16vec4 sA = clamp(fma(abs(dA), f16vec4(kA), f16vec4(4.5)), float16_t(0.0), float16_t(1.0));
const f16vec4 sB = clamp(fma(dB, f16vec4(kB), f16vec4(4.5)), float16_t(0.0), float16_t(1.0));
conf = fma(sA, f16vec4(e5), f16vec4(e6)) * sB;
}
const float16_t hv = float16_t(f343), hv2 = float16_t(0.1) + float16_t(0.9) * hv;
const float16_t lo = (hv - float16_t(p.c184)) * float16_t(p.c186), hi = (hv2 - float16_t(p.c184)) * float16_t(p.c186);
const float16_t avg = float16_t(0.25) * ((conf.x + conf.y) + (conf.z + conf.w));
const float alphaC = float(clamp(fma(avg, hv2 - hi, hi), float16_t(0.0), float16_t(1.0)));
const vec4 bl = vec4(clamp(fma(conf, f16vec4(hv - lo), f16vec4(lo)), float16_t(0.0), float16_t(1.0)));
if (!(abs(cur.x * F.x * F.y * F.z * F.w) < uintBitsToFloat(0x7F800000u))) { cur.xyz = vec3(0.0); F = vec4(0.0); }
// a single-rounded fma, as the reference: fma.rn.ftz.f32 N, bl, F - Hq, Hq
const vec4 N = fma(bl, F - Hq, Hq);
// Native averages the four f32 blend registers, before the encoded half-surface stores.
{
const vec4 enc = mix(fma(N - 0.99, vec4(1000.0), vec4(0.99)), N, lessThan(N, vec4(0.99)));
const u16vec4 b = float16BitsToUint16(f16vec4(enc));
imageStore(hiresOut, pix * 2, uvec4(b.x)); imageStore(hiresOut, pix * 2 + ivec2(1, 0), uvec4(b.y));
imageStore(hiresOut, pix * 2 + ivec2(0, 1), uvec4(b.z)); imageStore(hiresOut, pix * 2 + ivec2(1, 1), uvec4(b.w));
}
// the reference fuses the history product: fma(hc, 1 - alphaC, alphaC * cur)
precise const vec3 C = fma(hc, vec3(1.0 - alphaC), alphaC * cur.xyz);
precise const float Nm = 0.25 * (((N.x + N.y) + N.z) + N.w);
const vec2 chroma = C.yz / max(C.x, 1e-5) * Nm;
{
// history colour: chroma limited so that the RGB it stands for stays non-negative
const float Y = max(Nm, 0.0);
float lim = Y;
if (chroma.y - 1e-5 < -Y) lim = 1e-5 - chroma.y;
if (chroma.x - chroma.y - 1e-5 < -lim) lim = chroma.y - chroma.x + 1e-5;
if (-chroma.x - chroma.y - 1e-5 < -lim) lim = chroma.x + chroma.y + 1e-5;
const float sc = lim <= Y ? 1.0 : Y / lim;
imageStore(histOut, pix, vec4(Y, chroma * sc, clamp(max(f344 / (f344 + 1.0), p.minAlpha), 0.0, 1.0)));
}
if (pix.x < p.outSize.x && pix.y < p.outSize.y)
{
// the reference multiplies by ginv(Nm) first and divides by the exposure afterwards; the other
// association (ginv / exposure first) rounds differently before the RGB reconstruction
const vec3 ycc = vec3(Nm, chroma) * expandK(Nm) / p.expo;
imageStore(finalOut, p.outOrigin + pix, finalOutputValue(vec4(ycc.x + ycc.y - ycc.z, ycc.x + ycc.z, ycc.x - ycc.y - ycc.z, cur.a)));
}
// ---------------------------------------------------------------- per-token feature for the next frame's input stage
if (pix.x < p.tokW && pix.y < p.tokH)
{
const uint tb = uint(pix.y * p.tokW + pix.x) * 40u;
precise vec4 ms;
for (uint s = 0u; s < 4u; ++s)
ms[s] = 0.2 * (float(head.h[tb + s * 10u + 3u]) + float(head.h[tb + s * 10u + 4u]) + float(head.h[tb + s * 10u + 5u])
+ float(head.h[tb + s * 10u + 6u]) + float(head.h[tb + s * 10u + 7u]));
// Match the native scalar FMA chain; dot() permits a different association at half-store boundaries.
precise float norm = fma(ms.x, ms.x, ms.y * ms.y);
norm = fma(ms.z, ms.z, norm);
norm = fma(ms.w, ms.w, norm);
norm += 1.0;
precise const vec4 feature = ms * inversesqrt(norm);
imageStore(tokOut, pix, feature);
}
}