// Generated by gen_device_kernels_wgsl.sh from device_kernels.wgsl -- do not edit.
#pragma once
namespace std {
// The WGSL entry point running DK_* kernel k (device_kernels.h's ids).
static const char* device_kernel_wgsl_entry_(int k) {
if (k == 0) return "dk_ew";
if (k == 1) return "dk_ew";
if (k == 2) return "dk_ew";
if (k == 3) return "dk_ew";
if (k == 4) return "dk_ew";
if (k == 5) return "dk_ew";
if (k == 6) return "dk_ew";
if (k == 7) return "dk_ew";
if (k == 8) return "dk_ew";
if (k == 9) return "dk_ew";
if (k == 10) return "dk_ew";
if (k == 11) return "dk_ew";
if (k == 12) return "dk_ew";
if (k == 13) return "dk_ew";
if (k == 14) return "dk_ew";
if (k == 15) return "dk_ew";
if (k == 16) return "dk_ew";
if (k == 17) return "dk_ew";
if (k == 18) return "dk_ew";
if (k == 19) return "dk_ew";
if (k == 20) return "dk_ew";
if (k == 21) return "dk_ew";
if (k == 22) return "dk_ew";
if (k == 23) return "dk_ew";
if (k == 24) return "dk_ew";
if (k == 25) return "dk_ew";
if (k == 26) return "dk_ew";
if (k == 27) return "dk_ew";
if (k == 28) return "dk_ew";
if (k == 29) return "dk_ew";
if (k == 30) return "dk_ew";
if (k == 31) return "dk_ew";
if (k == 32) return "dk_ew";
if (k == 33) return "dk_ew";
if (k == 34) return "dk_ew";
if (k == 35) return "dk_ew";
if (k == 36) return "dk_ew";
if (k == 37) return "dk_ew";
if (k == 38) return "dk_ew";
if (k == 39) return "dk_ew";
if (k == 40) return "dk_ew";
if (k == 41) return "dk_ew";
if (k == 42) return "dk_sum";
if (k == 43) return "dk_bias_rows";
if (k == 44) return "dk_mul_row";
if (k == 45) return "dk_colsum_acc";
if (k == 46) return "dk_mul_row_acc";
if (k == 47) return "dk_rowsum";
if (k == 48) return "dk_rowbcast_acc";
if (k == 49) return "dk_concat";
if (k == 50) return "dk_slice";
if (k == 51) return "dk_gather_rows";
if (k == 52) return "dk_scatter_rows";
if (k == 53) return "dk_softmax_rows";
if (k == 54) return "dk_softmax_bwd";
if (k == 55) return "dk_transpose";
if (k == 56) return "dk_attn_fwd";
if (k == 57) return "dk_attn_bwd_q";
if (k == 58) return "dk_attn_bwd_kv";
if (k == 59) return "dk_ce_fwd";
if (k == 60) return "dk_ce_bwd";
if (k == 61) return "dk_ewloss_fwd";
if (k == 62) return "dk_ewloss_bwd";
if (k == 63) return "dk_layernorm";
if (k == 64) return "dk_layernorm_bwd";
if (k == 65) return "dk_opt";
if (k == 66) return "dk_opt_multi";
if (k == 67) return "dk_clip_scale";
if (k == 68) return "dk_im2col";
if (k == 69) return "dk_col2im";
if (k == 70) return "dk_swap01";
if (k == 71) return "dk_bias_channels";
if (k == 72) return "dk_bias_channels_bwd";
if (k == 73) return "dk_bn_partial";
if (k == 74) return "dk_bn_finish";
if (k == 75) return "dk_gn_stats";
if (k == 76) return "dk_norm_apply";
if (k == 77) return "dk_norm_sums";
if (k == 78) return "dk_bn_bwd_dx";
if (k == 79) return "dk_gn_bwd";
if (k == 80) return "dk_pool";
if (k == 81) return "dk_pool_bwd";
if (k == 82) return "dk_gpool";
if (k == 83) return "dk_upsample";
if (k == 84) return "dk_resize";
if (k == 85) return "dk_resize_bwd";
if (k == 86) return "dk_amp_unscale";
if (k == 87) return "dk_amp_unscale_multi";
if (k == 88) return "dk_amp_update";
if (k == 89) return "dk_spmm";
if (k == 90) return "dk_spmm_tgrad";
if (k == 91) return "dk_spmm_vgrad";
return "";
}
// The pipeline constant OPC kernel k takes (dk_ewop_<op>: op), else -1.
static int device_kernel_wgsl_opc_(int k) {
if (k == 1) return 0;
if (k == 2) return 1;
if (k == 3) return 2;
if (k == 4) return 3;
if (k == 5) return 4;
if (k == 6) return 5;
if (k == 7) return 6;
if (k == 8) return 7;
if (k == 9) return 8;
if (k == 10) return 9;
if (k == 11) return 10;
if (k == 12) return 11;
if (k == 13) return 12;
if (k == 14) return 13;
if (k == 15) return 14;
if (k == 16) return 20;
if (k == 17) return 21;
if (k == 18) return 22;
if (k == 19) return 23;
if (k == 20) return 24;
if (k == 21) return 25;
if (k == 22) return 26;
if (k == 23) return 27;
if (k == 24) return 28;
if (k == 25) return 29;
if (k == 26) return 30;
if (k == 27) return 31;
if (k == 28) return 32;
if (k == 29) return 33;
if (k == 30) return 40;
if (k == 31) return 41;
if (k == 32) return 42;
if (k == 33) return 43;
if (k == 34) return 44;
if (k == 35) return 45;
if (k == 36) return 46;
if (k == 37) return 47;
if (k == 38) return 48;
if (k == 39) return 49;
if (k == 40) return 50;
if (k == 41) return 51;
return -1;
}
static const char* device_kernel_wgsl_source_() {
return
"// SafeC Standard Library — the device kernels behind std::ml's tensor ops on\n"
"// WebGPU (device_webgpu.sc): device_kernels.cu's DK_* kernels in WGSL, one\n"
"// entry point per kernel under the same name (dk_ewop_<op> are dk_ew with\n"
"// the pipeline constant OPC = op), with the same arguments and semantics.\n"
"// device_kernels.cu is the specification and device_ref.sc its CPU\n"
"// reference; gen_device_kernels_wgsl.sh embeds this file in\n"
"// device_kernels_wgsl.h and fails unless every DK_KERNEL of\n"
"// device_kernels.cu has its entry point here. Edit both together.\n"
"//\n"
"// What differs from the CUDA source, and why:\n"
"// - WGSL has no pointers and no 64-bit integers. All device memory is one\n"
"// storage buffer (device_webgpu.sc's heap), bound as H (32-bit words)\n"
"// and HA (the same words as atomics); a device address is a byte offset\n"
"// into it. The backend passes DevArgs as 32 words: p0..p7 as byte\n"
"// offsets (0 = none), their element types (DEV_TAG_*), u0..u7, f0..f7.\n"
"// The tables of handles DK_OPT_MULTI / DK_AMP_UNSCALE_MULTI read hold\n"
"// the same offsets in their 64-bit entries' low words.\n"
"// - No float atomics: atomicAdd on a float is a compare-exchange loop on\n"
"// its bits (atomic_addf).\n"
"// - No doubles: dk_double accumulators are f32, as on Metal (the\n"
"// statistics stay two-pass, which holds device_check's tolerance), and\n"
"// dk_bn_partial's partials are floats in the buffer sized for doubles.\n"
"// - No 16-bit stores: a 16-bit element is written by clearing and setting\n"
"// its half of the word atomically (two threads may own the two halves).\n"
"// The elementwise kernels with a 16-bit output instead take two\n"
"// elements per thread and write whole words, as device_kernels.cu's\n"
"// pair paths do. fp16 converts with pack2x16float / unpack2x16float\n"
"// (core WGSL, no shader-f16 feature needed), bf16 with integer\n"
"// arithmetic (the same rounding as f2b_bits).\n"
"// - The kernels compile twice through the pipeline constant TAGGED, as\n"
"// device_kernels.cu compiles with and without DK_TAGGED: the plain\n"
"// pipelines read every buffer as fp32 and pay nothing for 16 bits.\n"
"// - The fused attention kernels' tile (ATT_T x ATT_T threads) and largest\n"
"// head dim (ATT_MAXD) are pipeline constants: 16 x 16 tiles for head\n"
"// dims up to 64 and 8 x 8 up to 128, as Metal's two builds (32 KB of\n"
"// workgroup memory doesn't hold four 16-row tiles of 128 floats).\n"
"// - No subgroup operations (no adapter here reports them): reductions go\n"
"// through workgroup memory.\n"
"// Besides the DK_* kernels, the wk_* entry points do what rocBLAS and MIOpen\n"
"// do on ROCm and device_kernels.cu's mk_* on Metal: GEMMs, convolutions as\n"
"// implicit GEMMs, copies, and gpu_webgpu.h's few ops of its own.\n"
"\n"
"// Group 0 is the heap (set once per compute pass), group 1 the arguments\n"
"// (a dynamic offset per dispatch: only that group changes between them).\n"
"struct Args { w: array<vec4u, 8> }\n"
"@group(0) @binding(0) var<storage, read_write> H: array<u32>;\n"
"@group(0) @binding(1) var<storage, read_write> HA: array<atomic<u32>>;\n"
"@group(0) @binding(2) var<storage, read> D: array<u32>; // host data queued with the dispatches (wk_upload)\n"
"@group(1) @binding(0) var<uniform> A: Args;\n"
"\n"
"override TAGGED: bool = false; // the 16-bit build (some argument is fp16 / bf16)\n"
"override OPC: i32 = -1; // dk_ew's op: >= 0 for dk_ewop_<op>\n"
"\n"
"// ── Arguments and element access ─────────────────────────────────────────────\n"
"fn aw(i: u32) -> u32 { return A.w[i >> 2u][i & 3u]; }\n"
"fn P(k: u32) -> u32 { return aw(k); } // p<k>: byte offset, 0 = none\n"
"fn TG(k: u32) -> u32 { return aw(8u + k); } // p<k>'s element type: 0 fp32, 1 fp16, 2 bf16\n"
"fn U(k: u32) -> u32 { return aw(16u + k); }\n"
"fn F(k: u32) -> f32 { return bitcast<f32>(aw(24u + k)); }\n"
"fn N() -> u32 { return aw(23u); } // u7: the thread count\n"
"\n"
"fn h2f(h: u32, t: u32) -> f32 {\n"
" if (t == 1u) { return unpack2x16float(h).x; }\n"
" return bitcast<f32>(h << 16u);\n"
"}\n"
"fn f2h(v: f32, t: u32) -> u32 {\n"
" if (t == 1u) { return pack2x16float(vec2f(v, 0.0)) & 0xffffu; }\n"
" let x = bitcast<u32>(v);\n"
" if ((x & 0x7f800000u) == 0x7f800000u && (x & 0x7fffffu) != 0u) { return ((x >> 16u) | 0x40u) & 0xffffu; }\n"
" return ((x + 0x7fffu + ((x >> 16u) & 1u)) >> 16u) & 0xffffu;\n"
"}\n"
"fn pack16(lo: f32, hi: f32, t: u32) -> u32 { return f2h(lo, t) | (f2h(hi, t) << 16u); }\n"
"\n"
"// Element i of the buffer at byte offset b holding type t.\n"
"fn ldp(b: u32, t: u32, i: u32) -> f32 {\n"
" if (!TAGGED || t == 0u) { return bitcast<f32>(H[(b >> 2u) + i]); }\n"
" let a = b + 2u * i;\n"
" return h2f((H[a >> 2u] >> ((a & 2u) * 8u)) & 0xffffu, t);\n"
"}\n"
"fn stp(b: u32, t: u32, i: u32, v: f32) {\n"
" if (!TAGGED || t == 0u) { H[(b >> 2u) + i] = bitcast<u32>(v); return; }\n"
" let a = b + 2u * i;\n"
" let s = (a & 2u) * 8u;\n"
" atomicAnd(&HA[a >> 2u], ~(0xffffu << s));\n"
" atomicOr(&HA[a >> 2u], f2h(v, t) << s);\n"
"}\n"
"// F(k)[i] read / written in its type; ACCUM(F(k)[i], v)\n"
"fn ld(k: u32, i: u32) -> f32 { return ldp(P(k), TG(k), i); }\n"
"fn st(k: u32, i: u32, v: f32) { stp(P(k), TG(k), i, v); }\n"
"fn accum(k: u32, i: u32, v: f32) {\n"
" if (U(6u) != 0u) { st(k, i, v); } else { st(k, i, ld(k, i) + v); }\n"
"}\n"
"// FA(k): fp32 buffers (gradients, sums, statistics)\n"
"fn ldf(k: u32, i: u32) -> f32 { return bitcast<f32>(H[(P(k) >> 2u) + i]); }\n"
"fn stf(k: u32, i: u32, v: f32) { H[(P(k) >> 2u) + i] = bitcast<u32>(v); }\n"
"fn rdw(w: u32) -> f32 { return bitcast<f32>(H[w]); }\n"
"fn wrw(w: u32, v: f32) { H[w] = bitcast<u32>(v); }\n"
"fn atomic_addf(w: u32, v: f32) {\n"
" var old = atomicLoad(&HA[w]);\n"
" loop {\n"
" let r = atomicCompareExchangeWeak(&HA[w], old, bitcast<u32>(bitcast<f32>(old) + v));\n"
" if (r.exchanged) { break; }\n"
" old = r.old_value;\n"
" }\n"
"}\n"
"fn atadd(k: u32, i: u32, v: f32) { atomic_addf((P(k) >> 2u) + i, v); }\n"
"fn finite(v: f32) -> bool { return (bitcast<u32>(v) & 0x7f800000u) != 0x7f800000u; }\n"
"// Output k is 16-bit and word-aligned: elementwise kernels write it in pairs.\n"
"fn pairs_out(k: u32) -> bool { return TAGGED && TG(k) != 0u && (P(k) & 3u) == 0u; }\n"
"// Argument k holds type t word-aligned (device_kernels.cu's pair_tag), and\n"
"// its elements 2j, 2j + 1 from one word.\n"
"fn pair_ok(k: u32, t: u32) -> bool { return TG(k) == t && (P(k) & 3u) == 0u; }\n"
"fn ldw2(k: u32, j: u32) -> vec2f {\n"
" let w = H[(P(k) >> 2u) + j];\n"
" return vec2f(h2f(w & 0xffffu, TG(k)), h2f(w >> 16u, TG(k)));\n"
"}\n"
"\n"
"// ── Elementwise: op codes match tensor.h's EW_* ─────────────────────────────\n"
"// p0 = in a, p1 = in b, p2 = out, p3 = aux; f0 = k, f1 = k2; u1 = 1: the\n"
"// accumulating forms write (see device_kernels.cu).\n"
"const EW_ADD = 0;\n"
"const EW_SUB = 1;\n"
"const EW_MUL = 2;\n"
"const EW_SCALE = 3;\n"
"const EW_RELU = 4;\n"
"const EW_FILL = 5;\n"
"const EW_COPY = 6;\n"
"const EW_ACC = 7;\n"
"const EW_DEC = 8;\n"
"const EW_ACC_MUL = 9;\n"
"const EW_ACC_SCALE = 10;\n"
"const EW_ACC_RELU = 11;\n"
"const EW_ACC_K = 12;\n"
"const EW_ACC_SIGB = 13;\n"
"const EW_ACC_TANHB = 14;\n"
"const EW_EXP = 20;\n"
"const EW_LOG = 21;\n"
"const EW_ABS = 22;\n"
"const EW_LEAKY = 23;\n"
"const EW_CLAMP = 24;\n"
"const EW_SIGMOID = 25;\n"
"const EW_TANH = 26;\n"
"const EW_GELU = 27;\n"
"const EW_SILU = 28;\n"
"const EW_ADDK = 29;\n"
"const EW_MIN = 30;\n"
"const EW_ELU = 31;\n"
"const EW_LERP = 32;\n"
"const EW_MULS = 33;\n"
"const EW_ACC_PK = 40;\n"
"const EW_ACC_EXPB = 41;\n"
"const EW_ACC_LOGB = 42;\n"
"const EW_ACC_ABSB = 43;\n"
"const EW_ACC_LEAKYB = 44;\n"
"const EW_ACC_CLAMPB = 45;\n"
"const EW_ACC_GELUB = 46;\n"
"const EW_ACC_SILUB = 47;\n"
"const EW_ACC_MINA = 48;\n"
"const EW_ACC_MINB = 49;\n"
"const EW_ACC_ELUB = 50;\n"
"const EW_ACC_PSCALE = 51;\n"
"\n"
"fn gelu_f(x: f32) -> f32 {\n"
" let c = 0.7978845608028654;\n"
" return 0.5 * x * (1.0 + tanh(c * (x + 0.044715 * x * x * x)));\n"
"}\n"
"fn gelu_d(x: f32) -> f32 {\n"
" let c = 0.7978845608028654;\n"
" let t = tanh(c * (x + 0.044715 * x * x * x));\n"
" let du = c * (1.0 + 0.134145 * x * x);\n"
" return 0.5 * (1.0 + t) + 0.5 * x * (1.0 - t * t) * du;\n"
"}\n"
"fn sig_f(x: f32) -> f32 { return 1.0 / (1.0 + exp(-x)); }\n"
"\n"
"fn ew_val(op: i32, x: f32, y: f32, z: f32, o: f32, k: f32, k2: f32, ow: bool) -> f32 {\n"
" let ob = select(o, 0.0, ow);\n"
" switch op {\n"
" case EW_ADD: { return x + y; }\n"
" case EW_SUB: { return x - y; }\n"
" case EW_MUL: { return x * y; }\n"
" case EW_SCALE: { return x * k; }\n"
" case EW_RELU: { return select(0.0, x, x > 0.0); }\n"
" case EW_FILL: { return k; }\n"
" case EW_COPY: { return x; }\n"
" case EW_ACC: { return ob + x; }\n"
" case EW_DEC: { return ob - x; }\n"
" case EW_ACC_MUL: { return fma(x, y, ob); }\n"
" case EW_ACC_SCALE: { return fma(x, k, ob); }\n"
" case EW_ACC_RELU: { return ob + select(0.0, x, y > 0.0); }\n"
" case EW_ACC_K: { return ob + k; }\n"
" case EW_ACC_SIGB: { return fma(x * y, 1.0 - y, ob); }\n"
" case EW_ACC_TANHB: { return fma(x, 1.0 - y * y, ob); }\n"
" case EW_EXP: { return exp(x); }\n"
" case EW_LOG: { return log(x); }\n"
" case EW_ABS: { return abs(x); }\n"
" case EW_LEAKY: { return select(k * x, x, x > 0.0); }\n"
" case EW_CLAMP: { return min(max(x, k), k2); }\n"
" case EW_SIGMOID: { return sig_f(x); }\n"
" case EW_TANH: { return tanh(x); }\n"
" case EW_GELU: { return gelu_f(x); }\n"
" case EW_SILU: { return x * sig_f(x); }\n"
" case EW_ADDK: { return x + k; }\n"
" case EW_MIN: { return select(y, x, x <= y); }\n"
" case EW_ELU: { return select(k * (exp(x) - 1.0), x, x > 0.0); }\n"
" case EW_LERP: { return fma(k, x - o, o); }\n"
" case EW_MULS: { return x * y; }\n"
" case EW_ACC_PK: { return ob + y; }\n"
" case EW_ACC_EXPB: { return fma(x, y, ob); }\n"
" case EW_ACC_LOGB: { return ob + x / y; }\n"
" case EW_ACC_ABSB: { return ob + select(select(0.0, -x, y < 0.0), x, y > 0.0); }\n"
" case EW_ACC_LEAKYB: { return ob + x * select(k, 1.0, y > 0.0); }\n"
" case EW_ACC_CLAMPB: { return ob + select(0.0, x, y == z); }\n"
" case EW_ACC_GELUB: { return fma(x, gelu_d(y), ob); }\n"
" case EW_ACC_SILUB: { let s = sig_f(y); return fma(x, s + y * s * (1.0 - s), ob); }\n"
" case EW_ACC_MINA: { return ob + select(0.0, x, y <= z); }\n"
" case EW_ACC_MINB: { return ob + select(0.0, x, !(z <= y)); }\n"
" case EW_ACC_ELUB: { return ob + x * select(z + k, 1.0, y > 0.0); }\n"
" case EW_ACC_PSCALE: { return ob + k * y; }\n"
" default: { }\n"
" }\n"
" return 0.0;\n"
"}\n"
"fn ew_uses_x(op: i32) -> bool { return op != EW_FILL && op != EW_ACC_K && op != EW_ACC_PK && op != EW_ACC_PSCALE; }\n"
"fn ew_y_scalar(op: i32) -> bool { return op == EW_MULS || op == EW_ACC_PK || op == EW_ACC_PSCALE; }\n"
"fn ew_uses_y(op: i32) -> bool {\n"
" return op == EW_ADD || op == EW_SUB || op == EW_MUL || op == EW_ACC_MUL || op == EW_ACC_RELU || op == EW_ACC_SIGB ||\n"
" op == EW_ACC_TANHB || op == EW_MIN || (op >= EW_ACC_EXPB && op <= EW_ACC_ELUB);\n"
"}\n"
"fn ew_uses_z(op: i32) -> bool { return op == EW_ACC_CLAMPB || op == EW_ACC_MINA || op == EW_ACC_MINB || op == EW_ACC_ELUB; }\n"
"fn ew_reads_o(op: i32, ow: bool) -> bool {\n"
" return op == EW_LERP || (!ow && ((op >= EW_ACC && op <= EW_ACC_TANHB) || (op >= EW_ACC_PK && op <= EW_ACC_PSCALE)));\n"
"}\n"
"fn ew_op() -> i32 { return select(i32(U(0u)), OPC, OPC >= 0); }\n"
"// Element i's result.\n"
"fn ew_at(op: i32, i: u32, y0: f32) -> f32 {\n"
" let ow = U(1u) != 0u;\n"
" var x = 0.0;\n"
" if (ew_uses_x(op)) { x = ld(0u, i); }\n"
" var y = y0;\n"
" if (ew_uses_y(op)) { y = ld(1u, i); }\n"
" var z = 0.0;\n"
" if (ew_uses_z(op)) { z = ld(3u, i); }\n"
" var o = 0.0;\n"
" if (ew_reads_o(op, ow)) { o = ld(2u, i); }\n"
" return ew_val(op, x, y, z, o, F(0u), F(1u), ow);\n"
"}\n"
"\n"
"// Any op (u0); dk_ewop_<code> are this kernel with OPC = code.\n"
"@compute @workgroup_size(256)\n"
"fn dk_ew(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let op = ew_op();\n"
" var y0 = 0.0;\n"
" if (ew_y_scalar(op)) { y0 = ld(1u, 0u); }\n"
" let n = N();\n"
" if (pairs_out(2u)) {\n"
" let o = P(2u) >> 2u;\n"
" let t = TG(2u);\n"
" for (var j = gid.x; j < n / 2u; j += nwg.x * 256u) {\n"
" H[o + j] = pack16(ew_at(op, 2u * j, y0), ew_at(op, 2u * j + 1u, y0), t);\n"
" }\n"
" if ((n & 1u) != 0u && gid.x == 0u) { st(2u, n - 1u, ew_at(op, n - 1u, y0)); }\n"
" return;\n"
" }\n"
" for (var i = gid.x; i < n; i += nwg.x * 256u) { st(2u, i, ew_at(op, i, y0)); }\n"
"}\n"
"\n"
"// ── Reductions (256 threads per workgroup, one atomic per workgroup) ─────────\n"
"var<workgroup> red: array<f32, 256>;\n"
"// The workgroup's total of v, returned to every thread (block_sum_d).\n"
"fn block_sum(v: f32, li: u32) -> f32 {\n"
" workgroupBarrier(); // the previous call's readers are done with red\n"
" red[li] = v;\n"
" workgroupBarrier();\n"
" for (var w = 128u; w > 0u; w >>= 1u) {\n"
" if (li < w) { red[li] += red[li + w]; }\n"
" workgroupBarrier();\n"
" }\n"
" return red[0];\n"
"}\n"
"fn block_add(v: f32, k: u32, li: u32) {\n"
" let s = block_sum(v, li);\n"
" if (li == 0u) { atadd(k, 0u, s); }\n"
"}\n"
"\n"
"// The row kernels (row sums, softmax, cross-entropy, layer norm) give each\n"
"// row RL lanes that split its columns (device_kernels.cu's GSR on Metal):\n"
"// RL = 32 for rows of 32 columns on, 1 (a thread per row, as on ROCm) for\n"
"// narrower ones. A workgroup takes 256 / RL rows at a time; every lane\n"
"// joins the reductions, those past the last row with neutral values.\n"
"override RL: u32 = 1u;\n"
"// The total (maximum) over the RL lanes of v's row, to each of them.\n"
"fn seg_sum(v: f32, li: u32) -> f32 {\n"
" if (RL == 1u) { return v; }\n"
" workgroupBarrier();\n"
" red[li] = v;\n"
" workgroupBarrier();\n"
" for (var w = RL / 2u; w > 0u; w >>= 1u) {\n"
" if (li % RL < w) { red[li] += red[li + w]; }\n"
" workgroupBarrier();\n"
" }\n"
" return red[li - li % RL];\n"
"}\n"
"fn seg_max(v: f32, li: u32) -> f32 {\n"
" if (RL == 1u) { return v; }\n"
" workgroupBarrier();\n"
" red[li] = v;\n"
" workgroupBarrier();\n"
" for (var w = RL / 2u; w > 0u; w >>= 1u) {\n"
" if (li % RL < w) { red[li] = max(red[li], red[li + w]); }\n"
" workgroupBarrier();\n"
" }\n"
" return red[li - li % RL];\n"
"}\n"
"\n"
"// p0 = in, p1 = out[0] (zeroed by the caller); u0: 0 sum, 1 sum of squares\n"
"@compute @workgroup_size(256)\n"
"fn dk_sum(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u,\n"
" @builtin(local_invocation_index) li: u32) {\n"
" var acc = 0.0;\n"
" let sq = U(0u) != 0u;\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let v = ld(0u, i);\n"
" acc += select(v, v * v, sq);\n"
" }\n"
" block_add(acc, 1u, li);\n"
"}\n"
"\n"
"// ── Rows and columns: u0 = cols ──────────────────────────────────────────────\n"
"// out[r,c] = x[r,c] + b[c] n = rows*cols\n"
"@compute @workgroup_size(256)\n"
"fn dk_bias_rows(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) { st(2u, i, ld(0u, i) + ld(1u, i % C)); }\n"
"}\n"
"// out[r,c] = x[r,c] * g[c] n = rows*cols\n"
"@compute @workgroup_size(256)\n"
"fn dk_mul_row(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) { st(2u, i, ld(0u, i) * ld(1u, i % C)); }\n"
"}\n"
"// gb[c] += sum_r d[r,c] (* x[r,c] when p2 != 0); u0 = cols, u1 = rows, u2 = S\n"
"// row slices; a thread per column n = ceil(cols/256)*S*256\n"
"@compute @workgroup_size(256)\n"
"fn dk_colsum_acc(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" let R = U(1u);\n"
" let S = max(U(2u), 1u);\n"
" let hx = P(2u) != 0u;\n"
" for (var u = wg.x; u < ((C + 255u) / 256u) * S; u += nwg.x) {\n"
" let c = (u / S) * 256u + li;\n"
" let sl = u % S;\n"
" if (c < C) {\n"
" var s = 0.0;\n"
" for (var r = sl; r < R; r += S) {\n"
" var v = ld(0u, r * C + c);\n"
" if (hx) { v *= ld(2u, r * C + c); }\n"
" s += v;\n"
" }\n"
" if (S > 1u) { atadd(1u, c, s); } else { stf(1u, c, ldf(1u, c) + s); }\n"
" }\n"
" }\n"
"}\n"
"// dx[r,c] += d[r,c] * g[c] n = rows*cols\n"
"@compute @workgroup_size(256)\n"
"fn dk_mul_row_acc(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) { accum(2u, i, ld(0u, i) * ld(1u, i % C)); }\n"
"}\n"
"// out[r] = sum_c x[r,c] n = rows\n"
"@compute @workgroup_size(256)\n"
"fn dk_rowsum(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" for (var rb = wg.x * (256u / RL); rb < N(); rb += nwg.x * (256u / RL)) {\n"
" let r = rb + li / RL;\n"
" let live = r < N();\n"
" var s = 0.0;\n"
" for (var c = li % RL; live && c < C; c += RL) { s += ld(0u, r * C + c); }\n"
" s = seg_sum(s, li);\n"
" if (live && li % RL == 0u) { st(1u, r, s); }\n"
" }\n"
"}\n"
"// dx[r,c] += g[r] n = rows*cols\n"
"@compute @workgroup_size(256)\n"
"fn dk_rowbcast_acc(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) { accum(1u, i, ld(0u, i / C)); }\n"
"}\n"
"// out [rows, p+q] = x [rows, p] ++ y [rows, q]; u0 = p, u1 = q n = rows*(p+q)\n"
"@compute @workgroup_size(256)\n"
"fn dk_concat(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let Pp = U(0u);\n"
" let Q = U(1u);\n"
" let W = Pp + Q;\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let r = i / W;\n"
" let c = i % W;\n"
" if (c < Pp) { st(2u, i, ld(0u, r * Pp + c)); } else { st(2u, i, ld(1u, r * Q + c - Pp)); }\n"
" }\n"
"}\n"
"// Columns [start, start+len) of a wide [rows, cols] tensor and a narrow one;\n"
"// u0 = cols, u1 = start, u2 = len, u3 = mode (0 slice, 1 slice backward,\n"
"// 2 concat backward) n = rows*len\n"
"@compute @workgroup_size(256)\n"
"fn dk_slice(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" let S = U(1u);\n"
" let L = U(2u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let r = i / L;\n"
" let c = i % L;\n"
" let w = r * C + S + c;\n"
" if (U(3u) == 1u) { st(1u, w, ld(1u, w) + ld(0u, i)); }\n"
" else if (U(3u) == 2u) { accum(1u, i, ld(0u, w)); }\n"
" else { st(1u, i, ld(0u, w)); }\n"
" }\n"
"}\n"
"// out[i, :] = table[idx[i], :]; u0 = d n = m*d\n"
"@compute @workgroup_size(256)\n"
"fn dk_gather_rows(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let D = U(0u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let r = i / D;\n"
" let c = i % D;\n"
" st(2u, i, ld(0u, u32(ld(1u, r)) * D + c));\n"
" }\n"
"}\n"
"// table.grad[idx[i], :] += g[i, :] n = m*d\n"
"@compute @workgroup_size(256)\n"
"fn dk_scatter_rows(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let D = U(0u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let r = i / D;\n"
" let c = i % D;\n"
" atadd(2u, u32(ld(1u, r)) * D + c, ld(0u, i));\n"
" }\n"
"}\n"
"\n"
"// ── Attention masks (attention.h) ────────────────────────────────────────────\n"
"fn attn_vis(flags: u32, w: u32, pos: i32, j: i32) -> bool {\n"
" if ((flags & 1u) != 0u && j > pos) { return false; }\n"
" if ((flags & 2u) != 0u && j <= pos - i32(w)) { return false; }\n"
" return true;\n"
"}\n"
"\n"
"// ── Softmax family: a row per RL lanes, u0 = cols ────────────────────────────\n"
"// u1: 0 softmax, 1 log-softmax. A masked softmax: u2 = mask flags, u3 =\n"
"// window, u4 = (int) keys - queries.\n"
"@compute @workgroup_size(256)\n"
"fn dk_softmax_rows(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" let fl = U(2u);\n"
" for (var rb = wg.x * (256u / RL); rb < N(); rb += nwg.x * (256u / RL)) {\n"
" let r = rb + li / RL;\n"
" let live = r < N();\n"
" let b = r * C;\n"
" let pos = i32(r) + i32(U(4u));\n"
" var mx = -3.0e38;\n"
" for (var c = li % RL; live && c < C; c += RL) {\n"
" if (fl == 0u || attn_vis(fl, U(3u), pos, i32(c))) { mx = max(mx, ld(0u, b + c)); }\n"
" }\n"
" mx = seg_max(mx, li);\n"
" var s = 0.0;\n"
" for (var c = li % RL; live && c < C; c += RL) {\n"
" if (fl == 0u || attn_vis(fl, U(3u), pos, i32(c))) { s += exp(ld(0u, b + c) - mx); }\n"
" }\n"
" s = seg_sum(s, li);\n"
" if (!live) { continue; }\n"
" if (fl != 0u) {\n"
" let inv = select(0.0, 1.0 / s, s > 0.0);\n"
" for (var c = li % RL; c < C; c += RL) {\n"
" var v = 0.0;\n"
" if (attn_vis(fl, U(3u), pos, i32(c))) { v = exp(ld(0u, b + c) - mx) * inv; }\n"
" st(1u, b + c, v);\n"
" }\n"
" } else if (U(1u) != 0u) {\n"
" let l = mx + log(s);\n"
" for (var c = li % RL; c < C; c += RL) { st(1u, b + c, ld(0u, b + c) - l); }\n"
" } else {\n"
" let inv = 1.0 / s;\n"
" for (var c = li % RL; c < C; c += RL) { st(1u, b + c, exp(ld(0u, b + c) - mx) * inv); }\n"
" }\n"
" }\n"
"}\n"
"// dx += y * (g - sum(g*y)) (softmax), dx += g - exp(y) * sum(g) (log-softmax)\n"
"// p0 = g, p1 = y, p2 = dx\n"
"@compute @workgroup_size(256)\n"
"fn dk_softmax_bwd(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" let lg = U(1u) != 0u;\n"
" for (var rb = wg.x * (256u / RL); rb < N(); rb += nwg.x * (256u / RL)) {\n"
" let r = rb + li / RL;\n"
" let live = r < N();\n"
" let b = r * C;\n"
" var s = 0.0;\n"
" for (var c = li % RL; live && c < C; c += RL) {\n"
" if (lg) { s += ld(0u, b + c); } else { s += ld(0u, b + c) * ld(1u, b + c); }\n"
" }\n"
" s = seg_sum(s, li);\n"
" if (!live) { continue; }\n"
" for (var c = li % RL; c < C; c += RL) {\n"
" if (lg) { st(2u, b + c, ld(2u, b + c) + ld(0u, b + c) - exp(ld(1u, b + c)) * s); }\n"
" else { st(2u, b + c, ld(2u, b + c) + ld(1u, b + c) * (ld(0u, b + c) - s)); }\n"
" }\n"
" }\n"
"}\n"
"\n"
"// ── Transpose ────────────────────────────────────────────────────────────────\n"
"// out [C, R] = in [R, C]^T, or with u2 = 1 out += it; u0 = R, u1 = C n = R*C\n"
"@compute @workgroup_size(256)\n"
"fn dk_transpose(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let R = U(0u);\n"
" let C = U(1u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let r = i / C;\n"
" let c = i % C;\n"
" let v = ld(0u, i);\n"
" if (U(2u) != 0u) { accum(1u, c * R + r, v); } else { st(1u, c * R + r, v); }\n"
" }\n"
"}\n"
"\n"
"// ── Fused multi-head attention (attention.h) ─────────────────────────────────\n"
"// FlashAttention-2's tiling as device_kernels.cu's dk_attn_*: a workgroup of\n"
"// ATT_T x ATT_T threads takes one head's tile of ATT_T rows and walks the\n"
"// other side in tiles of ATT_T through workgroup memory. Arguments: u0 = S,\n"
"// u1 = L, u2 = H, u3 = D, u4 = E, u5 = flags, u6 = window, f0 = scale;\n"
"// p4 = lse, p5 = dout, p7 = D (rowsum(dout * out)).\n"
"override ATT_T: u32 = 8u;\n"
"override ATT_MAXD: u32 = 128u;\n"
"const ATT_HID = -3.0e38;\n"
"const ATT_PER_MAX = 16u; // ATT_MAXD / ATT_T for either build\n"
"var<workgroup> aQ: array<f32, ATT_T * (ATT_MAXD + 1u)>;\n"
"var<workgroup> aK: array<f32, ATT_T * (ATT_MAXD + 1u)>;\n"
"var<workgroup> aV: array<f32, ATT_T * (ATT_MAXD + 1u)>;\n"
"var<workgroup> aG: array<f32, ATT_T * (ATT_MAXD + 1u)>;\n"
"var<workgroup> aP: array<f32, ATT_T * (ATT_T + 1u)>;\n"
"var<workgroup> aD: array<f32, ATT_T * (ATT_T + 1u)>;\n"
"var<workgroup> aRd: array<f32, ATT_T>;\n"
"var<workgroup> aRl: array<f32, ATT_T>;\n"
"\n"
"// The key range a tile of rows [i0, i0 + ATT_T) can see: [lo, hi).\n"
"fn attn_range(i0: u32) -> vec2i {\n"
" let off = i32(U(1u)) - i32(U(0u));\n"
" var lo = 0;\n"
" var hi = i32(U(1u));\n"
" if ((U(5u) & 1u) != 0u) { let m = i32(i0) + i32(ATT_T) - 1 + off + 1; if (m < hi) { hi = m; } }\n"
" if ((U(5u) & 2u) != 0u) { let m = i32(i0) + off - i32(U(6u)) + 1; if (m > lo) { lo = m; } }\n"
" if (lo < 0) { lo = 0; }\n"
" return vec2i(lo, hi);\n"
"}\n"
"// rows [r0, r0 + ATT_T) of argument k ([rows, H*W]), head h -> tile 'which'\n"
"// (0 aQ, 1 aK, 2 aV, 3 aG), 0 past the end\n"
"fn attn_load(which: u32, k: u32, r0: u32, rows: u32, Hh: u32, h: u32, W: u32, li: u32) {\n"
" let LD = ATT_MAXD + 1u;\n"
" for (var q = li; q < ATT_T * W; q += ATT_T * ATT_T) {\n"
" let r = q / W;\n"
" let c = q % W;\n"
" var v = 0.0;\n"
" if (r0 + r < rows) { v = ld(k, ((r0 + r) * Hh + h) * W + c); }\n"
" let at = r * LD + c;\n"
" switch which {\n"
" case 0u: { aQ[at] = v; }\n"
" case 1u: { aK[at] = v; }\n"
" case 2u: { aV[at] = v; }\n"
" default: { aG[at] = v; }\n"
" }\n"
" }\n"
"}\n"
"// p0 q, p1 k, p2 v, p3 out, p4 lse (written)\n"
"@compute @workgroup_size(ATT_T * ATT_T)\n"
"fn dk_attn_fwd(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let S = U(0u);\n"
" let L = U(1u);\n"
" let Hh = U(2u);\n"
" let D = U(3u);\n"
" let E = U(4u);\n"
" let fl = U(5u) & 3u;\n"
" let off = i32(L) - i32(S);\n"
" let tiles = (S + ATT_T - 1u) / ATT_T;\n"
" let r = li / ATT_T;\n"
" let ln = li % ATT_T;\n"
" let LD = ATT_MAXD + 1u;\n"
" let PL = ATT_T + 1u;\n"
" let PER = ATT_MAXD / ATT_T;\n"
" for (var u = wg.x; u < Hh * tiles; u += nwg.x) {\n"
" let h = u / tiles;\n"
" let i0 = (u % tiles) * ATT_T;\n"
" let i = i0 + r;\n"
" attn_load(0u, 0u, i0, S, Hh, h, D, li);\n"
" var acc: array<f32, ATT_PER_MAX>;\n"
" var m = ATT_HID;\n"
" var l = 0.0;\n"
" let rg = attn_range(i0);\n"
" for (var j0 = (rg.x / i32(ATT_T)) * i32(ATT_T); j0 < rg.y; j0 += i32(ATT_T)) {\n"
" workgroupBarrier();\n"
" attn_load(1u, 1u, u32(j0), L, Hh, h, D, li);\n"
" attn_load(2u, 2u, u32(j0), L, Hh, h, E, li);\n"
" workgroupBarrier();\n"
" let j = u32(j0) + ln;\n"
" var sc = ATT_HID;\n"
" if (i < S && j < L && attn_vis(fl, U(6u), i32(i) + off, i32(j))) {\n"
" var s = 0.0;\n"
" for (var d = 0u; d < D; d++) { s += aQ[r * LD + d] * aK[ln * LD + d]; }\n"
" sc = s * F(0u);\n"
" }\n"
" aP[r * PL + ln] = sc;\n"
" workgroupBarrier();\n"
" var mt = ATT_HID;\n"
" for (var c = 0u; c < ATT_T; c++) { mt = max(mt, aP[r * PL + c]); }\n"
" if (mt > ATT_HID) {\n"
" let mn = max(m, mt);\n"
" let cr = exp(m - mn);\n"
" var pv: array<f32, ATT_PER_MAX>;\n"
" var ps = 0.0;\n"
" for (var c = 0u; c < ATT_T; c++) {\n"
" let sv = aP[r * PL + c];\n"
" var p = 0.0;\n"
" if (sv > ATT_HID) { p = exp(sv - mn); }\n"
" pv[c] = p;\n"
" ps += p;\n"
" }\n"
" l = l * cr + ps;\n"
" for (var t = 0u; t < PER; t++) {\n"
" let e = ln + t * ATT_T;\n"
" if (e < E) {\n"
" var x = acc[t] * cr;\n"
" for (var c = 0u; c < ATT_T; c++) { x += pv[c] * aV[c * LD + e]; }\n"
" acc[t] = x;\n"
" }\n"
" }\n"
" m = mn;\n"
" }\n"
" }\n"
" if (i < S) {\n"
" var inv = 0.0;\n"
" if (l > 0.0) { inv = 1.0 / l; }\n"
" let ob = (i * Hh + h) * E;\n"
" for (var t = 0u; t < PER; t++) { let e = ln + t * ATT_T; if (e < E) { st(3u, ob + e, acc[t] * inv); } }\n"
" if (ln == 0u) {\n"
" var lse = 3.0e38; // no visible key: every p is 0\n"
" if (l > 0.0) { lse = m + log(l); }\n"
" stf(4u, h * S + i, lse);\n"
" }\n"
" }\n"
" workgroupBarrier();\n"
" }\n"
"}\n"
"// p0 q, p1 k, p2 v, p3 out, p6 dq (+)= scale * sum_j ds_ij k_j (0: none),\n"
"// ds_ij = p_ij (dout_i . v_j - D_i)\n"
"@compute @workgroup_size(ATT_T * ATT_T)\n"
"fn dk_attn_bwd_q(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let S = U(0u);\n"
" let L = U(1u);\n"
" let Hh = U(2u);\n"
" let D = U(3u);\n"
" let E = U(4u);\n"
" let fl = U(5u) & 3u;\n"
" let off = i32(L) - i32(S);\n"
" let tiles = (S + ATT_T - 1u) / ATT_T;\n"
" let r = li / ATT_T;\n"
" let ln = li % ATT_T;\n"
" let LD = ATT_MAXD + 1u;\n"
" let PL = ATT_T + 1u;\n"
" let PER = ATT_MAXD / ATT_T;\n"
" for (var u = wg.x; u < Hh * tiles; u += nwg.x) {\n"
" let h = u / tiles;\n"
" let i0 = (u % tiles) * ATT_T;\n"
" let i = i0 + r;\n"
" attn_load(3u, 5u, i0, S, Hh, h, E, li);\n"
" attn_load(1u, 3u, i0, S, Hh, h, E, li); // the output rows, for D (aK is free until the key tiles)\n"
" workgroupBarrier();\n"
" var part = 0.0;\n"
" for (var e = ln; e < E; e += ATT_T) { part += aG[r * LD + e] * aK[r * LD + e]; }\n"
" aD[r * PL + ln] = part;\n"
" workgroupBarrier();\n"
" if (ln == 0u) {\n"
" var Di = 0.0;\n"
" for (var c = 0u; c < ATT_T; c++) { Di += aD[r * PL + c]; }\n"
" aRd[r] = Di;\n"
" var rl = 0.0;\n"
" if (i < S) { rl = ldf(4u, h * S + i); stf(7u, h * S + i, Di); }\n"
" aRl[r] = rl;\n"
" }\n"
" if (P(6u) == 0u) { workgroupBarrier(); continue; }\n"
" attn_load(0u, 0u, i0, S, Hh, h, D, li);\n"
" var dq: array<f32, ATT_PER_MAX>;\n"
" let rg = attn_range(i0);\n"
" for (var j0 = (rg.x / i32(ATT_T)) * i32(ATT_T); j0 < rg.y; j0 += i32(ATT_T)) {\n"
" workgroupBarrier();\n"
" attn_load(1u, 1u, u32(j0), L, Hh, h, D, li);\n"
" attn_load(2u, 2u, u32(j0), L, Hh, h, E, li);\n"
" workgroupBarrier();\n"
" let j = u32(j0) + ln;\n"
" var ds = 0.0;\n"
" if (i < S && j < L && attn_vis(fl, U(6u), i32(i) + off, i32(j))) {\n"
" var s = 0.0;\n"
" var dp = 0.0;\n"
" for (var d = 0u; d < D; d++) { s += aQ[r * LD + d] * aK[ln * LD + d]; }\n"
" for (var e = 0u; e < E; e++) { dp += aG[r * LD + e] * aV[ln * LD + e]; }\n"
" ds = exp(s * F(0u) - aRl[r]) * (dp - aRd[r]);\n"
" }\n"
" aD[r * PL + ln] = ds;\n"
" workgroupBarrier();\n"
" for (var t = 0u; t < PER; t++) {\n"
" let d = ln + t * ATT_T;\n"
" if (d < D) {\n"
" var x = dq[t];\n"
" for (var c = 0u; c < ATT_T; c++) { x += aD[r * PL + c] * aK[c * LD + d]; }\n"
" dq[t] = x;\n"
" }\n"
" }\n"
" }\n"
" if (i < S) {\n"
" let ob = (i * Hh + h) * D;\n"
" for (var t = 0u; t < PER; t++) {\n"
" let d = ln + t * ATT_T;\n"
" if (d < D) {\n"
" if ((U(5u) & 4u) != 0u) { st(6u, ob + d, F(0u) * dq[t]); }\n"
" else { st(6u, ob + d, ld(6u, ob + d) + F(0u) * dq[t]); }\n"
" }\n"
" }\n"
" }\n"
" workgroupBarrier();\n"
" }\n"
"}\n"
"// p0 q, p1 k, p2 v, p3 dk (+)= scale * sum_i ds_ij q_i, p6 dv (+)= sum_i p_ij\n"
"// dout_i (either 0: none); after dk_attn_bwd_q. A workgroup takes ATT_T keys.\n"
"@compute @workgroup_size(ATT_T * ATT_T)\n"
"fn dk_attn_bwd_kv(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let S = U(0u);\n"
" let L = U(1u);\n"
" let Hh = U(2u);\n"
" let D = U(3u);\n"
" let E = U(4u);\n"
" let fl = U(5u) & 3u;\n"
" let off = i32(L) - i32(S);\n"
" let tiles = (L + ATT_T - 1u) / ATT_T;\n"
" let r = li / ATT_T;\n"
" let ln = li % ATT_T;\n"
" let LD = ATT_MAXD + 1u;\n"
" let PL = ATT_T + 1u;\n"
" let PER = ATT_MAXD / ATT_T;\n"
" for (var u = wg.x; u < Hh * tiles; u += nwg.x) {\n"
" let h = u / tiles;\n"
" let j0 = (u % tiles) * ATT_T;\n"
" let j = j0 + r;\n"
" attn_load(1u, 1u, j0, L, Hh, h, D, li);\n"
" attn_load(2u, 2u, j0, L, Hh, h, E, li);\n"
" var dk: array<f32, ATT_PER_MAX>;\n"
" var dv: array<f32, ATT_PER_MAX>;\n"
" // the queries that can see keys [j0, j0 + ATT_T): [qlo, qhi)\n"
" var qlo = 0;\n"
" var qhi = i32(S);\n"
" if ((fl & 1u) != 0u) { let m2 = i32(j0) - off; if (m2 > qlo) { qlo = m2; } }\n"
" if ((fl & 2u) != 0u) { let m2 = i32(j0) + i32(ATT_T) - 1 - off + i32(U(6u)); if (m2 < qhi) { qhi = m2; } }\n"
" if (qlo < 0) { qlo = 0; }\n"
" for (var i0 = (qlo / i32(ATT_T)) * i32(ATT_T); i0 < qhi; i0 += i32(ATT_T)) {\n"
" workgroupBarrier();\n"
" attn_load(0u, 0u, u32(i0), S, Hh, h, D, li);\n"
" attn_load(3u, 5u, u32(i0), S, Hh, h, E, li);\n"
" if (li < ATT_T) {\n"
" let iq = u32(i0) + li;\n"
" var rl = 0.0;\n"
" var rd = 0.0;\n"
" if (iq < S) { rl = ldf(4u, h * S + iq); rd = ldf(7u, h * S + iq); }\n"
" aRl[li] = rl;\n"
" aRd[li] = rd;\n"
" }\n"
" workgroupBarrier();\n"
" // thread (key r, query ln)\n"
" let i = u32(i0) + ln;\n"
" var p = 0.0;\n"
" var ds = 0.0;\n"
" if (i < S && j < L && attn_vis(fl, U(6u), i32(i) + off, i32(j))) {\n"
" var s = 0.0;\n"
" var dp = 0.0;\n"
" for (var d = 0u; d < D; d++) { s += aQ[ln * LD + d] * aK[r * LD + d]; }\n"
" for (var e = 0u; e < E; e++) { dp += aG[ln * LD + e] * aV[r * LD + e]; }\n"
" p = exp(s * F(0u) - aRl[ln]);\n"
" ds = p * (dp - aRd[ln]);\n"
" }\n"
" aP[r * PL + ln] = p;\n"
" aD[r * PL + ln] = ds;\n"
" workgroupBarrier();\n"
" for (var t = 0u; t < PER; t++) {\n"
" let d = ln + t * ATT_T;\n"
" if (d < E) {\n"
" var x = dv[t];\n"
" for (var c = 0u; c < ATT_T; c++) { x += aP[r * PL + c] * aG[c * LD + d]; }\n"
" dv[t] = x;\n"
" }\n"
" if (d < D) {\n"
" var x = dk[t];\n"
" for (var c = 0u; c < ATT_T; c++) { x += aD[r * PL + c] * aQ[c * LD + d]; }\n"
" dk[t] = x;\n"
" }\n"
" }\n"
" }\n"
" if (j < L) {\n"
" if (P(3u) != 0u) {\n"
" let ob = (j * Hh + h) * D;\n"
" for (var t = 0u; t < PER; t++) {\n"
" let d = ln + t * ATT_T;\n"
" if (d < D) {\n"
" if ((U(5u) & 8u) != 0u) { st(3u, ob + d, F(0u) * dk[t]); }\n"
" else { st(3u, ob + d, ld(3u, ob + d) + F(0u) * dk[t]); }\n"
" }\n"
" }\n"
" }\n"
" if (P(6u) != 0u) {\n"
" let ob = (j * Hh + h) * E;\n"
" for (var t = 0u; t < PER; t++) {\n"
" let e = ln + t * ATT_T;\n"
" if (e < E) {\n"
" if ((U(5u) & 16u) != 0u) { st(6u, ob + e, dv[t]); }\n"
" else { st(6u, ob + e, ld(6u, ob + e) + dv[t]); }\n"
" }\n"
" }\n"
" }\n"
" }\n"
" workgroupBarrier();\n"
" }\n"
"}\n"
"\n"
"// ── Fused losses -> loss[0] (zeroed by the caller), averaged ─────────────────\n"
"// Cross-entropy on class indices: p0 = logits [rows, C], p1 = labels, p2 =\n"
"// loss[0]; u0 = C; f0 = 1/rows n = rows\n"
"@compute @workgroup_size(256)\n"
"fn dk_ce_fwd(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" var acc = 0.0;\n"
" for (var rb = wg.x * (256u / RL); rb < N(); rb += nwg.x * (256u / RL)) {\n"
" let r = rb + li / RL;\n"
" let live = r < N();\n"
" let b = r * C;\n"
" var mx = -3.0e38;\n"
" for (var c = li % RL; live && c < C; c += RL) { mx = max(mx, ld(0u, b + c)); }\n"
" mx = seg_max(mx, li);\n"
" var s = 0.0;\n"
" for (var c = li % RL; live && c < C; c += RL) { s += exp(ld(0u, b + c) - mx); }\n"
" s = seg_sum(s, li);\n"
" if (live && li % RL == 0u) {\n"
" let lab = u32(ld(1u, r));\n"
" acc += (mx + log(s) - ld(0u, b + lab)) * F(0u);\n"
" }\n"
" }\n"
" block_add(acc, 2u, li);\n"
"}\n"
"// dlogits += seed[0] * (softmax - onehot) / rows; p3 = seed\n"
"@compute @workgroup_size(256)\n"
"fn dk_ce_bwd(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" let sc = ld(3u, 0u) * F(0u);\n"
" for (var rb = wg.x * (256u / RL); rb < N(); rb += nwg.x * (256u / RL)) {\n"
" let r = rb + li / RL;\n"
" let live = r < N();\n"
" let b = r * C;\n"
" var mx = -3.0e38;\n"
" for (var c = li % RL; live && c < C; c += RL) { mx = max(mx, ld(0u, b + c)); }\n"
" mx = seg_max(mx, li);\n"
" var s = 0.0;\n"
" for (var c = li % RL; live && c < C; c += RL) { s += exp(ld(0u, b + c) - mx); }\n"
" s = seg_sum(s, li);\n"
" if (!live) { continue; }\n"
" let lab = u32(ld(1u, r));\n"
" let inv = 1.0 / s;\n"
" for (var c = li % RL; c < C; c += RL) { accum(2u, b + c, sc * (exp(ld(0u, b + c) - mx) * inv - select(0.0, 1.0, c == lab))); }\n"
" }\n"
"}\n"
"// Elementwise losses: u0 = 0 BCE on logits, 1 Huber (f1 = delta); p0 = x,\n"
"// p1 = target, p2 = loss[0], f0 = 1/n n = elements\n"
"@compute @workgroup_size(256)\n"
"fn dk_ewloss_fwd(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u,\n"
" @builtin(local_invocation_index) li: u32) {\n"
" var acc = 0.0;\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let x = ld(0u, i);\n"
" let t = ld(1u, i);\n"
" var v = 0.0;\n"
" if (U(0u) == 0u) { v = max(x, 0.0) - x * t + log(1.0 + exp(-abs(x))); }\n"
" else {\n"
" let d = x - t;\n"
" let ad = abs(d);\n"
" let k = F(1u);\n"
" v = select(k * (ad - 0.5 * k), 0.5 * d * d, ad <= k);\n"
" }\n"
" acc += v * F(0u);\n"
" }\n"
" block_add(acc, 2u, li);\n"
"}\n"
"// dx += seed[0] / n * (sigmoid(x) - t) | clamp(x - t, -delta, delta); p2 = dx, p3 = seed\n"
"@compute @workgroup_size(256)\n"
"fn dk_ewloss_bwd(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let sc = ld(3u, 0u) * F(0u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let x = ld(0u, i);\n"
" let t = ld(1u, i);\n"
" var g = 0.0;\n"
" if (U(0u) == 0u) { g = sig_f(x) - t; } else { let k = F(1u); g = min(max(x - t, -k), k); }\n"
" accum(2u, i, sc * g);\n"
" }\n"
"}\n"
"\n"
"// ── Layer norm: a row per RL lanes, u0 = cols, f0 = eps ──────────────────────\n"
"// p0 = x, p1 = xhat out, p2 = inv std [rows]\n"
"@compute @workgroup_size(256)\n"
"fn dk_layernorm(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" for (var rb = wg.x * (256u / RL); rb < N(); rb += nwg.x * (256u / RL)) {\n"
" let r = rb + li / RL;\n"
" let live = r < N();\n"
" let b = r * C;\n"
" var m = 0.0;\n"
" for (var c = li % RL; live && c < C; c += RL) { m += ld(0u, b + c); }\n"
" m = seg_sum(m, li) / f32(C);\n"
" var v = 0.0;\n"
" for (var c = li % RL; live && c < C; c += RL) { let d = ld(0u, b + c) - m; v += d * d; }\n"
" let is = inverseSqrt(seg_sum(v, li) / f32(C) + F(0u));\n"
" if (!live) { continue; }\n"
" if (li % RL == 0u) { st(2u, r, is); }\n"
" for (var c = li % RL; c < C; c += RL) { st(1u, b + c, (ld(0u, b + c) - m) * is); }\n"
" }\n"
"}\n"
"// dx += inv * (g - mean(g) - xhat * mean(g * xhat)); p0 = g, p1 = xhat, p2 = inv, p3 = dx\n"
"@compute @workgroup_size(256)\n"
"fn dk_layernorm_bwd(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" for (var rb = wg.x * (256u / RL); rb < N(); rb += nwg.x * (256u / RL)) {\n"
" let r = rb + li / RL;\n"
" let live = r < N();\n"
" let b = r * C;\n"
" var mg = 0.0;\n"
" var mgy = 0.0;\n"
" for (var c = li % RL; live && c < C; c += RL) { let g = ld(0u, b + c); mg += g; mgy += g * ld(1u, b + c); }\n"
" mg = seg_sum(mg, li) / f32(C);\n"
" mgy = seg_sum(mgy, li) / f32(C);\n"
" if (!live) { continue; }\n"
" let is = ld(2u, r);\n"
" for (var c = li % RL; c < C; c += RL) {\n"
" st(3u, b + c, ld(3u, b + c) + is * (ld(0u, b + c) - mg - ld(1u, b + c) * mgy));\n"
" }\n"
" }\n"
"}\n"
"\n"
"// ── Optimizer step over one parameter (nn.h's OPT_* arithmetic) ─────────────\n"
"// p0 = w, p1 = g, p2 = m, p3 = v, p4 = scale[0] (gradient clip factor, or 0 =\n"
"// none); u0 = kind; f0 lr, f1 wd, f2 momentum, f3 beta1, f4 beta2, f5 eps,\n"
"// f6 1/bc1, f7 1/sqrt(bc2). w, g, m, v are word indices of fp32 buffers.\n"
"const OPT_SGD = 0;\n"
"const OPT_MOMENTUM = 1;\n"
"const OPT_NESTEROV = 2;\n"
"const OPT_RMSPROP = 3;\n"
"const OPT_ADAM = 4;\n"
"const OPT_ADAMW = 5;\n"
"const OPT_ADAGRAD = 6;\n"
"const OPT_LION = 7;\n"
"fn opt_one(w: u32, gr: u32, m: u32, v: u32, i: u32, cs: f32) {\n"
" let kind = i32(U(0u));\n"
" let lr = F(0u);\n"
" let wd = F(1u);\n"
" let mu = F(2u);\n"
" let b1 = F(3u);\n"
" let b2 = F(4u);\n"
" let eps = F(5u);\n"
" let g0 = rdw(gr + i) * cs;\n"
" let wi = rdw(w + i);\n"
" if (kind == OPT_SGD) { wrw(w + i, wi - lr * (g0 + wd * wi)); }\n"
" else if (kind == OPT_MOMENTUM || kind == OPT_NESTEROV) {\n"
" let g = g0 + wd * wi;\n"
" let mi = mu * rdw(m + i) + g;\n"
" wrw(m + i, mi);\n"
" wrw(w + i, wi - lr * select(mi, g + mu * mi, kind == OPT_NESTEROV));\n"
" } else if (kind == OPT_ADAGRAD) {\n"
" let g = g0 + wd * wi;\n"
" let vi = rdw(v + i) + g * g;\n"
" wrw(v + i, vi);\n"
" wrw(w + i, wi - lr * g / (sqrt(vi) + eps));\n"
" } else if (kind == OPT_LION) {\n"
" let mi = rdw(m + i);\n"
" let c = b1 * mi + (1.0 - b1) * g0;\n"
" let sg = select(select(0.0, -1.0, c < 0.0), 1.0, c > 0.0);\n"
" wrw(w + i, wi - lr * (sg + wd * wi));\n"
" wrw(m + i, b2 * mi + (1.0 - b2) * g0);\n"
" } else if (kind == OPT_RMSPROP) {\n"
" let g = g0 + wd * wi;\n"
" let vi = b2 * rdw(v + i) + (1.0 - b2) * g * g;\n"
" wrw(v + i, vi);\n"
" wrw(w + i, wi - lr * g / (sqrt(vi) + eps));\n"
" } else {\n"
" let l2 = select(wd, 0.0, kind == OPT_ADAMW);\n"
" let dec = select(0.0, wd, kind == OPT_ADAMW);\n"
" let g = g0 + l2 * wi;\n"
" let mi = b1 * rdw(m + i) + (1.0 - b1) * g;\n"
" let vi = b2 * rdw(v + i) + (1.0 - b2) * g * g;\n"
" wrw(m + i, mi);\n"
" wrw(v + i, vi);\n"
" wrw(w + i, wi - lr * ((mi * F(6u)) / (sqrt(vi) * F(7u) + eps) + dec * wi));\n"
" }\n"
"}\n"
"// p5 (when set): the loss scale's state; an overflowed step changes nothing.\n"
"@compute @workgroup_size(256)\n"
"fn dk_opt(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" if (P(5u) != 0u && ld(5u, 2u) != 0.0) { return; }\n"
" var cs = 1.0;\n"
" if (P(4u) != 0u) { cs = ld(4u, 0u); }\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" opt_one(P(0u) >> 2u, P(1u) >> 2u, P(2u) >> 2u, P(3u) >> 2u, i, cs);\n"
" }\n"
"}\n"
"// Every parameter's step in one launch: p0 = a table of u1 entries {w, g, m,\n"
"// v, n, start} (64-bit; their low words: byte offsets and counts)\n"
"@compute @workgroup_size(256)\n"
"fn dk_opt_multi(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let T = P(0u) >> 2u;\n"
" let cnt = U(1u);\n"
" if (P(5u) != 0u && ld(5u, 2u) != 0.0) { return; }\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" var lo = 0u;\n"
" var hi = cnt - 1u;\n"
" while (lo < hi) {\n"
" let mid = (lo + hi + 1u) / 2u;\n"
" if (H[T + mid * 12u + 10u] <= i) { lo = mid; } else { hi = mid - 1u; }\n"
" }\n"
" let e = T + lo * 12u;\n"
" opt_one(H[e] >> 2u, H[e + 2u] >> 2u, H[e + 4u] >> 2u, H[e + 6u] >> 2u, i - H[e + 10u], 1.0);\n"
" }\n"
"}\n"
"// scale[0] = sumsq[0] > maxNorm^2 ? maxNorm / sqrt(sumsq[0]) : 1; p0 = sumsq,\n"
"// p1 = scale, f0 = maxNorm\n"
"@compute @workgroup_size(256)\n"
"fn dk_clip_scale(@builtin(global_invocation_id) gid: vec3u) {\n"
" if (gid.x == 0u) {\n"
" let n = sqrt(ld(0u, 0u));\n"
" st(1u, 0u, select(1.0, F(0u) / n, n > F(0u) && n > 0.0));\n"
" }\n"
"}\n"
"\n"
"// ── Convolution helpers (convnet.h's layout: [batch, C*H*W] rows, NCHW) ─────\n"
"// u0 C, u1 H, u2 W, u3 kH, u4 kW, u5 OH, u6 OW; f0 strideH, f1 strideW, f2\n"
"// padH, f3 padW, f4 dilH, f5 dilW, f6 = B (integers carried as floats).\n"
"// col [C*kH*kW, B*OH*OW]; p0 = x [B, C*H*W], p1 = col n = C*kH*kW*B*OH*OW\n"
"struct ConvGeo { C: u32, Hh: u32, W: u32, KH: u32, KW: u32, OH: u32, OW: u32, B: u32,\n"
" sh: i32, sw: i32, ph: i32, pw: i32, dh: i32, dw: i32 }\n"
"fn conv_geo() -> ConvGeo {\n"
" return ConvGeo(U(0u), U(1u), U(2u), U(3u), U(4u), U(5u), U(6u), u32(F(6u)),\n"
" i32(F(0u)), i32(F(1u)), i32(F(2u)), i32(F(3u)), i32(F(4u)), i32(F(5u)));\n"
"}\n"
"// The input element of col entry i, or -1 outside the image.\n"
"fn col_src(g: ConvGeo, i: u32) -> i32 {\n"
" let OP = g.OH * g.OW;\n"
" let BP = g.B * OP;\n"
" let q = i % BP;\n"
" let r = i / BP;\n"
" let b = q / OP;\n"
" let p = q % OP;\n"
" let oh = p / g.OW;\n"
" let ow = p % g.OW;\n"
" let kw = r % g.KW;\n"
" let kh = (r / g.KW) % g.KH;\n"
" let c = r / (g.KW * g.KH);\n"
" let y = i32(oh) * g.sh - g.ph + i32(kh) * g.dh;\n"
" let x = i32(ow) * g.sw - g.pw + i32(kw) * g.dw;\n"
" if (y < 0 || y >= i32(g.Hh) || x < 0 || x >= i32(g.W)) { return -1; }\n"
" return i32(((b * g.C + c) * g.Hh + u32(y)) * g.W + u32(x));\n"
"}\n"
"@compute @workgroup_size(256)\n"
"fn dk_im2col(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let g = conv_geo();\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let s = col_src(g, i);\n"
" var v = 0.0;\n"
" if (s >= 0) { v = ld(0u, u32(s)); }\n"
" st(1u, i, v);\n"
" }\n"
"}\n"
"// x [B, C*H*W] += col2im(col): the adjoint of dk_im2col (atomics: windows overlap)\n"
"@compute @workgroup_size(256)\n"
"fn dk_col2im(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let g = conv_geo();\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let s = col_src(g, i);\n"
" if (s >= 0) { atadd(1u, u32(s), ld(0u, i)); }\n"
" }\n"
"}\n"
"// out [B2, A, P] (+)= in [A, B2, P] with the first two axes swapped; u0 A,\n"
"// u1 B2, u2 P, u3 = 1 to accumulate n = A*B2*P\n"
"@compute @workgroup_size(256)\n"
"fn dk_swap01(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let Aa = U(0u);\n"
" let B2 = U(1u);\n"
" let Pp = U(2u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let p = i % Pp;\n"
" let t = i / Pp;\n"
" let bi = t % B2;\n"
" let ai = t / B2;\n"
" let o = (bi * Aa + ai) * Pp + p;\n"
" if (U(3u) != 0u) { st(1u, o, ld(1u, o) + ld(0u, i)); } else { st(1u, o, ld(0u, i)); }\n"
" }\n"
"}\n"
"// The channel of value i in [batch, C, P].\n"
"fn chan_of(i: u32, C: u32, Pp: u32) -> u32 { return (i / Pp) % C; }\n"
"// out[b, c, p] += bias[c]; u0 = channels, u1 = plane n = batch*C*plane\n"
"@compute @workgroup_size(256)\n"
"fn dk_bias_channels(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" let Pp = U(1u);\n"
" let n = N();\n"
" if (pairs_out(0u) && (Pp & 1u) == 0u && (n & 1u) == 0u) {\n"
" // a pair shares its channel (even planes)\n"
" let o = P(0u) >> 2u;\n"
" let t = TG(0u);\n"
" for (var j = gid.x; j < n / 2u; j += nwg.x * 256u) {\n"
" let b = ldf(1u, chan_of(2u * j, C, Pp));\n"
" let w = H[o + j];\n"
" H[o + j] = pack16(h2f(w & 0xffffu, t) + b, h2f(w >> 16u, t) + b, t);\n"
" }\n"
" return;\n"
" }\n"
" for (var i = gid.x; i < n; i += nwg.x * 256u) { st(0u, i, ld(0u, i) + ld(1u, chan_of(i, C, Pp))); }\n"
"}\n"
"// The index of the q-th value of channel c.\n"
"fn chan_at(q: u32, c: u32, C: u32, Pp: u32) -> u32 {\n"
" let b = q / Pp;\n"
" return (b * C + c) * Pp + (q - b * Pp);\n"
"}\n"
"// gb[c] += sum over batch and plane of d; u0 = C, u1 = plane, u2 = batch,\n"
"// u3 = S slices: slice s takes the 256-value runs r with r % S == s n = C*S*256\n"
"@compute @workgroup_size(256)\n"
"fn dk_bias_channels_bwd(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" let Pp = U(1u);\n"
" let M = U(2u) * U(1u);\n"
" let S = U(3u);\n"
" for (var u = wg.x; u < C * S; u += nwg.x) {\n"
" let c = u / S;\n"
" let sl = u % S;\n"
" var s = 0.0;\n"
" for (var q = sl * 256u + li; q < M; q += S * 256u) { s += ld(0u, chan_at(q, c, C, Pp)); }\n"
" s = block_sum(s, li);\n"
" if (li == 0u) { atadd(1u, c, s); }\n"
" }\n"
"}\n"
"\n"
"// ── Batch / group normalization (convnet.h's layout [batch, C*plane]) ───────\n"
"// Batch norm statistics, pass 1: per (channel, slice) the count, mean and\n"
"// sum of squared deviations in p1 [C*S*3] (floats). p0 x; u0 C, u1 plane,\n"
"// u2 batch, u3 S n = C*S*256\n"
"@compute @workgroup_size(256)\n"
"fn dk_bn_partial(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" let Pp = U(1u);\n"
" let M = U(2u) * U(1u);\n"
" let S = U(3u);\n"
" for (var u = wg.x; u < C * S; u += nwg.x) {\n"
" let c = u / S;\n"
" let sl = u % S;\n"
" var s = 0.0;\n"
" var cnt = 0.0;\n"
" for (var q = sl * 256u + li; q < M; q += S * 256u) { s += ld(0u, chan_at(q, c, C, Pp)); cnt += 1.0; }\n"
" cnt = block_sum(cnt, li);\n"
" let tot = block_sum(s, li);\n"
" var mu = 0.0;\n"
" if (cnt > 0.0) { mu = tot / cnt; }\n"
" var sq = 0.0;\n"
" for (var q = sl * 256u + li; q < M; q += S * 256u) { let d = ld(0u, chan_at(q, c, C, Pp)) - mu; sq += d * d; }\n"
" sq = block_sum(sq, li);\n"
" if (li == 0u) { stf(1u, u * 3u, cnt); stf(1u, u * 3u + 1u, mu); stf(1u, u * 3u + 2u, sq); }\n"
" }\n"
"}\n"
"// Pass 2, one thread per channel: the slices combined (Chan et al.), the\n"
"// running statistics updated (training), mean and invstd. p0 partials, p1\n"
"// mean[C], p2 invstd[C], p3 running mean, p4 running var; u0 C, u3 S,\n"
"// u4 training; f0 momentum, f1 eps n = C\n"
"@compute @workgroup_size(256)\n"
"fn dk_bn_finish(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let S = U(3u);\n"
" for (var c = gid.x; c < N(); c += nwg.x * 256u) {\n"
" var mean = 0.0;\n"
" var vr = 0.0;\n"
" if (U(4u) != 0u) {\n"
" var n = 0.0;\n"
" var mu = 0.0;\n"
" var m2 = 0.0;\n"
" for (var s = 0u; s < S; s++) {\n"
" let p = (c * S + s) * 3u;\n"
" let pn = ldf(0u, p);\n"
" if (pn == 0.0) { continue; }\n"
" let nn = n + pn;\n"
" let d = ldf(0u, p + 1u) - mu;\n"
" mu += d * pn / nn;\n"
" m2 += ldf(0u, p + 2u) + d * d * n * pn / nn;\n"
" n = nn;\n"
" }\n"
" mean = mu;\n"
" vr = m2 / n;\n"
" var unb = vr;\n"
" if (n > 1.0) { unb = m2 / (n - 1.0); }\n"
" st(3u, c, (1.0 - F(0u)) * ld(3u, c) + F(0u) * mean);\n"
" st(4u, c, (1.0 - F(0u)) * ld(4u, c) + F(0u) * unb);\n"
" } else {\n"
" mean = ld(3u, c);\n"
" vr = ld(4u, c);\n"
" }\n"
" st(1u, c, mean);\n"
" st(2u, c, 1.0 / sqrt(vr + F(1u)));\n"
" }\n"
"}\n"
"// Group norm statistics per (image, group) unit of M contiguous values:\n"
"// p0 x, p1 mean[U], p2 invstd[U]; u0 M, u1 U; f0 eps n = U*256\n"
"@compute @workgroup_size(256)\n"
"fn dk_gn_stats(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let M = U(0u);\n"
" for (var u = wg.x; u < U(1u); u += nwg.x) {\n"
" let b = u * M;\n"
" var s = 0.0;\n"
" for (var q = li; q < M; q += 256u) { s += ld(0u, b + q); }\n"
" let mu = block_sum(s, li) / f32(M);\n"
" var sq = 0.0;\n"
" for (var q = li; q < M; q += 256u) { let d = ld(0u, b + q) - mu; sq += d * d; }\n"
" sq = block_sum(sq, li);\n"
" if (li == 0u) { st(1u, u, mu); st(2u, u, 1.0 / sqrt(sq / f32(M) + F(0u))); }\n"
" }\n"
"}\n"
"// The statistic of value i: its channel (batch norm, G = 0) or its (image, group).\n"
"fn norm_stat(i: u32, C: u32, Pp: u32, G: u32) -> u32 {\n"
" let t = i / Pp;\n"
" let img = t / C;\n"
" let c = t - img * C;\n"
" if (G == 0u) { return c; }\n"
" return img * G + c / (C / G);\n"
"}\n"
"// A normalization fused with the activation after it: u5 = 0 none, 1 ReLU,\n"
"// 2 leaky ReLU (slope f1); the backward kernels take its output in p7.\n"
"fn fused_act(v: f32) -> f32 {\n"
" if (U(5u) == 0u || v > 0.0) { return v; }\n"
" return select(F(1u) * v, 0.0, U(5u) == 1u);\n"
"}\n"
"fn fused_dy(i: u32) -> f32 {\n"
" let dy = ld(0u, i);\n"
" if (U(5u) == 0u || ld(7u, i) > 0.0) { return dy; }\n"
" return select(F(1u) * dy, 0.0, U(5u) == 1u);\n"
"}\n"
"// out = gamma[c] (x - mean) invstd + beta[c]: p0 x, p1 out, p2 mean, p3 invstd,\n"
"// p4 gamma, p5 beta; u0 C, u1 plane, u2 groups (0: batch norm) n = values\n"
"fn norm_apply_at(i: u32) -> f32 {\n"
" let C = U(0u);\n"
" let Pp = U(1u);\n"
" let c = chan_of(i, C, Pp);\n"
" let s = norm_stat(i, C, Pp, U(2u));\n"
" return fused_act(ld(4u, c) * (ld(0u, i) - ld(2u, s)) * ld(3u, s) + ld(5u, c));\n"
"}\n"
"@compute @workgroup_size(256)\n"
"fn dk_norm_apply(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let n = N();\n"
" let C = U(0u);\n"
" let Pp = U(1u);\n"
" if (pairs_out(1u) && pair_ok(0u, TG(1u)) && (Pp & 1u) == 0u && (n & 1u) == 0u) {\n"
" // both of a word's elements in one channel (even planes): its\n"
" // statistics read once, the input word once\n"
" let o = P(1u) >> 2u;\n"
" let t = TG(1u);\n"
" for (var j = gid.x; j < n / 2u; j += nwg.x * 256u) {\n"
" let c = chan_of(2u * j, C, Pp);\n"
" let sn = norm_stat(2u * j, C, Pp, U(2u));\n"
" let g = ld(4u, c) * ld(3u, sn);\n"
" let m = ld(2u, sn);\n"
" let b = ld(5u, c);\n"
" let x = ldw2(0u, j);\n"
" H[o + j] = pack16(fused_act(g * (x.x - m) + b), fused_act(g * (x.y - m) + b), t);\n"
" }\n"
" return;\n"
" }\n"
" if (pairs_out(1u)) {\n"
" let o = P(1u) >> 2u;\n"
" let t = TG(1u);\n"
" for (var j = gid.x; j < n / 2u; j += nwg.x * 256u) { H[o + j] = pack16(norm_apply_at(2u * j), norm_apply_at(2u * j + 1u), t); }\n"
" if ((n & 1u) != 0u && gid.x == 0u) { st(1u, n - 1u, norm_apply_at(n - 1u)); }\n"
" return;\n"
" }\n"
" for (var i = gid.x; i < n; i += nwg.x * 256u) { st(1u, i, norm_apply_at(i)); }\n"
"}\n"
"// Per channel: dgamma[c] += sum dy xh, dbeta[c] += sum dy, and (p6, when\n"
"// set) sums[c] += sum dy, sums[C + c] += sum dy xh. p0 dy, p1 x, p2 mean,\n"
"// p3 invstd, p4 dgamma, p5 dbeta (either may be 0); u0 C, u1 plane,\n"
"// u2 groups, u3 batch, u4 S slices n = C*S*256\n"
"@compute @workgroup_size(256)\n"
"fn dk_norm_sums(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" let Pp = U(1u);\n"
" let G = U(2u);\n"
" let M = U(3u) * U(1u);\n"
" let S = U(4u);\n"
" for (var u = wg.x; u < C * S; u += nwg.x) {\n"
" let c = u / S;\n"
" let sl = u % S;\n"
" var sd = 0.0;\n"
" var sx = 0.0;\n"
" for (var q = sl * 256u + li; q < M; q += S * 256u) {\n"
" let i = chan_at(q, c, C, Pp);\n"
" var s = c;\n"
" if (G != 0u) { s = norm_stat(i, C, Pp, G); }\n"
" let dy = fused_dy(i);\n"
" sd += dy;\n"
" sx += dy * (ld(1u, i) - ld(2u, s)) * ld(3u, s);\n"
" }\n"
" sd = block_sum(sd, li);\n"
" sx = block_sum(sx, li);\n"
" if (li == 0u) {\n"
" if (P(4u) != 0u) { atadd(4u, c, sx); }\n"
" if (P(5u) != 0u) { atadd(5u, c, sd); }\n"
" if (P(6u) != 0u) { atadd(6u, c, sd); atadd(6u, C + c, sx); }\n"
" }\n"
" }\n"
"}\n"
"// Batch norm: dx += gamma invstd (dy - m1 - xh m2), m1 = sum dy / N, m2 =\n"
"// sum dy xh / N in training (0 in eval). p0 dy, p1 x, p2 mean, p3 invstd,\n"
"// p4 gamma, p5 sums, p6 dx; u0 C, u1 plane, u2 training; f0 1/N n = values\n"
"fn bn_bwd_dx_at(i: u32) -> f32 {\n"
" let C = U(0u);\n"
" let c = chan_of(i, C, U(1u));\n"
" let inv = ld(3u, c);\n"
" var m1 = 0.0;\n"
" var m2 = 0.0;\n"
" if (U(2u) != 0u) { m1 = ld(5u, c) * F(0u); m2 = ld(5u, C + c) * F(0u); }\n"
" let xh = (ld(1u, i) - ld(2u, c)) * inv;\n"
" var v = ld(4u, c) * inv * (fused_dy(i) - m1 - xh * m2);\n"
" if (U(6u) == 0u) { v += ld(6u, i); }\n"
" return v;\n"
"}\n"
"@compute @workgroup_size(256)\n"
"fn dk_bn_bwd_dx(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let n = N();\n"
" let C = U(0u);\n"
" let Pp = U(1u);\n"
" let t = TG(6u);\n"
" if (pairs_out(6u) && pair_ok(0u, t) && pair_ok(1u, t) && (U(5u) == 0u || pair_ok(7u, t)) &&\n"
" (Pp & 1u) == 0u && (n & 1u) == 0u) {\n"
" // both of a word's elements in one channel (even planes): the\n"
" // channel's values read once, each input word once\n"
" let o = P(6u) >> 2u;\n"
" for (var j = gid.x; j < n / 2u; j += nwg.x * 256u) {\n"
" let c = chan_of(2u * j, C, Pp);\n"
" let inv = ld(3u, c);\n"
" var m1 = 0.0;\n"
" var m2 = 0.0;\n"
" if (U(2u) != 0u) { m1 = ld(5u, c) * F(0u); m2 = ld(5u, C + c) * F(0u); }\n"
" let g = ld(4u, c) * inv;\n"
" var dy = ldw2(0u, j);\n"
" if (U(5u) != 0u) {\n"
" let y = ldw2(7u, j);\n"
" let sl = select(F(1u), 0.0, U(5u) == 1u);\n"
" if (!(y.x > 0.0)) { dy.x *= sl; }\n"
" if (!(y.y > 0.0)) { dy.y *= sl; }\n"
" }\n"
" let xh = (ldw2(1u, j) - vec2f(ld(2u, c))) * inv;\n"
" var v = g * (dy - vec2f(m1) - xh * m2);\n"
" if (U(6u) == 0u) { v += ldw2(6u, j); }\n"
" H[o + j] = pack16(v.x, v.y, t);\n"
" }\n"
" return;\n"
" }\n"
" if (pairs_out(6u)) {\n"
" let o = P(6u) >> 2u;\n"
" let t = TG(6u);\n"
" for (var j = gid.x; j < n / 2u; j += nwg.x * 256u) { H[o + j] = pack16(bn_bwd_dx_at(2u * j), bn_bwd_dx_at(2u * j + 1u), t); }\n"
" if ((n & 1u) != 0u && gid.x == 0u) { st(6u, n - 1u, bn_bwd_dx_at(n - 1u)); }\n"
" return;\n"
" }\n"
" for (var i = gid.x; i < n; i += nwg.x * 256u) { st(6u, i, bn_bwd_dx_at(i)); }\n"
"}\n"
"// Group norm: per (image, group) unit, dxh = dy gamma[c],\n"
"// dx += invstd (dxh - mean(dxh) - xh mean(dxh xh)). p0 dy, p1 x, p2 mean,\n"
"// p3 invstd, p4 gamma, p5 dx; u0 C, u1 plane, u2 groups, u3 units n = U*256\n"
"@compute @workgroup_size(256)\n"
"fn dk_gn_bwd(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let C = U(0u);\n"
" let Pp = U(1u);\n"
" let G = U(2u);\n"
" let Cg = C / G;\n"
" let M = Cg * Pp;\n"
" for (var u = wg.x; u < U(3u); u += nwg.x) {\n"
" let b = u * M;\n"
" let g = u % G;\n"
" let mean = ld(2u, u);\n"
" let inv = ld(3u, u);\n"
" var s1 = 0.0;\n"
" var s2 = 0.0;\n"
" for (var q = li; q < M; q += 256u) {\n"
" let dxh = fused_dy(b + q) * ld(4u, g * Cg + q / Pp);\n"
" s1 += dxh;\n"
" s2 += dxh * (ld(1u, b + q) - mean) * inv;\n"
" }\n"
" let m1 = block_sum(s1, li) / f32(M);\n"
" let m2 = block_sum(s2, li) / f32(M);\n"
" for (var q = li; q < M; q += 256u) {\n"
" let dxh = fused_dy(b + q) * ld(4u, g * Cg + q / Pp);\n"
" let xh = (ld(1u, b + q) - mean) * inv;\n"
" st(5u, b + q, ld(5u, b + q) + inv * (dxh - m1 - xh * m2));\n"
" }\n"
" }\n"
"}\n"
"\n"
"// ── Pooling ([batch, C*H*W] planes) ──────────────────────────────────────────\n"
"// u0 C, u1 H, u2 W, u3 oH, u4 oW, u5 mode (0 max, 1 avg over the padded\n"
"// window, 2 avg over the part inside the image); f0..f5 kH, kW, sH, sW, pH, pW.\n"
"struct PoolWin { h0: u32, h1: u32, w0: u32, w1: u32, padded: u32 }\n"
"fn pool_win(oh: u32, ow: u32) -> PoolWin {\n"
" let Hh = i32(U(1u));\n"
" let W = i32(U(2u));\n"
" let kH = i32(F(0u));\n"
" let kW = i32(F(1u));\n"
" let sH = i32(F(2u));\n"
" let sW = i32(F(3u));\n"
" let pH = i32(F(4u));\n"
" let pW = i32(F(5u));\n"
" let h0 = i32(oh) * sH - pH;\n"
" let w0 = i32(ow) * sW - pW;\n"
" let h1 = h0 + kH;\n"
" let w1 = w0 + kW;\n"
" let hp = min(h1, Hh + pH);\n"
" let wp = min(w1, W + pW);\n"
" return PoolWin(u32(max(h0, 0)), u32(min(h1, Hh)), u32(max(w0, 0)), u32(min(w1, W)), u32((hp - h0) * (wp - w0)));\n"
"}\n"
"// The (first) largest value's index in window r of plane base (element offset).\n"
"fn pool_argmax(base: u32, r: PoolWin) -> u32 {\n"
" let W = U(2u);\n"
" var best = r.h0 * W + r.w0;\n"
" var bv = ld(0u, base + best);\n"
" for (var y = r.h0; y < r.h1; y++) {\n"
" for (var x = r.w0; x < r.w1; x++) {\n"
" let v = ld(0u, base + y * W + x);\n"
" if (v > bv) { bv = v; best = y * W + x; }\n"
" }\n"
" }\n"
" return best;\n"
"}\n"
"fn pool_div(r: PoolWin) -> f32 {\n"
" if (U(5u) == 1u) { return f32(r.padded); }\n"
" return f32((r.h1 - r.h0) * (r.w1 - r.w0));\n"
"}\n"
"// p0 x, p1 out [batch, C*oH*oW] (forward); p0 x, p1 dout, p2 dx (backward:\n"
"// u6 = 1, atomics) n = batch*C*oH*oW\n"
"@compute @workgroup_size(256)\n"
"fn dk_pool(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let Hh = U(1u);\n"
" let W = U(2u);\n"
" let oH = U(3u);\n"
" let oW = U(4u);\n"
" let mode = U(5u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let t = i / oW;\n"
" let ow = i - t * oW;\n"
" let pl = t / oH;\n"
" let oh = t - pl * oH;\n"
" let r = pool_win(oh, ow);\n"
" let base = pl * Hh * W;\n"
" if (mode == 0u) {\n"
" let best = pool_argmax(base, r);\n"
" if (U(6u) != 0u) { atadd(2u, base + best, ld(1u, i)); } else { st(1u, i, ld(0u, base + best)); }\n"
" } else {\n"
" let div = pool_div(r);\n"
" if (U(6u) != 0u) {\n"
" let share = ld(1u, i) / div;\n"
" for (var y = r.h0; y < r.h1; y++) { for (var x = r.w0; x < r.w1; x++) { atadd(2u, base + y * W + x, share); } }\n"
" } else {\n"
" var acc = 0.0;\n"
" for (var y = r.h0; y < r.h1; y++) { for (var x = r.w0; x < r.w1; x++) { acc += ld(0u, base + y * W + x); } }\n"
" st(1u, i, acc / div);\n"
" }\n"
" }\n"
" }\n"
"}\n"
"// The gradient input i gathers from the windows that hold it -- for max\n"
"// pooling, those whose (first) largest value it is.\n"
"fn pool_bwd_at(i: u32) -> f32 {\n"
" let Hh = U(1u);\n"
" let W = U(2u);\n"
" let oH = U(3u);\n"
" let oW = U(4u);\n"
" let kH = i32(F(0u));\n"
" let kW = i32(F(1u));\n"
" let sH = i32(F(2u));\n"
" let sW = i32(F(3u));\n"
" let pH = i32(F(4u));\n"
" let pW = i32(F(5u));\n"
" let t = i / W;\n"
" let x = i - t * W;\n"
" let pl = t / Hh;\n"
" let y = t - pl * Hh;\n"
" let base = pl * Hh * W;\n"
" let dbase = pl * oH * oW;\n"
" var oh0: i32;\n"
" var oh1: i32;\n"
" var ow0: i32;\n"
" var ow1: i32;\n"
" if (kH == sH && kW == sW && pH == 0 && pW == 0) {\n"
" // non-overlapping windows: exactly one holds (y, x)\n"
" oh0 = i32(y / u32(kH)); oh1 = oh0;\n"
" ow0 = i32(x / u32(kW)); ow1 = ow0;\n"
" } else {\n"
" oh0 = (i32(y) + pH - kH + sH) / sH; oh1 = (i32(y) + pH) / sH; // windows over row y\n"
" ow0 = (i32(x) + pW - kW + sW) / sW; ow1 = (i32(x) + pW) / sW;\n"
" if (i32(y) + pH - kH + 1 < 0) { oh0 = 0; }\n"
" if (i32(x) + pW - kW + 1 < 0) { ow0 = 0; }\n"
" }\n"
" if (oh1 > i32(oH) - 1) { oh1 = i32(oH) - 1; }\n"
" if (ow1 > i32(oW) - 1) { ow1 = i32(oW) - 1; }\n"
" var s = 0.0;\n"
" for (var oh = oh0; oh <= oh1; oh++) {\n"
" for (var ow = ow0; ow <= ow1; ow++) {\n"
" let r = pool_win(u32(oh), u32(ow));\n"
" if (y < r.h0 || y >= r.h1 || x < r.w0 || x >= r.w1) { continue; }\n"
" let d = ld(1u, dbase + u32(oh) * oW + u32(ow));\n"
" if (U(5u) == 0u) {\n"
" if (pool_argmax(base, r) == y * W + x) { s += d; }\n"
" } else {\n"
" s += d / pool_div(r);\n"
" }\n"
" }\n"
" }\n"
" if (U(6u) == 0u) { s += ld(2u, i); }\n"
" return s;\n"
"}\n"
"// The pooling's backward as a gather (no atomics, any element type). p0 x,\n"
"// p1 dout, p2 dx; geometry as dk_pool; u6 = 1: dx was an unwritten zero\n"
"// gradient n = planes*H*W\n"
"@compute @workgroup_size(256)\n"
"fn dk_pool_bwd(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let Hh = U(1u);\n"
" let W = U(2u);\n"
" let oH = U(3u);\n"
" let oW = U(4u);\n"
" let mode = U(5u);\n"
" let kH = i32(F(0u));\n"
" let kW = i32(F(1u));\n"
" let sH = i32(F(2u));\n"
" let sW = i32(F(3u));\n"
" let pH = i32(F(4u));\n"
" let pW = i32(F(5u));\n"
" let n = N();\n"
" let tiled = kH == sH && kW == sW && pH == 0 && pW == 0;\n"
" if (tiled && U(6u) != 0u && oH * u32(kH) == Hh && oW * u32(kW) == W && pairs_out(2u) && (kW & 1) == 0 && (W & 1u) == 0u) {\n"
" // exact tiling into a 16-bit gradient, windows an even number of\n"
" // columns wide: a thread per window writes its rows as whole words\n"
" let o = P(2u) >> 2u;\n"
" let t = TG(2u);\n"
" let planes = n / (Hh * W);\n"
" for (var i = gid.x; i < planes * oH * oW; i += nwg.x * 256u) {\n"
" let q = i / oW;\n"
" let ow = i - q * oW;\n"
" let pl = q / oH;\n"
" let oh = q - pl * oH;\n"
" let r = pool_win(oh, ow);\n"
" let base = pl * Hh * W;\n"
" let d = ld(1u, i);\n"
" var best = 0xffffffffu;\n"
" var share = 0.0;\n"
" if (mode == 0u) { best = pool_argmax(base, r); } else { share = d / pool_div(r); }\n"
" for (var yy = r.h0; yy < r.h1; yy++) {\n"
" for (var xx = r.w0; xx < r.w1; xx += 2u) {\n"
" let e = yy * W + xx;\n"
" var lo = share;\n"
" var hi = share;\n"
" if (mode == 0u) { lo = select(0.0, d, e == best); hi = select(0.0, d, e + 1u == best); }\n"
" H[o + (base + e) / 2u] = pack16(lo, hi, t);\n"
" }\n"
" }\n"
" }\n"
" return;\n"
" }\n"
" if (tiled && (U(6u) == 0u || (oH * u32(kH) == Hh && oW * u32(kW) == W))) {\n"
" // non-overlapping windows partition the inputs: a thread per window\n"
" let planes = n / (Hh * W);\n"
" for (var i = gid.x; i < planes * oH * oW; i += nwg.x * 256u) {\n"
" let t = i / oW;\n"
" let ow = i - t * oW;\n"
" let pl = t / oH;\n"
" let oh = t - pl * oH;\n"
" let r = pool_win(oh, ow);\n"
" let base = pl * Hh * W;\n"
" let d = ld(1u, i);\n"
" if (mode == 0u) {\n"
" let best = pool_argmax(base, r);\n"
" if (U(6u) != 0u) {\n"
" for (var yy = r.h0; yy < r.h1; yy++) {\n"
" for (var xx = r.w0; xx < r.w1; xx++) { st(2u, base + yy * W + xx, select(0.0, d, yy * W + xx == best)); }\n"
" }\n"
" } else { st(2u, base + best, ld(2u, base + best) + d); }\n"
" } else {\n"
" let share = d / pool_div(r);\n"
" for (var yy = r.h0; yy < r.h1; yy++) { for (var xx = r.w0; xx < r.w1; xx++) { accum(2u, base + yy * W + xx, share); } }\n"
" }\n"
" }\n"
" return;\n"
" }\n"
" if (pairs_out(2u)) {\n"
" // a 16-bit gradient otherwise: two inputs per thread, whole words\n"
" let o = P(2u) >> 2u;\n"
" let t = TG(2u);\n"
" for (var j = gid.x; j < n / 2u; j += nwg.x * 256u) { H[o + j] = pack16(pool_bwd_at(2u * j), pool_bwd_at(2u * j + 1u), t); }\n"
" if ((n & 1u) != 0u && gid.x == 0u) { st(2u, n - 1u, pool_bwd_at(n - 1u)); }\n"
" return;\n"
" }\n"
" for (var i = gid.x; i < n; i += nwg.x * 256u) { st(2u, i, pool_bwd_at(i)); }\n"
"}\n"
"// Global pooling, one plane of u0 values per unit: p0 x, p1 out [planes];\n"
"// backward (u2 = 1): p1 dout, p2 dx. u1 = 0 max, 1 avg n = planes*256\n"
"@compute @workgroup_size(256)\n"
"fn dk_gpool(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u,\n"
" @builtin(num_workgroups) nwg: vec3u) {\n"
" let Pp = U(0u);\n"
" for (var pl = wg.x; pl < N() / 256u; pl += nwg.x) {\n"
" let b = pl * Pp;\n"
" if (U(1u) != 0u) {\n"
" if (U(2u) != 0u) {\n"
" let share = ld(1u, pl) / f32(Pp);\n"
" for (var q = li; q < Pp; q += 256u) { accum(2u, b + q, share); }\n"
" } else {\n"
" var s = 0.0;\n"
" for (var q = li; q < Pp; q += 256u) { s += ld(0u, b + q); }\n"
" s = block_sum(s, li);\n"
" if (li == 0u) { st(1u, pl, s / f32(Pp)); }\n"
" }\n"
" } else if (li == 0u) {\n"
" var best = 0u;\n"
" for (var q = 1u; q < Pp; q++) { if (ld(0u, b + q) > ld(0u, b + best)) { best = q; } }\n"
" if (U(2u) != 0u) { st(2u, b + best, ld(2u, b + best) + ld(1u, pl)); } else { st(1u, pl, ld(0u, b + best)); }\n"
" }\n"
" }\n"
"}\n"
"\n"
"// ── Upsampling and resizing (planes of H x W -> oH x oW) ─────────────────────\n"
"// Nearest by an integer factor u2: forward (u3 = 0) p0 x, p1 out, one thread\n"
"// per output; backward (u3 = 1) p0 dout, p1 dx, one per input. u0 W, u1 H\n"
"@compute @workgroup_size(256)\n"
"fn dk_upsample(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let W = U(0u);\n"
" let Hh = U(1u);\n"
" let f = U(2u);\n"
" let oW = W * f;\n"
" let oH = Hh * f;\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" if (U(3u) == 0u) {\n"
" let ow = i % oW;\n"
" let oh = (i / oW) % oH;\n"
" let pl = i / (oW * oH);\n"
" st(1u, i, ld(0u, (pl * Hh + oh / f) * W + ow / f));\n"
" } else {\n"
" let w = i % W;\n"
" let h = (i / W) % Hh;\n"
" let pl = i / (W * Hh);\n"
" var s = 0.0;\n"
" for (var y = 0u; y < f; y++) { for (var x = 0u; x < f; x++) { s += ld(0u, (pl * oH + h * f + y) * oW + w * f + x); } }\n"
" accum(1u, i, s);\n"
" }\n"
" }\n"
"}\n"
"// Up to four taps per axis (unused ones weigh 0).\n"
"struct Taps { i0: u32, i1: u32, i2: u32, i3: u32, w0: f32, w1: f32, w2: f32, w3: f32 }\n"
"fn rs_src(o: u32, inn: u32, out: u32, align: u32, clampZero: bool) -> f32 {\n"
" if (align != 0u) {\n"
" if (out <= 1u) { return 0.0; }\n"
" return f32(o) * f32(inn - 1u) / f32(out - 1u);\n"
" }\n"
" var s = (f32(o) + 0.5) * f32(inn) / f32(out) - 0.5;\n"
" if (clampZero && s < 0.0) { s = 0.0; }\n"
" return s;\n"
"}\n"
"fn rs_clamp(j: i32, inn: u32) -> u32 {\n"
" if (j < 0) { return 0u; }\n"
" if (j >= i32(inn)) { return inn - 1u; }\n"
" return u32(j);\n"
"}\n"
"fn rs_taps(o: u32, inn: u32, out: u32, mode: u32, align: u32) -> Taps {\n"
" var t = Taps(0u, 0u, 0u, 0u, 0.0, 0.0, 0.0, 0.0);\n"
" if (mode == 0u) {\n"
" let i = u32(floor(f32(o) * (f32(inn) / f32(out))));\n"
" let c = min(i, inn - 1u);\n"
" t.i0 = c; t.i1 = c; t.i2 = c; t.i3 = c; t.w0 = 1.0;\n"
" } else if (mode == 1u) {\n"
" let s = rs_src(o, inn, out, align, true);\n"
" let i0 = min(u32(s), inn - 1u);\n"
" let l = s - f32(i0);\n"
" let i1 = select(inn - 1u, i0 + 1u, i0 + 1u < inn);\n"
" t.i0 = i0; t.i1 = i1; t.i2 = i1; t.i3 = i1;\n"
" t.w0 = 1.0 - l; t.w1 = l;\n"
" } else {\n"
" let s = rs_src(o, inn, out, align, false);\n"
" let f = floor(s);\n"
" let x = s - f;\n"
" let Aa = -0.75;\n"
" let x1 = x + 1.0;\n"
" let x2 = 1.0 - x;\n"
" let x3 = 2.0 - x;\n"
" t.w0 = ((Aa * x1 - 5.0 * Aa) * x1 + 8.0 * Aa) * x1 - 4.0 * Aa;\n"
" t.w1 = ((Aa + 2.0) * x - (Aa + 3.0)) * x * x + 1.0;\n"
" t.w2 = ((Aa + 2.0) * x2 - (Aa + 3.0)) * x2 * x2 + 1.0;\n"
" t.w3 = ((Aa * x3 - 5.0 * Aa) * x3 + 8.0 * Aa) * x3 - 4.0 * Aa;\n"
" let i = i32(f);\n"
" t.i0 = rs_clamp(i - 1, inn); t.i1 = rs_clamp(i, inn); t.i2 = rs_clamp(i + 1, inn); t.i3 = rs_clamp(i + 2, inn);\n"
" }\n"
" return t;\n"
"}\n"
"// sum_k w_k row[i_k] (row: an element offset into argument 0)\n"
"fn rs_dot(t: Taps, row: u32) -> f32 {\n"
" var s = 0.0;\n"
" if (t.w0 != 0.0) { s += t.w0 * ld(0u, row + t.i0); }\n"
" if (t.w1 != 0.0) { s += t.w1 * ld(0u, row + t.i1); }\n"
" if (t.w2 != 0.0) { s += t.w2 * ld(0u, row + t.i2); }\n"
" if (t.w3 != 0.0) { s += t.w3 * ld(0u, row + t.i3); }\n"
" return s;\n"
"}\n"
"// convnet.h's RESIZE_* (u4), alignCorners u5: p0 x, p1 out. u0 H, u1 W, u2 oH,\n"
"// u3 oW n = planes*oH*oW\n"
"@compute @workgroup_size(256)\n"
"fn dk_resize(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let Hh = U(0u);\n"
" let W = U(1u);\n"
" let oH = U(2u);\n"
" let oW = U(3u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let ow = i % oW;\n"
" let oh = (i / oW) % oH;\n"
" let pl = i / (oW * oH);\n"
" let ty = rs_taps(oh, Hh, oH, U(4u), U(5u));\n"
" let tx = rs_taps(ow, W, oW, U(4u), U(5u));\n"
" let img = pl * Hh * W;\n"
" var acc = 0.0;\n"
" if (ty.w0 != 0.0) { acc += ty.w0 * rs_dot(tx, img + ty.i0 * W); }\n"
" if (ty.w1 != 0.0) { acc += ty.w1 * rs_dot(tx, img + ty.i1 * W); }\n"
" if (ty.w2 != 0.0) { acc += ty.w2 * rs_dot(tx, img + ty.i2 * W); }\n"
" if (ty.w3 != 0.0) { acc += ty.w3 * rs_dot(tx, img + ty.i3 * W); }\n"
" st(1u, i, acc);\n"
" }\n"
"}\n"
"// The weight with which output o's taps read input i along one axis (0: none).\n"
"fn rs_weight(o: u32, i: u32, inn: u32, out: u32, mode: u32, align: u32) -> f32 {\n"
" let t = rs_taps(o, inn, out, mode, align);\n"
" return select(0.0, t.w0, t.i0 == i) + select(0.0, t.w1, t.i1 == i) + select(0.0, t.w2, t.i2 == i) +\n"
" select(0.0, t.w3, t.i3 == i);\n"
"}\n"
"// The outputs that can read input i: [lo, hi] (see device_kernels.cu).\n"
"fn rs_range(i: u32, inn: u32, out: u32, mode: u32, align: u32) -> vec2i {\n"
" var a0: f32;\n"
" var a1: f32;\n"
" let lo_s = select(f32(i) - 2.0, f32(i) - 1.0, mode == 1u);\n"
" let hi_s = select(f32(i) + 2.0, f32(i) + 1.0, mode == 1u);\n"
" if (mode == 0u) {\n"
" let r = f32(out) / f32(inn);\n"
" a0 = f32(i) * r;\n"
" a1 = f32(i + 1u) * r;\n"
" } else if (align != 0u) {\n"
" var r = f32(out);\n"
" if (inn > 1u) { r = f32(out - 1u) / f32(inn - 1u); }\n"
" a0 = lo_s * r;\n"
" a1 = hi_s * r;\n"
" } else {\n"
" let r = f32(out) / f32(inn);\n"
" a0 = (lo_s + 0.5) * r - 0.5;\n"
" a1 = (hi_s + 0.5) * r - 0.5;\n"
" }\n"
" let l = i32(floor(a0)) - 1;\n"
" let h = i32(ceil(a1)) + 1;\n"
" return vec2i(max(l, 0), min(h, i32(out) - 1));\n"
"}\n"
"// The resize's backward as two gathers (no atomics), one axis at a time:\n"
"// pass 0 (f0 = 0): t[pl, oy, x] = sum over ox reading x of dout[pl, oy, ox] wx;\n"
"// pass 1 (f0 = 1): dx[pl, y, x] (+)= sum over oy reading y of t[pl, oy, x] wy.\n"
"// p0 source, p1 destination; u0 H, u1 W, u2 oH, u3 oW, u4 mode, u5 align\n"
"@compute @workgroup_size(256)\n"
"fn dk_resize_bwd(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let Hh = U(0u);\n"
" let W = U(1u);\n"
" let oH = U(2u);\n"
" let oW = U(3u);\n"
" let mode = U(4u);\n"
" let align = U(5u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let x = i % W;\n"
" let r = i / W;\n"
" var s = 0.0;\n"
" if (F(0u) == 0.0) {\n"
" // r = pl * oH + oy: a row of dout\n"
" let q = rs_range(x, W, oW, mode, align);\n"
" for (var ox = q.x; ox <= q.y; ox++) {\n"
" let w = rs_weight(u32(ox), x, W, oW, mode, align);\n"
" if (w != 0.0) { s += w * ld(0u, r * oW + u32(ox)); }\n"
" }\n"
" st(1u, i, s);\n"
" } else {\n"
" let y = r % Hh;\n"
" let pl = r / Hh;\n"
" let q = rs_range(y, Hh, oH, mode, align);\n"
" for (var oy = q.x; oy <= q.y; oy++) {\n"
" let w = rs_weight(u32(oy), y, Hh, oH, mode, align);\n"
" if (w != 0.0) { s += w * ld(0u, pl * oH * W + x + u32(oy) * W); }\n"
" }\n"
" accum(1u, i, s);\n"
" }\n"
" }\n"
"}\n"
"\n"
"// ── The dynamic loss scale ───────────────────────────────────────────────────\n"
"// State (4 floats): [scale, 1/scale, overflow flag, good steps].\n"
"// g *= 1/scale, the flag raised on a non-finite value: p0 g, p1 state n = values\n"
"@compute @workgroup_size(256)\n"
"fn dk_amp_unscale(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let inv = ld(1u, 1u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let v = ld(0u, i) * inv;\n"
" if (!finite(v)) { st(1u, 2u, 1.0); }\n"
" st(0u, i, v);\n"
" }\n"
"}\n"
"// Every gradient of a table at once: p0 = u1 entries {g, n, start} (64-bit;\n"
"// low words), p1 state n = total values\n"
"@compute @workgroup_size(256)\n"
"fn dk_amp_unscale_multi(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let T = P(0u) >> 2u;\n"
" let cnt = U(1u);\n"
" let inv = ld(1u, 1u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" var lo = 0u;\n"
" var hi = cnt - 1u;\n"
" while (lo < hi) {\n"
" let mid = (lo + hi + 1u) / 2u;\n"
" if (H[T + mid * 6u + 4u] <= i) { lo = mid; } else { hi = mid - 1u; }\n"
" }\n"
" let g = (H[T + lo * 6u] >> 2u) + i - H[T + lo * 6u + 4u];\n"
" let v = rdw(g) * inv;\n"
" if (!finite(v)) { st(1u, 2u, 1.0); }\n"
" wrw(g, v);\n"
" }\n"
"}\n"
"// After a step: an overflow halves the scale, u0 good steps in a row double\n"
"// it; the flag is cleared. p0 state n = 1\n"
"@compute @workgroup_size(256)\n"
"fn dk_amp_update(@builtin(global_invocation_id) gid: vec3u) {\n"
" if (gid.x == 0u) {\n"
" if (ld(0u, 2u) != 0.0) { st(0u, 0u, max(ld(0u, 0u) * 0.5, 1.0)); st(0u, 3u, 0.0); }\n"
" else {\n"
" st(0u, 3u, ld(0u, 3u) + 1.0);\n"
" if (ld(0u, 3u) >= f32(U(0u))) { st(0u, 0u, min(ld(0u, 0u) * 2.0, 16777216.0)); st(0u, 3u, 0.0); }\n"
" }\n"
" st(0u, 1u, 1.0 / ld(0u, 0u));\n"
" st(0u, 2u, 0.0);\n"
" }\n"
"}\n"
"\n"
"// ── Sparse (CSR) x dense: sparse.h ───────────────────────────────────────────\n"
"// out [rows, n] (+)= S . X; p0 rowPtr, p1 colIdx (float indices), p2 values,\n"
"// p3 X, p4 out; u0 = n, u1 = 1 to accumulate n = rows*ncols\n"
"@compute @workgroup_size(256)\n"
"fn dk_spmm(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let NC = U(0u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let r = i / NC;\n"
" let j = i % NC;\n"
" let lo = u32(ld(0u, r));\n"
" let hi = u32(ld(0u, r + 1u));\n"
" var s = 0.0;\n"
" for (var e = lo; e < hi; e++) { s += ld(2u, e) * ld(3u, u32(ld(1u, e)) * NC + j); }\n"
" if (U(1u) != 0u) { st(4u, i, ld(4u, i) + s); } else { st(4u, i, s); }\n"
" }\n"
"}\n"
"// dX[cols, n] += S^T . G: one thread per (nonzero, column); p5 = row of\n"
"// each nonzero n = nnz*ncols\n"
"@compute @workgroup_size(256)\n"
"fn dk_spmm_tgrad(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let NC = U(0u);\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let e = i / NC;\n"
" let j = i % NC;\n"
" let r = u32(ld(5u, e));\n"
" let col = u32(ld(1u, e));\n"
" atadd(4u, col * NC + j, ld(2u, e) * ld(3u, r * NC + j));\n"
" }\n"
"}\n"
"// dvalues[e] += G[row(e), :] . X[col(e), :]; p1 colIdx, p3 X, p4 G, p5 row\n"
"// of each nonzero, p6 dvalues n = nnz\n"
"@compute @workgroup_size(256)\n"
"fn dk_spmm_vgrad(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let NC = U(0u);\n"
" for (var e = gid.x; e < N(); e += nwg.x * 256u) {\n"
" let r = u32(ld(5u, e));\n"
" let col = u32(ld(1u, e));\n"
" var s = 0.0;\n"
" for (var j = 0u; j < NC; j++) { s += ld(4u, r * NC + j) * ld(3u, col * NC + j); }\n"
" st(6u, e, ld(6u, e) + s);\n"
" }\n"
"}\n"
"\n"
"// ── WebGPU only: copies, GEMMs, convolutions ─────────────────────────────────\n"
"// dst = src for n 32-bit words (bit-exact): p0 src, p2 dst, u7 = n. Either\n"
"// may start mid-word (a 16-bit view): then it moves 2n halves.\n"
"@compute @workgroup_size(256)\n"
"fn wk_copy(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let s = P(0u);\n"
" let d = P(2u);\n"
" if (((s | d) & 3u) == 0u) {\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) { H[(d >> 2u) + i] = H[(s >> 2u) + i]; }\n"
" return;\n"
" }\n"
" for (var i = gid.x; i < 2u * N(); i += nwg.x * 256u) {\n"
" let a = s + 2u * i;\n"
" let h = (H[a >> 2u] >> ((a & 2u) * 8u)) & 0xffffu;\n"
" let b = d + 2u * i;\n"
" let sh = (b & 2u) * 8u;\n"
" atomicAnd(&HA[b >> 2u], ~(0xffffu << sh));\n"
" atomicOr(&HA[b >> 2u], h << sh);\n"
" }\n"
"}\n"
"\n"
"// dst = n words the host queued with this command buffer: an upload in\n"
"// order with the queued work (Metal's mk_upload). p0 = their byte offset in\n"
"// D, p2 = dst, u7 = n.\n"
"@compute @workgroup_size(256)\n"
"fn wk_upload(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let s = P(0u) >> 2u;\n"
" let d = P(2u) >> 2u;\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) { H[d + i] = D[s + i]; }\n"
"}\n"
"\n"
"// gpu_webgpu.h's elementwise ops with no DK_* kernel: p0 a, p1 b, p2 out,\n"
"// u0 = 0 a / b, 1 pow(a, b), 2 sqrt(a) n = elements\n"
"@compute @workgroup_size(256)\n"
"fn wk_map(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" for (var i = gid.x; i < N(); i += nwg.x * 256u) {\n"
" let a = ld(0u, i);\n"
" var v = 0.0;\n"
" if (U(0u) == 0u) { v = a / ld(1u, i); }\n"
" else if (U(0u) == 1u) { v = pow(a, ld(1u, i)); }\n"
" else { v = sqrt(a); }\n"
" st(2u, i, v);\n"
" }\n"
"}\n"
"\n"
"// A strided batch of c_i = alpha op(a_i) op(b_i) + beta c_i (device.h's\n"
"// DevGemm), row-major. Words: 0 a, 1 b, 2 c (byte offsets), 3 m, 4 n, 5 k,\n"
"// 6 lda, 7 ldb, 8 ldc, 9 sa, 10 sb, 11 sc (batch strides in elements),\n"
"// 12 alpha, 13 beta, 14 the first batch index of this dispatch, 15 the\n"
"// slices of k (split-K, wk_gemm: 1 = none), 16 k per slice. The pipeline\n"
"// constants fix the transposes and the element types (GAT: a and b's, GCT:\n"
"// c's), accumulation is fp32. A workgroup of 256 threads takes a 64 x 64\n"
"// tile of c, k in steps of 16 through workgroup memory, each thread a 4 x 4\n"
"// block of it; the grid is (n / 64, m / 64, batch x slices). The tiles are\n"
"// double-buffered: a step's global loads are in flight while the previous\n"
"// step's tile is multiplied, one barrier per step. A product with too few\n"
"// tiles to fill the GPU splits k over the grid, each slice adding its part\n"
"// atomically into c, which the backend has scaled by beta (fp32 c only).\n"
"// The tiles are scalar arrays: threads storing different components of one\n"
"// vec4 in workgroup memory race on Metal (the store rewrites the vector).\n"
"override GTA: bool = false;\n"
"override GTB: bool = false;\n"
"override GAT: u32 = 0u;\n"
"override GCT: u32 = 0u;\n"
"var<workgroup> gA: array<f32, 2048>; // two [k 16][m 64]\n"
"var<workgroup> gB: array<f32, 2048>; // two [k 16][n 64]\n"
"fn gld(base: u32, t: u32, i: u32) -> f32 {\n"
" if (t == 0u) { return bitcast<f32>(H[(base >> 2u) + i]); }\n"
" let a = base + 2u * i;\n"
" return h2f((H[a >> 2u] >> ((a & 2u) * 8u)) & 0xffffu, t);\n"
"}\n"
"fn gst1(base: u32, t: u32, i: u32, v: f32) {\n"
" if (t == 0u) { H[(base >> 2u) + i] = bitcast<u32>(v); return; }\n"
" let a = base + 2u * i;\n"
" let s = (a & 2u) * 8u;\n"
" atomicAnd(&HA[a >> 2u], ~(0xffffu << s));\n"
" atomicOr(&HA[a >> 2u], f2h(v, t) << s);\n"
"}\n"
"// Elements i .. i + 3 = v (the first n of them): a 16-bit run starting on a\n"
"// word boundary as two whole words (no atomics: the thread owns them), else\n"
"// element by element.\n"
"fn gst4(base: u32, t: u32, i: u32, v: vec4f, n: u32) {\n"
" if (t != 0u && n == 4u && (((base >> 1u) + i) & 1u) == 0u) {\n"
" let w = (base + 2u * i) >> 2u;\n"
" H[w] = pack16(v.x, v.y, t);\n"
" H[w + 1u] = pack16(v.z, v.w, t);\n"
" return;\n"
" }\n"
" if (n > 0u) { gst1(base, t, i, v.x); }\n"
" if (n > 1u) { gst1(base, t, i + 1u, v.y); }\n"
" if (n > 2u) { gst1(base, t, i + 2u, v.z); }\n"
" if (n > 3u) { gst1(base, t, i + 3u, v.w); }\n"
"}\n"
"// c's elements i .. i + 3 (n of them) = alpha acc + beta c (beta 0: c is not read).\n"
"fn gemm_out4(Cb: u32, i: u32, acc: vec4f, n: u32, alpha: f32, beta: f32) {\n"
" var v = alpha * acc;\n"
" if (beta != 0.0) {\n"
" var c = vec4f(0.0);\n"
" if (n > 0u) { c.x = gld(Cb, GCT, i); }\n"
" if (n > 1u) { c.y = gld(Cb, GCT, i + 1u); }\n"
" if (n > 2u) { c.z = gld(Cb, GCT, i + 2u); }\n"
" if (n > 3u) { c.w = gld(Cb, GCT, i + 3u); }\n"
" v += beta * c;\n"
" }\n"
" gst4(Cb, GCT, i, v, n);\n"
"}\n"
"// One element of the A (B) tile at step k0 for e = li + 256 q: row (column)\n"
"// e / 16 or e % 64, as the transposes make neighbouring threads read\n"
"// neighbouring addresses; 0 outside the matrix.\n"
"fn gemm_ea(Ab: u32, M: u32, lda: u32, m0: u32, k0: u32, kEnd: u32, e: u32) -> f32 {\n"
" var mm = e / 16u;\n"
" var kk = e % 16u;\n"
" if (GTA) { mm = e % 64u; kk = e / 64u; }\n"
" let gm = m0 + mm;\n"
" let gk = k0 + kk;\n"
" if (gm >= M || gk >= kEnd) { return 0.0; }\n"
" return gld(Ab, GAT, select(gm * lda + gk, gk * lda + gm, GTA));\n"
"}\n"
"fn gemm_eb(Bb: u32, Nn: u32, ldb: u32, n0: u32, k0: u32, kEnd: u32, e: u32) -> f32 {\n"
" var nn = e % 64u;\n"
" var kk = e / 64u;\n"
" if (GTB) { kk = e % 16u; nn = e / 16u; }\n"
" let gn = n0 + nn;\n"
" let gk = k0 + kk;\n"
" if (gn >= Nn || gk >= kEnd) { return 0.0; }\n"
" return gld(Bb, GAT, select(gk * ldb + gn, gn * ldb + gk, GTB));\n"
"}\n"
"// The thread's four elements of each tile (written out: a loop over them\n"
"// isn't unrolled, and a vector indexed by its counter costs more)\n"
"fn gemm_ld_a(Ab: u32, M: u32, lda: u32, m0: u32, k0: u32, kEnd: u32, li: u32) -> vec4f {\n"
" return vec4f(gemm_ea(Ab, M, lda, m0, k0, kEnd, li), gemm_ea(Ab, M, lda, m0, k0, kEnd, li + 256u),\n"
" gemm_ea(Ab, M, lda, m0, k0, kEnd, li + 512u), gemm_ea(Ab, M, lda, m0, k0, kEnd, li + 768u));\n"
"}\n"
"fn gemm_ld_b(Bb: u32, Nn: u32, ldb: u32, n0: u32, k0: u32, kEnd: u32, li: u32) -> vec4f {\n"
" return vec4f(gemm_eb(Bb, Nn, ldb, n0, k0, kEnd, li), gemm_eb(Bb, Nn, ldb, n0, k0, kEnd, li + 256u),\n"
" gemm_eb(Bb, Nn, ldb, n0, k0, kEnd, li + 512u), gemm_eb(Bb, Nn, ldb, n0, k0, kEnd, li + 768u));\n"
"}\n"
"// The index of e = li + 256 q's element in a tile: A (k-major), B.\n"
"fn gemm_ia(e: u32) -> u32 {\n"
" if (GTA) { return (e / 64u) * 64u + e % 64u; }\n"
" return (e % 16u) * 64u + e / 16u;\n"
"}\n"
"fn gemm_ib(e: u32) -> u32 {\n"
" if (GTB) { return (e % 16u) * 64u + e / 16u; }\n"
" return (e / 64u) * 64u + e % 64u;\n"
"}\n"
"@compute @workgroup_size(256)\n"
"fn wk_gemm(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u) {\n"
" let M = aw(3u);\n"
" let Nn = aw(4u);\n"
" let K = aw(5u);\n"
" let lda = aw(6u);\n"
" let ldb = aw(7u);\n"
" let ldc = aw(8u);\n"
" let S = aw(15u);\n"
" let z = wg.z / S + aw(14u);\n"
" let kBeg = (wg.z % S) * aw(16u);\n"
" let kEnd = min(K, kBeg + aw(16u));\n"
" let ea = select(2u, 4u, GAT == 0u);\n"
" let ec = select(2u, 4u, GCT == 0u);\n"
" let Ab = aw(0u) + z * aw(9u) * ea;\n"
" let Bb = aw(1u) + z * aw(10u) * ea;\n"
" let Cb = aw(2u) + z * aw(11u) * ec;\n"
" let m0 = wg.y * 64u;\n"
" let n0 = wg.x * 64u;\n"
" let tx = li % 16u;\n"
" let ty = li / 16u;\n"
" var acc: array<vec4f, 4>;\n"
" var ra = gemm_ld_a(Ab, M, lda, m0, kBeg, kEnd, li);\n"
" var rb = gemm_ld_b(Bb, Nn, ldb, n0, kBeg, kEnd, li);\n"
" let ia0 = gemm_ia(li);\n"
" let ia1 = gemm_ia(li + 256u);\n"
" let ia2 = gemm_ia(li + 512u);\n"
" let ia3 = gemm_ia(li + 768u);\n"
" let ib0 = gemm_ib(li);\n"
" let ib1 = gemm_ib(li + 256u);\n"
" let ib2 = gemm_ib(li + 512u);\n"
" let ib3 = gemm_ib(li + 768u);\n"
" var buf = 0u;\n"
" for (var k0 = kBeg; k0 < kEnd; k0 += 16u) {\n"
" gA[buf + ia0] = ra.x; gA[buf + ia1] = ra.y; gA[buf + ia2] = ra.z; gA[buf + ia3] = ra.w;\n"
" gB[buf + ib0] = rb.x; gB[buf + ib1] = rb.y; gB[buf + ib2] = rb.z; gB[buf + ib3] = rb.w;\n"
" workgroupBarrier();\n"
" if (k0 + 16u < kEnd) {\n"
" ra = gemm_ld_a(Ab, M, lda, m0, k0 + 16u, kEnd, li);\n"
" rb = gemm_ld_b(Bb, Nn, ldb, n0, k0 + 16u, kEnd, li);\n"
" }\n"
" // two k at a time (naga's loops aren't unrolled: each step would wait for its loads)\n"
" for (var kk = 0u; kk < 16u; kk += 2u) {\n"
" let ai = buf + kk * 64u + ty * 4u;\n"
" let bi = buf + kk * 64u + tx * 4u;\n"
" let a = vec4f(gA[ai], gA[ai + 1u], gA[ai + 2u], gA[ai + 3u]);\n"
" let b = vec4f(gB[bi], gB[bi + 1u], gB[bi + 2u], gB[bi + 3u]);\n"
" let a2 = vec4f(gA[ai + 64u], gA[ai + 65u], gA[ai + 66u], gA[ai + 67u]);\n"
" let b2 = vec4f(gB[bi + 64u], gB[bi + 65u], gB[bi + 66u], gB[bi + 67u]);\n"
" acc[0] += a.x * b + a2.x * b2;\n"
" acc[1] += a.y * b + a2.y * b2;\n"
" acc[2] += a.z * b + a2.z * b2;\n"
" acc[3] += a.w * b + a2.w * b2;\n"
" }\n"
" buf = 1024u - buf;\n"
" }\n"
" let alpha = bitcast<f32>(aw(12u));\n"
" let beta = bitcast<f32>(aw(13u));\n"
" let gn = n0 + tx * 4u;\n"
" let nc = select(0u, min(Nn - gn, 4u), gn < Nn);\n"
" for (var i = 0u; i < 4u; i++) {\n"
" let gm = m0 + ty * 4u + i;\n"
" if (gm >= M) { break; }\n"
" if (S > 1u) {\n"
" for (var j = 0u; j < nc; j++) { atomic_addf((Cb >> 2u) + gm * ldc + gn + j, alpha * acc[i][j]); }\n"
" continue;\n"
" }\n"
" gemm_out4(Cb, gm * ldc + gn, acc[i], nc, alpha, beta);\n"
" }\n"
"}\n"
"\n"
"// wk_gemm for small products (few outputs, short k): a thread per element of\n"
"// c, k in registers over loads the cache serves (no tiles, no barriers:\n"
"// the lowest latency). Same words and constants; the grid covers m n batch.\n"
"@compute @workgroup_size(256)\n"
"fn wk_gemm1(@builtin(global_invocation_id) gid: vec3u, @builtin(num_workgroups) nwg: vec3u) {\n"
" let M = aw(3u);\n"
" let Nn = aw(4u);\n"
" let K = aw(5u);\n"
" let lda = aw(6u);\n"
" let ldb = aw(7u);\n"
" let ldc = aw(8u);\n"
" let ea = select(2u, 4u, GAT == 0u);\n"
" let ec = select(2u, 4u, GCT == 0u);\n"
" let alpha = bitcast<f32>(aw(12u));\n"
" let beta = bitcast<f32>(aw(13u));\n"
" let mn = M * Nn;\n"
" let total = mn * aw(17u);\n"
" for (var t = gid.x; t < total; t += nwg.x * 256u) {\n"
" let zl = t / mn;\n"
" let r = t - zl * mn;\n"
" let i = r / Nn;\n"
" let j = r - i * Nn;\n"
" let z = zl + aw(14u);\n"
" let Ab = aw(0u) + z * aw(9u) * ea;\n"
" let Bb = aw(1u) + z * aw(10u) * ea;\n"
" let Cb = aw(2u) + z * aw(11u) * ec;\n"
" // eight k at a time: naga's loops keep Metal's compiler from\n"
" // unrolling them, and each step waits for its loads otherwise\n"
" let sak = select(1u, lda, GTA);\n"
" let sbk = select(ldb, 1u, GTB);\n"
" var ia = select(i * lda, i, GTA);\n"
" var ib = select(j, j * ldb, GTB);\n"
" var a0 = 0.0;\n"
" var a1 = 0.0;\n"
" var a2 = 0.0;\n"
" var a3 = 0.0;\n"
" var k = 0u;\n"
" for (; k + 8u <= K; k += 8u) {\n"
" a0 += gld(Ab, GAT, ia) * gld(Bb, GAT, ib);\n"
" a1 += gld(Ab, GAT, ia + sak) * gld(Bb, GAT, ib + sbk);\n"
" a2 += gld(Ab, GAT, ia + 2u * sak) * gld(Bb, GAT, ib + 2u * sbk);\n"
" a3 += gld(Ab, GAT, ia + 3u * sak) * gld(Bb, GAT, ib + 3u * sbk);\n"
" a0 += gld(Ab, GAT, ia + 4u * sak) * gld(Bb, GAT, ib + 4u * sbk);\n"
" a1 += gld(Ab, GAT, ia + 5u * sak) * gld(Bb, GAT, ib + 5u * sbk);\n"
" a2 += gld(Ab, GAT, ia + 6u * sak) * gld(Bb, GAT, ib + 6u * sbk);\n"
" a3 += gld(Ab, GAT, ia + 7u * sak) * gld(Bb, GAT, ib + 7u * sbk);\n"
" ia += 8u * sak;\n"
" ib += 8u * sbk;\n"
" }\n"
" for (; k + 4u <= K; k += 4u) {\n"
" a0 += gld(Ab, GAT, ia) * gld(Bb, GAT, ib);\n"
" a1 += gld(Ab, GAT, ia + sak) * gld(Bb, GAT, ib + sbk);\n"
" a2 += gld(Ab, GAT, ia + 2u * sak) * gld(Bb, GAT, ib + 2u * sbk);\n"
" a3 += gld(Ab, GAT, ia + 3u * sak) * gld(Bb, GAT, ib + 3u * sbk);\n"
" ia += 4u * sak;\n"
" ib += 4u * sbk;\n"
" }\n"
" for (; k < K; k++) {\n"
" a0 += gld(Ab, GAT, ia) * gld(Bb, GAT, ib);\n"
" ia += sak;\n"
" ib += sbk;\n"
" }\n"
" let acc = (a0 + a1) + (a2 + a3);\n"
" let idx = i * ldc + j;\n"
" var v = alpha * acc;\n"
" if (beta != 0.0) { v += beta * gld(Cb, GCT, idx); }\n"
" gst1(Cb, GCT, idx, v);\n"
" }\n"
"}\n"
"\n"
"// wk_gemm's larger tile for products with enough of them to fill the GPU:\n"
"// 128 x 128 per workgroup, each thread an 8 x 8 block (four FMAs per\n"
"// workgroup-memory load instead of two), k in steps of 16. Every thread\n"
"// stages whole vec4s (4 consecutive m of A, n of B) from loads neighbouring\n"
"// threads make at neighbouring addresses. Same words and constants.\n"
"var<workgroup> hA: array<vec4f, 512>; // [k 16][m 128 as 32 vec4]\n"
"var<workgroup> hB: array<vec4f, 512>; // [k 16][n 128 as 32 vec4]\n"
"fn g4(base: u32, t: u32, i0: u32, step: u32, n: u32) -> vec4f {\n"
" // elements i0, i0 + step, ... (n of them; the rest 0)\n"
" var v = vec4f(0.0);\n"
" if (n > 0u) { v.x = gld(base, t, i0); }\n"
" if (n > 1u) { v.y = gld(base, t, i0 + step); }\n"
" if (n > 2u) { v.z = gld(base, t, i0 + 2u * step); }\n"
" if (n > 3u) { v.w = gld(base, t, i0 + 3u * step); }\n"
" return v;\n"
"}\n"
"@compute @workgroup_size(256)\n"
"fn wk_gemm8(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u) {\n"
" let M = aw(3u);\n"
" let Nn = aw(4u);\n"
" let K = aw(5u);\n"
" let lda = aw(6u);\n"
" let ldb = aw(7u);\n"
" let ldc = aw(8u);\n"
" let z = wg.z + aw(14u);\n"
" let ea = select(2u, 4u, GAT == 0u);\n"
" let ec = select(2u, 4u, GCT == 0u);\n"
" let Ab = aw(0u) + z * aw(9u) * ea;\n"
" let Bb = aw(1u) + z * aw(10u) * ea;\n"
" let Cb = aw(2u) + z * aw(11u) * ec;\n"
" let m0 = wg.y * 128u;\n"
" let n0 = wg.x * 128u;\n"
" let tx = li % 16u;\n"
" let ty = li / 16u;\n"
" var acc: array<vec4f, 16>; // rows ty*8 + i, columns tx*8 + (0..3 | 4..7)\n"
" for (var k0 = 0u; k0 < K; k0 += 16u) {\n"
" for (var q = 0u; q < 2u; q++) {\n"
" let e = li + q * 256u;\n"
" var m4: u32;\n"
" var kk: u32;\n"
" if (GTA) { m4 = e % 32u; kk = e / 32u; } else { kk = e % 16u; m4 = e / 16u; }\n"
" let gm = m0 + m4 * 4u;\n"
" let gk = k0 + kk;\n"
" var v = vec4f(0.0);\n"
" if (gk < K && gm < M) {\n"
" let cnt = min(M - gm, 4u);\n"
" if (GTA) { v = g4(Ab, GAT, gk * lda + gm, 1u, cnt); } else { v = g4(Ab, GAT, gm * lda + gk, lda, cnt); }\n"
" }\n"
" hA[kk * 32u + m4] = v;\n"
" }\n"
" for (var q = 0u; q < 2u; q++) {\n"
" let e = li + q * 256u;\n"
" var n4: u32;\n"
" var kk: u32;\n"
" if (GTB) { kk = e % 16u; n4 = e / 16u; } else { n4 = e % 32u; kk = e / 32u; }\n"
" let gn = n0 + n4 * 4u;\n"
" let gk = k0 + kk;\n"
" var v = vec4f(0.0);\n"
" if (gk < K && gn < Nn) {\n"
" let cnt = min(Nn - gn, 4u);\n"
" if (GTB) { v = g4(Bb, GAT, gn * ldb + gk, ldb, cnt); } else { v = g4(Bb, GAT, gk * ldb + gn, 1u, cnt); }\n"
" }\n"
" hB[kk * 32u + n4] = v;\n"
" }\n"
" workgroupBarrier();\n"
" for (var kk = 0u; kk < 16u; kk++) {\n"
" let a0 = hA[kk * 32u + ty * 2u];\n"
" let a1 = hA[kk * 32u + ty * 2u + 1u];\n"
" let b0 = hB[kk * 32u + tx * 2u];\n"
" let b1 = hB[kk * 32u + tx * 2u + 1u];\n"
" acc[0] += a0.x * b0; acc[1] += a0.x * b1;\n"
" acc[2] += a0.y * b0; acc[3] += a0.y * b1;\n"
" acc[4] += a0.z * b0; acc[5] += a0.z * b1;\n"
" acc[6] += a0.w * b0; acc[7] += a0.w * b1;\n"
" acc[8] += a1.x * b0; acc[9] += a1.x * b1;\n"
" acc[10] += a1.y * b0; acc[11] += a1.y * b1;\n"
" acc[12] += a1.z * b0; acc[13] += a1.z * b1;\n"
" acc[14] += a1.w * b0; acc[15] += a1.w * b1;\n"
" }\n"
" workgroupBarrier();\n"
" }\n"
" let alpha = bitcast<f32>(aw(12u));\n"
" let beta = bitcast<f32>(aw(13u));\n"
" let gn = n0 + tx * 8u;\n"
" let n0c = select(0u, min(Nn - gn, 4u), gn < Nn);\n"
" let n1c = select(0u, min(Nn - gn - 4u, 4u), gn + 4u < Nn);\n"
" for (var i = 0u; i < 8u; i++) {\n"
" let gm = m0 + ty * 8u + i;\n"
" if (gm >= M) { break; }\n"
" gemm_out4(Cb, gm * ldc + gn, acc[i * 2u], n0c, alpha, beta);\n"
" gemm_out4(Cb, gm * ldc + gn + 4u, acc[i * 2u + 1u], n1c, alpha, beta);\n"
" }\n"
"}\n"
"\n"
"// Convolutions as implicit GEMMs (WebGpuBackend::conv: MIOpen / cuDNN's\n"
"// role), the column matrix gathered tile by tile instead of written out by\n"
"// im2col. In a plain convolution's terms -- input I [B, Ci, Hi, Wi], weights\n"
"// W [Co, Ci, kH, kW], output O [B, Co, Ho, Wo]; a transposed convolution is\n"
"// one with I and O swapped --\n"
"// COP 0 O = conv(I, W): m = co, n = (b, oy, ox), k = (ci, ky, kx)\n"
"// COP 1 I = conv^T(O, W): m = ci, n = (b, y, x), k = (co, ky, kx),\n"
"// a gather over the windows holding (y, x), no atomics\n"
"// COP 2 W += O . I-columns: m = co, n = (ci, ky, kx), k = (b, oy, ox),\n"
"// split over k (grid z, kSplit each), added atomically (fp32)\n"
"// Words: 0 x (I), 1 w, 2 y (O) byte offsets, 3 B, 4 Ci, 5 Hi, 6 Wi, 7 Co,\n"
"// 8 Ho, 9 Wo, 10 kH, 11 kW, 12 sH, 13 sW, 14 pH, 15 pW, 16 dH, 17 dW,\n"
"// 18 kSplit. A workgroup of 256 threads takes a CBM x CBN tile of the\n"
"// product (64 x 64, 32 x 128 or 16 x 256: few output channels take the\n"
"// flat ones), k in steps of 16, each thread a 4 x 4 block (8 x 8 blocks in\n"
"// larger tiles measured no faster once the gathers were unrolled). Each\n"
"// thread loads the same column of the B tile every step, so its pixel is\n"
"// decoded once. CT is I, W and O's element type (COP 2: fp32).\n"
"override COP: u32 = 0u;\n"
"override CBM: u32 = 64u;\n"
"override CT: u32 = 0u;\n"
"override CBN: u32 = 4096u / CBM; // the tile's columns\n"
"var<workgroup> cA: array<f32, 16u * CBM>; // [k 16][m CBM]\n"
"var<workgroup> cB: array<f32, 16u * CBN>; // [k 16][n CBN]\n"
"struct Cv { x: u32, w: u32, y: u32, B: u32, Ci: u32, Hi: u32, Wi: u32, Co: u32, Ho: u32, Wo: u32,\n"
" kH: u32, kW: u32, sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, kSplit: u32 }\n"
"fn cv_args() -> Cv {\n"
" return Cv(aw(0u), aw(1u), aw(2u), aw(3u), aw(4u), aw(5u), aw(6u), aw(7u), aw(8u), aw(9u), aw(10u), aw(11u),\n"
" aw(12u), aw(13u), aw(14u), aw(15u), aw(16u), aw(17u), aw(18u));\n"
"}\n"
"// The k step's 16 rows, decoded once per workgroup (cK): COP 0, 1 (the\n"
"// channel, ky dH, kx dW, the tap ky kW + kx; by stride phase: the channel,\n"
"// iy, ix, the tap), COP 2 (the image, oy sH, ox sW, the output pixel).\n"
"var<workgroup> cK: array<vec4i, 16>;\n"
"// A[m][k]: the weights (COP 0, 1) or the output gradient (COP 2)\n"
"fn cv_a(c: Cv, m: u32, k: u32, t: vec4i) -> f32 {\n"
" let KK = c.kH * c.kW;\n"
" if (COP == 0u) { return gld(c.w, CT, m * c.Ci * KK + k); }\n"
" if (COP == 1u) { return gld(c.w, CT, (u32(t.x) * c.Ci + m) * KK + u32(t.w)); }\n"
" return gld(c.y, 0u, (u32(t.x) * c.Co + m) * c.Ho * c.Wo + u32(t.w));\n"
"}\n"
"// B[k][n] for k's decoded row t and the thread's column: its decoded pixel\n"
"// or tap (p0, p1, p2) and image base (ib, an element offset)\n"
"fn cv_b(c: Cv, t: vec4i, p0: i32, p1: i32, p2: u32, ib: u32, phased: bool) -> f32 {\n"
" if (COP == 0u) {\n"
" // p0, p1: the window's corner (oy sH - pH, ox sW - pW)\n"
" let y = p0 + t.y;\n"
" let x = p1 + t.z;\n"
" if (y < 0 || y >= i32(c.Hi) || x < 0 || x >= i32(c.Wi)) { return 0.0; }\n"
" return gld(c.x, CT, ib + (u32(t.x) * c.Hi + u32(y)) * c.Wi + u32(x));\n"
" }\n"
" if (COP == 1u && phased) {\n"
" // p0, p1: the phase's first window row / column plus the pixel's\n"
" // (jy, jx); t: (co, iy, ix): the window (p0 - iy, p1 - ix)\n"
" let oy = p0 - t.y;\n"
" let ox = p1 - t.z;\n"
" if (oy < 0 || ox < 0 || oy >= i32(c.Ho) || ox >= i32(c.Wo)) { return 0.0; }\n"
" return gld(c.y, CT, ib + (u32(t.x) * c.Ho + u32(oy)) * c.Wo + u32(ox));\n"
" }\n"
" if (COP == 1u) {\n"
" // p0, p1: (y + pH, x + pW); the window at (oy, ox) holds (y, x) through the tap\n"
" let ny = p0 - t.y;\n"
" let nx = p1 - t.z;\n"
" if (ny < 0 || nx < 0) { return 0.0; }\n"
" var oy = u32(ny);\n"
" var ox = u32(nx);\n"
" if (c.sH != 1u || c.sW != 1u) {\n"
" oy = oy / c.sH;\n"
" ox = ox / c.sW;\n"
" if (oy * c.sH != u32(ny) || ox * c.sW != u32(nx)) { return 0.0; }\n"
" }\n"
" if (oy >= c.Ho || ox >= c.Wo) { return 0.0; }\n"
" return gld(c.y, CT, ib + (u32(t.x) * c.Ho + oy) * c.Wo + ox);\n"
" }\n"
" // COP 2: p0, p1: the tap's offset (ky dH - pH, kx dW - pW), p2 = ci\n"
" let y = t.y + p0;\n"
" let x = t.z + p1;\n"
" if (y < 0 || y >= i32(c.Hi) || x < 0 || x >= i32(c.Wi)) { return 0.0; }\n"
" return gld(c.x, 0u, ((u32(t.x) * c.Ci + p2) * c.Hi + u32(y)) * c.Wi + u32(x));\n"
"}\n"
"@compute @workgroup_size(256)\n"
"fn wk_conv(@builtin(local_invocation_index) li: u32, @builtin(workgroup_id) wg: vec3u) {\n"
" let c = cv_args();\n"
" let BN = CBN;\n"
" let KK = c.kH * c.kW;\n"
" let Po = c.Ho * c.Wo;\n"
" let Pi = c.Hi * c.Wi;\n"
" var M = c.Co;\n"
" var Ntot = c.B * Po;\n"
" var kBeg = 0u;\n"
" var kEnd = c.Ci * KK;\n"
" // the input gradient by stride phase (without dilation): grid z =\n"
" // (ry, rx); its pixels y = yb + jy sH (likewise x) take the taps\n"
" // ky = ry + iy sH only, through the window oyb + jy - iy: no tap misses\n"
" // and no division (3 in 4 of the taps miss for stride 2 otherwise)\n"
" let phased = COP == 1u && c.dH == 1u && c.dW == 1u;\n"
" var ry = 0u;\n"
" var rx = 0u;\n"
" var yb = 0u;\n"
" var xb = 0u;\n"
" var py = c.Hi;\n"
" var px = c.Wi;\n"
" var tyc = c.kH;\n"
" var txc = c.kW;\n"
" var oyb = 0;\n"
" var oxb = 0;\n"
" if (phased) {\n"
" ry = wg.z / c.sW;\n"
" rx = wg.z - ry * c.sW;\n"
" yb = (ry + c.sH - c.pH % c.sH) % c.sH;\n"
" xb = (rx + c.sW - c.pW % c.sW) % c.sW;\n"
" py = select(0u, (c.Hi - yb + c.sH - 1u) / c.sH, yb < c.Hi);\n"
" px = select(0u, (c.Wi - xb + c.sW - 1u) / c.sW, xb < c.Wi);\n"
" tyc = select(0u, (c.kH - ry + c.sH - 1u) / c.sH, ry < c.kH);\n"
" txc = select(0u, (c.kW - rx + c.sW - 1u) / c.sW, rx < c.kW);\n"
" oyb = (i32(yb) + i32(c.pH) - i32(ry)) / i32(c.sH);\n"
" oxb = (i32(xb) + i32(c.pW) - i32(rx)) / i32(c.sW);\n"
" }\n"
" if (COP == 1u) { M = c.Ci; Ntot = c.B * py * px; kEnd = c.Co * tyc * txc; }\n"
" if (COP == 2u) { Ntot = c.Ci * KK; kBeg = wg.z * c.kSplit; kEnd = min(c.B * Po, kBeg + c.kSplit); }\n"
" let m0 = wg.y * CBM;\n"
" let n0 = wg.x * BN;\n"
" if (n0 >= Ntot) { return; } // a phase with fewer pixels: the whole workgroup\n"
" // this thread's column of the B tile, decoded once\n"
" let bc = li % BN;\n"
" let n = n0 + bc;\n"
" let nIn = n < Ntot;\n"
" let nc = min(n, Ntot - 1u);\n"
" var p0 = 0;\n"
" var p1 = 0;\n"
" var p2 = 0u;\n"
" var ib = 0u;\n"
" if (COP == 0u) {\n"
" let b = nc / Po;\n"
" let op = nc - b * Po;\n"
" let oy = op / c.Wo;\n"
" p0 = i32(oy * c.sH) - i32(c.pH);\n"
" p1 = i32((op - oy * c.Wo) * c.sW) - i32(c.pW);\n"
" ib = b * c.Ci * Pi;\n"
" } else if (phased) {\n"
" let b = nc / (py * px);\n"
" let ip = nc - b * py * px;\n"
" let jy = ip / px;\n"
" p0 = oyb + i32(jy);\n"
" p1 = oxb + i32(ip - jy * px);\n"
" ib = b * c.Co * Po;\n"
" } else if (COP == 1u) {\n"
" let b = nc / Pi;\n"
" let ip = nc - b * Pi;\n"
" let y = ip / c.Wi;\n"
" p0 = i32(y + c.pH);\n"
" p1 = i32(ip - y * c.Wi + c.pW);\n"
" ib = b * c.Co * Po;\n"
" } else {\n"
" p2 = nc / KK;\n"
" let r = nc - p2 * KK;\n"
" let ky = r / c.kW;\n"
" p0 = i32(ky * c.dH) - i32(c.pH);\n"
" p1 = i32((r - ky * c.kW) * c.dW) - i32(c.pW);\n"
" }\n"
" let tx = li % (BN / 4u);\n"
" let ty = li / (BN / 4u);\n"
" var acc: array<vec4f, 4>;\n"
" for (var k0 = kBeg; k0 < kEnd; k0 += 16u) {\n"
" if (li < 16u) {\n"
" let k = min(k0 + li, kEnd - 1u);\n"
" if (COP == 2u) {\n"
" let b = k / Po;\n"
" let op = k - b * Po;\n"
" let oy = op / c.Wo;\n"
" cK[li] = vec4i(i32(b), i32(oy * c.sH), i32((op - oy * c.Wo) * c.sW), i32(op));\n"
" } else if (phased) {\n"
" let co = k / (tyc * txc);\n"
" let r = k - co * tyc * txc;\n"
" let iy = r / txc;\n"
" let ix = r - iy * txc;\n"
" cK[li] = vec4i(i32(co), i32(iy), i32(ix), i32((ry + iy * c.sH) * c.kW + rx + ix * c.sW));\n"
" } else {\n"
" let q = k / KK;\n"
" let r = k - q * KK;\n"
" let ky = r / c.kW;\n"
" cK[li] = vec4i(i32(q), i32(ky * c.dH), i32((r - ky * c.kW) * c.dW), i32(r));\n"
" }\n"
" }\n"
" workgroupBarrier();\n"
" // the loads four at a time where the tile allows (naga's loops\n"
" // aren't unrolled: each load would wait for the one before)\n"
" let kk = li % 16u;\n"
" let ka = k0 + kk < kEnd;\n"
" let ta = cK[kk];\n"
" if (CBM >= 64u) {\n"
" for (var q = 0u; q < CBM / 16u; q += 4u) {\n"
" let mm = li / 16u + q * 16u;\n"
" var v0 = 0.0;\n"
" var v1 = 0.0;\n"
" var v2 = 0.0;\n"
" var v3 = 0.0;\n"
" if (ka && m0 + mm < M) { v0 = cv_a(c, m0 + mm, k0 + kk, ta); }\n"
" if (ka && m0 + mm + 16u < M) { v1 = cv_a(c, m0 + mm + 16u, k0 + kk, ta); }\n"
" if (ka && m0 + mm + 32u < M) { v2 = cv_a(c, m0 + mm + 32u, k0 + kk, ta); }\n"
" if (ka && m0 + mm + 48u < M) { v3 = cv_a(c, m0 + mm + 48u, k0 + kk, ta); }\n"
" cA[kk * CBM + mm] = v0;\n"
" cA[kk * CBM + mm + 16u] = v1;\n"
" cA[kk * CBM + mm + 32u] = v2;\n"
" cA[kk * CBM + mm + 48u] = v3;\n"
" }\n"
" } else {\n"
" for (var q = 0u; q < CBM / 16u; q++) {\n"
" let mm = li / 16u + q * 16u;\n"
" var v = 0.0;\n"
" if (ka && m0 + mm < M) { v = cv_a(c, m0 + mm, k0 + kk, ta); }\n"
" cA[kk * CBM + mm] = v;\n"
" }\n"
" }\n"
" let rs = 256u / BN; // B rows a pass of the workgroup loads; BN / 16 passes, a multiple of 4\n"
" for (var r = li / BN; r < 16u; r += 4u * rs) {\n"
" var v0 = 0.0;\n"
" var v1 = 0.0;\n"
" var v2 = 0.0;\n"
" var v3 = 0.0;\n"
" if (nIn && k0 + r < kEnd) { v0 = cv_b(c, cK[r], p0, p1, p2, ib, phased); }\n"
" if (nIn && k0 + r + rs < kEnd) { v1 = cv_b(c, cK[r + rs], p0, p1, p2, ib, phased); }\n"
" if (nIn && k0 + r + 2u * rs < kEnd) { v2 = cv_b(c, cK[r + 2u * rs], p0, p1, p2, ib, phased); }\n"
" if (nIn && k0 + r + 3u * rs < kEnd) { v3 = cv_b(c, cK[r + 3u * rs], p0, p1, p2, ib, phased); }\n"
" cB[r * BN + bc] = v0;\n"
" cB[(r + rs) * BN + bc] = v1;\n"
" cB[(r + 2u * rs) * BN + bc] = v2;\n"
" cB[(r + 3u * rs) * BN + bc] = v3;\n"
" }\n"
" workgroupBarrier();\n"
" for (var kk = 0u; kk < 16u; kk++) {\n"
" let ai = kk * CBM + ty * 4u;\n"
" let bi = kk * BN + tx * 4u;\n"
" let a = vec4f(cA[ai], cA[ai + 1u], cA[ai + 2u], cA[ai + 3u]);\n"
" let b = vec4f(cB[bi], cB[bi + 1u], cB[bi + 2u], cB[bi + 3u]);\n"
" acc[0] += a.x * b;\n"
" acc[1] += a.y * b;\n"
" acc[2] += a.z * b;\n"
" acc[3] += a.w * b;\n"
" }\n"
" workgroupBarrier();\n"
" }\n"
" for (var i = 0u; i < 4u; i++) {\n"
" let gm = m0 + ty * 4u + i;\n"
" if (gm >= M) { break; }\n"
" // the four outputs' indices: whole 16-bit words where they run on\n"
" var o = vec4u(0u);\n"
" var cnt = 0u;\n"
" for (var j = 0u; j < 4u; j++) {\n"
" let gn = n0 + tx * 4u + j;\n"
" if (gn >= Ntot) { break; }\n"
" if (COP == 0u) {\n"
" let b = gn / Po;\n"
" o[j] = (b * c.Co + gm) * Po + gn - b * Po;\n"
" } else if (phased) {\n"
" let b = gn / (py * px);\n"
" let ip = gn - b * py * px;\n"
" let jy = ip / px;\n"
" o[j] = (b * c.Ci + gm) * Pi + (yb + jy * c.sH) * c.Wi + xb + (ip - jy * px) * c.sW;\n"
" } else if (COP == 1u) {\n"
" let b = gn / Pi;\n"
" o[j] = (b * c.Ci + gm) * Pi + gn - b * Pi;\n"
" } else {\n"
" atomic_addf((c.w >> 2u) + gm * Ntot + gn, acc[i][j]);\n"
" }\n"
" cnt = j + 1u;\n"
" }\n"
" if (COP == 2u) { continue; }\n"
" let ob = select(c.x, c.y, COP == 0u);\n"
" if (cnt == 4u && o.y == o.x + 1u && o.z == o.x + 2u && o.w == o.x + 3u) { gst4(ob, CT, o.x, acc[i], 4u); }\n"
" else {\n"
" for (var j = 0u; j < cnt; j++) { gst1(ob, CT, o[j], acc[i][j]); }\n"
" }\n"
" }\n"
"}\n"
"\n";
}
} // namespace std