Mirrors/totk-pc-recomp
1
0 Fork 0
mirror of https://github.com/valentderah/totk-pc-recomp.git synced 2026-11-11 01:41:02 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
dev
totk-pc-recomp/runtime/shaders/astc_decode.comp
T

430 lines
18 KiB
Plaintext
Raw Permalink Blame History

// ASTC LDR 2D decode to RGBA8 on the GPU, one invocation per texel. Ported from runtime/tex_astc.c (itself from LAS
// Recomp runtime/nvn/astc.c @ da35b83a, MIT, Copyright (c) 2026 Peiikai), so the output is the same bytes as
// astc_decode_ldr and astcenc's 8-bit decode: endpoints expand to 16 bits (linear v * 257, sRGB v << 8 | 0x80),
// texel = (e0 * (64 - w) + e1 * w + 32) >> 6, top 8 bits. Error blocks are magenta; HDR endpoint modes give magenta
// endpoints for their partition.
//
// A workgroup of 64 lanes takes up to four blocks of one block row (main below): every lane decodes its block's header,
// the block's lanes share the integer-sequence decode of its weights and colour values through shared memory, and
// each lane then interpolates its own texels.
// Compiled to runtime/astc_decode_spv.h by tools/build_shaders.py; vk_astc.c records the dispatch.
#version 450
layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
layout(std430, set = 0, binding = 0) readonly buffer Blocks { uvec4 blk[]; };
layout(std430, set = 0, binding = 1) writeonly buffer Texels { uint px[]; };
// uint 0..11: byte offsets of the weight unquantization rows (QUANT_2..QUANT_32), 12..27: of the colour rows
// (QUANT_6..QUANT_256 minus one); the bytes follow (vk_astc.c astc_tables).
layout(std430, set = 0, binding = 2) readonly buffer Tables { uint tab[]; };
layout(push_constant) uniform Params {
uint in_base; // first block of slice 0 (uvec4 index into blk)
uint out_base; // first texel of slice 0 (index into px)
uint width, height;
uint bw, bh; // block footprint
uint bx; // blocks per row
uint bslice; // blocks per slice
uint srgb;
} pc;
uint tab_byte(uint i) { return (tab[i >> 2] >> ((i & 3u) * 8u)) & 0xFFu; }
// n (<= 32) bits at pos; bits past 127 read as zero.
uint rbits(uvec4 b, int pos, int n) {
if (n == 0 || pos >= 128) return 0u;
int wi = pos >> 5, sh = pos & 31;
uint v = b[wi] >> uint(sh);
if (sh != 0 && wi < 3) v |= b[wi + 1] << uint(32 - sh);
return n >= 32 ? v : v & ((1u << uint(n)) - 1u);
}
// Quantization levels QUANT_2 .. QUANT_256 (0..20): bits, trits, quints of tex_astc.c q_bits / q_trits / q_quints.
int q_bits(int q) {
if (q < 2) return q == 0 ? 1 : 0;
int g = (q - 2) / 3, r = (q - 2) - 3 * g;
return r == 0 ? g + 2 : r == 1 ? g : g + 1;
}
bool q_trit(int q) { return q == 1 || (q >= 2 && (q - 2) % 3 == 2); }
bool q_quint(int q) { return q >= 2 && (q - 2) % 3 == 1; }
int ise_bitcount(int count, int q) {
int n = count * q_bits(q);
if (q_trit(q)) n += (8 * count + 4) / 5;
if (q_quint(q)) n += (7 * count + 2) / 3;
return n;
}
int trit_at(uint T, int k) {
uint C;
int t0, t1, t2, t3, t4;
if (((T >> 2) & 7u) == 7u) {
C = ((T >> 5) & 7u) << 2 | (T & 3u);
t4 = 2; t3 = 2;
} else {
C = T & 0x1Fu;
if (((T >> 5) & 3u) == 3u) { t4 = 2; t3 = int((T >> 7) & 1u); }
else { t4 = int((T >> 7) & 1u); t3 = int((T >> 5) & 3u); }
}
if ((C & 3u) == 3u) {
t2 = 2; t1 = int((C >> 4) & 1u);
t0 = int(((C >> 3) & 1u) << 1 | (((C >> 2) & 1u) & ~(C >> 3) & 1u));
} else if (((C >> 2) & 3u) == 3u) {
t2 = 2; t1 = 2; t0 = int(C & 3u);
} else {
t2 = int((C >> 4) & 1u); t1 = int((C >> 2) & 3u);
t0 = int(((C >> 1) & 1u) << 1 | ((C & 1u) & ~(C >> 1) & 1u));
}
return k == 0 ? t0 : k == 1 ? t1 : k == 2 ? t2 : k == 3 ? t3 : t4;
}
int quint_at(uint Q, int k) {
int q0, q1, q2;
if (((Q >> 1) & 3u) == 3u && ((Q >> 5) & 3u) == 0u) {
uint b0 = Q & 1u;
q2 = int(b0 << 2 | (((Q >> 4) & 1u) & ~b0 & 1u) << 1 | (((Q >> 3) & 1u) & ~b0 & 1u));
q1 = 4; q0 = 4;
} else {
uint C;
if (((Q >> 1) & 3u) == 3u) { q2 = 4; C = ((Q >> 3) & 3u) << 3 | (~(Q >> 5) & 3u) << 1 | (Q & 1u); }
else { q2 = int((Q >> 5) & 3u); C = Q & 0x1Fu; }
if ((C & 7u) == 5u) { q1 = 4; q0 = int((C >> 3) & 3u); }
else { q1 = int((C >> 3) & 3u); q0 = int(C & 7u); }
}
return k == 0 ? q0 : k == 1 ? q1 : q2;
}
// Raw ISE value idx (trit/quint << bits | bits) of a sequence of count values at bit start; tex_astc.c ise_decode
// reads the sequence front to back, this reads one value's group only. Bits of values past count are not read.
uint ise_at(uvec4 b, int start, int count, int q, int idx) {
int nb = q_bits(q);
if (q_trit(q)) {
// value k of a group: nb bits at k * nb + ts[k], then tb[k] trit bits; ts = {0, 2, 4, 5, 7}, tb = {2, 2, 1, 2, 1}
int g = idx / 5, k = idx - 5 * g, gs = start + g * (5 * nb + 8);
uint m = rbits(b, gs + k * nb + (k * 7 + 2) / 4, nb);
uint T = 0u;
for (int j = 0; j < 5; j++) {
int ts = (j * 7 + 2) / 4, tb = j == 2 || j == 4 ? 1 : 2;
if (5 * g + j < count) T |= rbits(b, gs + (j + 1) * nb + ts, tb) << uint(ts);
}
return uint(trit_at(T, k)) << uint(nb) | m;
}
if (q_quint(q)) {
// qs = {0, 3, 5}, qb = {3, 2, 2}
int g = idx / 3, k = idx - 3 * g, gs = start + g * (3 * nb + 7);
uint m = rbits(b, gs + k * nb + (k == 0 ? 0 : k == 1 ? 3 : 5), nb);
uint Q = 0u;
for (int j = 0; j < 3; j++) {
int qs = j == 0 ? 0 : j == 1 ? 3 : 5, qb = j == 0 ? 3 : 2;
if (3 * g + j < count) Q |= rbits(b, gs + (j + 1) * nb + qs, qb) << uint(qs);
}
return uint(quint_at(Q, k)) << uint(nb) | m;
}
return rbits(b, start + idx * nb, nb);
}
uint hash52(uint p) {
p ^= p >> 15; p -= p << 17; p += p << 7; p += p << 4;
p ^= p >> 5; p += p << 16; p ^= p >> 7; p ^= p >> 3;
p ^= p << 6; p ^= p >> 17;
return p;
}
int part_of(int seed, int count, int x, int y) {
seed += (count - 1) * 1024;
uint rnum = hash52(uint(seed));
int sh1, sh2;
if ((seed & 1) != 0) { sh1 = (seed & 2) != 0 ? 4 : 5; sh2 = count == 3 ? 6 : 5; }
else { sh1 = count == 3 ? 6 : 5; sh2 = (seed & 2) != 0 ? 4 : 5; }
int s[8];
for (int i = 0; i < 8; i++) {
int v = int((rnum >> uint(4 * i)) & 0xFu);
s[i] = (v * v) >> ((i & 1) != 0 ? sh2 : sh1);
}
int a = int((uint(s[0] * x + s[1] * y) + (rnum >> 14)) & 0x3Fu);
int b = int((uint(s[2] * x + s[3] * y) + (rnum >> 10)) & 0x3Fu);
int c = count >= 3 ? int((uint(s[4] * x + s[5] * y) + (rnum >> 6)) & 0x3Fu) : 0;
int d = count >= 4 ? int((uint(s[6] * x + s[7] * y) + (rnum >> 2)) & 0x3Fu) : 0;
if (a >= b && a >= c && a >= d) return 0;
if (b >= c && b >= d) return 1;
if (c >= d) return 2;
return 3;
}
int clamp255(int v) { return clamp(v, 0, 255); }
// v: the partition's unquantized colour values. False for an HDR mode.
bool unpack_endpoints(int cem, int v[8], out ivec4 e0, out ivec4 e1) {
e0 = ivec4(0); e1 = ivec4(0);
switch (cem) {
case 0:
e0 = ivec4(v[0], v[0], v[0], 255); e1 = ivec4(v[1], v[1], v[1], 255);
return true;
case 1: {
int l0 = v[0] >> 2 | (v[1] & 0xC0);
int l1 = min(l0 + (v[1] & 0x3F), 255);
e0 = ivec4(l0, l0, l0, 255); e1 = ivec4(l1, l1, l1, 255);
return true;
}
case 4:
e0 = ivec4(v[0], v[0], v[0], v[2]); e1 = ivec4(v[1], v[1], v[1], v[3]);
return true;
case 5: {
int l0 = v[0], l1 = v[1], a0 = v[2], a1 = v[3];
l0 |= (l1 & 0x80) << 1; a0 |= (a1 & 0x80) << 1;
l1 &= 0x7F; a1 &= 0x7F;
if ((l1 & 0x40) != 0) l1 -= 0x80;
if ((a1 & 0x40) != 0) a1 -= 0x80;
l0 >>= 1; a0 >>= 1; l1 >>= 1; a1 >>= 1;
l1 = clamp255(l1 + l0); a1 = clamp255(a1 + a0);
e0 = ivec4(l0, l0, l0, a0); e1 = ivec4(l1, l1, l1, a1);
return true;
}
case 6:
case 10: {
int s = v[3];
e1 = ivec4(v[0], v[1], v[2], 255);
e0 = ivec4(v[0] * s >> 8, v[1] * s >> 8, v[2] * s >> 8, 255);
if (cem == 10) { e0.w = v[4]; e1.w = v[5]; }
return true;
}
case 8:
case 12: {
ivec4 a = ivec4(v[0], v[2], v[4], cem == 12 ? v[6] : 0);
ivec4 b = ivec4(v[1], v[3], v[5], cem == 12 ? v[7] : 0);
if (a.x + a.y + a.z > b.x + b.y + b.z) {
e0 = b; e1 = a;
e0.x = (e0.x + e0.z) >> 1; e0.y = (e0.y + e0.z) >> 1;
e1.x = (e1.x + e1.z) >> 1; e1.y = (e1.y + e1.z) >> 1;
} else {
e0 = a; e1 = b;
}
if (cem == 8) { e0.w = 255; e1.w = 255; }
return true;
}
case 9:
case 13: {
ivec4 a = ivec4(0), b = ivec4(0);
for (int c = 0; c < 4; c++) {
if (c == 3 && cem != 13) break;
int base = v[2 * c], delta = v[2 * c + 1];
a[c] = base >> 1 | (delta & 0x80);
int d = (delta >> 1) & 0x3F;
if ((d & 0x20) != 0) d -= 0x40;
b[c] = d;
}
int sum = b.x + b.y + b.z;
b += a;
if (sum < 0) {
e0 = b; e1 = a;
e0.x = (e0.x + e0.z) >> 1; e0.y = (e0.y + e0.z) >> 1;
e1.x = (e1.x + e1.z) >> 1; e1.y = (e1.y + e1.z) >> 1;
} else {
e0 = a; e1 = b;
}
e0 = clamp(e0, ivec4(0), ivec4(255)); e1 = clamp(e1, ivec4(0), ivec4(255));
if (cem == 9) { e0.w = 255; e1.w = 255; }
return true;
}
default:
return false;
}
}
const uint MAGENTA = 0xFFFF00FFu;
// What every texel of a block needs from its header (tex_astc.c astc_decode_block up to the weight infill).
struct Hdr {
uint konst; // the colour of a void-extent or error block
bool solid; // konst is the whole block
int wx, wy, dual, quant, wtotal, pcount, cem0, cem1, cem2, cem3, n0, n1, n2, ccount, cq, ccs;
};
Hdr header(uvec4 b, int bw, int bh) {
Hdr H;
H.solid = true; H.konst = MAGENTA;
H.wx = 0; H.wy = 0; H.dual = 0; H.quant = 0; H.wtotal = 0; H.pcount = 1;
H.cem0 = 0; H.cem1 = 0; H.cem2 = 0; H.cem3 = 0; H.n0 = 0; H.n1 = 0; H.n2 = 0; H.ccount = 0; H.cq = 20; H.ccs = -1;
uint mode = b.x & 0x7FFu;
if ((mode & 0x1FFu) == 0x1FCu) { // void extent
if ((mode & 0x200u) != 0u || ((b.x >> 10) & 3u) != 3u) return H;
uint s0 = rbits(b, 12, 13), s1 = rbits(b, 25, 13), t0 = rbits(b, 38, 13), t1 = rbits(b, 51, 13);
bool all_ones = s0 == 0x1FFFu && s1 == 0x1FFFu && t0 == 0x1FFFu && t1 == 0x1FFFu;
if ((s0 >= s1 || t0 >= t1) && !all_ones) return H;
H.konst = (b.z >> 8 & 0xFFu) | (b.z >> 24) << 8 | (b.w >> 8 & 0xFFu) << 16 | (b.w >> 24) << 24;
return H;
}
// block mode (tex_astc.c block_mode_2d)
uint m = mode;
uint r = (m >> 4) & 1u, h = (m >> 9) & 1u, d = (m >> 10) & 1u, a = (m >> 5) & 3u;
int wx, wy;
if ((m & 3u) != 0u) {
r |= (m & 3u) << 1;
uint bb = (m >> 7) & 3u;
uint sel = (m >> 2) & 3u;
if (sel == 0u) { wx = int(bb) + 4; wy = int(a) + 2; }
else if (sel == 1u) { wx = int(bb) + 8; wy = int(a) + 2; }
else if (sel == 2u) { wx = int(a) + 2; wy = int(bb) + 8; }
else {
bb &= 1u;
if ((m & 0x100u) != 0u) { wx = int(bb) + 2; wy = int(a) + 2; }
else { wx = int(a) + 2; wy = int(bb) + 6; }
}
} else {
r |= ((m >> 2) & 3u) << 1;
if (((m >> 2) & 3u) == 0u) return H;
uint bb = (m >> 9) & 3u;
uint sel = (m >> 7) & 3u;
if (sel == 0u) { wx = 12; wy = int(a) + 2; }
else if (sel == 1u) { wx = int(a) + 2; wy = 12; }
else if (sel == 2u) { wx = int(a) + 6; wy = int(bb) + 6; d = 0u; h = 0u; }
else {
if (a == 0u) { wx = 6; wy = 10; }
else if (a == 1u) { wx = 10; wy = 6; }
else return H;
}
}
int dual = int(d);
int wtotal = (wx * wy) << dual;
int quant = int(r) - 2 + 6 * int(h);
int wbits = ise_bitcount(wtotal, quant);
if (wtotal > 64 || wbits < 24 || wbits > 96 || wx > bw || wy > bh) return H;
int pcount = int((b.x >> 11) & 3u) + 1;
if (dual != 0 && pcount == 4) return H;
// colour endpoint modes
int cem0, cem1 = 0, cem2 = 0, cem3 = 0;
int below = 128 - wbits, high = 0;
if (pcount == 1) {
cem0 = int((b.x >> 13) & 0xFu);
} else {
high = 3 * pcount - 4;
below -= high;
uint enc = ((b.x >> 23) & 0x3Fu) | rbits(b, below, high) << 6;
uint base = enc & 3u;
if (base == 0u) {
cem0 = cem1 = cem2 = cem3 = int((enc >> 2) & 0xFu);
below += high;
high = 0;
} else {
base--;
cem0 = int((((enc >> 2) & 1u) + base) << 2 | ((enc >> uint(2 + pcount)) & 3u));
cem1 = int((((enc >> 3) & 1u) + base) << 2 | ((enc >> uint(4 + pcount)) & 3u));
cem2 = int((((enc >> 4) & 1u) + base) << 2 | ((enc >> uint(6 + pcount)) & 3u));
cem3 = int((((enc >> 5) & 1u) + base) << 2 | ((enc >> uint(8 + pcount)) & 3u));
}
}
int n0 = ((cem0 >> 2) + 1) * 2, n1 = pcount > 1 ? ((cem1 >> 2) + 1) * 2 : 0;
int n2 = pcount > 2 ? ((cem2 >> 2) + 1) * 2 : 0, n3 = pcount > 3 ? ((cem3 >> 2) + 1) * 2 : 0;
int ccount = n0 + n1 + n2 + n3;
if (ccount > 18) return H;
int cbits = max((pcount == 1 ? 111 : 99) - wbits - high - (dual != 0 ? 2 : 0), 0);
int cq = -1;
for (int q = 20; q >= 0; q--) {
if (ise_bitcount(ccount, q) <= cbits) { cq = q; break; }
}
if (cq < 4) return H; // QUANT_6
H.solid = false;
H.wx = wx; H.wy = wy; H.dual = dual; H.quant = quant; H.wtotal = wtotal; H.pcount = pcount;
H.cem0 = cem0; H.cem1 = cem1; H.cem2 = cem2; H.cem3 = cem3; H.n0 = n0; H.n1 = n1; H.n2 = n2;
H.ccount = ccount; H.cq = cq;
H.ccs = dual != 0 ? int(rbits(b, below - 2, 2)) : -1;
return H;
}
// A workgroup decodes BPW blocks of one block row when a block has 64 texels or fewer (64 / texels of them), else one
// block: S lanes per block (its texels, at most 64). The lanes of a block decode its weights and colour values into
// shared memory, the first pcount lanes then unpack the endpoints, and each lane writes its texels.
const int MAXB = 4; // 4x4: 4 blocks of 16 lanes
shared uint s_w[MAXB][64]; // unquantized weights 0..64 (dual plane: interleaved as in the bit stream)
shared int s_cv[MAXB][18]; // unquantized colour values
shared uvec2 s_e0[MAXB][4], s_e1[MAXB][4]; // 16-bit endpoints, two channels per uint
uvec2 pack16(ivec4 e) { return uvec2(uint(e.x) | uint(e.y) << 16, uint(e.z) | uint(e.w) << 16); }
int ch16(uvec2 p, int c) { return int(((c < 2 ? p.x : p.y) >> (16 * (c & 1))) & 0xFFFFu); }
void main() {
int bw = int(pc.bw), bh = int(pc.bh), G = bw * bh;
int bpw = G <= 64 ? 64 / G : 1, S = G <= 64 ? G : 64;
int l = int(gl_LocalInvocationID.x), lb = l / S, lane = l - lb * S;
bool live = lb < bpw;
int bxi = int(gl_WorkGroupID.x) * bpw + lb, byi = int(gl_WorkGroupID.y);
live = live && bxi < int(pc.bx);
lb = min(lb, MAXB - 1);
uvec4 b = live ? blk[pc.in_base + gl_WorkGroupID.z * pc.bslice + uint(byi) * pc.bx + uint(bxi)] : uvec4(0u);
Hdr H = header(b, bw, bh);
bool work = live && !H.solid;
if (work) {
uvec4 rev = uvec4(bitfieldReverse(b.w), bitfieldReverse(b.z), bitfieldReverse(b.y), bitfieldReverse(b.x));
uint wrow = tab[H.quant];
for (int j = lane; j < H.wtotal; j += S) s_w[lb][j] = tab_byte(wrow + ise_at(rev, 0, H.wtotal, H.quant, j));
uint crow = H.cq == 20 ? 0u : tab[12 + H.cq - 4];
int cstart = H.pcount == 1 ? 17 : 29;
for (int j = lane; j < H.ccount; j += S) {
uint raw = ise_at(b, cstart, H.ccount, H.cq, j);
s_cv[lb][j] = H.cq == 20 ? int(raw) : int(tab_byte(crow + raw));
}
}
barrier();
if (work && lane < H.pcount) {
int p = lane;
int cem = p == 0 ? H.cem0 : p == 1 ? H.cem1 : p == 2 ? H.cem2 : H.cem3;
int k = p == 0 ? 0 : p == 1 ? H.n0 : p == 2 ? H.n0 + H.n1 : H.n0 + H.n1 + H.n2;
int n = ((cem >> 2) + 1) * 2;
int v[8];
for (int j = 0; j < 8; j++) {
v[j] = 0;
if (j < n) v[j] = s_cv[lb][k + j];
}
ivec4 e0, e1;
if (!unpack_endpoints(cem, v, e0, e1)) { e0 = ivec4(0xFF, 0, 0xFF, 0xFF); e1 = e0; }
if (pc.srgb != 0u) { e0 = e0 << 8 | 0x80; e1 = e1 << 8 | 0x80; }
else { e0 *= 257; e1 *= 257; }
s_e0[lb][p] = pack16(e0);
s_e1[lb][p] = pack16(e1);
}
barrier();
if (!live) return;
bool small = G < 32;
int seed = int((b.x >> 13) & 0x3FFu);
int sx = int((1024 + bw / 2) / (bw - 1)), sy = int((1024 + bh / 2) / (bh - 1));
for (int t = lane; t < G; t += S) {
int x = t % bw, y = t / bw;
uint gx = uint(bxi * bw + x), gy = uint(byi * bh + y);
if (gx >= pc.width || gy >= pc.height) continue;
uint o = H.konst;
if (!H.solid) {
int p = H.pcount > 1 ? part_of(seed, H.pcount, small ? x << 1 : x, small ? y << 1 : y) : 0;
int vx = (sx * x * (H.wx - 1) + 32) >> 6, vy = (sy * y * (H.wy - 1) + 32) >> 6;
int fs = vx & 0xF, ft = vy & 0xF;
int i0 = (vy >> 4) * H.wx + (vx >> 4);
int w11 = (fs * ft + 8) >> 4, w10 = fs - w11, w01 = ft - w11, w00 = 16 - fs - ft + w11;
int wcount = H.wx * H.wy;
int w[2];
for (int pl = 0; pl < 2; pl++) {
w[pl] = 0;
if (pl > H.dual) continue;
int acc = 0;
for (int c = 0; c < 4; c++) {
int idx = i0 + (c == 0 ? 0 : c == 1 ? 1 : c == 2 ? H.wx : H.wx + 1);
int f = c == 0 ? w00 : c == 1 ? w10 : c == 2 ? w01 : w11;
if (f == 0 || idx >= wcount) continue;
acc += int(s_w[lb][H.dual != 0 ? 2 * idx + pl : idx]) * f;
}
w[pl] = (acc + 8) >> 4;
}
uvec2 q0 = s_e0[lb][p], q1 = s_e1[lb][p];
o = 0u;
for (int c = 0; c < 4; c++) {
int wc = c == H.ccs ? w[1] : w[0];
o |= uint((ch16(q0, c) * (64 - wc) + ch16(q1, c) * wc + 32) >> 14) << uint(8 * c);
}
}
px[pc.out_base + (gl_WorkGroupID.z * pc.height + gy) * pc.width + gx] = o;
}
}
Reference in new issue View Git Blame Copy Permalink