| |
| const PRELUDE = `enable f16;\n`; |
|
|
| const REDUCE8 = ` |
| if (tid < 128u) { red[tid] += red[tid + 128u]; } workgroupBarrier(); |
| if (tid < 64u) { red[tid] += red[tid + 64u]; } workgroupBarrier(); |
| if (tid < 32u) { red[tid] += red[tid + 32u]; } workgroupBarrier(); |
| if (tid < 16u) { red[tid] += red[tid + 16u]; } workgroupBarrier(); |
| if (tid < 8u) { red[tid] += red[tid + 8u]; } workgroupBarrier(); |
| if (tid < 4u) { red[tid] += red[tid + 4u]; } workgroupBarrier(); |
| if (tid < 2u) { red[tid] += red[tid + 2u]; } workgroupBarrier(); |
| if (tid < 1u) { red[tid] += red[tid + 1u]; } workgroupBarrier(); |
| `; |
| const REDUCE8MAX = ` |
| if (tid < 128u) { red[tid] = max(red[tid], red[tid + 128u]); } workgroupBarrier(); |
| if (tid < 64u) { red[tid] = max(red[tid], red[tid + 64u]); } workgroupBarrier(); |
| if (tid < 32u) { red[tid] = max(red[tid], red[tid + 32u]); } workgroupBarrier(); |
| if (tid < 16u) { red[tid] = max(red[tid], red[tid + 16u]); } workgroupBarrier(); |
| if (tid < 8u) { red[tid] = max(red[tid], red[tid + 8u]); } workgroupBarrier(); |
| if (tid < 4u) { red[tid] = max(red[tid], red[tid + 4u]); } workgroupBarrier(); |
| if (tid < 2u) { red[tid] = max(red[tid], red[tid + 2u]); } workgroupBarrier(); |
| if (tid < 1u) { red[tid] = max(red[tid], red[tid + 1u]); } workgroupBarrier(); |
| `; |
|
|
| const GEMM = ` |
| enable chromium_experimental_subgroup_matrix; |
| diagnostic(off, chromium.subgroup_matrix_uniformity); |
| struct P { M: u32, N: u32, K: u32, pad: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> A: array<f16>; |
| @group(0) @binding(2) var<storage, read> B: array<f16>; |
| @group(0) @binding(3) var<storage, read> bias: array<f16>; |
| @group(0) @binding(4) var<storage, read_write> C: array<f16>; |
| const TR: u32 = 128u; const TC: u32 = 64u; const TK: u32 = 32u; |
| var<workgroup> tile_A: array<f16, TR * TK>; |
| var<workgroup> tile_B: array<f16, TC * TK>; |
| var<workgroup> scratch: array<array<f32, 64>, 8>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) { |
| let bm = wg.x * TR; let bn = wg.y * TC; |
| let sg = tid / 32u; let lane = tid % 32u; |
| let sg_row = sg / 2u; let sg_col = sg % 2u; |
| var acc: array<subgroup_matrix_result<f32, 8, 8>, 16>; |
| for (var i = 0u; i < 16u; i++) { acc[i] = subgroup_matrix_result<f32, 8, 8>(); } |
| let nk = min((p.K + TK - 1u) / TK, 2048u); |
| for (var kt = 0u; kt < nk; kt++) { |
| let k0 = kt * TK; |
| for (var i = 0u; i < 16u; i++) { |
| let idx = i * 256u + tid; let r = idx / TK; let c = idx % TK; |
| let gr = bm + r; let gc = k0 + c; |
| var v: f16 = f16(0.0); |
| if (gr < p.M && gc < p.K) { v = A[gr * p.K + gc]; } |
| tile_A[idx] = v; |
| } |
| for (var i = 0u; i < 8u; i++) { |
| let idx = i * 256u + tid; let gn = idx % TC; let gk = idx / TC; |
| let gr = bn + gn; let gc = k0 + gk; |
| var v: f16 = f16(0.0); |
| if (gr < p.N && gc < p.K) { v = B[gc * p.N + gr]; } |
| tile_B[gn * TK + gk] = v; |
| } |
| workgroupBarrier(); |
| for (var step = 0u; step < TK; step += 8u) { |
| let a_off = (sg_row * 32u) * TK + step; |
| let m0 = subgroupMatrixLoad<subgroup_matrix_left<f16, 8, 8>>(&tile_A, a_off, false, TK); |
| let m1 = subgroupMatrixLoad<subgroup_matrix_left<f16, 8, 8>>(&tile_A, a_off + 8u * TK, false, TK); |
| let m2 = subgroupMatrixLoad<subgroup_matrix_left<f16, 8, 8>>(&tile_A, a_off + 16u * TK, false, TK); |
| let m3 = subgroupMatrixLoad<subgroup_matrix_left<f16, 8, 8>>(&tile_A, a_off + 24u * TK, false, TK); |
| let b_off = (sg_col * 32u) * TK + step; |
| let n0 = subgroupMatrixLoad<subgroup_matrix_right<f16, 8, 8>>(&tile_B, b_off, true, TK); |
| let n1 = subgroupMatrixLoad<subgroup_matrix_right<f16, 8, 8>>(&tile_B, b_off + 8u * TK, true, TK); |
| let n2 = subgroupMatrixLoad<subgroup_matrix_right<f16, 8, 8>>(&tile_B, b_off + 16u * TK, true, TK); |
| let n3 = subgroupMatrixLoad<subgroup_matrix_right<f16, 8, 8>>(&tile_B, b_off + 24u * TK, true, TK); |
| acc[0] = subgroupMatrixMultiplyAccumulate(m0, n0, acc[0]); |
| acc[1] = subgroupMatrixMultiplyAccumulate(m0, n1, acc[1]); |
| acc[2] = subgroupMatrixMultiplyAccumulate(m0, n2, acc[2]); |
| acc[3] = subgroupMatrixMultiplyAccumulate(m0, n3, acc[3]); |
| acc[4] = subgroupMatrixMultiplyAccumulate(m1, n0, acc[4]); |
| acc[5] = subgroupMatrixMultiplyAccumulate(m1, n1, acc[5]); |
| acc[6] = subgroupMatrixMultiplyAccumulate(m1, n2, acc[6]); |
| acc[7] = subgroupMatrixMultiplyAccumulate(m1, n3, acc[7]); |
| acc[8] = subgroupMatrixMultiplyAccumulate(m2, n0, acc[8]); |
| acc[9] = subgroupMatrixMultiplyAccumulate(m2, n1, acc[9]); |
| acc[10] = subgroupMatrixMultiplyAccumulate(m2, n2, acc[10]); |
| acc[11] = subgroupMatrixMultiplyAccumulate(m2, n3, acc[11]); |
| acc[12] = subgroupMatrixMultiplyAccumulate(m3, n0, acc[12]); |
| acc[13] = subgroupMatrixMultiplyAccumulate(m3, n1, acc[13]); |
| acc[14] = subgroupMatrixMultiplyAccumulate(m3, n2, acc[14]); |
| acc[15] = subgroupMatrixMultiplyAccumulate(m3, n3, acc[15]); |
| } |
| workgroupBarrier(); |
| } |
| let row0 = bm + sg_row * 32u; |
| let col0 = bn + sg_col * 32u; |
| for (var f = 0u; f < 16u; f++) { |
| subgroupMatrixStore(&scratch[sg], 0u, acc[f], false, 8u); |
| workgroupBarrier(); |
| let fr = f / 4u; let fc = f % 4u; |
| for (var e = lane; e < 64u; e += 32u) { |
| let i = e / 8u; let j = e % 8u; |
| let gr = row0 + fr * 8u + i; let gc = col0 + fc * 8u + j; |
| if (gr < p.M && gc < p.N) { C[gr * p.N + gc] = f16(clamp(scratch[sg][e] + f32(bias[gc]), -60000.0, 60000.0)); } |
| } |
| workgroupBarrier(); |
| } |
| } |
| `; |
|
|
| const GEMM_SIMT_I4 = ` |
| struct P { M: u32, N: u32, K: u32, group: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> A: array<f16>; |
| @group(0) @binding(2) var<storage, read> Pk: array<u32>; |
| @group(0) @binding(3) var<storage, read> Sc: array<f16>; |
| @group(0) @binding(4) var<storage, read> bias: array<f16>; |
| @group(0) @binding(5) var<storage, read_write> C: array<f16>; |
| const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 32u; |
| var<workgroup> As: array<f32, 2048>; |
| var<workgroup> Bs: array<f32, 2048>; |
| @compute @workgroup_size(16, 16) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) { |
| let bm = wg.x * BM; let bn = wg.y * BN; |
| let ty = tid / 16u; let tx = tid % 16u; |
| var acc: array<vec4<f32>, 4>; |
| for (var i = 0u; i < 4u; i++) { acc[i] = vec4<f32>(0.0); } |
| let nk = min((p.K + BK - 1u) / BK, 2048u); |
| for (var kt = 0u; kt < nk; kt++) { |
| let k0 = kt * BK; |
| for (var i = 0u; i < 8u; i++) { |
| let idx = i * 256u + tid; let r = idx / BK; let c = idx % BK; |
| let gr = bm + r; let gc = k0 + c; |
| var v: f32 = 0.0; |
| if (gr < p.M && gc < p.K) { v = f32(A[gr * p.K + gc]); } |
| As[r * BK + c] = v; |
| } |
| // dequant B tile: 256 words (16 kp x 16 words of 4 columns) |
| let kp = tid / 16u; let n0 = (tid % 16u) * 4u; |
| let gk2 = k0 / 2u + kp; let gn0 = bn + n0; |
| var word: u32 = 0u; |
| if (gk2 < p.K / 2u && gn0 < p.N) { word = Pk[(gk2 * p.N + gn0) / 4u]; } |
| let grp = (k0 + 2u * kp) / p.group; |
| for (var j = 0u; j < 4u; j++) { |
| let n = n0 + j; |
| if (bn + n < p.N) { |
| let byte = (word >> (8u * j)) & 0xFFu; |
| let s = f32(Sc[grp * p.N + bn + n]); |
| Bs[(2u * kp) * BN + n] = (f32(byte & 0xFu) - 8.0) * s; |
| Bs[(2u * kp + 1u) * BN + n] = (f32(byte >> 4u) - 8.0) * s; |
| } else { |
| Bs[(2u * kp) * BN + n] = 0.0; |
| Bs[(2u * kp + 1u) * BN + n] = 0.0; |
| } |
| } |
| workgroupBarrier(); |
| for (var k = 0u; k < BK; k++) { |
| let a0 = As[(ty * 4u + 0u) * BK + k]; |
| let a1 = As[(ty * 4u + 1u) * BK + k]; |
| let a2 = As[(ty * 4u + 2u) * BK + k]; |
| let a3 = As[(ty * 4u + 3u) * BK + k]; |
| let b0 = Bs[k * BN + tx * 4u + 0u]; |
| let b1 = Bs[k * BN + tx * 4u + 1u]; |
| let b2 = Bs[k * BN + tx * 4u + 2u]; |
| let b3 = Bs[k * BN + tx * 4u + 3u]; |
| let bv = vec4<f32>(b0, b1, b2, b3); |
| acc[0] += a0 * bv; acc[1] += a1 * bv; acc[2] += a2 * bv; acc[3] += a3 * bv; |
| } |
| workgroupBarrier(); |
| } |
| for (var r = 0u; r < 4u; r++) { |
| let gr = bm + ty * 4u + r; |
| if (gr < p.M) { |
| for (var c = 0u; c < 4u; c++) { |
| let gc = bn + tx * 4u + c; |
| if (gc < p.N) { C[gr * p.N + gc] = f16(clamp(acc[r][c] + f32(bias[gc]), -60000.0, 60000.0)); } |
| } |
| } |
| } |
| } |
| `; |
|
|
| const GEMM_SIMT = ` |
| struct P { M: u32, N: u32, K: u32, pad: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> A: array<f16>; |
| @group(0) @binding(2) var<storage, read> B: array<f16>; |
| @group(0) @binding(3) var<storage, read> bias: array<f16>; |
| @group(0) @binding(4) var<storage, read_write> C: array<f16>; |
| const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 32u; |
| var<workgroup> As: array<f32, 2048>; |
| var<workgroup> Bs: array<f32, 2048>; |
| @compute @workgroup_size(16, 16) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) { |
| let bm = wg.x * BM; let bn = wg.y * BN; |
| let ty = tid / 16u; let tx = tid % 16u; |
| var acc: array<vec4<f32>, 4>; |
| for (var i = 0u; i < 4u; i++) { acc[i] = vec4<f32>(0.0); } |
| let nk = min((p.K + BK - 1u) / BK, 2048u); |
| for (var kt = 0u; kt < nk; kt++) { |
| let k0 = kt * BK; |
| for (var i = 0u; i < 8u; i++) { |
| let idx = i * 256u + tid; let r = idx / BK; let c = idx % BK; |
| let gr = bm + r; let gc = k0 + c; |
| var v: f32 = 0.0; |
| if (gr < p.M && gc < p.K) { v = f32(A[gr * p.K + gc]); } |
| As[r * BK + c] = v; |
| } |
| for (var i = 0u; i < 8u; i++) { |
| let idx = i * 256u + tid; let r = idx / BN; let c = idx % BN; |
| let gr = k0 + r; let gc = bn + c; |
| var v: f32 = 0.0; |
| if (gr < p.K && gc < p.N) { v = f32(B[gr * p.N + gc]); } |
| Bs[r * BN + c] = v; |
| } |
| workgroupBarrier(); |
| for (var k = 0u; k < BK; k++) { |
| let a0 = As[(ty * 4u + 0u) * BK + k]; |
| let a1 = As[(ty * 4u + 1u) * BK + k]; |
| let a2 = As[(ty * 4u + 2u) * BK + k]; |
| let a3 = As[(ty * 4u + 3u) * BK + k]; |
| let b0 = Bs[k * BN + tx * 4u + 0u]; |
| let b1 = Bs[k * BN + tx * 4u + 1u]; |
| let b2 = Bs[k * BN + tx * 4u + 2u]; |
| let b3 = Bs[k * BN + tx * 4u + 3u]; |
| let bv = vec4<f32>(b0, b1, b2, b3); |
| acc[0] += a0 * bv; acc[1] += a1 * bv; acc[2] += a2 * bv; acc[3] += a3 * bv; |
| } |
| workgroupBarrier(); |
| } |
| for (var r = 0u; r < 4u; r++) { |
| let gr = bm + ty * 4u + r; |
| if (gr < p.M) { |
| for (var c = 0u; c < 4u; c++) { |
| let gc = bn + tx * 4u + c; |
| if (gc < p.N) { C[gr * p.N + gc] = f16(clamp(acc[r][c] + f32(bias[gc]), -60000.0, 60000.0)); } |
| } |
| } |
| } |
| } |
| `; |
|
|
| const GEMV = ` |
| struct P { N: u32, K: u32, pad0: u32, pad1: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> A: array<f16>; |
| @group(0) @binding(2) var<storage, read> B: array<f16>; |
| @group(0) @binding(3) var<storage, read> bias: array<f16>; |
| @group(0) @binding(4) var<storage, read_write> C: array<f16>; |
| var<workgroup> shA: array<f32, 1024>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) { |
| for (var i = tid; i < p.K; i += 256u) { shA[i] = f32(A[i]); } |
| workgroupBarrier(); |
| let base = wg.x * 1024u + tid * 4u; |
| var acc = vec4<f32>(0.0); |
| let KK = min(p.K, 65536u); |
| for (var k = 0u; k < KK; k++) { |
| let a = shA[k]; |
| let o = k * p.N + base; |
| acc += a * vec4<f32>(f32(B[o]), f32(B[o + 1u]), f32(B[o + 2u]), f32(B[o + 3u])); |
| } |
| for (var j = 0u; j < 4u; j++) { |
| let c = base + j; |
| if (c < p.N) { C[c] = f16(acc[j] + f32(bias[c])); } |
| } |
| } |
| `; |
|
|
| const GEMV_I4 = ` |
| struct P { N: u32, K: u32, group: u32, pad1: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> A: array<f16>; |
| @group(0) @binding(2) var<storage, read> Pk: array<u32>; |
| @group(0) @binding(3) var<storage, read> Sc: array<f16>; |
| @group(0) @binding(4) var<storage, read> bias: array<f16>; |
| @group(0) @binding(5) var<storage, read_write> C: array<f16>; |
| var<workgroup> shA: array<f32, 1024>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) { |
| for (var i = tid; i < p.K; i += 256u) { shA[i] = f32(A[i]); } |
| workgroupBarrier(); |
| let base = wg.x * 1024u + tid * 4u; |
| var acc = vec4<f32>(0.0); |
| let KK = min(p.K, 65536u); |
| for (var k = 0u; k < KK; k++) { |
| let a = shA[k]; |
| let word = Pk[((k / 2u) * p.N + base) / 4u]; |
| let sh = (k & 1u) * 4u; |
| let s = vec4<f32>(f32(Sc[(k / p.group) * p.N + base]), f32(Sc[(k / p.group) * p.N + base + 1u]), |
| f32(Sc[(k / p.group) * p.N + base + 2u]), f32(Sc[(k / p.group) * p.N + base + 3u])); |
| let q = vec4<f32>(f32((word >> sh) & 0xFu), f32((word >> (8u + sh)) & 0xFu), |
| f32((word >> (16u + sh)) & 0xFu), f32((word >> (24u + sh)) & 0xFu)); |
| acc += a * (q - vec4<f32>(8.0)) * s; |
| } |
| let bb = vec4<f32>(f32(bias[base]), f32(bias[base + 1u]), f32(bias[base + 2u]), f32(bias[base + 3u])); |
| for (var j = 0u; j < 4u; j++) { if (base + j < p.N) { C[base + j] = f16(acc[j] + bb[j]); } } |
| } |
| `; |
|
|
| const LNM = ` |
| struct P { M: u32, D: u32, offS: u32, offT: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> X: array<f16>; |
| @group(0) @binding(2) var<storage, read> MOD: array<f16>; |
| @group(0) @binding(3) var<storage, read_write> Y: array<f16>; |
| var<workgroup> red: array<f32, 256>; |
| var<workgroup> stats: array<f32, 2>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) { |
| let row = wg.x; let D = p.D; let off = row * D; |
| var s = 0.0; |
| for (var i = tid; i < D; i += 256u) { s += f32(X[off + i]); } |
| red[tid] = s; workgroupBarrier(); |
| ${REDUCE8} |
| if (tid == 0u) { stats[0] = red[0] / f32(D); } |
| workgroupBarrier(); |
| let mu = stats[0]; |
| var vs = 0.0; |
| for (var i = tid; i < D; i += 256u) { let d = f32(X[off + i]) - mu; vs += d * d; } |
| red[tid] = vs; workgroupBarrier(); |
| ${REDUCE8} |
| if (tid == 0u) { stats[1] = inverseSqrt(red[0] / f32(D) + 1e-6); } |
| workgroupBarrier(); |
| let inv = stats[1]; |
| for (var i = tid; i < D; i += 256u) { |
| let sc = 1.0 + f32(MOD[p.offS + i]); |
| let sh = f32(MOD[p.offT + i]); |
| Y[off + i] = f16((f32(X[off + i]) - mu) * inv * sc + sh); |
| } |
| } |
| `; |
|
|
| const LNA = ` |
| struct P { M: u32, D: u32, pad0: u32, pad1: u32, eps: f32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> X: array<f16>; |
| @group(0) @binding(2) var<storage, read> Wt: array<f16>; |
| @group(0) @binding(3) var<storage, read> Bs: array<f16>; |
| @group(0) @binding(4) var<storage, read_write> Y: array<f16>; |
| var<workgroup> red: array<f32, 256>; |
| var<workgroup> stats: array<f32, 2>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) { |
| let row = wg.x; let D = p.D; let off = row * D; |
| var s = 0.0; |
| for (var i = tid; i < D; i += 256u) { s += f32(X[off + i]); } |
| red[tid] = s; workgroupBarrier(); |
| ${REDUCE8} |
| if (tid == 0u) { stats[0] = red[0] / f32(D); } |
| workgroupBarrier(); |
| let mu = stats[0]; |
| var vs = 0.0; |
| for (var i = tid; i < D; i += 256u) { let d = f32(X[off + i]) - mu; vs += d * d; } |
| red[tid] = vs; workgroupBarrier(); |
| ${REDUCE8} |
| if (tid == 0u) { stats[1] = inverseSqrt(red[0] / f32(D) + p.eps); } |
| workgroupBarrier(); |
| let inv = stats[1]; |
| for (var i = tid; i < D; i += 256u) { |
| Y[off + i] = f16((f32(X[off + i]) - mu) * inv * f32(Wt[i]) + f32(Bs[i])); |
| } |
| } |
| `; |
|
|
| const ROPE = ` |
| struct P { M: u32, N: u32, total: u32, pad: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> FREQ: array<f32>; |
| @group(0) @binding(2) var<storage, read_write> X: array<f16>; |
| const SC: f32 = 0.35355339059327379; |
| @compute @workgroup_size(256) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| let i = gid.x; |
| if (i >= p.total) { return; } |
| let pair = i % 32u; let t = i / 32u; let h = t % 16u; let row = t / 16u; |
| let pos = row % p.N; |
| let theta = f32(pos) * FREQ[pair]; |
| let c = cos(theta); let s = sin(theta); |
| let base = row * 1024u + h * 64u + pair * 2u; |
| let x0 = f32(X[base]); let x1 = f32(X[base + 1u]); |
| X[base] = f16((x0 * c - x1 * s) * SC); |
| X[base + 1u] = f16((x0 * s + x1 * c) * SC); |
| } |
| `; |
|
|
| const QK = ` |
| struct P { N: u32, B: u32, pad0: u32, pad1: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> Q: array<f16>; |
| @group(0) @binding(2) var<storage, read> K: array<f16>; |
| @group(0) @binding(3) var<storage, read_write> S: array<f32>; |
| var<workgroup> Qs: array<f16, 4096>; |
| var<workgroup> Ks: array<f16, 4096>; |
| @compute @workgroup_size(16, 16) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) { |
| let b = wg.z / 16u; let h = wg.z % 16u; |
| let qi = wg.x * 64u; let kj = wg.y * 64u; |
| let N = min(p.N, 8192u); |
| for (var i = 0u; i < 16u; i++) { |
| let idx = i * 256u + tid; let r = idx / 64u; let c = idx % 64u; |
| let qr = qi + r; var v: f16 = f16(0.0); |
| if (qr < N) { v = Q[((b * N + qr) * 16u + h) * 64u + c]; } |
| Qs[idx] = v; |
| } |
| for (var i = 0u; i < 16u; i++) { |
| let idx = i * 256u + tid; let r = idx / 64u; let c = idx % 64u; |
| let kr = kj + r; var v: f16 = f16(0.0); |
| if (kr < N) { v = K[((b * N + kr) * 16u + h) * 64u + c]; } |
| Ks[idx] = v; |
| } |
| workgroupBarrier(); |
| let ty = tid / 16u; let tx = tid % 16u; |
| var acc: array<vec4<f32>, 4>; |
| for (var i = 0u; i < 4u; i++) { acc[i] = vec4<f32>(0.0); } |
| for (var k = 0u; k < 64u; k++) { |
| let a0 = f32(Qs[(ty * 4u + 0u) * 64u + k]); |
| let a1 = f32(Qs[(ty * 4u + 1u) * 64u + k]); |
| let a2 = f32(Qs[(ty * 4u + 2u) * 64u + k]); |
| let a3 = f32(Qs[(ty * 4u + 3u) * 64u + k]); |
| let b0 = f32(Ks[(tx * 4u + 0u) * 64u + k]); |
| let b1 = f32(Ks[(tx * 4u + 1u) * 64u + k]); |
| let b2 = f32(Ks[(tx * 4u + 2u) * 64u + k]); |
| let b3 = f32(Ks[(tx * 4u + 3u) * 64u + k]); |
| let bv = vec4<f32>(b0, b1, b2, b3); |
| acc[0] += a0 * bv; acc[1] += a1 * bv; acc[2] += a2 * bv; acc[3] += a3 * bv; |
| } |
| let bidx = wg.z; |
| for (var r = 0u; r < 4u; r++) { |
| let gr = qi + ty * 4u + r; |
| if (gr < N) { |
| for (var c = 0u; c < 4u; c++) { |
| let gc = kj + tx * 4u + c; |
| if (gc < N) { S[(bidx * N + gr) * N + gc] = acc[r][c]; } |
| } |
| } |
| } |
| } |
| `; |
|
|
| const SOFTMAX = ` |
| struct P { N: u32, rows: u32, pad0: u32, pad1: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read_write> S: array<f32>; |
| var<workgroup> red: array<f32, 256>; |
| var<workgroup> st: array<f32, 2>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) { |
| let off = wg.x * p.N; |
| var m = -3.4e38; |
| for (var i = tid; i < p.N; i += 256u) { m = max(m, S[off + i]); } |
| red[tid] = m; workgroupBarrier(); |
| ${REDUCE8MAX} |
| if (tid == 0u) { st[0] = red[0]; } |
| workgroupBarrier(); |
| let mx = st[0]; |
| var sum = 0.0; |
| for (var i = tid; i < p.N; i += 256u) { sum += exp(S[off + i] - mx); } |
| red[tid] = sum; workgroupBarrier(); |
| ${REDUCE8} |
| if (tid == 0u) { st[1] = 1.0 / red[0]; } |
| workgroupBarrier(); |
| let inv = st[1]; |
| for (var i = tid; i < p.N; i += 256u) { S[off + i] = exp(S[off + i] - mx) * inv; } |
| } |
| `; |
|
|
| const AV = ` |
| struct P { N: u32, pad0: u32, pad1: u32, pad2: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> S: array<f32>; |
| @group(0) @binding(2) var<storage, read> V: array<f16>; |
| @group(0) @binding(3) var<storage, read_write> O: array<f16>; |
| var<workgroup> Ss: array<f32, 4096>; |
| var<workgroup> Vs: array<f16, 4096>; |
| @compute @workgroup_size(16, 16) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) { |
| let b = wg.z / 16u; let h = wg.z % 16u; |
| let qi = wg.x * 64u; |
| let N = min(p.N, 8192u); let bidx = wg.z; |
| let ty = tid / 16u; let tx = tid % 16u; |
| var acc: array<vec4<f32>, 4>; |
| for (var i = 0u; i < 4u; i++) { acc[i] = vec4<f32>(0.0); } |
| let nk = min((N + 63u) / 64u, 256u); |
| for (var kt = 0u; kt < nk; kt++) { |
| let k0 = kt * 64u; |
| for (var i = 0u; i < 16u; i++) { |
| let idx = i * 256u + tid; let r = idx / 64u; let c = idx % 64u; |
| let gr = qi + r; let skey = k0 + c; let vkey = k0 + r; |
| var v0: f32 = 0.0; |
| if (gr < N && skey < N) { v0 = S[(bidx * N + gr) * N + skey]; } |
| Ss[r * 64u + c] = v0; |
| var v1: f16 = f16(0.0); |
| if (vkey < N) { v1 = V[((b * N + vkey) * 16u + h) * 64u + c]; } |
| Vs[r * 64u + c] = v1; |
| } |
| workgroupBarrier(); |
| for (var k = 0u; k < 64u; k++) { |
| let a0 = Ss[(ty * 4u + 0u) * 64u + k]; |
| let a1 = Ss[(ty * 4u + 1u) * 64u + k]; |
| let a2 = Ss[(ty * 4u + 2u) * 64u + k]; |
| let a3 = Ss[(ty * 4u + 3u) * 64u + k]; |
| let b0 = f32(Vs[k * 64u + tx * 4u + 0u]); |
| let b1 = f32(Vs[k * 64u + tx * 4u + 1u]); |
| let b2 = f32(Vs[k * 64u + tx * 4u + 2u]); |
| let b3 = f32(Vs[k * 64u + tx * 4u + 3u]); |
| let bv = vec4<f32>(b0, b1, b2, b3); |
| acc[0] += a0 * bv; acc[1] += a1 * bv; acc[2] += a2 * bv; acc[3] += a3 * bv; |
| } |
| workgroupBarrier(); |
| } |
| for (var r = 0u; r < 4u; r++) { |
| let gr = qi + ty * 4u + r; |
| if (gr < N) { |
| for (var c = 0u; c < 4u; c++) { |
| O[((b * N + gr) * 16u + h) * 64u + tx * 4u + c] = f16(acc[r][c]); |
| } |
| } |
| } |
| } |
| `; |
|
|
| const ADDGATE = ` |
| struct P { M: u32, D: u32, gOff: u32, pad: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read_write> X: array<f16>; |
| @group(0) @binding(2) var<storage, read> Z: array<f16>; |
| @group(0) @binding(3) var<storage, read> MOD: array<f16>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| let i = gid.x; |
| if (i >= p.M * p.D) { return; } |
| let d = i % p.D; |
| X[i] = f16(f32(X[i]) + f32(MOD[p.gOff + d]) * f32(Z[i])); |
| } |
| `; |
|
|
| const GELU = ` |
| struct P { total: u32, pad0: u32, pad1: u32, pad2: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read_write> X: array<f16>; |
| fn erf(x: f32) -> f32 { |
| let ax = abs(x); |
| let t = 1.0 / (1.0 + 0.3275911 * ax); |
| let y = 1.0 - (((((1.061405429 * t - 1.453152027) * t) + 1.421413741) * t - 0.284496736) * t + 0.254829592) * t * exp(-ax * ax); |
| return select(-y, y, x >= 0.0); |
| } |
| @compute @workgroup_size(256) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| if (gid.x >= p.total) { return; } |
| let x = f32(X[gid.x]); |
| X[gid.x] = f16(0.5 * x * (1.0 + erf(x * 0.70710678118654752))); |
| } |
| `; |
|
|
| const SILU = ` |
| struct P { total: u32, pad0: u32, pad1: u32, pad2: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read_write> X: array<f16>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| if (gid.x >= p.total) { return; } |
| let x = f32(X[gid.x]); |
| X[gid.x] = f16(x / (1.0 + exp(-x))); |
| } |
| `; |
|
|
| const MISH = ` |
| struct P { total: u32, pad0: u32, pad1: u32, pad2: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read_write> X: array<f16>; |
| fn tanh_safe(v: f32) -> f32 { |
| if (v > 15.0) { return 1.0; } |
| if (v < -15.0) { return -1.0; } |
| return tanh(v); |
| } |
| @compute @workgroup_size(256) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| if (gid.x >= p.total) { return; } |
| let x = f32(X[gid.x]); |
| var sp = x; |
| if (x < 20.0) { sp = log(1.0 + exp(x)); } |
| X[gid.x] = f16(x * tanh_safe(sp)); |
| } |
| `; |
|
|
| const CONCAT3 = ` |
| struct P { N: u32, B: u32, pad0: u32, pad1: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> X: array<f16>; |
| @group(0) @binding(2) var<storage, read> C: array<f16>; |
| @group(0) @binding(3) var<storage, read> T: array<f16>; |
| @group(0) @binding(4) var<storage, read_write> O: array<f16>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| let i = gid.x; |
| if (i >= p.B * p.N * 712u) { return; } |
| let col = i % 712u; let row = i / 712u; |
| let b = row / p.N; let f = row % p.N; |
| var v: f16 = f16(0.0); |
| if (col < 100u) { v = X[f * 100u + col]; } |
| else if (col < 200u) { if (b == 0u) { v = C[f * 100u + (col - 100u)]; } } |
| else { if (b == 0u) { v = T[f * 512u + (col - 200u)]; } } |
| O[i] = v; |
| } |
| `; |
|
|
| const GATHER = ` |
| struct P { N: u32, pad0: u32, pad1: u32, pad2: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> IDS: array<i32>; |
| @group(0) @binding(2) var<storage, read> EMB: array<f16>; |
| @group(0) @binding(3) var<storage, read> FRQ: array<f16>; |
| @group(0) @binding(4) var<storage, read_write> O: array<f16>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| let i = gid.x; |
| if (i >= p.N * 512u) { return; } |
| let f = i / 512u; let d = i % 512u; |
| let id = IDS[f] + 1; |
| if (id <= 0) { O[i] = f16(0.0); return; } |
| O[i] = f16(f32(EMB[u32(id) * 512u + d]) + f32(FRQ[f * 512u + d])); |
| } |
| `; |
|
|
| const DWCONV = ` |
| struct P { M: u32, D: u32, pad0: u32, pad1: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> X: array<f16>; |
| @group(0) @binding(2) var<storage, read> Wt: array<f16>; |
| @group(0) @binding(3) var<storage, read> Bs: array<f16>; |
| @group(0) @binding(4) var<storage, read_write> Y: array<f16>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| let i = gid.x; |
| if (i >= p.M * p.D) { return; } |
| let d = i32(i % p.D); let f = i32(i / p.D); |
| var acc = f32(Bs[u32(d)]); |
| for (var k = 0; k < 7; k++) { |
| let ff = f - 3 + k; |
| if (ff >= 0 && ff < i32(p.M)) { acc += f32(Wt[u32(d) * 7u + u32(k)]) * f32(X[u32(ff) * p.D + u32(d)]); } |
| } |
| Y[i] = f16(acc); |
| } |
| `; |
|
|
| const GRNRED = ` |
| struct P { M: u32, D: u32, off: u32, pad: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> X: array<f16>; |
| @group(0) @binding(2) var<storage, read_write> NM: array<f32>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>) { |
| let d = wg.x + p.off; |
| var s = 0.0; |
| for (var f = 0u; f < p.M; f++) { let v = f32(X[f * p.D + d]); s += v * v; } |
| NM[d] = sqrt(s); |
| } |
| `; |
|
|
| const GRNMEAN = ` |
| struct P { D: u32, pad0: u32, pad1: u32, pad2: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> NM: array<f32>; |
| @group(0) @binding(2) var<storage, read_write> MEAN: array<f32>; |
| var<workgroup> red: array<f32, 256>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(local_invocation_index) tid: u32) { |
| var acc = 0.0; |
| for (var c = tid; c < p.D; c += 256u) { acc += NM[c]; } |
| red[tid] = acc; workgroupBarrier(); |
| ${REDUCE8} |
| if (tid == 0u) { MEAN[0] = red[0] / f32(p.D); } |
| } |
| `; |
|
|
| const GRNAPPLY = ` |
| struct P { M: u32, D: u32, pad0: u32, pad1: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> X: array<f16>; |
| @group(0) @binding(2) var<storage, read_write> Y: array<f16>; |
| @group(0) @binding(3) var<storage, read> G: array<f16>; |
| @group(0) @binding(4) var<storage, read> B: array<f16>; |
| @group(0) @binding(5) var<storage, read> NM: array<f32>; |
| @group(0) @binding(6) var<storage, read> MEAN: array<f32>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| let i = gid.x; |
| if (i >= p.M * p.D) { return; } |
| let d = i % p.D; |
| let v = f32(X[i]); |
| Y[i] = f16(v * (NM[d] / (MEAN[0] + 1e-6)) * f32(G[d]) + f32(B[d]) + v); |
| } |
| `; |
|
|
| const ADDRES = ` |
| struct P { total: u32, pad0: u32, pad1: u32, pad2: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read_write> Y: array<f16>; |
| @group(0) @binding(2) var<storage, read> X: array<f16>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| if (gid.x >= p.total) { return; } |
| Y[gid.x] = f16(f32(Y[gid.x]) + f32(X[gid.x])); |
| } |
| `; |
|
|
| const ZEROFILL = ` |
| struct P { M: u32, D: u32, pad0: u32, pad1: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> IDS: array<i32>; |
| @group(0) @binding(2) var<storage, read_write> X: array<f16>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| let i = gid.x; |
| if (i >= p.M * p.D) { return; } |
| if (IDS[i / p.D] < 0) { X[i] = f16(0.0); } |
| } |
| `; |
|
|
| const CONVPOS = ` |
| struct P { N: u32, M: u32, pad0: u32, pad1: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> X: array<f16>; |
| @group(0) @binding(2) var<storage, read> Wt: array<f16>; |
| @group(0) @binding(3) var<storage, read> Bs: array<f16>; |
| @group(0) @binding(4) var<storage, read_write> Y: array<f16>; |
| var<workgroup> Ws: array<f16, 2 * 64 * 31>; |
| var<workgroup> Xs: array<f32, 2 * 94>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_index) tid: u32) { |
| let g = wg.x / 16u; let pt = wg.x % 16u; |
| let p0 = pt * 64u; |
| let b = wg.y; |
| let tx = tid % 16u; let ty = tid / 16u; |
| let c0 = g * 64u + ty * 4u; |
| var acc: array<f32, 16>; |
| for (var i = 0u; i < 16u; i++) { acc[i] = 0.0; } |
| for (var cc = 0u; cc < 32u; cc++) { |
| let ci0 = cc * 2u; |
| for (var i = tid; i < 2u * 64u * 31u; i += 256u) { |
| let k = i % 31u; let oc = (i / 31u) % 64u; let ic = i / (31u * 64u); |
| Ws[i] = Wt[((g * 64u + oc) * 64u + (ci0 + ic)) * 31u + k]; |
| } |
| for (var i = tid; i < 2u * 94u; i += 256u) { |
| let ic = i / 94u; let pp = i % 94u; |
| let pos = i32(p0) - 15 + i32(pp); |
| var v: f32 = 0.0; |
| if (pos >= 0 && pos < i32(p.N)) { v = f32(X[(b * p.N + u32(pos)) * 1024u + g * 64u + ci0 + ic]); } |
| Xs[i] = v; |
| } |
| workgroupBarrier(); |
| for (var ic = 0u; ic < 2u; ic++) { |
| for (var k = 0u; k < 31u; k++) { |
| let x0 = Xs[ic * 94u + tx * 4u + k]; |
| let x1 = Xs[ic * 94u + tx * 4u + 1u + k]; |
| let x2 = Xs[ic * 94u + tx * 4u + 2u + k]; |
| let x3 = Xs[ic * 94u + tx * 4u + 3u + k]; |
| let w0 = f32(Ws[(ic * 64u + ty * 4u + 0u) * 31u + k]); |
| let w1 = f32(Ws[(ic * 64u + ty * 4u + 1u) * 31u + k]); |
| let w2 = f32(Ws[(ic * 64u + ty * 4u + 2u) * 31u + k]); |
| let w3 = f32(Ws[(ic * 64u + ty * 4u + 3u) * 31u + k]); |
| acc[0] += x0 * w0; acc[1] += x0 * w1; acc[2] += x0 * w2; acc[3] += x0 * w3; |
| acc[4] += x1 * w0; acc[5] += x1 * w1; acc[6] += x1 * w2; acc[7] += x1 * w3; |
| acc[8] += x2 * w0; acc[9] += x2 * w1; acc[10] += x2 * w2; acc[11] += x2 * w3; |
| acc[12] += x3 * w0; acc[13] += x3 * w1; acc[14] += x3 * w2; acc[15] += x3 * w3; |
| } |
| } |
| workgroupBarrier(); |
| } |
| for (var oc = 0u; oc < 4u; oc++) { |
| for (var j = 0u; j < 4u; j++) { |
| let pos = p0 + tx * 4u + j; |
| if (pos < p.N) { |
| Y[(b * p.N + pos) * 1024u + c0 + oc] = f16(clamp(acc[j * 4u + oc] + f32(Bs[c0 + oc]), -60000.0, 60000.0)); |
| } |
| } |
| } |
| } |
| `; |
|
|
| const FLOW = ` |
| struct P { total: u32, N: u32, B: u32, pad1: u32, dt: f32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read_write> X: array<f16>; |
| @group(0) @binding(2) var<storage, read> V: array<f16>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| if (gid.x >= p.total) { return; } |
| let i = gid.x; |
| let d = i % 100u; let f = i / 100u; |
| let v0 = f32(V[f * 100u + d]); |
| var v = v0; |
| if (p.B == 2u) { let v1 = f32(V[(p.N + f) * 100u + d]); v = v0 + (v0 - v1) * 2.0; } |
| X[i] = f16(f32(X[i]) + p.dt * v); |
| } |
| `; |
|
|
| const TEB = ` |
| struct P { N: u32, B: u32, pad0: u32, pad1: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> T: array<f16>; |
| @group(0) @binding(2) var<storage, read_write> O: array<f16>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| let i = gid.x; |
| if (i >= p.B * p.N * 512u) { return; } |
| let row = i / 512u; let d = i % 512u; |
| if (row < p.N) { O[i] = T[row * 512u + d]; } else { O[i] = f16(0.0); } |
| } |
| `; |
|
|
| const COPY = ` |
| struct P { total: u32, pad0: u32, pad1: u32, pad2: u32 }; |
| @group(0) @binding(0) var<uniform> p: P; |
| @group(0) @binding(1) var<storage, read> X: array<f16>; |
| @group(0) @binding(2) var<storage, read_write> Y: array<f16>; |
| @compute @workgroup_size(256) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| if (gid.x >= p.total) { return; } |
| Y[gid.x] = X[gid.x]; |
| } |
| `; |
|
|
| export const KERNELS = { |
| gemm: GEMM, gemv: GEMV, lnm: LNM, lna: LNA, rope: ROPE, qk: QK, softmax: SOFTMAX, |
| av: AV, addgate: ADDGATE, gelu: GELU, silu: SILU, mish: MISH, concat3: CONCAT3, |
| gather: GATHER, dwconv: DWCONV, grnred: GRNRED, grnmean: GRNMEAN, grnapply: GRNAPPLY, |
| addres: ADDRES, zerofill: ZEROFILL, convpos: CONVPOS, flow: FLOW, teb: TEB, copy: COPY, |
| }; |
|
|
| function f16tof32(h) { |
| const sign = (h & 0x8000) ? -1 : 1, exp = (h >> 10) & 0x1f, frac = h & 0x3ff; |
| if (exp === 0) return sign * Math.pow(2, -14) * (frac / 1024); |
| if (exp === 31) return frac ? NaN : sign * Infinity; |
| return sign * Math.pow(2, exp - 15) * (1 + frac / 1024); |
| } |
|
|
| export class F5Engine { |
| static async create(device, opts = {}) { |
| const e = new F5Engine(device, opts); |
| await e._init(); |
| return e; |
| } |
| constructor(device, opts) { this.device = device; this.opts = opts; this.f32 = opts.dtype === 'f32'; this.es = this.f32 ? 4 : 2; } |
|
|
| async _init() { |
| const dev = this.device; |
| const [mr, br] = await Promise.all([ |
| fetch(this.opts.metaUrl || '/webgpu/f5w.json'), |
| fetch(this.opts.binUrl || '/webgpu/f5w.bin'), |
| ]); |
| this.meta = await mr.json(); |
| this.binSize = Number(br.headers.get('content-length')) || 0; |
| this.wsize = Math.ceil((this.f32 ? this.binSize * 2 : this.binSize) / 256) * 256; |
| this._planChunks(); |
| this.wbufs = this.chunkStarts.map((st, i) => { |
| const end = i + 1 < this.chunkStarts.length ? this.chunkStarts[i + 1] : this.wsize; |
| return dev.createBuffer({ size: Math.ceil((end - st) / 256) * 256, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST }); |
| }); |
| await this._streamWeights(br); |
| this.ubufs = []; |
| this.ubNext = 0; |
| this.compileErrors = []; |
| const cre = (name, code) => { |
| if (this.f32) code = code.replace(/enable f16;\n/g, '').replace(/\bf16\b/g, 'f32'); |
| const mod = dev.createShaderModule({ code: (this.f32 ? '' : PRELUDE) + code }); |
| mod.getCompilationInfo().then((ci) => { |
| for (const m of ci.messages) if (m.type === 'error') this.compileErrors.push(name + ': ' + m.message + ' @' + m.lineNum); |
| }); |
| return dev.createComputePipelineAsync({ layout: 'auto', compute: { module: mod, entryPoint: 'main' } }); |
| }; |
| this.gemmKind = this.opts.gemm || (this.device.features.has('chromium-experimental-subgroup-matrix') ? 'sgmat' : 'simt'); |
| console.log('DBG gemmKind', this.gemmKind, 'opts', JSON.stringify(Object.keys(this.opts || {}))); |
| const defs = { |
| gemm: this.gemmKind === 'sgmat' ? GEMM : GEMM_SIMT, gemm_i4: GEMM_SIMT_I4, gemv: GEMV, gemv_i4: GEMV_I4, lnm: LNM, lna: LNA, rope: ROPE, |
| qk: QK, softmax: SOFTMAX, av: AV, addgate: ADDGATE, gelu: GELU, |
| silu: SILU, mish: MISH, concat3: CONCAT3, gather: GATHER, |
| dwconv: DWCONV, grnred: GRNRED, grnmean: GRNMEAN, |
| grnapply: GRNAPPLY, addres: ADDRES, zerofill: ZEROFILL, |
| convpos: CONVPOS, flow: FLOW, teb: TEB, copy: COPY, |
| }; |
| this.pipe = {}; |
| for (const [name, code] of Object.entries(defs)) { |
| this.pipe[name] = await cre(name, code); |
| } |
| this.pipeNames = new Map(Object.entries(this.pipe).map(([k, v]) => [v, k])); |
| } |
|
|
| _planChunks() { |
| const mult = this.f32 ? 2 : 1; |
| const ents = Object.entries(this.meta).map(([n, t]) => { |
| if (t.q) return { n, off: t.packed * mult, len: (t.scales + t.scale_bytes) - t.packed }; |
| const len = Math.ceil(t.shape.reduce((a, b) => a * b, 1) * this.es / 4) * 4; |
| return { n, off: t.offset * mult, len }; |
| }).sort((a, b) => a.off - b.off); |
| const CHUNK = 640 << 20; |
| const starts = [0]; |
| let cur = 0; |
| this.bufOf = new Map(); |
| for (const e of ents) { |
| if (cur > 0 && cur + e.len > CHUNK) { starts.push(e.off); cur = 0; } |
| this.bufOf.set(e.n, starts.length - 1); |
| cur = e.off + e.len - starts[starts.length - 1]; |
| } |
| this.chunkStarts = starts; |
| } |
|
|
| _writeAt(globalOff, data, byteLen) { |
| let off = globalOff, src = 0, remain = byteLen; |
| while (remain > 0) { |
| let bi = this.chunkStarts.length - 1; |
| for (let i = 0; i < this.chunkStarts.length; i++) { |
| const end = i + 1 < this.chunkStarts.length ? this.chunkStarts[i + 1] : this.wsize; |
| if (off < end) { bi = i; break; } |
| } |
| const base = this.chunkStarts[bi]; |
| const bo = off - base; |
| const end = bi + 1 < this.chunkStarts.length ? this.chunkStarts[bi + 1] : this.wsize; |
| const n = Math.min(remain, end - off); |
| if (data instanceof Float32Array) { |
| this.device.queue.writeBuffer(this.wbufs[bi], bo, data.subarray(src / 4, (src + n) / 4)); |
| } else { |
| this.device.queue.writeBuffer(this.wbufs[bi], bo, data.subarray(src, src + n)); |
| } |
| off += n; src += n; remain -= n; |
| } |
| } |
|
|
| async _streamWeights(res) { |
| |
| const dev = this.device; |
| const reader = res.body.getReader(); |
| const CH = 32 << 20; |
| let stage = new Uint8Array(CH + 8); |
| let staged = 0, gpuOff = 0; |
| const flush = async (final) => { |
| if (!staged) return; |
| let n = staged; |
| const extra = final ? (4 - (n % 4)) % 4 : 0; |
| if (extra) { stage.fill(0, n, n + extra); n += extra; } |
| if (this.f32) { |
| const elems = n >> 1; |
| const out = new Float32Array(elems); |
| out.set(new Float16Array(stage.buffer, 0, elems)); |
| this._writeAt(gpuOff, out, elems * 4); |
| gpuOff += elems * 4; |
| } else { |
| this._writeAt(gpuOff, stage.subarray(0, n), n); |
| gpuOff += n; |
| } |
| staged = 0; |
| if (this.opts.onProgress) this.opts.onProgress(gpuOff / (this.f32 ? this.binSize * 2 : this.binSize)); |
| }; |
| while (true) { |
| const { done, value } = await reader.read(); |
| if (done) break; |
| let off = 0; |
| while (off < value.length) { |
| const take = Math.min(value.length - off, stage.length - staged); |
| stage.set(value.subarray(off, off + take), staged); |
| staged += take; off += take; |
| if (staged >= CH) await flush(false); |
| } |
| } |
| await flush(true); |
| const pad = Math.ceil(gpuOff / 256) * 256 - gpuOff; |
| if (pad > 0) this._writeAt(gpuOff, new Uint8Array(pad), pad); |
| } |
|
|
| _t(name) { const t = this.meta[name]; if (!t) throw new Error('missing ' + name); return t; } |
| _wslice(name) { |
| const t = this._t(name); |
| const n = t.shape.reduce((a, b) => a * b, 1); |
| const mult = this.f32 ? 2 : 1; |
| const bi = this.bufOf.get(name); |
| const base = this.chunkStarts[bi]; |
| return { buffer: this.wbufs[bi], offset: t.offset * mult - base, size: Math.ceil(n * this.es / 4) * 4 }; |
| } |
| _alc(bytes) { |
| const size = Math.ceil(bytes / 256) * 256; |
| return this.device.createBuffer({ size, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC }); |
| } |
| _ub(vals, floats) { |
| const ab = new ArrayBuffer(256); |
| const dv = new DataView(ab); |
| for (let i = 0; i < vals.length; i++) { |
| if (floats && floats.indexOf(i) >= 0) dv.setFloat32(i * 4, vals[i], true); |
| else dv.setUint32(i * 4, vals[i] >>> 0, true); |
| } |
| const buf = this.device.createBuffer({ size: 256, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST }); |
| this.device.queue.writeBuffer(buf, 0, ab); |
| this.ubNext++; |
| this.ubufs.push(buf); |
| return { buffer: buf, offset: 0, size: 256 }; |
| } |
| _bg(pipe, u, bufs) { |
| const wrap = (b) => (b && b.buffer) ? b : { buffer: b, offset: 0, size: b.size }; |
| return this.device.createBindGroup({ |
| layout: pipe.getBindGroupLayout(0), |
| entries: [{ binding: 0, resource: wrap(u) }, ...bufs.map((b, i) => ({ binding: i + 1, resource: wrap(b) }))], |
| }); |
| } |
|
|
| setup(N) { |
| const dev = this.device; |
| this.N = N; this.B = this.opts.branches || 2; this.M2 = this.B * N; |
| const A = (name, elems) => (this[name] = this._alc(elems * this.es)); |
| A('xBuf', N * 100); A('condBuf', N * 100); |
| A('teBuf', N * 512); A('tebBuf', 2 * N * 512); A('catBuf', 2 * N * 712); |
| A('hBuf', 2 * N * 1024); A('tmpBuf', 2 * N * 1024); A('qBuf', 2 * N * 1024); |
| A('kBuf', 2 * N * 1024); A('vBuf', 2 * N * 1024); A('attnBuf', 2 * N * 1024); |
| A('ffBuf', 2 * N * 2048); A('velBuf', 2 * N * 100); |
| A('timeIn', 256); A('tembBuf', 1024); A('temb2Buf', 1024); |
| A('modBuf', 137216); this.scoresBuf = this._alc(this.B * 16 * N * N * 4); |
| A('tBuf', N * 1024); A('tBuf2', N * 1024); A('tBuf3', N * 2048); |
| A('tnmBuf', 1024 * 4); A('tmeanBuf', 4); |
| this.idsBuf = this._alc(N * 4); |
| this.ubNext = 0; |
| this.ropeFreq = dev.createBuffer({ size: 256, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST }); |
| const f = new Float32Array(64); |
| for (let j = 0; j < 32; j++) f[j] = Math.pow(10000, -j / 32); |
| dev.queue.writeBuffer(this.ropeFreq, 0, f); |
| this.calls = []; |
| this._buildStepCalls(); |
| const stepCalls = this.calls; |
| this.calls = []; |
| this._buildTextCalls(); |
| this.textCalls = this.calls; |
| this.calls = stepCalls; |
| } |
|
|
| _add(pipe, bg, x, y = 1, z = 1, meta) { |
| const c = { pipe, bg, x: Math.ceil(x), y, z, name: this.pipeNames.get(pipe), meta }; |
| this.calls.push(c); |
| if (this.debug) this.debug.push(c); |
| return c; |
| } |
|
|
| _qdesc(name) { |
| const t = this._t(name); |
| const mult = this.f32 ? 2 : 1; |
| const bi = this.bufOf.get(name); |
| const base = this.chunkStarts[bi]; |
| return { |
| packed: { buffer: this.wbufs[bi], offset: t.packed * mult - base, size: t.packed_bytes }, |
| scales: { buffer: this.wbufs[bi], offset: t.scales * mult - base, size: t.scale_bytes }, |
| group: t.group, |
| }; |
| } |
|
|
| _gemm(A, wn, bn, C, M, K, N) { |
| const t = this._t(wn); |
| if (t.q) { |
| if (!(M > 0 && N > 0 && K > 0 && M < 1e5 && N < 1e6 && K < 1e6)) throw new Error(`bad gemm args ${M},${K},${N} (${wn})`); |
| const q = this._qdesc(wn); |
| if (q.packed.offset + q.packed.size > q.packed.buffer.size) throw new Error(`q4 packed OOB ${wn}`); |
| const u = this._ub([M, N, K, q.group]); |
| this._add(this.pipe.gemm_i4, this._bg(this.pipe.gemm_i4, u, [A, q.packed, q.scales, this._wslice(bn), C]), |
| Math.ceil(M / 64), Math.ceil(N / 64)); |
| return; |
| } |
| const wa = this._wslice(wn), ba = this._wslice(bn); |
| if (!(M > 0 && N > 0 && K > 0 && M < 1e5 && N < 1e6 && K < 1e6)) throw new Error(`bad gemm args ${M},${K},${N} (${wn})`); |
| if (wa.offset + wa.size > wa.buffer.size) throw new Error(`gemm weight OOB ${wn}`); |
| if (ba.offset + ba.size > ba.buffer.size) throw new Error(`gemm bias OOB ${bn}`); |
| const u = this._ub([M, N, K, 0]); |
| this._add(this.pipe.gemm, this._bg(this.pipe.gemm, u, [A, wa, ba, C]), Math.ceil(M / (this.gemmKind === 'sgmat' ? 128 : 64)), Math.ceil(N / 64), 1, |
| { vals: [M, N, K], wa: [wa.offset, wa.size, wa.buffer.size], ba: [ba.offset, ba.size, ba.buffer.size] }); |
| } |
| _gemv(A, wn, bn, C, K, N) { |
| const t = this._t(wn); |
| if (t.q) { |
| const q = this._qdesc(wn); |
| const u = this._ub([N, K, q.group, 0]); |
| this._add(this.pipe.gemv_i4, this._bg(this.pipe.gemv_i4, u, [A, q.packed, q.scales, this._wslice(bn), C]), Math.ceil(N / 1024)); |
| return; |
| } |
| const u = this._ub([N, K, 0, 0]); |
| this._add(this.pipe.gemv, this._bg(this.pipe.gemv, u, [A, this._wslice(wn), this._wslice(bn), C]), Math.ceil(N / 1024)); |
| } |
| _elem(pipe, X, total) { |
| const u = this._ub([total, 0, 0, 0]); |
| this._add(pipe, this._bg(pipe, u, [X]), Math.ceil(total / 256)); |
| } |
|
|
| _buildStepCalls() { |
| const N = this.N, M2 = this.M2; |
| this._gemm(this.timeIn, 'time.0', 'time.0.b', this.tembBuf, 1, 256, 1024); |
| this._elem(this.pipe.silu, this.tembBuf, 1024); |
| this._gemm(this.tembBuf, 'time.2', 'time.2.b', this.temb2Buf, 1, 1024, 1024); |
| this._elem(this.pipe.silu, this.temb2Buf, 1024); |
| this._gemv(this.temb2Buf, 'ada.w', 'ada.b', this.modBuf, 1024, 137216); |
| this._add(this.pipe.concat3, this._bg(this.pipe.concat3, this._ub([N, this.B, 0, 0]), [this.xBuf, this.condBuf, this.tebBuf, this.catBuf]), Math.ceil(this.B * N * 712 / 256)); |
| this._gemm(this.catBuf, 'in.proj', 'in.proj.b', this.hBuf, M2, 712, 1024); |
| const cp = (X, wn, bn, Y) => { |
| const u = this._ub([N, M2, 0, 0]); |
| this._add(this.pipe.convpos, this._bg(this.pipe.convpos, u, [X, this._wslice(wn), this._wslice(bn), Y]), 16 * 16, this.B); |
| }; |
| cp(this.hBuf, 'convpos.0.w', 'convpos.0.b', this.tmpBuf); |
| this._elem(this.pipe.mish, this.tmpBuf, M2 * 1024); |
| cp(this.tmpBuf, 'convpos.2.w', 'convpos.2.b', this.qBuf); |
| this._elem(this.pipe.mish, this.qBuf, M2 * 1024); |
| { |
| const u = this._ub([M2 * 1024, 0, 0, 0]); |
| this._add(this.pipe.addres, this._bg(this.pipe.addres, u, [this.hBuf, this.qBuf]), Math.ceil(M2 * 1024 / 256)); |
| } |
| for (let L = 0; L < 22; L++) { |
| const base = L * 6144; |
| { |
| const u = this._ub([M2, 1024, base + 1024, base]); |
| this._add(this.pipe.lnm, this._bg(this.pipe.lnm, u, [this.hBuf, this.modBuf, this.tmpBuf]), M2); |
| } |
| this._gemm(this.tmpBuf, `b${L}.q`, `b${L}.q.b`, this.qBuf, M2, 1024, 1024); |
| this._gemm(this.tmpBuf, `b${L}.k`, `b${L}.k.b`, this.kBuf, M2, 1024, 1024); |
| this._gemm(this.tmpBuf, `b${L}.v`, `b${L}.v.b`, this.vBuf, M2, 1024, 1024); |
| for (const buf of [this.qBuf, this.kBuf]) { |
| const u = this._ub([M2, N, M2 * 16 * 32, 0]); |
| this._add(this.pipe.rope, this._bg(this.pipe.rope, u, [this.ropeFreq, buf]), Math.ceil(M2 * 16 * 32 / 256)); |
| } |
| { |
| const u = this._ub([N, 0, 0, 0]); |
| this._add(this.pipe.qk, this._bg(this.pipe.qk, u, [this.qBuf, this.kBuf, this.scoresBuf]), Math.ceil(N / 64), Math.ceil(N / 64), this.B * 16); |
| const u2 = this._ub([N, this.B * 16 * N, 0, 0]); |
| this._add(this.pipe.softmax, this._bg(this.pipe.softmax, u2, [this.scoresBuf]), this.B * 16 * N); |
| const u3 = this._ub([N, 0, 0, 0]); |
| this._add(this.pipe.av, this._bg(this.pipe.av, u3, [this.scoresBuf, this.vBuf, this.attnBuf]), Math.ceil(N / 64), 1, this.B * 16); |
| } |
| this._gemm(this.attnBuf, `b${L}.o`, `b${L}.o.b`, this.tmpBuf, M2, 1024, 1024); |
| { |
| const u = this._ub([M2, 1024, base + 2048, 0]); |
| this._add(this.pipe.addgate, this._bg(this.pipe.addgate, u, [this.hBuf, this.tmpBuf, this.modBuf]), Math.ceil(M2 * 1024 / 256)); |
| } |
| { |
| const u = this._ub([M2, 1024, base + 4096, base + 3072]); |
| this._add(this.pipe.lnm, this._bg(this.pipe.lnm, u, [this.hBuf, this.modBuf, this.tmpBuf]), M2); |
| } |
| this._gemm(this.tmpBuf, `b${L}.up`, `b${L}.up.b`, this.ffBuf, M2, 1024, 2048); |
| this._elem(this.pipe.gelu, this.ffBuf, M2 * 2048); |
| this._gemm(this.ffBuf, `b${L}.down`, `b${L}.down.b`, this.tmpBuf, M2, 2048, 1024); |
| { |
| const u = this._ub([M2, 1024, base + 5120, 0]); |
| this._add(this.pipe.addgate, this._bg(this.pipe.addgate, u, [this.hBuf, this.tmpBuf, this.modBuf]), Math.ceil(M2 * 1024 / 256)); |
| } |
| } |
| const fbase = 22 * 6144; |
| { |
| const u = this._ub([M2, 1024, fbase, fbase + 1024]); |
| this._add(this.pipe.lnm, this._bg(this.pipe.lnm, u, [this.hBuf, this.modBuf, this.tmpBuf]), M2); |
| } |
| this._gemm(this.tmpBuf, 'out.w', 'out.b', this.velBuf, M2, 1024, 100); |
| this.flowUniform = this._ub([N * 100, N, this.B, 0]); |
| this._add(this.pipe.flow, this._bg(this.pipe.flow, this.flowUniform, [this.xBuf, this.velBuf]), Math.ceil(N * 100 / 256)); |
| } |
|
|
| _buildTextCalls() { |
| const N = this.N; |
| { |
| const u = this._ub([N, 0, 0, 0]); |
| this._add(this.pipe.gather, this._bg(this.pipe.gather, u, [this.idsBuf, this._wslice('txt.emb'), this._wslice('txt.freqs'), this.teBuf]), Math.ceil(N * 512 / 256)); |
| } |
| for (let b = 0; b < 4; b++) { |
| { |
| const u = this._ub([N, 512, 0, 0]); |
| this._add(this.pipe.dwconv, this._bg(this.pipe.dwconv, u, [this.teBuf, this._wslice(`tb${b}.dw.w`), this._wslice(`tb${b}.dw.b`), this.tBuf]), Math.ceil(N * 512 / 256)); |
| } |
| { |
| const u = this._ub([N, 512, 0, 0, 1e-5, 0, 0, 0], [4]); |
| this._add(this.pipe.lna, this._bg(this.pipe.lna, u, [this.tBuf, this._wslice(`tb${b}.ln.w`), this._wslice(`tb${b}.ln.b`), this.tBuf3]), N); |
| } |
| this._gemm(this.tBuf3, `tb${b}.p1`, `tb${b}.p1.b`, this.tBuf2, N, 512, 1024); |
| this._elem(this.pipe.gelu, this.tBuf2, N * 1024); |
| { |
| const u = this._ub([N, 1024, 0, 0]); |
| this._add(this.pipe.grnred, this._bg(this.pipe.grnred, u, [this.tBuf2, this.tnmBuf]), 1024); |
| const u2 = this._ub([1024, 0, 0, 0]); |
| this._add(this.pipe.grnmean, this._bg(this.pipe.grnmean, u2, [this.tnmBuf, this.tmeanBuf]), 1); |
| const u3 = this._ub([N, 1024, 0, 0]); |
| this._add(this.pipe.grnapply, this._bg(this.pipe.grnapply, u3, [this.tBuf2, this.tBuf, this._wslice(`tb${b}.grn.g`), this._wslice(`tb${b}.grn.b`), this.tnmBuf, this.tmeanBuf]), Math.ceil(N * 1024 / 256)); |
| } |
| this._gemm(this.tBuf, `tb${b}.p2`, `tb${b}.p2.b`, this.tBuf2, N, 1024, 512); |
| { |
| const u = this._ub([N * 512, 0, 0, 0]); |
| this._add(this.pipe.addres, this._bg(this.pipe.addres, u, [this.tBuf2, this.teBuf]), Math.ceil(N * 512 / 256)); |
| const u2 = this._ub([N, 512, 0, 0]); |
| this._add(this.pipe.zerofill, this._bg(this.pipe.zerofill, u2, [this.idsBuf, this.tBuf2]), Math.ceil(N * 512 / 256)); |
| const u3 = this._ub([N * 512, 0, 0, 0]); |
| this._add(this.pipe.copy, this._bg(this.pipe.copy, u3, [this.tBuf2, this.teBuf]), Math.ceil(N * 512 / 256)); |
| } |
| } |
| { |
| const u = this._ub([N, this.B, 0, 0]); |
| this._add(this.pipe.teb, this._bg(this.pipe.teb, u, [this.teBuf, this.tebBuf]), Math.ceil(this.B * N * 512 / 256)); |
| } |
| } |
|
|
| textNames() { return this.textCalls.map((c) => c.name); } |
| async testGemm2(M, K, N, mode) { |
| const dev = this.device; |
| const mk = (bytes) => dev.createBuffer({ size: Math.ceil(bytes / 4) * 4, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC }); |
| const A = mode === 'dedicated' ? mk(M * K * 2) : this.tBuf3; |
| const C = mode === 'dedicated' ? mk(M * N * 2) : this.tBuf2; |
| const W = mode === 'dedicated' ? { buffer: mk(K * N * 2), offset: 0, size: K * N * 2 } : this._wslice('tb0.p1'); |
| const Bs = mode === 'dedicated' ? { buffer: mk(N * 2), offset: 0, size: N * 2 } : this._wslice('tb0.p1.b'); |
| const u = this._ub([M, N, K, 0]); |
| const bg = this._bg(this.pipe.gemm, u, [A, W, Bs, C]); |
| this.encodeCalls([{ pipe: this.pipe.gemm, bg, x: Math.ceil(M / 128), y: Math.ceil(N / 64), z: 1 }]); |
| await this.device.queue.onSubmittedWorkDone(); |
| return mode + ' ok'; |
| } |
|
|
| async testGemm(M, K, N) { |
| const u = this._ub([M, N, K, 0]); |
| const bg = this._bg(this.pipe.gemm, u, [this.tBuf3, this._wslice('tb0.p1'), this._wslice('tb0.p1.b'), this.tBuf2]); |
| this.encodeCalls([{ pipe: this.pipe.gemm, bg, x: Math.ceil(M / 64), y: Math.ceil(N / 64), z: 1 }]); |
| await this.device.queue.onSubmittedWorkDone(); |
| return M + 'x' + K + 'x' + N; |
| } |
|
|
| debugText() { return this.textCalls.map((c) => ({ name: c.name, x: c.x, y: c.y, meta: c.meta })); } |
|
|
| async runTextPrefix(k) { |
| this.encodeCalls(this.textCalls.slice(0, k)); |
| await this.device.queue.onSubmittedWorkDone(); |
| return k; |
| } |
|
|
| async runOne(list, k) { |
| this.encodeCalls([list[k]]); |
| await this.device.queue.onSubmittedWorkDone(); |
| return list[k].name; |
| } |
|
|
| async setCond(condMelF32, ids, limit) { |
| const dev = this.device, N = this.N; |
| const ca = this.f32 ? new Float32Array(N * 100) : new Float16Array(N * 100); |
| ca.set(condMelF32.subarray(0, N * 100)); |
| dev.queue.writeBuffer(this.condBuf, 0, ca); |
| const idarr = new Int32Array(N).fill(-1); |
| for (let i = 0; i < Math.min(ids.length, N); i++) idarr[i] = ids[i]; |
| dev.queue.writeBuffer(this.idsBuf, 0, idarr); |
| const list = (limit === undefined) ? this.textCalls : this.textCalls.slice(0, limit); |
| if (list.length) this.encodeCalls(list); |
| if (limit !== 0) await dev.queue.onSubmittedWorkDone(); |
| } |
|
|
| setX(xF32) { |
| const f = this.f32 ? new Float32Array(this.N * 100) : new Float16Array(this.N * 100); |
| f.set(xF32.subarray(0, this.N * 100)); |
| this.device.queue.writeBuffer(this.xBuf, 0, f); |
| } |
|
|
| encodeCalls(list) { |
| const enc = this.device.createCommandEncoder(); |
| const p = enc.beginComputePass(); |
| for (const c of list) { |
| p.setPipeline(c.pipe); |
| p.setBindGroup(0, c.bg); |
| p.dispatchWorkgroups(c.x, c.y, c.z); |
| } |
| p.end(); |
| this.device.queue.submit([enc.finish()]); |
| } |
|
|
| writeTime(t) { |
| const half = 128; |
| const tb = this.f32 ? new Float32Array(256) : new Float16Array(256); |
| for (let j = 0; j < 128; j++) { |
| const fr = Math.exp(-Math.log(10000) / (half - 1) * j); |
| const a = t * 1000 * fr; |
| tb[j] = Math.sin(a); tb[128 + j] = Math.cos(a); |
| } |
| this.device.queue.writeBuffer(this.timeIn, 0, tb); |
| } |
|
|
| async step(t, dt) { |
| this.writeTime(t); |
| this.device.queue.writeBuffer(this.flowUniform.buffer, this.flowUniform.offset + 16, new Float32Array([dt])); |
| this.encodeCalls(this.calls); |
| await this.device.queue.onSubmittedWorkDone(); |
| } |
|
|
| async runStepPrefix(k, t) { |
| this.writeTime(t === undefined ? 0.5 : t); |
| this.encodeCalls(this.calls.slice(0, k)); |
| await this.device.queue.onSubmittedWorkDone(); |
| } |
|
|
| async readBuf(buf, elems) { |
| const bytes = elems * this.es; |
| const staging = this.device.createBuffer({ size: Math.ceil(bytes / 4) * 4, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST }); |
| const enc = this.device.createCommandEncoder(); |
| enc.copyBufferToBuffer(buf, 0, staging, 0, bytes); |
| this.device.queue.submit([enc.finish()]); |
| await staging.mapAsync(GPUMapMode.READ); |
| const ab = staging.getMappedRange().slice(0); |
| staging.unmap(); |
| return this.f32 ? new Float32Array(ab) : new Float32Array(new Float16Array(ab)); |
| } |
|
|
| async readVel() { |
| const bytes = this.M2 * 100 * this.es; |
| const buf = this.device.createBuffer({ size: Math.ceil(bytes / 4) * 4, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST }); |
| const enc = this.device.createCommandEncoder(); |
| enc.copyBufferToBuffer(this.velBuf, 0, buf, 0, bytes); |
| this.device.queue.submit([enc.finish()]); |
| await buf.mapAsync(GPUMapMode.READ); |
| const ab = buf.getMappedRange().slice(0); |
| const out = this.f32 ? new Float32Array(ab) : new Float32Array(new Float16Array(ab)); |
| buf.unmap(); |
| return out; |
| } |
|
|
| async readX() { |
| const bytes = this.N * 100 * this.es; |
| const buf = this.device.createBuffer({ size: Math.ceil(bytes / 4) * 4, usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST }); |
| const enc = this.device.createCommandEncoder(); |
| enc.copyBufferToBuffer(this.xBuf, 0, buf, 0, bytes); |
| this.device.queue.submit([enc.finish()]); |
| await buf.mapAsync(GPUMapMode.READ); |
| const ab = buf.getMappedRange().slice(0); |
| const out = this.f32 ? new Float32Array(ab) : new Float32Array(new Float16Array(ab)); |
| buf.unmap(); |
| return out; |
| } |
| } |
|
|
