150 lines
8.3 KiB
Plaintext
150 lines
8.3 KiB
Plaintext
#version 460
|
|
// d4r engine, preset M: the stage after the network. Expands the last decoder's tokens (64 channels, FP8) to 17
|
|
// values per render pixel (one matrix product per 4x4-token block) and turns them into the reconstruction's
|
|
// per-pixel inputs: a recurrent four-channel feature blended with its reprojected history, three gate values,
|
|
// three softmax weights and a 2x2 covariance. See M_NOTES.md.
|
|
#extension GL_KHR_cooperative_matrix : require
|
|
#extension GL_KHR_memory_scope_semantics : require
|
|
#extension GL_KHR_shader_subgroup_basic : require
|
|
#extension GL_EXT_shader_explicit_arithmetic_types : require
|
|
#extension GL_EXT_float_e4m3 : require
|
|
#extension GL_GOOGLE_include_directive : require
|
|
#include "round_half.glsl"
|
|
#include "normalized_store.glsl"
|
|
#include "activation_m.glsl"
|
|
layout(local_size_x = 32) in; // one wave: 16 tokens = 8x8 render pixels, two pixels per lane
|
|
layout(set = 0, binding = 0, std430) readonly buffer Params
|
|
{
|
|
ivec2 size; // render size
|
|
int tokW; uint reset;
|
|
vec2 rndOffs; // jitter
|
|
vec2 rnd2hi; // render pixels -> pixels of the 3x image (the motion vectors' unit)
|
|
vec2 hiSize;
|
|
vec2 histOffs; // added to the history uv, in render pixels
|
|
float sigma, rho;
|
|
vec2 renderInv; // host-rounded texture coordinates, distinct from GPU rcp.approx offsets
|
|
} p;
|
|
layout(set = 0, binding = 1) uniform sampler2D motionTex; // texelFetch
|
|
layout(set = 0, binding = 2) uniform sampler2D featTex; // the previous feature, bilinear
|
|
layout(set = 0, binding = 3, rgba8_snorm) uniform writeonly image2D featOut;
|
|
layout(set = 0, binding = 4, rgba8) uniform writeonly image2D gateOut;
|
|
layout(set = 0, binding = 5, rgba8) uniform writeonly image2D mixOut;
|
|
layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D covOut;
|
|
layout(set = 0, binding = 7, std430) readonly buffer Tokens { floate4m3_t a[]; } tok; // [tokH][tokW][64]
|
|
layout(set = 0, binding = 8, std430) readonly buffer W8 { floate4m3_t w[]; } w8; // [64][80]: 20 columns per sub-position
|
|
layout(set = 0, binding = 9, std430) readonly buffer W16 { f16vec4 w[]; } w16; // bias, 16 equal rows of 80
|
|
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
|
|
#define HACC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
|
|
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
|
|
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
|
|
#define RM gl_CooperativeMatrixLayoutRowMajor
|
|
shared floate4m3_t TK[16 * 72];
|
|
shared f16vec4 EO[16 * 88 / 4];
|
|
float16_t q8(float16_t v) { return float16_t(floate4m3_t(clamp(v, float16_t(-448.0), float16_t(448.0)))); }
|
|
f16vec4 q8(f16vec4 v) { return f16vec4(q8(v.x), q8(v.y), q8(v.z), q8(v.w)); }
|
|
// NVIDIA's expansion GEMM is chained m16n8k32 mma with an f16 accumulator: the bias is the initial C and every
|
|
// 32 products are summed with one f16 rounding (swin_model.py mma_chain; the same k32 step the native kernels'
|
|
// SWIN_EXACT build implements). K32_ROUND reproduces that: after each pair of k16 cooperative products the
|
|
// accumulator is rounded to f16 through the shared EO region (ROUND_HALF_STORE).
|
|
#ifndef K32_ROUND
|
|
#define K32_ROUND 1
|
|
#endif
|
|
#if K32_ROUND
|
|
#define R32(m, column) ROUND_HALF_STORE(m, EO, (column) * 4u, 22u)
|
|
#else
|
|
#define R32(m, column)
|
|
#endif
|
|
#ifndef PLANES
|
|
#define PLANES 0
|
|
#endif
|
|
void main()
|
|
{
|
|
const uint lane = gl_LocalInvocationID.x, row = lane & 15u, hf = lane >> 4;
|
|
const ivec2 t0 = 4 * ivec2(gl_WorkGroupID.xy);
|
|
#if PLANES
|
|
// NGX's layout, two planes of 32 channels (the HIP network's output); the dispatch covers the token grid
|
|
const uint src = (hf * 16u * gl_NumWorkGroups.x * gl_NumWorkGroups.y + uint((t0.y + int(row >> 2)) * p.tokW + t0.x + int(row & 3u))) * 32u;
|
|
#else
|
|
const uint src = uint((t0.y + int(row >> 2)) * p.tokW + t0.x + int(row & 3u)) * 64u + hf * 32u;
|
|
#endif
|
|
for (uint c = 0u; c < 32u; ++c) TK[row * 72u + hf * 32u + c] = tok.a[src + c];
|
|
memoryBarrierShared(); barrier();
|
|
// Share each activation tile across three output tiles, then the remaining two.
|
|
// Five live accumulators cost more registers than the saved loads on gfx1201.
|
|
{
|
|
HACC b0; coopMatLoad(b0, w16.w, 0u, 20u, RM);
|
|
ACC a0 = ACC(b0);
|
|
HACC b1; coopMatLoad(b1, w16.w, 4u, 20u, RM);
|
|
ACC a1 = ACC(b1);
|
|
HACC b2; coopMatLoad(b2, w16.w, 8u, 20u, RM);
|
|
ACC a2 = ACC(b2);
|
|
for (uint kt = 0u; kt < 4u; ++kt)
|
|
{
|
|
A8 a; coopMatLoad(a, TK, kt * 16u, 72u, RM);
|
|
B8 w0; coopMatLoad(w0, w8.w, kt * 16u * 80u, 80u, RM);
|
|
a0 = coopMatMulAdd(a, w0, a0);
|
|
B8 w1; coopMatLoad(w1, w8.w, kt * 16u * 80u + 16u, 80u, RM);
|
|
a1 = coopMatMulAdd(a, w1, a1);
|
|
B8 w2; coopMatLoad(w2, w8.w, kt * 16u * 80u + 32u, 80u, RM);
|
|
a2 = coopMatMulAdd(a, w2, a2);
|
|
if ((kt & 1u) != 0u) { R32(a0, 0u); R32(a1, 1u); R32(a2, 2u); }
|
|
}
|
|
coopMatStore(HACC(a0), EO, 0u, 22u, RM);
|
|
coopMatStore(HACC(a1), EO, 4u, 22u, RM);
|
|
coopMatStore(HACC(a2), EO, 8u, 22u, RM);
|
|
}
|
|
{
|
|
HACC b0; coopMatLoad(b0, w16.w, 12u, 20u, RM);
|
|
ACC a0 = ACC(b0);
|
|
HACC b1; coopMatLoad(b1, w16.w, 16u, 20u, RM);
|
|
ACC a1 = ACC(b1);
|
|
for (uint kt = 0u; kt < 4u; ++kt)
|
|
{
|
|
A8 a; coopMatLoad(a, TK, kt * 16u, 72u, RM);
|
|
B8 w0; coopMatLoad(w0, w8.w, kt * 16u * 80u + 48u, 80u, RM);
|
|
a0 = coopMatMulAdd(a, w0, a0);
|
|
B8 w1; coopMatLoad(w1, w8.w, kt * 16u * 80u + 64u, 80u, RM);
|
|
a1 = coopMatMulAdd(a, w1, a1);
|
|
if ((kt & 1u) != 0u) { R32(a0, 3u); R32(a1, 4u); }
|
|
}
|
|
coopMatStore(HACC(a0), EO, 12u, 22u, RM);
|
|
coopMatStore(HACC(a1), EO, 16u, 22u, RM);
|
|
}
|
|
memoryBarrierShared(); barrier();
|
|
for (uint pi = lane; pi < 64u; pi += 32u)
|
|
{
|
|
const uint px = pi & 7u, py = pi >> 3;
|
|
const ivec2 at = 2 * t0 + ivec2(px, py);
|
|
if (at.x >= p.size.x || at.y >= p.size.y) continue;
|
|
const uint e = (((py >> 1) * 4u + (px >> 1)) * 88u + ((py & 1u) * 2u + (px & 1u)) * 20u) / 4u;
|
|
const f16vec4 v0 = q8(EO[e]), v1 = q8(EO[e + 1u]), v2 = q8(EO[e + 2u]), v3 = q8(EO[e + 3u]), v4 = q8(EO[e + 4u]);
|
|
// values: 0, 1, 10 gates; 2..5 feature; 6..9 its blend; 11..13 mix; 14, 15, 16 covariance
|
|
imageStore(gateOut, at, unorm8Value(vec4(sig16(v0.x), sig16(v0.y), sig16(v2.z), 0.0)));
|
|
const vec3 m = vec3(v2.w, v3.x, v3.y), ex = exp2((m - max(max(m.x, m.y), m.z)) * 1.442695);
|
|
imageStore(mixOut, at, unorm8Value(vec4(ex * (1.0 / (ex.z + (ex.x + ex.y))), 0.0)));
|
|
const float a = p.sigma * float(sig16(v3.z));
|
|
const float b = p.sigma * float(sig16(v3.w));
|
|
const float t = float(v4.x) * p.rho, s = t * t, m_ = abs(t);
|
|
const float big = m_ >= 9.0109138 ? 1.0 : fma(1.0 / (exp2(m_ * 2.88539) + 1.0), -2.0, 1.0);
|
|
const float rho = m_ >= 0.6 ? (t < 0.0 ? -big : big) : fma(fma(fma(fma(fma(0.01573968306183815, s, -0.05230396240949631), s, 0.13315297663211823), s, -0.33332768082618713), s, 0.0), t, t);
|
|
imageStore(covOut, at, vec4(a * a, b * b, a * b * rho, 0.0));
|
|
precise vec2 f = fma(vec2(at) + 0.5 - p.rndOffs, p.rnd2hi, texelFetch(motionTex, at, 0).xy);
|
|
f16vec4 hist = f16vec4(0.0);
|
|
if (p.reset == 0u && all(greaterThanEqual(f, vec2(0.0))) && all(lessThanEqual(f, p.hiSize))) {
|
|
precise vec2 normalized = fma(f, 1.0 / p.hiSize, p.histOffs * (1.0 / vec2(p.size)));
|
|
precise vec2 texels = normalized * vec2(p.size);
|
|
if (any(notEqual(textureSize(featTex, 0), p.size)))
|
|
texels = clamp(texels, vec2(0.5), vec2(p.size) - 0.5);
|
|
precise vec2 uv = texels * p.renderInv;
|
|
hist = f16vec4(roundHalfValue(textureLod(featTex, uv, 0.0)));
|
|
}
|
|
// the recurrent feature: gated mix of its reprojected history and the new value. The reference's PTX
|
|
// computes it with four separate f16 operations (sub, mul, mul, add:
|
|
// build/gfx1201-perf/ptx/dlss-0162-00.ptx:1416-1447), so `precise` keeps the compiler from fusing a
|
|
// product into an fma and dropping a rounding.
|
|
const f16vec4 g = gate16(f16vec4(v1.z, v1.w, v2.x, v2.y));
|
|
precise f16vec4 blend = (f16vec4(1.0) - g) * hist + g * tan16(f16vec4(v0.z, v0.w, v1.x, v1.y));
|
|
imageStore(featOut, at, snorm8Value(vec4(blend)));
|
|
}
|
|
}
|