diff --git a/docs/backends.md b/docs/backends.md index 82f13bb..ffa55a1 100644 --- a/docs/backends.md +++ b/docs/backends.md @@ -100,11 +100,15 @@ That is what lets a backend realize a launch with no per-kernel host code: followed by the params' fields, four bytes apiece. Its kernels take their pointers first and 4-byte scalars after (116 of the 117 do; `tools/cuda_trace/gen_kernel_sigs.py` reads this off the `.cu`). -- **WebGPU**'s kernels predate the ABI: every entry point reads one 96-byte - uniform layout, a family picks its operation by number, and the bind group is - fixed (A and B read, C written, D and E read). Its core carries a `marshal_` - from the canonical params into that layout, per kernel id. This stays inside - `webgpu.h`; a backend whose kernels follow the ABI needs none. +- **WebGPU** binds view *i*, whole, at binding *i* and a uniform at binding + *n*: the views' element offsets, then the params. A binding offset has to be + 256-byte aligned, which a view's is not, so the offsets travel in the uniform + and each kernel folds them into its indexing. The WGSL declares each + kernel's bindings, read-only for an input and read-write for an output, and + each pipeline takes its layout from them. WGSL has no templates, so a family + of operations (add, sub, ...) is one entry point, and the operation is the + pipeline-overridable constant `OP`, set per kernel id from the kernel table + rather than carried in the params. Where two backends' kernels disagreed, the CUDA kernel's order is canonical, because the `.cu` is the one source no development machine here can run, and @@ -228,7 +232,7 @@ which never reaches the evaluator, is profiled from its first kernel. 1. If every backend can run it as one kernel with the same buffers: add its params struct to `gpu_abi.h`, the function to `gpu_ops.h`, the kernel to each backend's source with that layout, and the id to each backend's kernel table - (`kernel_name_` in `metal.h` and `cuda.h`, `marshal_` in `webgpu.h`). + (`kernel_name_` in `metal.h` and `cuda.h`, `kernel_` in `webgpu.h`). 2. Otherwise add the detecting wrapper to `gpu_ops.h` and the member to the `own` struct of each backend that implements it. 3. A backend that gets neither simply declines the op. Nothing else changes. @@ -266,11 +270,6 @@ What a test may ask of a backend is asked in code, not by platform macro: `gpu::caps`, `gpu::traits`. A test that needs to know whether a kernel exists asks; a new backend edits no test. -## What is not shared yet - -- WebGPU's kernels still predate the canonical ABI and go through `marshal_` - (see above). - ## Verifying a change All of these run on a development Mac. diff --git a/include/gpu_abi.h b/include/gpu_abi.h index d169dca..2549f83 100644 --- a/include/gpu_abi.h +++ b/include/gpu_abi.h @@ -14,7 +14,8 @@ // order the kernel takes its scalars. That is what lets each backend realize a // launch generically — Metal binds view i at buffer index i and the params // after them; CUDA builds cuLaunchKernel's argv as the view addresses followed -// by the params' fields, four bytes apiece. +// by the params' fields, four bytes apiece; WebGPU binds view i at binding i +// and a uniform after them holding the views' offsets and then the params. #include #include diff --git a/include/webgpu.h b/include/webgpu.h index 70b478f..49dc8d5 100644 --- a/include/webgpu.h +++ b/include/webgpu.h @@ -39,6 +39,7 @@ #include "gpu_abi.h" // the op vocabulary and the launch contract #include "profile.h" +#include "shape.h" // tl::contiguous_strides_into (concat_part's meta) #include "types.h" #if defined(TENSORLIB_WEBGPU) && defined(__EMSCRIPTEN__) @@ -49,9 +50,11 @@ // not in emscripten/html5_webgpu.h (that is the old built-in binding's home). #include +#include #include #include #include +#include #include #include #include @@ -71,70 +74,87 @@ inline const char* wgsl_source_() { return src; } -// Uniform params, laid out to match the WGSL Params struct. One struct serves -// every kernel family (see the comment on Params in the .wgsl): unused fields -// cost a few bytes of a 256-byte slot, and it keeps one uniform ring and one -// bind group layout for the whole backend. -struct params { - uint32_t M, N, K; - uint32_t lda, ldb, ldc; - uint32_t a_off, b_off, c_off; - uint32_t ta, tb; - uint32_t ars, acs, brs, bcs; - uint32_t op; - float scale, offset; - uint32_t pad0, pad1, pad2, pad3, pad4; - float arg; // a scalar operand (ew_scalar's s) +// The kernel table: a kernel id's WGSL entry point and the OP its family +// selects the operation by (-1: the entry point has no OP), or no entry for an +// id this backend has no kernel for. The counterpart of metal.h's and cuda.h's +// kernel_name_; the OP values are the case labels of the family's switch in +// the .wgsl. +struct kernel { + const char* entry = nullptr; + int op = -1; }; - -// WGSL gives a uniform-address-space struct align 16, so Params is 96 bytes -// there. This must agree: the bind group's minBindingSize comes from sizeof -// here, and a short one fails validation on every dispatch. -static_assert(sizeof(params) == 96, "params must match the WGSL Params size"); - -// Which operation within a family, matching the OP_* constants in the WGSL. -// Families have separate numbering, so this is only meaningful alongside the -// entry point it is passed to. -inline uint32_t kernel_op_(kop op) { - switch (op) { - case kop::add: case kop::badd: return 0; - case kop::sub: case kop::bsub: return 1; - case kop::mul: case kop::bmul: return 2; - case kop::div: case kop::bdiv: return 3; - case kop::pow_: case kop::bpow: return 4; - - case kop::exp_: return 0; - case kop::log_: return 1; - case kop::sqrt_: return 2; - case kop::sigmoid: return 3; - case kop::relu: return 4; - case kop::affine: return 5; - case kop::tanh_: return 6; - case kop::sin_: return 7; - case kop::cos_: return 8; - - case kop::row_sum: return 0; - case kop::row_max: return 1; - default: return 0; +inline kernel kernel_(kop k) { + switch (k) { + case kop::add: return {"ew_binary", 0}; + case kop::sub: return {"ew_binary", 1}; + case kop::mul: return {"ew_binary", 2}; + case kop::div: return {"ew_binary", 3}; + case kop::pow_: return {"ew_binary", 4}; + + case kop::exp_: return {"ew_unary", 0}; + case kop::log_: return {"ew_unary", 1}; + case kop::sqrt_: return {"ew_unary", 2}; + case kop::sigmoid: return {"ew_unary", 3}; + case kop::relu: return {"ew_unary", 4}; + case kop::affine: return {"ew_unary", 5}; + case kop::tanh_: return {"ew_unary", 6}; + case kop::sin_: return {"ew_unary", 7}; + case kop::cos_: return {"ew_unary", 8}; + + case kop::badd: return {"ew_bcast", 0}; + case kop::bsub: return {"ew_bcast", 1}; + case kop::bmul: return {"ew_bcast", 2}; + case kop::bdiv: return {"ew_bcast", 3}; + case kop::bpow: return {"ew_bcast", 4}; + + case kop::gt_: return {"cmp", 0}; + case kop::lt_: return {"cmp", 1}; + case kop::ge_: return {"cmp", 2}; + case kop::le_: return {"cmp", 3}; + case kop::eq_: return {"cmp", 4}; + case kop::ne_: return {"cmp", 5}; + + case kop::pow_s_: return {"ew_scalar", 0}; + case kop::gt_s_: return {"ew_scalar", 1}; + case kop::lt_s_: return {"ew_scalar", 2}; + case kop::ge_s_: return {"ew_scalar", 3}; + case kop::le_s_: return {"ew_scalar", 4}; + case kop::eq_s_: return {"ew_scalar", 5}; + case kop::ne_s_: return {"ew_scalar", 6}; + + case kop::row_sum: return {"row_reduce", 0}; + case kop::row_max: return {"row_reduce", 1}; + case kop::softmax: return {"softmax", -1}; + case kop::clamp_: return {"clamp_", -1}; + case kop::layer_norm_: return {"layer_norm", -1}; + case kop::index_select: return {"index_select", -1}; + default: return {}; } } +// The uniform a kernel reads (tensorlib_webgpu.wgsl's header): its views' +// element offsets, one slot a view, then its params struct. +constexpr size_t kMaxViews = 8; +constexpr size_t kViewsBytes = kMaxViews * sizeof(uint32_t); + // WebGPU guarantees maxComputeWorkgroupsPerDimension >= 65535. Anything past // that returns false and falls to CPU rather than silently truncating. constexpr int64_t kMaxWorkgroups = 65535; -// Dynamic uniform offsets must be a multiple of the adapter's +// A uniform binding's offset must be a multiple of the adapter's // minUniformBufferOffsetAlignment; 256 is the spec's guaranteed-safe maximum. constexpr uint64_t kUniformSlotBytes = 256; -// One flush can batch this many dispatches; past it, encode_ forces a blocking -// flush mid-graph. At 96 bytes of payload per 256-byte slot the ring is pure -// device memory (1 MB here) with no binding-size implication, so it is sized to -// put that forced stall well beyond any graph the backend is aimed at. +// One flush can batch this many launches; past it, launch_ forces a blocking +// flush mid-graph. The ring is pure device memory (1 MB here) with no +// binding-size implication, so it is sized to put that forced stall well +// beyond any graph the backend is aimed at. constexpr uint32_t kUniformSlotCount = 4096; -// Every WGSL entry point. The context prebuilds a pipeline for each and the -// browser harness asserts each one dispatched — both need the same list, and a -// new kernel missing from either loses a guarantee silently. +// Every WGSL entry point. The context prebuilds a pipeline for each (with its +// OP at the default; the family's other operations are the same code under +// another constant, built on first use) and the browser harness asserts each +// one dispatched — both need the same list, and a new kernel missing from +// either loses a guarantee silently. inline constexpr const char* kEntryPoints[] = { "sgemm", "ew_binary", "ew_unary", "ew_bcast", "softmax", "row_reduce", "pad", "fold", @@ -161,12 +181,19 @@ inline bool has_preinitialized_device_() { return EM_ASM_INT({ return Module["preinitializedWebGPUDevice"] ? 1 : 0; }) != 0; } +// Byte offsets must be 4-aligned to convert to the element offsets the +// kernels index with. They always are for f32 views; anything else falls to +// the CPU rather than silently truncating. +inline bool elem_off_(int64_t byte_off, uint32_t* out) { + if (byte_off % 4) return false; + *out = (uint32_t)(byte_off / 4); + return true; +} + struct context { wgpu::Instance instance; wgpu::Device device; wgpu::Queue queue; - wgpu::BindGroupLayout bgl; - wgpu::PipelineLayout play; wgpu::ShaderModule mod; wgpu::Buffer uniforms; // ring of kUniformSlots x kUniformSlot bytes bool ready = false; @@ -182,7 +209,7 @@ struct context { wgpu::ComputePassEncoder pass; uint32_t slot = 0; // next free uniform ring slot - // pad_/fold_'s per-call shape metadata ring — same idea as the uniform + // The N-D kernels' per-call shape metadata ring — same idea as the uniform // ring above (queue.WriteBuffer runs ahead of whatever is still sitting in // the unsubmitted encoder, so two calls sharing one buffer before a flush // would have the second's write stomp the first dispatch's not-yet- @@ -208,10 +235,19 @@ struct context { // kernel operand, only mapped. std::unordered_map> staging_pool; - // Compute pipelines, keyed by WGSL entry point. Every one is built in the - // constructor and they all share a layout, so by the time encode_ runs this - // is a pure lookup and pipeline_'s create branch is unreachable. - std::unordered_map pipelines; + // Compute pipelines, keyed by WGSL entry point and OP, each with the bind + // group layout it took from its entry point (an auto layout: the WGSL's + // declarations are the kernel's side of the ABI). + struct pipeline { + wgpu::ComputePipeline pipe; + wgpu::BindGroupLayout layout; + const char* entry; // for the census and the profile row + }; + std::unordered_map pipelines; + // Each kernel id's pipeline once dispatch has resolved it, so a shared op's + // launch indexes an array where the map lookup would build a key. + // unordered_map never moves its values, so the pointers stay good. + std::array kop_pipelines{}; // Per-entry-point dispatch census. Unlike the native backends, this one has // no test runner that fails when it is absent: the browser suite passes // whether or not the GPU engages, because every unported op falls back to @@ -274,41 +310,6 @@ struct context { } if (!compiled) return; - // Explicit layout rather than GetBindGroupLayout(0): the auto-generated - // one has no dynamic offset on the uniform binding, which the ring needs. - // Bindings 4 (D) and 5 (E) are a third and fourth read-only operand — - // binary_bcast_nd and where_nd need more real buffers (operands + N-D - // shape/stride meta) than A/B/C alone can hold; see tensorlib_webgpu.wgsl's - // comment on D/E for which kernel binds what there. - wgpu::BindGroupLayoutEntry be[6] = {}; - for (int i = 0; i < 3; ++i) { - be[i].binding = i; - be[i].visibility = wgpu::ShaderStage::Compute; - be[i].buffer.type = i == 2 ? wgpu::BufferBindingType::Storage - : wgpu::BufferBindingType::ReadOnlyStorage; - } - be[3].binding = 3; - be[3].visibility = wgpu::ShaderStage::Compute; - be[3].buffer.type = wgpu::BufferBindingType::Uniform; - be[3].buffer.hasDynamicOffset = true; - be[3].buffer.minBindingSize = sizeof(params); - for (int i = 4; i < 6; ++i) { - be[i].binding = i; - be[i].visibility = wgpu::ShaderStage::Compute; - be[i].buffer.type = wgpu::BufferBindingType::ReadOnlyStorage; - } - wgpu::BindGroupLayoutDescriptor bgld = {}; - bgld.entryCount = 6; - bgld.entries = be; - bgl = device.CreateBindGroupLayout(&bgld); - if (!bgl) return; - - wgpu::PipelineLayoutDescriptor pld = {}; - pld.bindGroupLayoutCount = 1; - pld.bindGroupLayouts = &bgl; - play = device.CreatePipelineLayout(&pld); - if (!play) return; - wgpu::BufferDescriptor ud = {}; ud.size = kUniformSlotBytes * kUniformSlotCount; ud.usage = wgpu::BufferUsage::Uniform | wgpu::BufferUsage::CopyDst; @@ -319,22 +320,34 @@ struct context { // error would surface as an op quietly falling back to CPU forever; here // it makes available() false, which the harness reports. for (const char* ep : kEntryPoints) { - if (!pipeline_(ep)) return; + if (!pipeline_(ep, -1)) return; } ready = true; } - wgpu::ComputePipeline pipeline_(const char* entry) { - auto it = pipelines.find(entry); - if (it != pipelines.end()) return it->second; + // `entry` with OP set to `op`. 0 is OP's default in the WGSL, so it and -1 + // (no OP) name the one pipeline the constructor prebuilt. No layout is + // given, so the pipeline takes the auto layout of what its entry point + // declares. + const pipeline* pipeline_(const char* entry, int op) { + std::string key = entry; + if (op > 0) key += '#' + std::to_string(op); + auto it = pipelines.find(key); + if (it != pipelines.end()) return &it->second; + wgpu::ConstantEntry constant = {}; + constant.key = "OP"; + constant.value = op; wgpu::ComputePipelineDescriptor pd = {}; - pd.layout = play; pd.compute.module = mod; pd.compute.entryPoint = entry; + if (op > 0) { + pd.compute.constantCount = 1; + pd.compute.constants = &constant; + } wgpu::ComputePipeline p = device.CreateComputePipeline(&pd); - if (p) pipelines[entry] = p; - return p; + if (!p) return nullptr; + return &(pipelines[key] = pipeline{p, p.GetBindGroupLayout(0), entry}); } // Blocking wait on a single future — the one place anything suspends. @@ -352,50 +365,60 @@ struct context { // live bytes are the host's has no encoded command touching it: a pending // kernel write would have made the device copy live, and a pending kernel // read would have come through here and uploaded already. - void before_kernel_(void* native, gpu::access a) { - mirror* m = mirror_(native); - if (m && m->live.before_kernel(a)) queue.WriteBuffer(m->dev, 0, m->host, m->bytes); + void before_kernel_(mirror& m, gpu::access a) { + if (m.live.before_kernel(a)) queue.WriteBuffer(m.dev, 0, m.host, m.bytes); } - void device_read_(void* native) { before_kernel_(native, gpu::access::in); } - void device_write_(void* native) { before_kernel_(native, gpu::access::out); } - // The one place a dispatch is encoded. Every op differs only in which - // pipeline, which params and what grid — keeping the bind group, uniform - // ring and encoder bookkeeping in a single copy is the same discipline - // metal_kernels.metal applies to its kernel bodies. + // The one place a kernel is launched, the shared ops' and this backend's own + // alike: view i bound whole at binding i, then the uniform — the views' + // element offsets and the params — at binding n (tensorlib_webgpu.wgsl's + // header). Nothing here knows a kernel. Declines, so the op falls back to + // the CPU, for a view with no device buffer or an offset that is not a + // whole float. // - // `b` may be the same mirror as `a` (a unary or reduce kernel binds its one - // input twice): two read-only bindings may alias. `out` is always a fresh - // allocation from the evaluator, so a writable binding never does. `d`/`e` - // are the optional third/fourth read-only operand (binary_bcast_nd/ - // where_nd's extra tensor operand and/or N-D shape/stride meta); every - // other kernel leaves them null, which binds them to `a` -- unused by that - // kernel's WGSL, but bind group validation requires every declared binding - // be present regardless of which ones the active entry point reads. - bool encode_(const char* entry, mirror* a, mirror* b, mirror* out, - const params& p, int64_t gx, int64_t gy, mirror* d = nullptr, - mirror* e = nullptr) { - if (gx <= 0 || gy <= 0) return false; - if (gx > kMaxWorkgroups || gy > kMaxWorkgroups) return false; - wgpu::ComputePipeline pipe = pipeline_(entry); - if (!pipe) return false; + // Two inputs may be the same buffer: read-only bindings may alias. A buffer + // bound as an output may not be bound again in the same launch — WebGPU + // invalidates the whole encoder for it, and with it every launch batched + // since the last flush — so no caller hands one over twice. + bool launch_(const pipeline* p, const gpu::arg* args, size_t n, + const void* params, size_t params_bytes, uint32_t gx, + uint32_t gy) { + if (!p || n == 0 || n > kMaxViews || + kViewsBytes + params_bytes > kUniformSlotBytes) { + return false; + } + if (gx == 0 || gy == 0 || gx > kMaxWorkgroups || gy > kMaxWorkgroups) { + return false; + } + mirror* views[kMaxViews]; + uint32_t offs[kMaxViews] = {}; + for (size_t i = 0; i < n; i++) { + views[i] = mirror_(args[i].s.buf); + if (!views[i] || !elem_off_(args[i].s.off, &offs[i])) return false; + } + if (slot >= kUniformSlotCount) flush_(); // ring exhausted; new batch + for (size_t i = 0; i < n; i++) before_kernel_(*views[i], args[i].a); const uint32_t off = slot++ * (uint32_t)kUniformSlotBytes; - queue.WriteBuffer(uniforms, off, &p, sizeof(p)); - - mirror* md = d ? d : a; - mirror* me = e ? e : a; - wgpu::BindGroupEntry bge[6] = {}; - bge[0].binding = 0; bge[0].buffer = a->dev; bge[0].size = a->bytes; - bge[1].binding = 1; bge[1].buffer = b->dev; bge[1].size = b->bytes; - bge[2].binding = 2; bge[2].buffer = out->dev; bge[2].size = out->bytes; - bge[3].binding = 3; bge[3].buffer = uniforms; bge[3].size = sizeof(params); - bge[4].binding = 4; bge[4].buffer = md->dev; bge[4].size = md->bytes; - bge[5].binding = 5; bge[5].buffer = me->dev; bge[5].size = me->bytes; + unsigned char u[kUniformSlotBytes]; + std::memcpy(u, offs, kViewsBytes); + std::memcpy(u + kViewsBytes, params, params_bytes); + queue.WriteBuffer(uniforms, off, u, kViewsBytes + params_bytes); + + wgpu::BindGroupEntry bge[kMaxViews + 1] = {}; + for (size_t i = 0; i < n; i++) { + bge[i].binding = static_cast(i); + bge[i].buffer = views[i]->dev; + bge[i].size = views[i]->bytes; + } + bge[n].binding = static_cast(n); + bge[n].buffer = uniforms; + bge[n].offset = off; + bge[n].size = kUniformSlotBytes; wgpu::BindGroupDescriptor bgd = {}; - bgd.layout = bgl; - bgd.entryCount = 6; + bgd.layout = p->layout; + bgd.entryCount = n + 1; bgd.entries = bge; wgpu::BindGroup bg = device.CreateBindGroup(&bgd); @@ -403,13 +426,13 @@ struct context { enc = device.CreateCommandEncoder(); pass = enc.BeginComputePass(); } - pass.SetPipeline(pipe); - pass.SetBindGroup(0, bg, 1, &off); - pass.DispatchWorkgroups((uint32_t)gx, (uint32_t)gy, 1); + pass.SetPipeline(p->pipe); + pass.SetBindGroup(0, bg); + pass.DispatchWorkgroups(gx, gy, 1); pending = true; - dispatch_counts[entry]++; - gpu::launched(entry); // counted, untimed + dispatch_counts[p->entry]++; + gpu::launched(p->entry); // counted, untimed return true; } @@ -574,243 +597,102 @@ inline void sync_to_host(void* native, bool for_write) { if (for_write) m->live.host_wrote(); } -// Shared host-side prologue for every op: resolve the operand mirrors and -// stage the lazy copies. Returns false when any operand is untracked — that -// means a heap storage with no device buffer, so the op belongs on the CPU. -// `b` may be null for one-input kernels, which then bind `a` twice. -inline bool operands_(context& c, void* a, void* b, void* out, - context::mirror** ma, context::mirror** mb, - context::mirror** mo) { - *ma = c.mirror_(a); - *mb = b ? c.mirror_(b) : *ma; - *mo = c.mirror_(out); - if (!*ma || !*mb || !*mo) return false; - c.device_read_(a); - if (b) c.device_read_(b); - c.device_write_(out); - return true; -} - -// Byte offsets must be 4-aligned to convert to the element offsets the -// kernels index with. They always are for f32 views; anything else falls to -// the CPU rather than silently truncating. -inline bool elem_off_(int64_t byte_off, uint32_t* out) { - if (byte_off % 4) return false; - *out = (uint32_t)(byte_off / 4); - return true; -} - -// out = op(a) @ op(b) * scale + offset. -inline bool own::gemm(gpu::span a, int64_t lda, bool ta, gpu::span b, - int64_t ldb, bool tb, gpu::span out, int64_t m, int64_t n, - int64_t k, float scale, float offset) { +// The device core's one way to run a kernel for the shared ops: the id's +// entry point and OP from the kernel table, and the rest as the shared op +// handed it — views, params, grid — with no per-kernel code. +inline bool dispatch(kop k, const gpu::arg* args, size_t n, const void* params, + size_t params_bytes, const gpu::grid& g) { auto& c = context::get(); - if (!c.ready || m <= 0 || n <= 0 || k <= 0) return false; - params p = {}; - if (!elem_off_(a.off, &p.a_off) || !elem_off_(b.off, &p.b_off) || - !elem_off_(out.off, &p.c_off)) { - return false; + if (!c.ready) return false; + const context::pipeline*& p = c.kop_pipelines[static_cast(k)]; + if (!p) { + const kernel kn = kernel_(k); + if (!kn.entry) return false; + p = c.pipeline_(kn.entry, kn.op); } - context::mirror *ma, *mb, *mo; - if (!operands_(c, a.buf, b.buf, out.buf, &ma, &mb, &mo)) return false; - - p.M = (uint32_t)m; - p.N = (uint32_t)n; - p.K = (uint32_t)k; - p.lda = (uint32_t)lda; - p.ldb = (uint32_t)ldb; - p.ldc = (uint32_t)n; // the eval seam always hands us a contiguous output - p.ta = ta ? 1u : 0u; - p.tb = tb ? 1u : 0u; - p.scale = scale; - p.offset = offset; - return c.encode_("sgemm", ma, mb, mo, p, (n + 63) / 64, (m + 63) / 64); + return c.launch_(p, args, n, params, params_bytes, g.gx, g.gy); } -// This backend's kernels predate the shared kernel ABI (gpu_abi.h): every WGSL -// entry point reads the one `params` layout above, and a family picks its -// operation by number. So each kernel id the shared ops may dispatch is -// marshalled here, from its canonical params into that layout; the entry -// point comes back, or null for an id this backend has no kernel for. -// -// `in_off` holds the element offsets of the views the kernel reads. A and B's -// are placed by dispatch; a kernel that reads D or E at an offset moves it -// into its own field and clears it here, so an offset nobody consumed is -// caught rather than dropped. -inline const char* marshal_(kop k, const void* canonical, uint32_t* in_off, - params& p) { - switch (k) { - case kop::badd: case kop::bsub: case kop::bmul: case kop::bdiv: - case kop::bpow: { - const auto& q = *static_cast(canonical); - p.M = q.m; - p.N = q.n; - p.ars = q.ars; - p.acs = q.acs; - p.brs = q.brs; - p.bcs = q.bcs; - p.op = kernel_op_(k); - p.scale = q.scale; - p.offset = q.offset; - return "ew_bcast"; - } - case kop::gt_: case kop::lt_: case kop::ge_: case kop::le_: case kop::eq_: - case kop::ne_: { - const auto& q = *static_cast(canonical); - if (q.n == 0) return nullptr; - p.M = q.n; - p.ars = q.bstride; - p.op = static_cast(k) - static_cast(kop::gt_); - return "cmp"; - } - case kop::clamp_: { // no epilogue: scale/offset carry lo/hi - const auto& q = *static_cast(canonical); - if (q.n == 0) return nullptr; - p.M = q.n; - p.scale = q.lo; - p.offset = q.hi; - return "clamp_"; - } - case kop::pow_s_: case kop::gt_s_: case kop::lt_s_: case kop::ge_s_: - case kop::le_s_: case kop::eq_s_: case kop::ne_s_: { - const auto& q = *static_cast(canonical); - if (q.n == 0) return nullptr; - p.M = q.n; - p.op = static_cast(k) - static_cast(kop::pow_s_); - p.arg = q.s; - p.scale = q.scale; - p.offset = q.offset; - return "ew_scalar"; - } - case kop::softmax: case kop::row_sum: case kop::row_max: { - const auto& q = *static_cast(canonical); - p.M = q.rows; - p.N = q.cols; - p.op = kernel_op_(k); - p.scale = q.scale; - p.offset = q.offset; - return k == kop::softmax ? "softmax" : "row_reduce"; - } - case kop::layer_norm_: { // A = x, B = g, D = b at pad3, arg = eps - const auto& q = *static_cast(canonical); - p.M = q.rows; - p.N = q.cols; - p.arg = q.eps; - p.scale = q.scale; - p.offset = q.offset; - p.pad3 = in_off[2]; - in_off[2] = 0; - return "layer_norm"; - } - case kop::index_select: { - const auto& q = *static_cast(canonical); - p.M = q.n; - p.pad0 = q.row_size; - return "index_select"; - } - case kop::add: case kop::sub: case kop::mul: case kop::div: case kop::pow_: - case kop::exp_: case kop::log_: case kop::sqrt_: case kop::sigmoid: - case kop::relu: case kop::affine: case kop::tanh_: case kop::sin_: - case kop::cos_: { - const auto& q = *static_cast(canonical); - if (q.n == 0) return nullptr; - p.M = q.n; - p.op = kernel_op_(k); - p.scale = q.scale; - p.offset = q.offset; - return k <= kop::pow_ ? "ew_binary" : "ew_unary"; - } - default: return nullptr; - } -} +namespace detail_ { + +// This backend's own kernels' params: 4-byte fields, which +// tensorlib_webgpu.wgsl declares in the same order after the views' offsets. +struct gemm_params { + uint32_t m, n, k, lda, ldb, ldc, ta, tb; + float scale, offset; +}; +struct pad_params { uint32_t n, rank, axis, before; }; +struct fold_params { uint32_t n, rank, axis, step; }; +struct sum_to_params { uint32_t n, rank, reduced_n; }; +struct concat_params { uint32_t n, rank, shift; }; +struct bcast_nd_params { + uint32_t n, rank; + float scale, offset; +}; +struct where_nd_params { uint32_t n, rank; }; +struct index_add_params { uint32_t n, row_size, k; }; +struct scatter_axis_params { uint32_t n, size; }; +struct rope_params { + uint32_t n, t, d, pos, half; + float base; +}; + +inline uint32_t u32(int64_t v) { return static_cast(v); } -// The device core's one way to run a kernel for the shared ops. The bind -// group is fixed — A and B read, C written, D and E read — so the view a -// kernel writes is C and the ones it reads fill A, B, D, E in order; a -// one-input kernel binds its input twice. Offsets ride in the uniform as -// element counts (A, B and C only: D and E are bound whole). -inline bool dispatch(kop k, const gpu::arg* args, size_t n, - const void* canonical, size_t /*params_bytes*/, - const gpu::grid& g) { +// An own op's launch, through the same launch_ as the shared ops. +template +inline bool launch_own_(const char* entry, int op, + std::initializer_list args, const P& p, + const gpu::grid& g) { auto& c = context::get(); if (!c.ready) return false; - params p = {}; - context::mirror* in[4] = {}; - uint32_t in_off[4] = {}; - context::mirror* out = nullptr; - size_t ins = 0; - for (size_t i = 0; i < n; i++) { - context::mirror* m = c.mirror_(args[i].s.buf); - uint32_t off = 0; - if (!m || !elem_off_(args[i].s.off, &off)) return false; // CPU's - if (args[i].a == gpu::access::in) { - if (ins == 4) return false; - in_off[ins] = off; - in[ins++] = m; - } else { - if (out) return false; // one writable binding - out = m; - p.c_off = off; - } - } - if (!out || ins == 0) return false; - if (ins == 1) { - in[1] = in[0]; - in_off[1] = in_off[0]; - } - p.a_off = in_off[0]; - p.b_off = in_off[1]; - const char* entry = marshal_(k, canonical, in_off, p); - if (!entry || in_off[2] || in_off[3]) return false; - for (size_t i = 0; i < n; i++) c.before_kernel_(args[i].s.buf, args[i].a); - return c.encode_(entry, in[0], in[1], out, p, g.gx, g.gy, in[2], in[3]); + return c.launch_(c.pipeline_(entry, op), args.begin(), args.size(), &p, + sizeof(P), g.gx, g.gy); } -// A one-input elementwise dispatch over n elements: the kernel binds its one -// input twice, and `fill` sets the family's own fields. -template -inline bool encode_one_input_(const char* entry, void* a, int64_t ao, void* out, - int64_t oo, int64_t n, Fill&& fill) { - auto& c = context::get(); - if (!c.ready || n <= 0) return false; - params p = {}; - if (!elem_off_(ao, &p.a_off) || !elem_off_(oo, &p.c_off)) return false; - context::mirror *ma, *mb, *mo; - if (!operands_(c, a, nullptr, out, &ma, &mb, &mo)) return false; - p.b_off = p.a_off; - p.M = static_cast(n); - fill(p); - return c.encode_(entry, ma, mb, mo, p, (n + 255) / 256, 1); +} // namespace detail_ + +// out = op(a) @ op(b) * scale + offset. +inline bool own::gemm(gpu::span a, int64_t lda, bool ta, gpu::span b, + int64_t ldb, bool tb, gpu::span out, int64_t m, int64_t n, + int64_t k, float scale, float offset) { + using detail_::u32; + if (m <= 0 || n <= 0 || k <= 0) return false; + // ldc = n: the eval seam always hands us a contiguous output. + detail_::gemm_params p{u32(m), u32(n), u32(k), u32(lda), u32(ldb), u32(n), + ta ? 1u : 0u, tb ? 1u : 0u, scale, offset}; + return detail_::launch_own_("sgemm", -1, + {gpu::in(a), gpu::in(b), gpu::out(out)}, p, + gpu::grid{u32((n + 63) / 64), u32((m + 63) / 64)}); } // A ring, not one reused buffer: queue.WriteBuffer runs ahead of whatever is -// still sitting in the unsubmitted encoder (see device_read_'s comment -// above), so two pad/fold calls batched into the same unflushed pass would -// have the second call's metadata write stomp the first dispatch's -// not-yet-executed read of the same buffer — the exact hazard the uniform -// ring above (kUniformSlotCount) already exists to avoid for Params, just for -// a second resource. One slot comfortably covers the rank-8 cap (pad needs -// 2*rank <= 16 words, fold (rank-1)+rank <= 15, binary_bcast_nd 3*rank <= 24, -// where_nd's own [out_shape, cond_strides, a_strides, b_strides] 4*rank <= -// 32 -- the largest of the four, which is what sizes this), which is why -// this state lives on `context` (meta_ring_tok/meta_ring_host/meta_slot) -// right beside `slot` instead of as a second, independent ring: flush() -// resets both counters together, the way it already resets `slot`. +// still sitting in the unsubmitted encoder (see before_kernel_'s comment +// above), so two N-D calls batched into the same unflushed pass would have +// the second call's metadata write stomp the first launch's not-yet-executed +// read of the same buffer — the exact hazard the uniform ring above +// (kUniformSlotCount) already exists to avoid, just for a second resource. +// One slot comfortably covers the rank-8 cap (pad needs 2*rank <= 16 words, +// fold (rank-1)+rank <= 15, binary_bcast_nd 3*rank <= 24, where_nd's own +// [out_shape, cond_strides, a_strides, b_strides] 4*rank <= 32 -- the largest, +// which is what sizes this), which is why this state lives on `context` +// (meta_ring_tok/meta_ring_host/meta_slot) right beside `slot` instead of as +// a second, independent ring: flush() resets both counters together, the way +// it already resets `slot`. inline constexpr size_t kMetaSlotWords = 32; inline constexpr size_t kMetaSlotCount = 4096; -// Rank cap for pad_/fold_'s GPU dispatch — matches cuda.h's own -// kPadFoldMaxRank (not unified with it: the two backends derive their caps -// from different physical constraints, kMetaSlotWords here vs a fixed-size -// on-stack index array there) and kernels/tensorlib_webgpu.wgsl's -// kPadFoldMaxRank, which the WGSL side needs as its own `const` since a -// shader can't see a host-side C++ constant. +// Rank cap for the N-D kernels — matches cuda.h's own kPadFoldMaxRank (not +// unified with it: the two backends derive their caps from different physical +// constraints, kMetaSlotWords here vs a fixed-size on-stack index array +// there) and kernels/tensorlib_webgpu.wgsl's kPadFoldMaxRank, which the WGSL +// side needs as its own `const` since a shader can't see a host-side C++ +// constant. inline constexpr int kPadFoldMaxRank = 8; // Allocated once, sized for the whole ring — through the same alloc() pool -// every tensor buffer uses (by the time pad()/fold() below can run, alloc() -// is already defined above). The token is opaque to eval_one's storage -// layer, so nothing else could mistake it for a live array. +// every tensor buffer uses. The token is opaque to eval_one's storage layer, +// so nothing else could mistake it for a live array. inline void* context::meta_ring_(float** host_out) { if (!meta_ring_tok) { meta_ring_tok = alloc( @@ -823,45 +705,49 @@ inline void* context::meta_ring_(float** host_out) { } // This call's word offset into the ring, advancing like the uniform ring's -// `slot` and forcing the same flush-then-reset on wraparound. +// `slot`. A full ring — either ring: the launch this slot is for comes next, +// and were launch_ to flush between them, later calls in the new batch would +// be handed this slot again and rewrite it before the launch had read it — +// flushes first, so a slot and its launch always share a batch. inline uint32_t context::meta_reserve_slot_() { - if (meta_slot >= kMetaSlotCount) { + if (meta_slot >= kMetaSlotCount || slot >= kUniformSlotCount) { flush(); meta_slot = 0; } return (meta_slot++) * static_cast(kMetaSlotWords); } -// Reserve one meta-ring slot: returns the u32 write pointer for this call's -// slot (already offset into the ring) via `host_words`, and this slot's word -// offset via `word_off_out`; the return value is the ring's own opaque -// token, or null if the ring couldn't be allocated. Shared by pad()/fold()/ -// binary_bcast_nd()/where_nd() below, which differ only in how many words -// they fill in and with what. -inline void* reserve_meta_(context& c, uint32_t* word_off_out, - uint32_t** host_words) { +// A run of `len` shape or stride values, one u32 word each in a meta slot. +struct meta_run { + const int64_t* p; + int len; +}; + +// `runs`, back to back, in a fresh slot of the meta ring, as a view an N-D +// kernel reads like any other input; a null span when they overflow a slot +// or the ring could not be allocated. +inline gpu::span meta_(std::initializer_list runs) { + uint32_t words[kMetaSlotWords]; + size_t count = 0; + for (const meta_run& r : runs) { + if (r.len < 0 || count + r.len > kMetaSlotWords) return {}; + for (int d = 0; d < r.len; d++) words[count++] = detail_::u32(r.p[d]); + } + auto& c = context::get(); float* ring_host = nullptr; void* ring_tok = c.meta_ring_(&ring_host); - if (!ring_tok) return nullptr; - *word_off_out = c.meta_reserve_slot_(); - *host_words = reinterpret_cast(ring_host) + *word_off_out; - return ring_tok; -} - -// Upload a filled meta-ring slot and mark it live -- the other half of -// reserve_meta_ above, split from it so the caller can fill `host_words` (the -// pointer reserve_meta_ handed back) in between. -inline context::mirror* commit_meta_(context& c, void* ring_tok, - uint32_t word_off, const uint32_t* words, - size_t word_count) { + if (!ring_tok) return {}; + const uint32_t word_off = c.meta_reserve_slot_(); + std::memcpy(reinterpret_cast(ring_host) + word_off, words, + count * 4); + // This slot only, at its own byte offset, and unconditionally: the ring's + // residency never settles into one live copy (each slot is written once, + // read once, never again), so the host and device copies are both live + // after it rather than one or the other. context::mirror* mm = c.mirror_(ring_tok); - // This call's slot only, at its own byte offset — an unconditional - // WriteBuffer, not device_read_'s dirty-flag check, since the ring's - // mirror never legitimately settles into a single steady HOST/DEVICE state - // (each slot is written once, read once, never again). - c.queue.WriteBuffer(mm->dev, word_off * 4, words, word_count * 4); + c.queue.WriteBuffer(mm->dev, word_off * 4, words, count * 4); mm->live.uploaded(); - return mm; + return {ring_tok, static_cast(word_off) * 4}; } // Gather-style pad/fold (M11): unlike CUDA's scatter+atomicAdd, WGSL has no @@ -869,38 +755,19 @@ inline context::mirror* commit_meta_(context& c, void* ring_tok, // have it read (pad) or sum (fold) whatever cells of `a` map to it — no // output cell is ever written by two invocations, so unlike cuda.h's pad/fold // neither needs a pre-zeroed buffer. `a_shape`/`out_shape` (length `rank`, -// `rank-1` for fold's `out_shape`) ride the otherwise-unused second storage -// binding as bit-reinterpreted u32 — WGSL's fixed Params uniform (used by -// every other kernel here) has no room for a variable-length array, and -// WriteBuffer is a raw byte copy regardless of the binding's declared type. +// `rank-1` for fold's `out_shape`) ride the meta ring. inline bool own::pad(gpu::span a, gpu::span out, const int64_t* a_shape, const int64_t* out_shape, int rank, int axis, int64_t before, int64_t n, int64_t out_n) { + using detail_::u32; (void)n; - auto& c = context::get(); - if (!c.ready || rank <= 0 || rank > kPadFoldMaxRank) return false; - params p = {}; - if (!elem_off_(a.off, &p.a_off) || !elem_off_(out.off, &p.c_off)) return false; - context::mirror* ma = c.mirror_(a.buf); - context::mirror* mo = c.mirror_(out.buf); - if (!ma || !mo) return false; - c.device_read_(a.buf); - c.device_write_(out.buf); - - uint32_t word_off, *raw; - void* ring_tok = reserve_meta_(c, &word_off, &raw); - if (!ring_tok) return false; - for (int d = 0; d < rank; d++) raw[d] = static_cast(out_shape[d]); - for (int d = 0; d < rank; d++) raw[rank + d] = static_cast(a_shape[d]); - context::mirror* mm = - commit_meta_(c, ring_tok, word_off, raw, 2 * static_cast(rank)); - - p.M = static_cast(out_n); - p.b_off = word_off; - p.pad0 = static_cast(rank); - p.pad1 = static_cast(axis); - p.pad2 = static_cast(before); - return c.encode_("pad", ma, mm, mo, p, (out_n + 255) / 256, 1); + if (rank <= 0 || rank > kPadFoldMaxRank || out_n <= 0) return false; + gpu::span meta = meta_({{out_shape, rank}, {a_shape, rank}}); + if (!meta) return false; + return detail_::launch_own_( + "pad", -1, {gpu::in(a), gpu::in(meta), gpu::out(out)}, + detail_::pad_params{u32(out_n), u32(rank), u32(axis), u32(before)}, + gpu::policy::flat(out_n)); } // unfold's inverse. Each output element sums over the bounded range of @@ -912,217 +779,101 @@ inline bool own::pad(gpu::span a, gpu::span out, const int64_t* a_shape, inline bool own::fold(gpu::span a, gpu::span out, const int64_t* a_shape, const int64_t* out_shape, int rank, int axis, int64_t step, int64_t n, int64_t out_n) { + using detail_::u32; (void)n; - auto& c = context::get(); - if (!c.ready || rank <= 0 || rank > kPadFoldMaxRank) return false; - params p = {}; - if (!elem_off_(a.off, &p.a_off) || !elem_off_(out.off, &p.c_off)) return false; - context::mirror* ma = c.mirror_(a.buf); - context::mirror* mo = c.mirror_(out.buf); - if (!ma || !mo) return false; - c.device_read_(a.buf); - c.device_write_(out.buf); - - int out_rank = rank - 1; - uint32_t word_off, *raw; - void* ring_tok = reserve_meta_(c, &word_off, &raw); - if (!ring_tok) return false; - for (int d = 0; d < out_rank; d++) raw[d] = static_cast(out_shape[d]); - for (int d = 0; d < rank; d++) raw[out_rank + d] = static_cast(a_shape[d]); - context::mirror* mm = - commit_meta_(c, ring_tok, word_off, raw, - static_cast(out_rank) + static_cast(rank)); - - p.M = static_cast(out_n); - p.b_off = word_off; - p.pad0 = static_cast(rank); - p.pad1 = static_cast(axis); - p.pad2 = static_cast(step); - return c.encode_("fold", ma, mm, mo, p, (out_n + 255) / 256, 1); + if (rank <= 0 || rank > kPadFoldMaxRank || out_n <= 0) return false; + const int out_rank = rank - 1; + gpu::span meta = meta_({{out_shape, out_rank}, {a_shape, rank}}); + if (!meta) return false; + return detail_::launch_own_( + "fold", -1, {gpu::in(a), gpu::in(meta), gpu::out(out)}, + detail_::fold_params{u32(out_n), u32(rank), u32(axis), u32(step)}, + gpu::policy::flat(out_n)); } // index_select's dual, rewritten as a gather: WGSL has no float atomicAdd, // the same gap pad/fold above work around, so this sums over every source // row matching each OUTPUT row instead of scattering into a pre-zeroed -// buffer -- no zeroing needed. A = idx, B = values, C = out; p.pad0 = -// row_size, p.pad1 = k (number of source rows to scan). +// buffer -- no zeroing needed. inline bool own::index_add(gpu::span idx, gpu::span values, gpu::span out, int64_t row_size, int64_t k, int64_t out_n) { - auto& c = context::get(); - if (!c.ready || out_n <= 0) return false; - params p = {}; - if (!elem_off_(idx.off, &p.a_off) || !elem_off_(values.off, &p.b_off) || - !elem_off_(out.off, &p.c_off)) { - return false; - } - context::mirror *ma, *mb, *mo; - if (!operands_(c, idx.buf, values.buf, out.buf, &ma, &mb, &mo)) { - return false; - } - p.M = static_cast(out_n); - p.pad0 = static_cast(row_size); - p.pad1 = static_cast(k); - return c.encode_("index_add", ma, mb, mo, p, (out_n + 255) / 256, 1); + using detail_::u32; + if (out_n <= 0) return false; + return detail_::launch_own_( + "index_add", -1, {gpu::in(idx), gpu::in(values), gpu::out(out)}, + detail_::index_add_params{u32(out_n), u32(row_size), u32(k)}, + gpu::policy::flat(out_n)); } // One-hot scatter into a new trailing axis, as a gather: out[pos,k] = // values[pos] where indices[pos] == k, else 0. Every output element reads, -// never writes twice, so -- like index_select above -- no zeroing needed. -// A = idx, B = values, C = out; p.pad0 = size. +// never writes twice, so -- like index_select -- no zeroing needed. inline bool own::scatter_to_axis(gpu::span idx, gpu::span values, gpu::span out, int64_t n, int64_t size) { - auto& c = context::get(); - int64_t out_n = n * size; - if (!c.ready || out_n <= 0) return false; - params p = {}; - if (!elem_off_(idx.off, &p.a_off) || !elem_off_(values.off, &p.b_off) || - !elem_off_(out.off, &p.c_off)) { - return false; - } - context::mirror *ma, *mb, *mo; - if (!operands_(c, idx.buf, values.buf, out.buf, &ma, &mb, &mo)) { - return false; - } - p.M = static_cast(out_n); - p.pad0 = static_cast(size); - return c.encode_("scatter_axis", ma, mb, mo, p, (out_n + 255) / 256, 1); + using detail_::u32; + const int64_t out_n = n * size; + if (out_n <= 0) return false; + return detail_::launch_own_( + "scatter_axis", -1, {gpu::in(idx), gpu::in(values), gpu::out(out)}, + detail_::scatter_axis_params{u32(out_n), u32(size)}, + gpu::policy::flat(out_n)); } -// N-D broadcast binary: generalizes binary_bcast() above to any rank (a +// N-D broadcast binary: generalizes binary_bcast() to any rank (a // Transformer's [N,S,D] LayerNorm broadcasting a [N,S,1] mean, rank 3). // a_strides/b_strides are the broadcast strides (0 on a broadcast axis) -// array.h computes host-side via the same broadcast_strides() the CPU -// oracle uses. A = a, B = b, D = meta [out_shape(rank), a_strides(rank), -// b_strides(rank)] (a and b already fill A/B, unlike pad/fold where B was -// free for this); p.pad0 = rank, p.pad3 = meta's word offset into D. +// array.h computes host-side via the same broadcast_strides() the CPU oracle +// uses. `op` is the rank-2 broadcast id (badd..bpow), whose OP this kernel +// shares. inline bool own::binary_bcast_nd(kop op, gpu::span a, const int64_t* a_strides, gpu::span b, const int64_t* b_strides, gpu::span out, const int64_t* out_shape, int rank, int64_t n, float scale, float offset) { - auto& c = context::get(); - if (!c.ready || rank <= 0 || rank > kPadFoldMaxRank || n <= 0) return false; - params p = {}; - if (!elem_off_(a.off, &p.a_off) || !elem_off_(b.off, &p.b_off) || - !elem_off_(out.off, &p.c_off)) { - return false; - } - context::mirror *ma, *mb, *mo; - if (!operands_(c, a.buf, b.buf, out.buf, &ma, &mb, &mo)) { - return false; - } - - uint32_t word_off, *raw; - void* ring_tok = reserve_meta_(c, &word_off, &raw); - if (!ring_tok) return false; - for (int d = 0; d < rank; d++) raw[d] = static_cast(out_shape[d]); - for (int d = 0; d < rank; d++) { - raw[rank + d] = static_cast(a_strides[d]); - } - for (int d = 0; d < rank; d++) { - raw[2 * rank + d] = static_cast(b_strides[d]); - } - context::mirror* mm = - commit_meta_(c, ring_tok, word_off, raw, 3 * static_cast(rank)); - - p.M = static_cast(n); - p.op = kernel_op_(op); - p.pad0 = static_cast(rank); - p.pad3 = word_off; - p.scale = scale; - p.offset = offset; - return c.encode_("ew_bcast_nd", ma, mb, mo, p, (n + 255) / 256, 1, mm); + using detail_::u32; + if (op < kop::badd || op > kop::bpow) return false; + if (rank <= 0 || rank > kPadFoldMaxRank || n <= 0) return false; + gpu::span meta = + meta_({{out_shape, rank}, {a_strides, rank}, {b_strides, rank}}); + if (!meta) return false; + return detail_::launch_own_( + "ew_bcast_nd", kernel_(op).op, + {gpu::in(a), gpu::in(b), gpu::in(meta), gpu::out(out)}, + detail_::bcast_nd_params{u32(n), u32(rank), scale, offset}, + gpu::policy::flat(n)); } -// N-D broadcast ternary select: Tensor.where's GPU dispatch. A = cond, B = a, -// D = b, E = meta [out_shape(rank), cond_strides(rank), a_strides(rank), -// b_strides(rank)] (cond/a/b fill A/B/D, leaving E free for meta); p.pad0 = -// rank, p.pad3 = b's element offset into D, p.pad4 = meta's word offset -// into E. Bypasses operands_() (built for two real operands) since this one -// needs three, the same way pad()/fold() above do their own mirror lookups. +// N-D broadcast ternary select: Tensor.where's GPU dispatch, each operand +// through its own broadcast strides. inline bool own::where_nd(gpu::span cond, const int64_t* c_strides, gpu::span a, const int64_t* a_strides, gpu::span b, const int64_t* b_strides, gpu::span out, const int64_t* out_shape, int rank, int64_t n) { - auto& c = context::get(); - if (!c.ready || rank <= 0 || rank > kPadFoldMaxRank || n <= 0) return false; - params p = {}; - uint32_t b_elem_off; - if (!elem_off_(cond.off, &p.a_off) || !elem_off_(a.off, &p.b_off) || - !elem_off_(b.off, &b_elem_off) || !elem_off_(out.off, &p.c_off)) { - return false; - } - context::mirror* mcond = c.mirror_(cond.buf); - context::mirror* ma = c.mirror_(a.buf); - context::mirror* mb = c.mirror_(b.buf); - context::mirror* mo = c.mirror_(out.buf); - if (!mcond || !ma || !mb || !mo) return false; - c.device_read_(cond.buf); - c.device_read_(a.buf); - c.device_read_(b.buf); - c.device_write_(out.buf); - - uint32_t word_off, *raw; - void* ring_tok = reserve_meta_(c, &word_off, &raw); - if (!ring_tok) return false; - for (int d = 0; d < rank; d++) raw[d] = static_cast(out_shape[d]); - for (int d = 0; d < rank; d++) { - raw[rank + d] = static_cast(c_strides[d]); - } - for (int d = 0; d < rank; d++) { - raw[2 * rank + d] = static_cast(a_strides[d]); - } - for (int d = 0; d < rank; d++) { - raw[3 * rank + d] = static_cast(b_strides[d]); - } - context::mirror* mm = - commit_meta_(c, ring_tok, word_off, raw, 4 * static_cast(rank)); - - p.M = static_cast(n); - p.pad0 = static_cast(rank); - p.pad3 = b_elem_off; - p.pad4 = word_off; - return c.encode_("where_nd", mcond, ma, mo, p, (n + 255) / 256, 1, mb, mm); + using detail_::u32; + if (rank <= 0 || rank > kPadFoldMaxRank || n <= 0) return false; + gpu::span meta = meta_({{out_shape, rank}, {c_strides, rank}, + {a_strides, rank}, {b_strides, rank}}); + if (!meta) return false; + return detail_::launch_own_( + "where_nd", -1, + {gpu::in(cond), gpu::in(a), gpu::in(b), gpu::in(meta), gpu::out(out)}, + detail_::where_nd_params{u32(n), u32(rank)}, gpu::policy::flat(n)); } // sum_to (un-broadcast a gradient): gather, mirrors cuda.h's tl_sum_to and // metal.h's own sum_to -- one invocation per OUTPUT element sums every `a` -// element that broadcasts onto it, so no atomics (unlike index_add). Only -// one real tensor operand (`a`), so -- like pad/fold above -- B is free for -// the meta ring: [a_shape(rank), a_strides(rank), acc(rank)]. +// element that broadcasts onto it, so no atomics (unlike index_add). inline bool own::sum_to(gpu::span a, const int64_t* a_shape, const int64_t* a_strides, const int64_t* acc, int rank, int64_t out_n, int64_t reduced_n, gpu::span out) { - auto& c = context::get(); - if (!c.ready || rank <= 0 || rank > kPadFoldMaxRank) return false; - params p = {}; - if (!elem_off_(a.off, &p.a_off) || !elem_off_(out.off, &p.c_off)) return false; - context::mirror* ma = c.mirror_(a.buf); - context::mirror* mo = c.mirror_(out.buf); - if (!ma || !mo) return false; - c.device_read_(a.buf); - c.device_write_(out.buf); - - uint32_t word_off, *raw; - void* ring_tok = reserve_meta_(c, &word_off, &raw); - if (!ring_tok) return false; - for (int d = 0; d < rank; d++) { - raw[d] = static_cast(a_shape[d]); - } - for (int d = 0; d < rank; d++) { - raw[rank + d] = static_cast(a_strides[d]); - } - for (int d = 0; d < rank; d++) { - raw[2 * rank + d] = static_cast(acc[d]); - } - context::mirror* mm = - commit_meta_(c, ring_tok, word_off, raw, 3 * static_cast(rank)); - - p.M = static_cast(out_n); - p.b_off = word_off; - p.pad0 = static_cast(rank); - p.pad1 = static_cast(reduced_n); - return c.encode_("sum_to", ma, mm, mo, p, (out_n + 255) / 256, 1); + using detail_::u32; + if (rank <= 0 || rank > kPadFoldMaxRank || out_n <= 0) return false; + gpu::span meta = meta_({{a_shape, rank}, {a_strides, rank}, {acc, rank}}); + if (!meta) return false; + return detail_::launch_own_( + "sum_to", -1, {gpu::in(a), gpu::in(meta), gpu::out(out)}, + detail_::sum_to_params{u32(out_n), u32(rank), u32(reduced_n)}, + gpu::policy::flat(out_n)); } // concat_part (Tensor.concat along an arbitrary axis, KV-cache append): @@ -1130,79 +881,40 @@ inline bool own::sum_to(gpu::span a, const int64_t* a_shape, // axis -- see kernels/tensorlib_webgpu.wgsl's own concat_part for why this // is its own entry rather than reusing pad's (that one dispatches over // OUTPUT elements for its zero border; concat has no border and wants the -// much smaller SOURCE element count instead). Only one real tensor operand -// (`a`), so -- like pad/fold/sum_to above -- B is free for the meta ring: -// [a_shape(rank), out_strides(rank)] (out_strides computed host-side, -// mirrors cuda.h's own upload_pad_fold_meta_). +// much smaller SOURCE element count instead). out_strides are computed +// host-side, as cuda.h's upload_pad_fold_meta_ does. inline bool own::concat_part(gpu::span a, gpu::span out, const int64_t* a_shape, const int64_t* out_shape, int rank, int axis, int64_t before, int64_t n) { - auto& c = context::get(); - if (!c.ready || rank <= 0 || rank > kPadFoldMaxRank || n <= 0) { - return false; - } - params p = {}; - if (!elem_off_(a.off, &p.a_off) || !elem_off_(out.off, &p.c_off)) return false; - context::mirror* ma = c.mirror_(a.buf); - context::mirror* mo = c.mirror_(out.buf); - if (!ma || !mo) return false; - c.device_read_(a.buf); - c.device_write_(out.buf); - + using detail_::u32; + if (rank <= 0 || rank > kPadFoldMaxRank || n <= 0) return false; int64_t out_strides[kPadFoldMaxRank]; - int64_t acc = 1; - for (int d = rank - 1; d >= 0; d--) { - out_strides[d] = acc; - acc *= out_shape[d]; - } - - uint32_t word_off, *raw; - void* ring_tok = reserve_meta_(c, &word_off, &raw); - if (!ring_tok) return false; - for (int d = 0; d < rank; d++) raw[d] = static_cast(a_shape[d]); - for (int d = 0; d < rank; d++) { - raw[rank + d] = static_cast(out_strides[d]); - } - context::mirror* mm = - commit_meta_(c, ring_tok, word_off, raw, 2 * static_cast(rank)); - - p.M = static_cast(n); - p.b_off = word_off; - p.pad0 = static_cast(rank); - p.pad1 = static_cast(before * out_strides[axis]); - return c.encode_("concat_part", ma, mm, mo, p, (n + 255) / 256, 1); + contiguous_strides_into(out_shape, rank, out_strides); + gpu::span meta = meta_({{a_shape, rank}, {out_strides, rank}}); + if (!meta) return false; + return detail_::launch_own_( + "concat_part", -1, {gpu::in(a), gpu::in(meta), gpu::out(out)}, + detail_::concat_params{u32(n), u32(rank), u32(before * out_strides[axis])}, + gpu::policy::flat(n)); } // RoPE (rotary position embedding), half-split (GPT-NeoX / HF-llama) // convention -- mirrors cuda.h's/metal.h's own rope. x is [rows, D] -// contiguous (rows = H*T); dispatched flat over rows*(D/2), matching -// this file's other flat-1D kernels. No second real tensor operand, so -// -- like unary()/clamp() above -- B binds x a second time (unused by -// the WGSL). x/out always view at offset 0 (array.h's gpu_rope_ requires -// x.offset_ == 0 and hands a fresh allocation for out), so there is no -// ao/oo in this signature to convert. +// contiguous (rows = H*T); dispatched flat over rows*(D/2), one invocation a +// rotated pair. inline bool own::rope(gpu::span x, gpu::span out, int64_t rows, int64_t T, int64_t D, int64_t pos, float base, gpu::span bias) { - auto& c = context::get(); - if (!c.ready || D <= 0 || (D & 1) || bias.buf) return false; // no fused bias - int64_t half = D / 2; - int64_t n = rows * half; + using detail_::u32; + if (D <= 0 || (D & 1) || bias.buf) return false; // no fused bias + const int64_t half = D / 2; + const int64_t n = rows * half; if (n <= 0) return false; - params p = {}; - if (!elem_off_(x.off, &p.a_off) || !elem_off_(out.off, &p.c_off)) return false; - context::mirror *ma, *mb, *mo; - if (!operands_(c, x.buf, nullptr, out.buf, &ma, &mb, &mo)) return false; - p.b_off = p.a_off; - p.M = static_cast(n); - p.N = static_cast(T); - p.K = static_cast(D); - p.pad0 = static_cast(pos); - p.pad1 = static_cast(half); - p.scale = base; - return c.encode_("rope", ma, mb, mo, p, (n + 255) / 256, 1); + return detail_::launch_own_( + "rope", -1, {gpu::in(x), gpu::out(out)}, + detail_::rope_params{u32(n), u32(T), u32(D), u32(pos), u32(half), base}, + gpu::policy::flat(n)); } - // What a model may ask of this backend beyond the kernel contract (gpu.h), // and the graph-capture group it names: none of it here, so each answers // false or does nothing and a decoder takes its host-position path. diff --git a/kernels/tensorlib_webgpu.wgsl b/kernels/tensorlib_webgpu.wgsl index 0a9f409..1efcf14 100644 --- a/kernels/tensorlib_webgpu.wgsl +++ b/kernels/tensorlib_webgpu.wgsl @@ -6,70 +6,67 @@ // committed because the wasm build is a flat emcc line that never runs CMake // (the CUDA backend's PTX goes through bin2c for the same reason). // -// View offsets arrive as ELEMENT offsets in the params block and are folded -// into the indexing here, rather than as bind-group binding offsets: WebGPU -// requires those to be 256-byte aligned, which an arbitrary view offset is -// not. So every binding covers its whole buffer. (CUDA instead folds offsets -// host-side into the pointer it passes, which WebGPU has no equivalent of.) - -// One Params struct and one bind group layout serve every kernel here, rather -// than the per-family structs metal_kernels.metal uses. WebGPU's bind group -// ceremony is heavy enough that a second layout would buy nothing: the fields -// each family ignores cost 4 bytes of a 256-byte uniform slot. Kernels that -// take one input (unary, the row reductions) get A bound to B as well — two -// read-only bindings may alias, and the output is always a fresh allocation, -// so no writable binding ever aliases a readable one. -struct Params { - M : u32, // gemm rows | elementwise element count | reduce rows - N : u32, // gemm cols | reduce cols - K : u32, - lda : u32, - ldb : u32, - ldc : u32, - a_off : u32, - b_off : u32, - c_off : u32, - ta : u32, - tb : u32, - ars : u32, // broadcast: per-operand row/col strides, in elements - acs : u32, - brs : u32, - bcs : u32, - op : u32, // which operation, within the entry point's family - scale : f32, - offset : f32, - // A uniform-address-space struct has align 16, so its size rounds up to a - // multiple of 16. Pad explicitly to 96 bytes so the host struct (which the - // bind group's minBindingSize comes from) matches exactly — a short - // minBindingSize fails bind group validation for every dispatch. - _pad0 : u32, - _pad1 : u32, - _pad2 : u32, - _pad3 : u32, - _pad4 : u32, - arg : f32, // a scalar operand (ew_scalar's s) +// The kernel ABI (gpu_abi.h), as WGSL states it. A kernel binds view i at +// binding i, in the order gpu_ops.h lists its views, and its uniform at the +// binding after the last view. The uniform is the views' element offsets +// (`Views`, by view index) followed by the kernel's params, the same fields +// in the same order as its gpu_abi.h struct. webgpu.h fills both the same way +// for every kernel. +// +// View offsets ride the uniform rather than the bind group: a binding offset +// has to be 256-byte aligned, which an arbitrary view offset is not. So every +// binding covers its whole buffer, and each kernel folds its views' offsets +// into its indexing. (Metal binds at the offset and CUDA adds it to the +// pointer; WebGPU has neither.) +// +// A binding's type follows its view: read-only storage for an input, +// read_write for an output. Binding i is an input in one kernel and an output +// in another, so the storage variables are declared once per binding and +// access, in0.. and out1.., and each entry point uses the ones its views +// need. WGSL checks binding clashes per entry point, and each pipeline takes +// its layout from its entry point (webgpu.h builds them with an auto layout), +// so what a kernel declares here is its layout. +// +// A family's operation — add or sub, exp or log — is not a params field but +// the pipeline-overridable constant OP: webgpu.h builds a pipeline per kernel +// id with OP set from its kernel table, the counterpart of the kernel name +// Metal and CUDA look an id up by. Within a pipeline OP is a constant, so the +// switch on it takes the same arm in every invocation. + +override OP : u32 = 0u; + +// Each view's element offset into its binding, by view index. +struct Views { + v0 : u32, v1 : u32, v2 : u32, v3 : u32, + v4 : u32, v5 : u32, v6 : u32, v7 : u32, }; -@group(0) @binding(0) var A : array; -@group(0) @binding(1) var B : array; -@group(0) @binding(2) var C : array; -@group(0) @binding(3) var p : Params; -// A third and fourth read-only operand, for kernels A/B alone can't cover: -// binary_bcast_nd's two real tensor operands (a, b) already fill A and B, so -// its shape/stride meta needs D; where_nd's three real operands (cond, a, b) -// fill A, B and D, so its meta needs E. Every other entry point ignores -// these (webgpu.h's encode_ binds them to A when a kernel has no use for -// them — bind group validation requires every declared binding be present -// regardless of which bindings the active entry point actually reads). -@group(0) @binding(4) var D : array; -@group(0) @binding(5) var E : array; - -// ---- sgemm: C(M,N) = op(A)(M,K) @ op(B)(K,N) * scale + offset +@group(0) @binding(0) var in0 : array; +@group(0) @binding(1) var in1 : array; +@group(0) @binding(2) var in2 : array; +@group(0) @binding(3) var in3 : array; +@group(0) @binding(1) var out1 : array; +@group(0) @binding(2) var out2 : array; +@group(0) @binding(3) var out3 : array; +@group(0) @binding(4) var out4 : array; + +// ---- sgemm: out(M,N) = op(a)(M,K) @ op(b)(K,N) * scale + offset // // 64x64 workgroup tile, 16x16 = 256 invocations, each holding a 4x4 register // accumulator. Mirrors the shape of the Metal sgemm_64_ kernel; MMA intrinsics // have no WGSL equivalent, so the inner product is plain FMA over registers. // Measured at ~580-630 GF/s for n=1024 on an M1 Pro (see spike/webgpu). +// +// Views: a, b -> out. Params: webgpu.h's gemm_params. + +struct GemmArgs { + v : Views, + m : u32, n : u32, k : u32, + lda : u32, ldb : u32, ldc : u32, + ta : u32, tb : u32, + scale : f32, offset : f32, +}; +@group(0) @binding(3) var gm : GemmArgs; const BM : u32 = 64u; const BN : u32 = 64u; @@ -84,15 +81,15 @@ var Bs : array; // BK * BN // Row/column strides for a possibly-transposed operand: transposing swaps // which axis walks by the leading dimension (cf. metal_kernels.metal:119-120). fn a_index(m : u32, k : u32) -> u32 { - let rs = select(p.lda, 1u, p.ta == 1u); - let cs = select(1u, p.lda, p.ta == 1u); - return p.a_off + m * rs + k * cs; + let rs = select(gm.lda, 1u, gm.ta == 1u); + let cs = select(1u, gm.lda, gm.ta == 1u); + return gm.v.v0 + m * rs + k * cs; } fn b_index(k : u32, n : u32) -> u32 { - let rs = select(p.ldb, 1u, p.tb == 1u); - let cs = select(1u, p.ldb, p.tb == 1u); - return p.b_off + k * rs + n * cs; + let rs = select(gm.ldb, 1u, gm.tb == 1u); + let cs = select(1u, gm.ldb, gm.tb == 1u); + return gm.v.v1 + k * rs + n * cs; } @compute @workgroup_size(16, 16, 1) @@ -104,7 +101,7 @@ fn sgemm(@builtin(workgroup_id) wg : vec3, var acc : array; // TM * TN, zero-initialized - let n_tiles = (p.K + BK - 1u) / BK; + let n_tiles = (gm.k + BK - 1u) / BK; for (var kt : u32 = 0u; kt < n_tiles; kt = kt + 1u) { let k0 = kt * BK; @@ -116,13 +113,13 @@ fn sgemm(@builtin(workgroup_id) wg : vec3, let am = m_base + i / BK; let ak = k0 + i % BK; - let a_ok = am < p.M && ak < p.K; - As[i] = select(0.0, A[a_index(am, ak)], a_ok); + let a_ok = am < gm.m && ak < gm.k; + As[i] = select(0.0, in0[a_index(am, ak)], a_ok); let bk = k0 + i / BN; let bn = n_base + i % BN; - let b_ok = bk < p.K && bn < p.N; - Bs[i] = select(0.0, B[b_index(bk, bn)], b_ok); + let b_ok = bk < gm.k && bn < gm.n; + Bs[i] = select(0.0, in1[b_index(bk, bn)], b_ok); } workgroupBarrier(); @@ -148,11 +145,11 @@ fn sgemm(@builtin(workgroup_id) wg : vec3, for (var i : u32 = 0u; i < TM; i = i + 1u) { let m = m_base + lid.y * TM + i; - if (m >= p.M) { continue; } + if (m >= gm.m) { continue; } for (var j : u32 = 0u; j < TN; j = j + 1u) { let n = n_base + lid.x * TN + j; - if (n >= p.N) { continue; } - C[p.c_off + m * p.ldc + n] = acc[i * TN + j] * p.scale + p.offset; + if (n >= gm.n) { continue; } + out2[gm.v.v2 + m * gm.ldc + n] = acc[i * TN + j] * gm.scale + gm.offset; } } } @@ -162,12 +159,10 @@ fn sgemm(@builtin(workgroup_id) wg : vec3, // WGSL has neither templates nor a preprocessor, so the per-op variants that // metal_kernels.metal generates from a macro would have to be copy-pasted // here — exactly the edge-tile bug class that file's header warns against. -// Instead the operation is a uniform field and each family is ONE entry point -// that switches on it. The branch is uniform across the dispatch and these -// kernels are memory-bound, so it costs nothing measurable; what it buys is a -// single copy of every bounds check and epilogue. +// Instead each family is ONE entry point that switches on OP. What that buys +// is a single copy of every bounds check and epilogue. // -// Op codes are assigned by kernel_op_() in webgpu.h. +// OP's values per family are assigned by kernel_() in webgpu.h. const OP_ADD : u32 = 0u; const OP_SUB : u32 = 1u; const OP_MUL : u32 = 2u; @@ -200,7 +195,7 @@ const OP_NE : u32 = 5u; // it. Threads whose row is shorter than the workgroup contribute this. const NEG_HUGE : f32 = -3.4e38; -// Shared by the contiguous and the broadcast binary: same five operations, +// Shared by the contiguous and the broadcast binaries: same five operations, // only the addressing differs. fn binary_op(op : u32, av : f32, bv : f32) -> f32 { switch (op) { @@ -226,10 +221,9 @@ fn unary_op(op : u32, v : f32) -> f32 { } } -// Elementwise comparison: out = (a OP b) ? 1.0 : 0.0. No scale/offset -- masks -// don't compose with the affine epilogue, so ew_cmp below skips it (unlike -// ew_binary/ew_unary). Own family/entry point, not folded into binary_op, -// since it returns a bool-as-float rather than composing with `bv`. +// Elementwise comparison: (a OP b) ? 1.0 : 0.0. No scale/offset -- masks +// don't compose with the affine epilogue. Not folded into binary_op, since it +// returns a bool-as-float rather than composing with `bv`. fn cmp_op(op : u32, av : f32, bv : f32) -> f32 { switch (op) { case 1u: { return select(0.0, 1.0, av < bv); } @@ -241,50 +235,57 @@ fn cmp_op(op : u32, av : f32, bv : f32) -> f32 { } } -// Contiguous elementwise binary over p.M elements. +// Views: a, b -> out (binary) or a -> out (unary). Params: ew_params. +struct EwArgs { v : Views, n : u32, scale : f32, offset : f32 }; +@group(0) @binding(3) var ewb : EwArgs; +@group(0) @binding(2) var ewu : EwArgs; + @compute @workgroup_size(256, 1, 1) fn ew_binary(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let v = binary_op(p.op, A[p.a_off + i], B[p.b_off + i]); - C[p.c_off + i] = fma(v, p.scale, p.offset); + if (i >= ewb.n) { return; } + let v = binary_op(OP, in0[ewb.v.v0 + i], in1[ewb.v.v1 + i]); + out2[ewb.v.v2 + i] = fma(v, ewb.scale, ewb.offset); } @compute @workgroup_size(256, 1, 1) fn ew_unary(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let v = unary_op(p.op, A[p.a_off + i]); - C[p.c_off + i] = fma(v, p.scale, p.offset); + if (i >= ewu.n) { return; } + let v = unary_op(OP, in0[ewu.v.v0 + i]); + out1[ewu.v.v1 + i] = fma(v, ewu.scale, ewu.offset); } -// Elementwise comparison over p.M elements: out = (a OP b) ? 1.0 : 0.0 -// (no epilogue -- masks don't compose with scale/offset). p.ars carries the -// bstride webgpu.h's compare() receives: 1 for a same-shape b, 0 for an -// explicit size-1 b (`x > s` itself is ew_scalar) -- an unused field for -// this family, repurposed rather than widening Params. +// out[i] = (a[i] OP b[i * bstride]) ? 1.0 : 0.0: bstride 1 for a same-shape +// b, 0 for a size-1 one (`x > s` itself is ew_scalar). +// Views: a, b -> out. Params: cmp_params. +struct CmpArgs { v : Views, n : u32, bstride : u32 }; +@group(0) @binding(3) var cm : CmpArgs; + @compute @workgroup_size(256, 1, 1) fn cmp(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let bv = B[p.b_off + i * p.ars]; - C[p.c_off + i] = cmp_op(p.op, A[p.a_off + i], bv); + if (i >= cm.n) { return; } + let bv = in1[cm.v.v1 + i * cm.bstride]; + out2[cm.v.v2 + i] = cmp_op(OP, in0[cm.v.v0 + i], bv); } -// clamp(x, lo, hi): Clip's forward. No affine epilogue -- p.scale/p.offset -// carry lo/hi instead (a dedicated entry, same as this family's clamp_ in -// metal_kernels.metal/tl_clamp in tensorlib_cuda.cu). Named clamp_, not -// clamp, so the entry point doesn't shadow WGSL's builtin of that name. +// clamp(x, lo, hi): Clip's forward, no epilogue. Named clamp_, not clamp, so +// the entry point doesn't shadow WGSL's builtin of that name. +// Views: a -> out. Params: clamp_params. +struct ClampArgs { v : Views, n : u32, lo : f32, hi : f32 }; +@group(0) @binding(2) var cl : ClampArgs; + @compute @workgroup_size(256, 1, 1) fn clamp_(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - C[p.c_off + i] = clamp(A[p.a_off + i], p.scale, p.offset); + if (i >= cl.n) { return; } + out1[cl.v.v1 + i] = clamp(in0[cl.v.v0 + i], cl.lo, cl.hi); } -// Tensor-scalar over p.M elements: out = f(a, s) * scale + offset, s in p.arg. -// Mirrors tensorlib_cuda.cu's TL_EW_SCALAR; op 0 is pow, 1.. the comparisons -// in cmp_op's order. +// Tensor-scalar: out = f(a, s) * scale + offset. Mirrors tensorlib_cuda.cu's +// TL_EW_SCALAR; OP 0 is pow, 1.. the comparisons in cmp_op's order. +// Views: a -> out. Params: scalar_params. fn scalar_op(op : u32, av : f32, s : f32) -> f32 { switch (op) { case 0u: { return pow(av, s); } @@ -292,31 +293,43 @@ fn scalar_op(op : u32, av : f32, s : f32) -> f32 { } } +struct ScalarArgs { v : Views, n : u32, s : f32, scale : f32, offset : f32 }; +@group(0) @binding(2) var sc : ScalarArgs; + @compute @workgroup_size(256, 1, 1) fn ew_scalar(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - C[p.c_off + i] = fma(scalar_op(p.op, A[p.a_off + i], p.arg), p.scale, p.offset); + if (i >= sc.n) { return; } + out1[sc.v.v1 + i] = fma(scalar_op(OP, in0[sc.v.v0 + i], sc.s), sc.scale, sc.offset); } // Rank-2 broadcast binary: out[r,c] = f(a[r*ars + c*acs], b[r*brs + c*bcs]) -// into a contiguous [M,N] output. Per-operand strides express every rank-2 +// into a contiguous [m,n] output. Per-operand strides express every rank-2 // broadcast (row vector, column vector, scalar) in one kernel, which keeps // bias/gamma/beta chains on the GPU — a CPU fallback mid-graph costs a full // submit-and-wait, and that is dearer here than on Metal. +// Views: a, b -> out. Params: bcast_params. +struct BcastArgs { + v : Views, + m : u32, n : u32, + ars : u32, acs : u32, brs : u32, bcs : u32, + scale : f32, offset : f32, +}; +@group(0) @binding(3) var bc : BcastArgs; + @compute @workgroup_size(32, 8, 1) fn ew_bcast(@builtin(global_invocation_id) gid : vec3) { let c = gid.x; let r = gid.y; - if (c >= p.N || r >= p.M) { return; } - let av = A[p.a_off + r * p.ars + c * p.acs]; - let bv = B[p.b_off + r * p.brs + c * p.bcs]; - C[p.c_off + r * p.N + c] = fma(binary_op(p.op, av, bv), p.scale, p.offset); + if (c >= bc.n || r >= bc.m) { return; } + let av = in0[bc.v.v0 + r * bc.ars + c * bc.acs]; + let bv = in1[bc.v.v1 + r * bc.brs + c * bc.bcs]; + out2[bc.v.v2 + r * bc.n + c] = fma(binary_op(OP, av, bv), bc.scale, bc.offset); } // ---- Row reductions over the last axis: one workgroup per row, 256 -// invocations, workgroup-scratch tree reduction. p.N (cols) may exceed the -// invocation count, so each thread strides over the row first. +// invocations, workgroup-scratch tree reduction. The row (cols) may exceed the +// invocation count, so each thread strides over it first. const T : u32 = 256u; var scratch : array; @@ -344,6 +357,11 @@ fn tree_reduce(op : u32, lid : u32, v : f32) -> f32 { return r; } +// Views: a -> out. Params: reduce_params. softmax writes rows x cols, +// row_reduce one value a row. +struct ReduceArgs { v : Views, rows : u32, cols : u32, scale : f32, offset : f32 }; +@group(0) @binding(2) var rd : ReduceArgs; + // Numerically stable softmax (subtract the row max). Applying an affine // epilogue to a softmax is not meaningful, so scale/offset are ignored here, // as they are on Metal. @@ -351,133 +369,142 @@ fn tree_reduce(op : u32, lid : u32, v : f32) -> f32 { fn softmax(@builtin(workgroup_id) wg : vec3, @builtin(local_invocation_index) lid : u32) { let row = wg.x; - let src = p.a_off + row * p.N; - let dst = p.c_off + row * p.N; + let src = rd.v.v0 + row * rd.cols; + let dst = rd.v.v1 + row * rd.cols; var m = NEG_HUGE; - for (var c : u32 = lid; c < p.N; c = c + T) { m = max(m, A[src + c]); } + for (var c : u32 = lid; c < rd.cols; c = c + T) { m = max(m, in0[src + c]); } let row_max = tree_reduce(OP_ROW_MAX, lid, m); var sum = 0.0; - for (var c : u32 = lid; c < p.N; c = c + T) { sum = sum + exp(A[src + c] - row_max); } + for (var c : u32 = lid; c < rd.cols; c = c + T) { sum = sum + exp(in0[src + c] - row_max); } let inv = 1.0 / tree_reduce(OP_ROW_SUM, lid, sum); - for (var c : u32 = lid; c < p.N; c = c + T) { - C[dst + c] = exp(A[src + c] - row_max) * inv; + for (var c : u32 = lid; c < rd.cols; c = c + T) { + out1[dst + c] = exp(in0[src + c] - row_max) * inv; + } +} + +// One value per row, with the affine epilogue. +@compute @workgroup_size(256, 1, 1) +fn row_reduce(@builtin(workgroup_id) wg : vec3, + @builtin(local_invocation_index) lid : u32) { + let row = wg.x; + let src = rd.v.v0 + row * rd.cols; + + var acc = select(0.0, NEG_HUGE, OP == OP_ROW_MAX); + for (var c : u32 = lid; c < rd.cols; c = c + T) { + acc = reduce_op(OP, acc, in0[src + c]); } + let r = tree_reduce(OP, lid, acc); + if (lid == 0u) { out1[rd.v.v1 + row] = r * rd.scale + rd.offset; } } // Layer norm over the last axis with the affine epilogue: (x - mu) / -// sqrt(var + eps) * g + b. A = x, B = g, D = b (p._pad3 its element offset), -// p.arg = eps. Two tree sums (x, then the squared deviations), each scaled by -// 1/N after the tree as row_reduce's mean is. +// sqrt(var + eps) * g + b. Two tree sums (x, then the squared deviations), +// each scaled by 1/cols after the tree as row_reduce's mean is. +// Views: x, g, b -> out. Params: layer_norm_params. +struct LayerNormArgs { + v : Views, + rows : u32, cols : u32, + eps : f32, scale : f32, offset : f32, +}; +@group(0) @binding(4) var ln : LayerNormArgs; + @compute @workgroup_size(256, 1, 1) fn layer_norm(@builtin(workgroup_id) wg : vec3, @builtin(local_invocation_index) lid : u32) { let row = wg.x; - let src = p.a_off + row * p.N; - let dst = p.c_off + row * p.N; - let inv_n = 1.0 / f32(p.N); + let src = ln.v.v0 + row * ln.cols; + let dst = ln.v.v3 + row * ln.cols; + let inv_n = 1.0 / f32(ln.cols); var sum = 0.0; - for (var c : u32 = lid; c < p.N; c = c + T) { sum = sum + A[src + c]; } + for (var c : u32 = lid; c < ln.cols; c = c + T) { sum = sum + in0[src + c]; } let mu = tree_reduce(OP_ROW_SUM, lid, sum) * inv_n; var ss = 0.0; - for (var c : u32 = lid; c < p.N; c = c + T) { - let v = A[src + c] - mu; + for (var c : u32 = lid; c < ln.cols; c = c + T) { + let v = in0[src + c] - mu; ss = ss + v * v; } - let inv = 1.0 / sqrt(tree_reduce(OP_ROW_SUM, lid, ss) * inv_n + p.arg); + let inv = 1.0 / sqrt(tree_reduce(OP_ROW_SUM, lid, ss) * inv_n + ln.eps); - for (var c : u32 = lid; c < p.N; c = c + T) { - let y = (A[src + c] - mu) * inv * B[p.b_off + c] + D[p._pad3 + c]; - C[dst + c] = y * p.scale + p.offset; + for (var c : u32 = lid; c < ln.cols; c = c + T) { + let y = (in0[src + c] - mu) * inv * in1[ln.v.v1 + c] + in2[ln.v.v2 + c]; + out3[dst + c] = y * ln.scale + ln.offset; } } -// One value per row, with the affine epilogue. -@compute @workgroup_size(256, 1, 1) -fn row_reduce(@builtin(workgroup_id) wg : vec3, - @builtin(local_invocation_index) lid : u32) { - let row = wg.x; - let src = p.a_off + row * p.N; - - var acc = select(0.0, NEG_HUGE, p.op == OP_ROW_MAX); - for (var c : u32 = lid; c < p.N; c = c + T) { - acc = reduce_op(p.op, acc, A[src + c]); - } - let r = tree_reduce(p.op, lid, acc); - if (lid == 0u) { C[p.c_off + row] = r * p.scale + p.offset; } -} - -// ---- im2col's pad/fold (M11) -// -// Both dispatch one invocation per OUTPUT element and gather from A, rather -// than CUDA's scatter+atomicAdd: WGSL has no float atomicAdd. A gather needs -// no pre-zeroed output (every C cell is written exactly once, by exactly one -// invocation) and no atomics — the tradeoff is fold's small bounded loop over -// the window indices that could cover a given output cell, in place of one -// atomicAdd per source element. -// -// Neither family's shape metadata fits the fixed Params uniform (a -// variable-length array has no home there), so it rides B as -// bit-reinterpreted u32 — B's declared type is array, but WriteBuffer on -// the host side is a raw byte copy regardless, and bitcast reads it back -// correctly. p._pad0/_pad1/_pad2 (otherwise-unused Params padding) carry -// rank/axis/before-or-step; p.M is the output element count (webgpu.h's -// pad()/fold() dispatch over out_n, not a's own size). +// ---- Shape metadata. The N-D kernels below take a variable-length shape +// or stride list, which has no home in a fixed params struct, so it rides a +// view of its own (webgpu.h's meta ring), one slot per call: an input like +// any other, whose f32 words are u32 bit patterns — WriteBuffer is a raw +// byte copy, and bitcast reads them back. // // Rank cap — matches webgpu.h's own kPadFoldMaxRank (a shader can't see a // host-side C++ constant, so this is its own copy) and bounds every // fixed-size local array below. const kPadFoldMaxRank : u32 = 8u; -// Shared by both: row-major decode of a dispatch-global thread id `i` against -// a shape held in B at word offset `base`, into a fixed-size local array. -// `rank8` bounds the loop so callers can pass either a full-rank or a -// (rank-1)-length shape. +// Word d of the meta view bound at in1. +fn meta1(base : u32, d : u32) -> u32 { return bitcast(in1[base + d]); } + +// Row-major decode of a dispatch-global thread id `i` against a shape held in +// in1 from `base`, into a fixed-size local array. `rank8` bounds the loop so +// callers can pass either a full-rank or a (rank-1)-length shape. fn decode_idx(i : u32, base : u32, rank8 : u32, - out_idx : ptr>) { + out_idx : ptr>) { var rem = i; for (var d : i32 = i32(rank8) - 1; d >= 0; d = d - 1) { - let dim = bitcast(B[base + u32(d)]); + let dim = meta1(base, u32(d)); (*out_idx)[u32(d)] = rem % dim; rem = rem / dim; } } -// Row-major strides of a contiguous tensor whose shape is held in B at word -// offset `base`, length `rank` — used to address `A`, which pad_/fold_'s GPU +// Row-major strides of a contiguous tensor whose shape is held in in1 from +// `base`, length `rank` — used to address the source, which pad_/fold_'s GPU // dispatch (array.h's gpu_pad_/gpu_fold_) requires to be contiguous. fn a_strides_from_shape(base : u32, rank : u32, a_shape : ptr>, out_strides : ptr>) { var acc : u32 = 1u; for (var d : i32 = i32(rank) - 1; d >= 0; d = d - 1) { - let dim = bitcast(B[base + u32(d)]); + let dim = meta1(base, u32(d)); (*a_shape)[u32(d)] = dim; (*out_strides)[u32(d)] = acc; acc = acc * dim; } } -// B layout at word offset p.b_off: [out_shape(rank), a_shape(rank)] — webgpu.h -// reserves a fresh ring slot per call (see meta_reserve_slot_), so two pad/ -// fold calls batched into the same unflushed pass never share one offset. +// ---- im2col's pad/fold (M11) +// +// Both dispatch one invocation per OUTPUT element and gather from the source, +// rather than CUDA's scatter+atomicAdd: WGSL has no float atomicAdd. A gather +// needs no pre-zeroed output (every cell is written exactly once, by exactly +// one invocation) and no atomics — the tradeoff is fold's small bounded loop +// over the window indices that could cover a given output cell, in place of +// one atomicAdd per source element. + +// Views: a, meta [out_shape(rank), a_shape(rank)] -> out. +// Params: webgpu.h's pad_params (n = the output's element count). +struct PadArgs { v : Views, n : u32, rank : u32, axis : u32, before : u32 }; +@group(0) @binding(3) var pd : PadArgs; + @compute @workgroup_size(256, 1, 1) fn pad(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let rank = p._pad0; - let axis = p._pad1; - let before = i32(p._pad2); + if (i >= pd.n) { return; } + let rank = pd.rank; + let axis = pd.axis; + let before = i32(pd.before); var out_idx : array; - decode_idx(i, p.b_off, rank, &out_idx); + decode_idx(i, pd.v.v1, rank, &out_idx); var a_shape : array; var a_strides : array; - a_strides_from_shape(p.b_off + rank, rank, &a_shape, &a_strides); + a_strides_from_shape(pd.v.v1 + rank, rank, &a_shape, &a_strides); var src : u32 = 0u; var in_bounds = true; @@ -487,32 +514,36 @@ fn pad(@builtin(global_invocation_id) gid : vec3) { c = c - before; if (c < 0 || c >= i32(a_shape[d])) { in_bounds = false; } } - // Clamped even out of range: A[src] is still read below (WGSL's select - // evaluates both operands), so src must stay in bounds regardless of - // in_bounds — only the select, not the address, decides the result. + // Clamped even out of range: the source is still read below (WGSL's + // select evaluates both operands), so src must stay in bounds regardless + // of in_bounds — only the select, not the address, decides the result. let cc = clamp(c, 0, i32(a_shape[d]) - 1); src = src + u32(cc) * a_strides[d]; } - C[p.c_off + i] = select(0.0, A[p.a_off + src], in_bounds); + out2[pd.v.v2 + i] = select(0.0, in0[pd.v.v0 + src], in_bounds); } -// B layout at word offset p.b_off: [out_shape(rank-1), a_shape(rank)] — same -// per-call ring slot as pad above. a's last dim is the sliding window (size -// a_shape[rank-1]); a's `axis` dim (size a_shape[axis]) is the window count. +// Views: a, meta [out_shape(rank-1), a_shape(rank)] -> out. a's last dim is +// the sliding window (size a_shape[rank-1]); a's `axis` dim (size +// a_shape[axis]) is the window count. +// Params: webgpu.h's fold_params (n = the output's element count). +struct FoldArgs { v : Views, n : u32, rank : u32, axis : u32, step : u32 }; +@group(0) @binding(3) var fd : FoldArgs; + @compute @workgroup_size(256, 1, 1) fn fold(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let rank = p._pad0; - let axis = p._pad1; - let step = i32(p._pad2); + if (i >= fd.n) { return; } + let rank = fd.rank; + let axis = fd.axis; + let step = i32(fd.step); let out_rank = rank - 1u; var out_idx : array; - decode_idx(i, p.b_off, out_rank, &out_idx); + decode_idx(i, fd.v.v1, out_rank, &out_idx); var a_shape : array; var a_strides : array; - a_strides_from_shape(p.b_off + out_rank, rank, &a_shape, &a_strides); + a_strides_from_shape(fd.v.v1 + out_rank, rank, &a_shape, &a_strides); let win : i32 = i32(a_shape[rank - 1u]); let nwin : i32 = i32(a_shape[axis]); @@ -533,116 +564,187 @@ fn fold(@builtin(global_invocation_id) gid : vec3) { src = src + coord * a_strides[d]; } src = src + u32(k) * a_strides[rank - 1u]; - sum = sum + A[p.a_off + src]; + sum = sum + in0[fd.v.v0 + src]; } - C[p.c_off + i] = sum; + out2[fd.v.v2 + i] = sum; } // sum_to: sum `a` down to a smaller broadcast-target shape (the dual of // broadcast_to that every arithmetic op's backward uses to un-broadcast a // gradient). Gather, not scatter: one invocation per OUTPUT element sums -// every `a` element that broadcasts onto it, so -- like pad/fold above -- -// no write conflict and no atomics. B layout at word offset p.b_off: -// [a_shape(rank), a_strides(rank), acc(rank)] -- acc is a's shape -// broadcast-aligned against the output's own strides (array.h's -// broadcast_strides(target, out.strides(), a.shape())), 0 on a reduced -// axis. p._pad0 = rank, p._pad1 = reduced_n (product of a_shape over -// exactly the zero-acc axes; 1 if there are none). +// every `a` element that broadcasts onto it, so -- like pad/fold above -- no +// write conflict and no atomics. +// Views: a, meta [a_shape(rank), a_strides(rank), acc(rank)] -> out. acc is +// a's shape broadcast-aligned against the output's own strides (array.h's +// broadcast_strides(target, out.strides(), a.shape())), 0 on a reduced axis. +// Params: webgpu.h's sum_to_params (n = the output's element count; +// reduced_n = the product of a_shape over exactly the zero-acc axes, 1 if +// there are none). +struct SumToArgs { v : Views, n : u32, rank : u32, reduced_n : u32 }; +@group(0) @binding(3) var st : SumToArgs; + @compute @workgroup_size(256, 1, 1) fn sum_to(@builtin(global_invocation_id) gid : vec3) { let t = gid.x; - if (t >= p.M) { return; } - let rank = p._pad0; - let reduced_n = p._pad1; + if (t >= st.n) { return; } + let rank = st.rank; + let m = st.v.v1; var base : u32 = 0u; var red_axis : array; var red_count : u32 = 0u; for (var d : u32 = 0u; d < rank; d = d + 1u) { - let a_shape_d = bitcast(B[p.b_off + d]); - let acc_d = bitcast(B[p.b_off + 2u * rank + d]); + let a_shape_d = meta1(m, d); + let acc_d = meta1(m, 2u * rank + d); if (acc_d != 0u) { - let a_strides_d = bitcast(B[p.b_off + rank + d]); let idx = (t / acc_d) % a_shape_d; - base = base + idx * a_strides_d; + base = base + idx * meta1(m, rank + d); } else { red_axis[red_count] = d; red_count = red_count + 1u; } } var sum : f32 = 0.0; - for (var r : u32 = 0u; r < reduced_n; r = r + 1u) { + for (var r : u32 = 0u; r < st.reduced_n; r = r + 1u) { var rem = r; var off = base; for (var k : i32 = i32(red_count) - 1; k >= 0; k = k - 1) { let d = red_axis[u32(k)]; - let dim = bitcast(B[p.b_off + d]); + let dim = meta1(m, d); let coord = rem % dim; rem = rem / dim; - let stride = bitcast(B[p.b_off + rank + d]); - off = off + coord * stride; + off = off + coord * meta1(m, rank + d); } - sum = sum + A[p.a_off + off]; + sum = sum + in0[st.v.v0 + off]; } - C[p.c_off + t] = sum; + out2[st.v.v2 + t] = sum; } // concat_part: writes `a` (contiguous, one part of an N-ary Tensor.concat) -// into `out` at `p._pad1` (already before*out_strides[axis], a flat -// element offset) along one axis. One invocation per SOURCE element (this -// part's own count, p.M) -- unlike pad above, there is no zero border and -// no bounds check: concat's parts exhaustively and disjointly cover `out`, -// so this is a plain scatter that never collides across the separate -// per-part dispatches building up one `out`. B layout at word offset -// p.b_off: [a_shape(rank), out_strides(rank)] -- concat's own meta -// convention (mirrors cuda.h's upload_pad_fold_meta_), different from -// pad/fold's [out_shape, a_shape] above since this dispatch walks the -// SOURCE part's own shape, not the output's. p._pad0 = rank. +// into `out` along one axis, `shift` elements in (already before * +// out_strides[axis]). One invocation per SOURCE element -- unlike pad above, +// there is no zero border and no bounds check: concat's parts exhaustively +// and disjointly cover `out`, so this is a plain scatter that never collides +// across the separate per-part dispatches building up one `out`. +// Views: a, meta [a_shape(rank), out_strides(rank)] -> out -- walks the +// SOURCE part's own shape, not the output's, unlike pad/fold's meta. +// Params: webgpu.h's concat_params (n = this part's element count). +struct ConcatArgs { v : Views, n : u32, rank : u32, shift : u32 }; +@group(0) @binding(3) var cp : ConcatArgs; + @compute @workgroup_size(256, 1, 1) fn concat_part(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let rank = p._pad0; - let shift = p._pad1; + if (i >= cp.n) { return; } + let rank = cp.rank; var rem = i; var dst : u32 = 0u; for (var d : i32 = i32(rank) - 1; d >= 0; d = d - 1) { - let dim = bitcast(B[p.b_off + u32(d)]); + let dim = meta1(cp.v.v1, u32(d)); + let coord = rem % dim; + rem = rem / dim; + dst = dst + coord * meta1(cp.v.v1, rank + u32(d)); + } + out2[cp.v.v2 + dst + cp.shift] = in0[cp.v.v0 + i]; +} + +// ---- N-D broadcast binary (any rank) and N-D broadcast ternary select +// (Tensor.where's masking) -- the WGSL counterparts of tl_b*_nd/tl_where_nd +// in kernels/tensorlib_cuda.cu. Same flat-index decode as pad/fold above, +// against strides the host supplies (broadcast_strides(), 0 on a broadcast +// axis) rather than derived from a shape. Their meta sits at in2 and in3, +// after the tensor operands, so they decode it inline rather than through +// meta1 (WGSL has no templates to parameterize over which binding to read). + +// Views: a, b, meta [out_shape(rank), a_strides(rank), b_strides(rank)] -> +// out. Params: webgpu.h's bcast_nd_params; OP as ew_binary's. +struct BcastNdArgs { v : Views, n : u32, rank : u32, scale : f32, offset : f32 }; +@group(0) @binding(4) var bn : BcastNdArgs; + +@compute @workgroup_size(256, 1, 1) +fn ew_bcast_nd(@builtin(global_invocation_id) gid : vec3) { + let i = gid.x; + if (i >= bn.n) { return; } + let rank = bn.rank; + let base = bn.v.v2; + var rem = i; + var a_off : u32 = 0u; + var b_off : u32 = 0u; + for (var d : i32 = i32(rank) - 1; d >= 0; d = d - 1) { + let dim = bitcast(in2[base + u32(d)]); + let coord = rem % dim; + rem = rem / dim; + a_off = a_off + coord * bitcast(in2[base + rank + u32(d)]); + b_off = b_off + coord * bitcast(in2[base + 2u * rank + u32(d)]); + } + let av = in0[bn.v.v0 + a_off]; + let bv = in1[bn.v.v1 + b_off]; + out3[bn.v.v3 + i] = fma(binary_op(OP, av, bv), bn.scale, bn.offset); +} + +// Views: cond, a, b, meta [out_shape(rank), cond_strides(rank), +// a_strides(rank), b_strides(rank)] -> out. Params: webgpu.h's +// where_nd_params. +struct WhereNdArgs { v : Views, n : u32, rank : u32 }; +@group(0) @binding(5) var wn : WhereNdArgs; + +@compute @workgroup_size(256, 1, 1) +fn where_nd(@builtin(global_invocation_id) gid : vec3) { + let i = gid.x; + if (i >= wn.n) { return; } + let rank = wn.rank; + let base = wn.v.v3; + var rem = i; + var c_off : u32 = 0u; + var a_off : u32 = 0u; + var b_off : u32 = 0u; + for (var d : i32 = i32(rank) - 1; d >= 0; d = d - 1) { + let dim = bitcast(in3[base + u32(d)]); let coord = rem % dim; rem = rem / dim; - let os = bitcast(B[p.b_off + rank + u32(d)]); - dst = dst + coord * os; + c_off = c_off + coord * bitcast(in3[base + rank + u32(d)]); + a_off = a_off + coord * bitcast(in3[base + 2u * rank + u32(d)]); + b_off = b_off + coord * bitcast(in3[base + 3u * rank + u32(d)]); } - C[p.c_off + dst + shift] = A[p.a_off + i]; + let cv = in0[wn.v.v0 + c_off]; + let av = in1[wn.v.v1 + a_off]; + let bv = in2[wn.v.v2 + b_off]; + out4[wn.v.v4 + i] = select(bv, av, cv != 0.0); } -// RoPE (rotary position embedding), half-split (GPT-NeoX / HF-llama) -// convention -- mirrors cuda.h's/metal.h's own rope. A is [rows, D] -// contiguous (rows = H*T: a [H,T,D] tensor flattened, or [H,D] with -// T=1); row r's head-dim vector sits at position `pos + (r % T)`; pairs -// (j, j+D/2) rotate by angle = position * base^(-2j/D). Dispatched flat -// over rows*(D/2), one invocation per (r, j) pair. p.N = T, p.K = D, -// p._pad0 = pos, p._pad1 = D/2, p.scale = base (repurposed -- rope has -// no affine epilogue to compose with, same idea as clamp_'s lo/hi above). +// ---- RoPE (rotary position embedding), half-split (GPT-NeoX / HF-llama) +// convention -- mirrors cuda.h's/metal.h's own rope. x is [rows, D] +// contiguous (rows = H*T: a [H,T,D] tensor flattened, or [H,D] with T=1); +// row r's head-dim vector sits at position `pos + (r % T)`; pairs (j, j+D/2) +// rotate by angle = position * base^(-2j/D). Dispatched flat over +// rows*(D/2), one invocation per (r, j) pair. +// Views: x -> out. Params: webgpu.h's rope_params (n = rows * D/2). +struct RopeArgs { + v : Views, + n : u32, t : u32, d : u32, pos : u32, half : u32, + base : f32, +}; +@group(0) @binding(2) var rp : RopeArgs; + @compute @workgroup_size(256, 1, 1) fn rope(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let half = p._pad1; + if (i >= rp.n) { return; } + let half = rp.half; let r = i / half; let j = i % half; var t : u32 = 0u; - if (p.N > 0u) { t = r % p.N; } - let position = f32(p._pad0 + t); - let theta = pow(p.scale, -2.0 * f32(j) / f32(p.K)); + if (rp.t > 0u) { t = r % rp.t; } + let position = f32(rp.pos + t); + let theta = pow(rp.base, -2.0 * f32(j) / f32(rp.d)); let ang = position * theta; let c = cos(ang); let s = sin(ang); - let bi = r * p.K; - let x0 = A[p.a_off + bi + j]; - let x1 = A[p.a_off + bi + j + half]; - C[p.c_off + bi + j] = x0 * c - x1 * s; - C[p.c_off + bi + j + half] = x0 * s + x1 * c; + let bi = r * rp.d; + let x0 = in0[rp.v.v0 + bi + j]; + let x1 = in0[rp.v.v0 + bi + j + half]; + out1[rp.v.v1 + bi + j] = x0 * c - x1 * s; + out1[rp.v.v1 + bi + j + half] = x0 * s + x1 * c; } // ---- Embedding-table lookup (index_select/index_add) and pooling-style @@ -657,106 +759,54 @@ fn rope(@builtin(global_invocation_id) gid : vec3) { // per OUTPUT element, summing over every source row whose index matches it, // in place of scattering into a pre-zeroed buffer. -// A = a (table), B = idx, C = out. p._pad0 = row_size. +// Views: a (table), idx -> out. Params: gather_params. +struct GatherArgs { v : Views, row_size : u32, n : u32 }; +@group(0) @binding(3) var gs : GatherArgs; + @compute @workgroup_size(256, 1, 1) fn index_select(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let row_size = p._pad0; - let row = i / row_size; - let col = i % row_size; - let src_row = u32(B[p.b_off + row] + 0.5); - C[p.c_off + i] = A[p.a_off + src_row * row_size + col]; + if (i >= gs.n) { return; } + let row = i / gs.row_size; + let col = i % gs.row_size; + let src_row = u32(in1[gs.v.v1 + row] + 0.5); + out2[gs.v.v2 + i] = in0[gs.v.v0 + src_row * gs.row_size + col]; } -// A = idx, B = values, C = out. p._pad0 = row_size, p._pad1 = k (source rows). +// Views: idx, values -> out. Params: webgpu.h's index_add_params (n = the +// output's element count, k = the source rows to scan). +struct IndexAddArgs { v : Views, n : u32, row_size : u32, k : u32 }; +@group(0) @binding(3) var ia : IndexAddArgs; + @compute @workgroup_size(256, 1, 1) fn index_add(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let row_size = p._pad0; - let k_count = p._pad1; - let row = i / row_size; - let col = i % row_size; + if (i >= ia.n) { return; } + let row = i / ia.row_size; + let col = i % ia.row_size; var sum : f32 = 0.0; - for (var k : u32 = 0u; k < k_count; k = k + 1u) { - let idx_row = u32(A[p.a_off + k] + 0.5); + for (var k : u32 = 0u; k < ia.k; k = k + 1u) { + let idx_row = u32(in0[ia.v.v0 + k] + 0.5); if (idx_row == row) { - sum = sum + B[p.b_off + k * row_size + col]; + sum = sum + in1[ia.v.v1 + k * ia.row_size + col]; } } - C[p.c_off + i] = sum; + out2[ia.v.v2 + i] = sum; } -// A = idx, B = values, C = out. p._pad0 = size (the new trailing axis). // out[pos, k] = values[pos] where idx[pos] == k, else 0 -- gather, so no // zeroing needed, unlike CUDA's pre-zeroed scatter. -@compute @workgroup_size(256, 1, 1) -fn scatter_axis(@builtin(global_invocation_id) gid : vec3) { - let i = gid.x; - if (i >= p.M) { return; } - let size = p._pad0; - let pos = i / size; - let k = i % size; - let dst_k = u32(A[p.a_off + pos] + 0.5); - C[p.c_off + i] = select(0.0, B[p.b_off + pos], dst_k == k); -} +// Views: idx, values -> out. Params: webgpu.h's scatter_axis_params (n = the +// output's element count, size = the new trailing axis). +struct ScatterAxisArgs { v : Views, n : u32, size : u32 }; +@group(0) @binding(3) var sa : ScatterAxisArgs; -// ---- N-D broadcast binary (any rank) and N-D broadcast ternary select -// (Tensor.where's masking) -- the WGSL counterparts of tl_b*_nd/tl_where_nd -// in kernels/tensorlib_cuda.cu. Same flat-index decode as pad/fold above, -// against strides the host supplies (broadcast_strides(), 0 on a broadcast -// axis) rather than derived from a shape -- inlined per kernel rather than -// shared via decode_idx, since that helper always reads from B and here the -// meta buffer is D or E (WGSL has no templates to parameterize over which -// storage binding to read). - -// A = a, B = b, D = meta [out_shape(rank), a_strides(rank), b_strides(rank)], -// E unused. p._pad0 = rank, p._pad3 = meta's word offset into D. @compute @workgroup_size(256, 1, 1) -fn ew_bcast_nd(@builtin(global_invocation_id) gid : vec3) { - let i = gid.x; - if (i >= p.M) { return; } - let rank = p._pad0; - let base = p._pad3; - var rem = i; - var a_off : u32 = 0u; - var b_off : u32 = 0u; - for (var d : i32 = i32(rank) - 1; d >= 0; d = d - 1) { - let dim = bitcast(D[base + u32(d)]); - let coord = rem % dim; - rem = rem / dim; - a_off = a_off + coord * bitcast(D[base + rank + u32(d)]); - b_off = b_off + coord * bitcast(D[base + 2u * rank + u32(d)]); - } - let av = A[p.a_off + a_off]; - let bv = B[p.b_off + b_off]; - C[p.c_off + i] = fma(binary_op(p.op, av, bv), p.scale, p.offset); -} - -// A = cond, B = a, D = b, E = meta [out_shape(rank), cond_strides(rank), -// a_strides(rank), b_strides(rank)]. p._pad0 = rank, p._pad3 = b's element -// offset into D, p._pad4 = meta's word offset into E. -@compute @workgroup_size(256, 1, 1) -fn where_nd(@builtin(global_invocation_id) gid : vec3) { +fn scatter_axis(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let rank = p._pad0; - let base = p._pad4; - var rem = i; - var c_off : u32 = 0u; - var a_off : u32 = 0u; - var b_off : u32 = 0u; - for (var d : i32 = i32(rank) - 1; d >= 0; d = d - 1) { - let dim = bitcast(E[base + u32(d)]); - let coord = rem % dim; - rem = rem / dim; - c_off = c_off + coord * bitcast(E[base + rank + u32(d)]); - a_off = a_off + coord * bitcast(E[base + 2u * rank + u32(d)]); - b_off = b_off + coord * bitcast(E[base + 3u * rank + u32(d)]); - } - let cv = A[p.a_off + c_off]; - let av = B[p.b_off + a_off]; - let bv = D[p._pad3 + b_off]; - C[p.c_off + i] = select(bv, av, cv != 0.0); + if (i >= sa.n) { return; } + let pos = i / sa.size; + let k = i % sa.size; + let dst_k = u32(in0[sa.v.v0 + pos] + 0.5); + out2[sa.v.v2 + i] = select(0.0, in1[sa.v.v1 + pos], dst_k == k); } diff --git a/kernels/tensorlib_webgpu_wgsl.inc b/kernels/tensorlib_webgpu_wgsl.inc index 0f1598b..e72e7e2 100644 --- a/kernels/tensorlib_webgpu_wgsl.inc +++ b/kernels/tensorlib_webgpu_wgsl.inc @@ -8,70 +8,67 @@ R"WGSL( // committed because the wasm build is a flat emcc line that never runs CMake // (the CUDA backend's PTX goes through bin2c for the same reason). // -// View offsets arrive as ELEMENT offsets in the params block and are folded -// into the indexing here, rather than as bind-group binding offsets: WebGPU -// requires those to be 256-byte aligned, which an arbitrary view offset is -// not. So every binding covers its whole buffer. (CUDA instead folds offsets -// host-side into the pointer it passes, which WebGPU has no equivalent of.) - -// One Params struct and one bind group layout serve every kernel here, rather -// than the per-family structs metal_kernels.metal uses. WebGPU's bind group -// ceremony is heavy enough that a second layout would buy nothing: the fields -// each family ignores cost 4 bytes of a 256-byte uniform slot. Kernels that -// take one input (unary, the row reductions) get A bound to B as well — two -// read-only bindings may alias, and the output is always a fresh allocation, -// so no writable binding ever aliases a readable one. -struct Params { - M : u32, // gemm rows | elementwise element count | reduce rows - N : u32, // gemm cols | reduce cols - K : u32, - lda : u32, - ldb : u32, - ldc : u32, - a_off : u32, - b_off : u32, - c_off : u32, - ta : u32, - tb : u32, - ars : u32, // broadcast: per-operand row/col strides, in elements - acs : u32, - brs : u32, - bcs : u32, - op : u32, // which operation, within the entry point's family - scale : f32, - offset : f32, - // A uniform-address-space struct has align 16, so its size rounds up to a - // multiple of 16. Pad explicitly to 96 bytes so the host struct (which the - // bind group's minBindingSize comes from) matches exactly — a short - // minBindingSize fails bind group validation for every dispatch. - _pad0 : u32, - _pad1 : u32, - _pad2 : u32, - _pad3 : u32, - _pad4 : u32, - arg : f32, // a scalar operand (ew_scalar's s) +// The kernel ABI (gpu_abi.h), as WGSL states it. A kernel binds view i at +// binding i, in the order gpu_ops.h lists its views, and its uniform at the +// binding after the last view. The uniform is the views' element offsets +// (`Views`, by view index) followed by the kernel's params, the same fields +// in the same order as its gpu_abi.h struct. webgpu.h fills both the same way +// for every kernel. +// +// View offsets ride the uniform rather than the bind group: a binding offset +// has to be 256-byte aligned, which an arbitrary view offset is not. So every +// binding covers its whole buffer, and each kernel folds its views' offsets +// into its indexing. (Metal binds at the offset and CUDA adds it to the +// pointer; WebGPU has neither.) +// +// A binding's type follows its view: read-only storage for an input, +// read_write for an output. Binding i is an input in one kernel and an output +// in another, so the storage variables are declared once per binding and +// access, in0.. and out1.., and each entry point uses the ones its views +// need. WGSL checks binding clashes per entry point, and each pipeline takes +// its layout from its entry point (webgpu.h builds them with an auto layout), +// so what a kernel declares here is its layout. +// +// A family's operation — add or sub, exp or log — is not a params field but +// the pipeline-overridable constant OP: webgpu.h builds a pipeline per kernel +// id with OP set from its kernel table, the counterpart of the kernel name +// Metal and CUDA look an id up by. Within a pipeline OP is a constant, so the +// switch on it takes the same arm in every invocation. + +override OP : u32 = 0u; + +// Each view's element offset into its binding, by view index. +struct Views { + v0 : u32, v1 : u32, v2 : u32, v3 : u32, + v4 : u32, v5 : u32, v6 : u32, v7 : u32, }; -@group(0) @binding(0) var A : array; -@group(0) @binding(1) var B : array; -@group(0) @binding(2) var C : array; -@group(0) @binding(3) var p : Params; -// A third and fourth read-only operand, for kernels A/B alone can't cover: -// binary_bcast_nd's two real tensor operands (a, b) already fill A and B, so -// its shape/stride meta needs D; where_nd's three real operands (cond, a, b) -// fill A, B and D, so its meta needs E. Every other entry point ignores -// these (webgpu.h's encode_ binds them to A when a kernel has no use for -// them — bind group validation requires every declared binding be present -// regardless of which bindings the active entry point actually reads). -@group(0) @binding(4) var D : array; -@group(0) @binding(5) var E : array; - -// ---- sgemm: C(M,N) = op(A)(M,K) @ op(B)(K,N) * scale + offset +@group(0) @binding(0) var in0 : array; +@group(0) @binding(1) var in1 : array; +@group(0) @binding(2) var in2 : array; +@group(0) @binding(3) var in3 : array; +@group(0) @binding(1) var out1 : array; +@group(0) @binding(2) var out2 : array; +@group(0) @binding(3) var out3 : array; +@group(0) @binding(4) var out4 : array; + +// ---- sgemm: out(M,N) = op(a)(M,K) @ op(b)(K,N) * scale + offset // // 64x64 workgroup tile, 16x16 = 256 invocations, each holding a 4x4 register // accumulator. Mirrors the shape of the Metal sgemm_64_ kernel; MMA intrinsics // have no WGSL equivalent, so the inner product is plain FMA over registers. // Measured at ~580-630 GF/s for n=1024 on an M1 Pro (see spike/webgpu). +// +// Views: a, b -> out. Params: webgpu.h's gemm_params. + +struct GemmArgs { + v : Views, + m : u32, n : u32, k : u32, + lda : u32, ldb : u32, ldc : u32, + ta : u32, tb : u32, + scale : f32, offset : f32, +}; +@group(0) @binding(3) var gm : GemmArgs; const BM : u32 = 64u; const BN : u32 = 64u; @@ -86,15 +83,15 @@ var Bs : array; // BK * BN // Row/column strides for a possibly-transposed operand: transposing swaps // which axis walks by the leading dimension (cf. metal_kernels.metal:119-120). fn a_index(m : u32, k : u32) -> u32 { - let rs = select(p.lda, 1u, p.ta == 1u); - let cs = select(1u, p.lda, p.ta == 1u); - return p.a_off + m * rs + k * cs; + let rs = select(gm.lda, 1u, gm.ta == 1u); + let cs = select(1u, gm.lda, gm.ta == 1u); + return gm.v.v0 + m * rs + k * cs; } fn b_index(k : u32, n : u32) -> u32 { - let rs = select(p.ldb, 1u, p.tb == 1u); - let cs = select(1u, p.ldb, p.tb == 1u); - return p.b_off + k * rs + n * cs; + let rs = select(gm.ldb, 1u, gm.tb == 1u); + let cs = select(1u, gm.ldb, gm.tb == 1u); + return gm.v.v1 + k * rs + n * cs; } @compute @workgroup_size(16, 16, 1) @@ -106,7 +103,7 @@ fn sgemm(@builtin(workgroup_id) wg : vec3, var acc : array; // TM * TN, zero-initialized - let n_tiles = (p.K + BK - 1u) / BK; + let n_tiles = (gm.k + BK - 1u) / BK; for (var kt : u32 = 0u; kt < n_tiles; kt = kt + 1u) { let k0 = kt * BK; @@ -118,13 +115,13 @@ fn sgemm(@builtin(workgroup_id) wg : vec3, let am = m_base + i / BK; let ak = k0 + i % BK; - let a_ok = am < p.M && ak < p.K; - As[i] = select(0.0, A[a_index(am, ak)], a_ok); + let a_ok = am < gm.m && ak < gm.k; + As[i] = select(0.0, in0[a_index(am, ak)], a_ok); let bk = k0 + i / BN; let bn = n_base + i % BN; - let b_ok = bk < p.K && bn < p.N; - Bs[i] = select(0.0, B[b_index(bk, bn)], b_ok); + let b_ok = bk < gm.k && bn < gm.n; + Bs[i] = select(0.0, in1[b_index(bk, bn)], b_ok); } workgroupBarrier(); @@ -150,11 +147,11 @@ fn sgemm(@builtin(workgroup_id) wg : vec3, for (var i : u32 = 0u; i < TM; i = i + 1u) { let m = m_base + lid.y * TM + i; - if (m >= p.M) { continue; } + if (m >= gm.m) { continue; } for (var j : u32 = 0u; j < TN; j = j + 1u) { let n = n_base + lid.x * TN + j; - if (n >= p.N) { continue; } - C[p.c_off + m * p.ldc + n] = acc[i * TN + j] * p.scale + p.offset; + if (n >= gm.n) { continue; } + out2[gm.v.v2 + m * gm.ldc + n] = acc[i * TN + j] * gm.scale + gm.offset; } } } @@ -164,12 +161,10 @@ fn sgemm(@builtin(workgroup_id) wg : vec3, // WGSL has neither templates nor a preprocessor, so the per-op variants that // metal_kernels.metal generates from a macro would have to be copy-pasted // here — exactly the edge-tile bug class that file's header warns against. -// Instead the operation is a uniform field and each family is ONE entry point -// that switches on it. The branch is uniform across the dispatch and these -// kernels are memory-bound, so it costs nothing measurable; what it buys is a -// single copy of every bounds check and epilogue. +// Instead each family is ONE entry point that switches on OP. What that buys +// is a single copy of every bounds check and epilogue. // -// Op codes are assigned by kernel_op_() in webgpu.h. +// OP's values per family are assigned by kernel_() in webgpu.h. const OP_ADD : u32 = 0u; const OP_SUB : u32 = 1u; const OP_MUL : u32 = 2u; @@ -202,7 +197,7 @@ const OP_NE : u32 = 5u; // it. Threads whose row is shorter than the workgroup contribute this. const NEG_HUGE : f32 = -3.4e38; -// Shared by the contiguous and the broadcast binary: same five operations, +// Shared by the contiguous and the broadcast binaries: same five operations, // only the addressing differs. fn binary_op(op : u32, av : f32, bv : f32) -> f32 { switch (op) { @@ -228,10 +223,9 @@ fn unary_op(op : u32, v : f32) -> f32 { } } -// Elementwise comparison: out = (a OP b) ? 1.0 : 0.0. No scale/offset -- masks -// don't compose with the affine epilogue, so ew_cmp below skips it (unlike -// ew_binary/ew_unary). Own family/entry point, not folded into binary_op, -// since it returns a bool-as-float rather than composing with `bv`. +// Elementwise comparison: (a OP b) ? 1.0 : 0.0. No scale/offset -- masks +// don't compose with the affine epilogue. Not folded into binary_op, since it +// returns a bool-as-float rather than composing with `bv`. fn cmp_op(op : u32, av : f32, bv : f32) -> f32 { switch (op) { case 1u: { return select(0.0, 1.0, av < bv); } @@ -243,50 +237,57 @@ fn cmp_op(op : u32, av : f32, bv : f32) -> f32 { } } -// Contiguous elementwise binary over p.M elements. +// Views: a, b -> out (binary) or a -> out (unary). Params: ew_params. +struct EwArgs { v : Views, n : u32, scale : f32, offset : f32 }; +@group(0) @binding(3) var ewb : EwArgs; +@group(0) @binding(2) var ewu : EwArgs; + @compute @workgroup_size(256, 1, 1) fn ew_binary(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let v = binary_op(p.op, A[p.a_off + i], B[p.b_off + i]); - C[p.c_off + i] = fma(v, p.scale, p.offset); + if (i >= ewb.n) { return; } + let v = binary_op(OP, in0[ewb.v.v0 + i], in1[ewb.v.v1 + i]); + out2[ewb.v.v2 + i] = fma(v, ewb.scale, ewb.offset); } @compute @workgroup_size(256, 1, 1) fn ew_unary(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let v = unary_op(p.op, A[p.a_off + i]); - C[p.c_off + i] = fma(v, p.scale, p.offset); + if (i >= ewu.n) { return; } + let v = unary_op(OP, in0[ewu.v.v0 + i]); + out1[ewu.v.v1 + i] = fma(v, ewu.scale, ewu.offset); } -// Elementwise comparison over p.M elements: out = (a OP b) ? 1.0 : 0.0 -// (no epilogue -- masks don't compose with scale/offset). p.ars carries the -// bstride webgpu.h's compare() receives: 1 for a same-shape b, 0 for an -// explicit size-1 b (`x > s` itself is ew_scalar) -- an unused field for -// this family, repurposed rather than widening Params. +// out[i] = (a[i] OP b[i * bstride]) ? 1.0 : 0.0: bstride 1 for a same-shape +// b, 0 for a size-1 one (`x > s` itself is ew_scalar). +// Views: a, b -> out. Params: cmp_params. +struct CmpArgs { v : Views, n : u32, bstride : u32 }; +@group(0) @binding(3) var cm : CmpArgs; + @compute @workgroup_size(256, 1, 1) fn cmp(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let bv = B[p.b_off + i * p.ars]; - C[p.c_off + i] = cmp_op(p.op, A[p.a_off + i], bv); + if (i >= cm.n) { return; } + let bv = in1[cm.v.v1 + i * cm.bstride]; + out2[cm.v.v2 + i] = cmp_op(OP, in0[cm.v.v0 + i], bv); } -// clamp(x, lo, hi): Clip's forward. No affine epilogue -- p.scale/p.offset -// carry lo/hi instead (a dedicated entry, same as this family's clamp_ in -// metal_kernels.metal/tl_clamp in tensorlib_cuda.cu). Named clamp_, not -// clamp, so the entry point doesn't shadow WGSL's builtin of that name. +// clamp(x, lo, hi): Clip's forward, no epilogue. Named clamp_, not clamp, so +// the entry point doesn't shadow WGSL's builtin of that name. +// Views: a -> out. Params: clamp_params. +struct ClampArgs { v : Views, n : u32, lo : f32, hi : f32 }; +@group(0) @binding(2) var cl : ClampArgs; + @compute @workgroup_size(256, 1, 1) fn clamp_(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - C[p.c_off + i] = clamp(A[p.a_off + i], p.scale, p.offset); + if (i >= cl.n) { return; } + out1[cl.v.v1 + i] = clamp(in0[cl.v.v0 + i], cl.lo, cl.hi); } -// Tensor-scalar over p.M elements: out = f(a, s) * scale + offset, s in p.arg. -// Mirrors tensorlib_cuda.cu's TL_EW_SCALAR; op 0 is pow, 1.. the comparisons -// in cmp_op's order. +// Tensor-scalar: out = f(a, s) * scale + offset. Mirrors tensorlib_cuda.cu's +// TL_EW_SCALAR; OP 0 is pow, 1.. the comparisons in cmp_op's order. +// Views: a -> out. Params: scalar_params. fn scalar_op(op : u32, av : f32, s : f32) -> f32 { switch (op) { case 0u: { return pow(av, s); } @@ -294,31 +295,43 @@ fn scalar_op(op : u32, av : f32, s : f32) -> f32 { } } +struct ScalarArgs { v : Views, n : u32, s : f32, scale : f32, offset : f32 }; +@group(0) @binding(2) var sc : ScalarArgs; + @compute @workgroup_size(256, 1, 1) fn ew_scalar(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - C[p.c_off + i] = fma(scalar_op(p.op, A[p.a_off + i], p.arg), p.scale, p.offset); + if (i >= sc.n) { return; } + out1[sc.v.v1 + i] = fma(scalar_op(OP, in0[sc.v.v0 + i], sc.s), sc.scale, sc.offset); } // Rank-2 broadcast binary: out[r,c] = f(a[r*ars + c*acs], b[r*brs + c*bcs]) -// into a contiguous [M,N] output. Per-operand strides express every rank-2 +// into a contiguous [m,n] output. Per-operand strides express every rank-2 // broadcast (row vector, column vector, scalar) in one kernel, which keeps // bias/gamma/beta chains on the GPU — a CPU fallback mid-graph costs a full // submit-and-wait, and that is dearer here than on Metal. +// Views: a, b -> out. Params: bcast_params. +struct BcastArgs { + v : Views, + m : u32, n : u32, + ars : u32, acs : u32, brs : u32, bcs : u32, + scale : f32, offset : f32, +}; +@group(0) @binding(3) var bc : BcastArgs; + @compute @workgroup_size(32, 8, 1) fn ew_bcast(@builtin(global_invocation_id) gid : vec3) { let c = gid.x; let r = gid.y; - if (c >= p.N || r >= p.M) { return; } - let av = A[p.a_off + r * p.ars + c * p.acs]; - let bv = B[p.b_off + r * p.brs + c * p.bcs]; - C[p.c_off + r * p.N + c] = fma(binary_op(p.op, av, bv), p.scale, p.offset); + if (c >= bc.n || r >= bc.m) { return; } + let av = in0[bc.v.v0 + r * bc.ars + c * bc.acs]; + let bv = in1[bc.v.v1 + r * bc.brs + c * bc.bcs]; + out2[bc.v.v2 + r * bc.n + c] = fma(binary_op(OP, av, bv), bc.scale, bc.offset); } // ---- Row reductions over the last axis: one workgroup per row, 256 -// invocations, workgroup-scratch tree reduction. p.N (cols) may exceed the -// invocation count, so each thread strides over the row first. +// invocations, workgroup-scratch tree reduction. The row (cols) may exceed the +// invocation count, so each thread strides over it first. const T : u32 = 256u; var scratch : array; @@ -346,6 +359,11 @@ fn tree_reduce(op : u32, lid : u32, v : f32) -> f32 { return r; } +// Views: a -> out. Params: reduce_params. softmax writes rows x cols, +// row_reduce one value a row. +struct ReduceArgs { v : Views, rows : u32, cols : u32, scale : f32, offset : f32 }; +@group(0) @binding(2) var rd : ReduceArgs; + // Numerically stable softmax (subtract the row max). Applying an affine // epilogue to a softmax is not meaningful, so scale/offset are ignored here, // as they are on Metal. @@ -353,133 +371,142 @@ fn tree_reduce(op : u32, lid : u32, v : f32) -> f32 { fn softmax(@builtin(workgroup_id) wg : vec3, @builtin(local_invocation_index) lid : u32) { let row = wg.x; - let src = p.a_off + row * p.N; - let dst = p.c_off + row * p.N; + let src = rd.v.v0 + row * rd.cols; + let dst = rd.v.v1 + row * rd.cols; var m = NEG_HUGE; - for (var c : u32 = lid; c < p.N; c = c + T) { m = max(m, A[src + c]); } + for (var c : u32 = lid; c < rd.cols; c = c + T) { m = max(m, in0[src + c]); } let row_max = tree_reduce(OP_ROW_MAX, lid, m); var sum = 0.0; - for (var c : u32 = lid; c < p.N; c = c + T) { sum = sum + exp(A[src + c] - row_max); } + for (var c : u32 = lid; c < rd.cols; c = c + T) { sum = sum + exp(in0[src + c] - row_max); } let inv = 1.0 / tree_reduce(OP_ROW_SUM, lid, sum); - for (var c : u32 = lid; c < p.N; c = c + T) { - C[dst + c] = exp(A[src + c] - row_max) * inv; + for (var c : u32 = lid; c < rd.cols; c = c + T) { + out1[dst + c] = exp(in0[src + c] - row_max) * inv; + } +} + +// One value per row, with the affine epilogue. +@compute @workgroup_size(256, 1, 1) +fn row_reduce(@builtin(workgroup_id) wg : vec3, + @builtin(local_invocation_index) lid : u32) { + let row = wg.x; + let src = rd.v.v0 + row * rd.cols; + + var acc = select(0.0, NEG_HUGE, OP == OP_ROW_MAX); + for (var c : u32 = lid; c < rd.cols; c = c + T) { + acc = reduce_op(OP, acc, in0[src + c]); } + let r = tree_reduce(OP, lid, acc); + if (lid == 0u) { out1[rd.v.v1 + row] = r * rd.scale + rd.offset; } } // Layer norm over the last axis with the affine epilogue: (x - mu) / -// sqrt(var + eps) * g + b. A = x, B = g, D = b (p._pad3 its element offset), -// p.arg = eps. Two tree sums (x, then the squared deviations), each scaled by -// 1/N after the tree as row_reduce's mean is. +// sqrt(var + eps) * g + b. Two tree sums (x, then the squared deviations), +// each scaled by 1/cols after the tree as row_reduce's mean is. +// Views: x, g, b -> out. Params: layer_norm_params. +struct LayerNormArgs { + v : Views, + rows : u32, cols : u32, + eps : f32, scale : f32, offset : f32, +}; +@group(0) @binding(4) var ln : LayerNormArgs; + @compute @workgroup_size(256, 1, 1) fn layer_norm(@builtin(workgroup_id) wg : vec3, @builtin(local_invocation_index) lid : u32) { let row = wg.x; - let src = p.a_off + row * p.N; - let dst = p.c_off + row * p.N; - let inv_n = 1.0 / f32(p.N); + let src = ln.v.v0 + row * ln.cols; + let dst = ln.v.v3 + row * ln.cols; + let inv_n = 1.0 / f32(ln.cols); var sum = 0.0; - for (var c : u32 = lid; c < p.N; c = c + T) { sum = sum + A[src + c]; } + for (var c : u32 = lid; c < ln.cols; c = c + T) { sum = sum + in0[src + c]; } let mu = tree_reduce(OP_ROW_SUM, lid, sum) * inv_n; var ss = 0.0; - for (var c : u32 = lid; c < p.N; c = c + T) { - let v = A[src + c] - mu; + for (var c : u32 = lid; c < ln.cols; c = c + T) { + let v = in0[src + c] - mu; ss = ss + v * v; } - let inv = 1.0 / sqrt(tree_reduce(OP_ROW_SUM, lid, ss) * inv_n + p.arg); + let inv = 1.0 / sqrt(tree_reduce(OP_ROW_SUM, lid, ss) * inv_n + ln.eps); - for (var c : u32 = lid; c < p.N; c = c + T) { - let y = (A[src + c] - mu) * inv * B[p.b_off + c] + D[p._pad3 + c]; - C[dst + c] = y * p.scale + p.offset; + for (var c : u32 = lid; c < ln.cols; c = c + T) { + let y = (in0[src + c] - mu) * inv * in1[ln.v.v1 + c] + in2[ln.v.v2 + c]; + out3[dst + c] = y * ln.scale + ln.offset; } } -// One value per row, with the affine epilogue. -@compute @workgroup_size(256, 1, 1) -fn row_reduce(@builtin(workgroup_id) wg : vec3, - @builtin(local_invocation_index) lid : u32) { - let row = wg.x; - let src = p.a_off + row * p.N; - - var acc = select(0.0, NEG_HUGE, p.op == OP_ROW_MAX); - for (var c : u32 = lid; c < p.N; c = c + T) { - acc = reduce_op(p.op, acc, A[src + c]); - } - let r = tree_reduce(p.op, lid, acc); - if (lid == 0u) { C[p.c_off + row] = r * p.scale + p.offset; } -} - -// ---- im2col's pad/fold (M11) -// -// Both dispatch one invocation per OUTPUT element and gather from A, rather -// than CUDA's scatter+atomicAdd: WGSL has no float atomicAdd. A gather needs -// no pre-zeroed output (every C cell is written exactly once, by exactly one -// invocation) and no atomics — the tradeoff is fold's small bounded loop over -// the window indices that could cover a given output cell, in place of one -// atomicAdd per source element. -// -// Neither family's shape metadata fits the fixed Params uniform (a -// variable-length array has no home there), so it rides B as -// bit-reinterpreted u32 — B's declared type is array, but WriteBuffer on -// the host side is a raw byte copy regardless, and bitcast reads it back -// correctly. p._pad0/_pad1/_pad2 (otherwise-unused Params padding) carry -// rank/axis/before-or-step; p.M is the output element count (webgpu.h's -// pad()/fold() dispatch over out_n, not a's own size). +// ---- Shape metadata. The N-D kernels below take a variable-length shape +// or stride list, which has no home in a fixed params struct, so it rides a +// view of its own (webgpu.h's meta ring), one slot per call: an input like +// any other, whose f32 words are u32 bit patterns — WriteBuffer is a raw +// byte copy, and bitcast reads them back. // // Rank cap — matches webgpu.h's own kPadFoldMaxRank (a shader can't see a // host-side C++ constant, so this is its own copy) and bounds every // fixed-size local array below. const kPadFoldMaxRank : u32 = 8u; -// Shared by both: row-major decode of a dispatch-global thread id `i` against -// a shape held in B at word offset `base`, into a fixed-size local array. -// `rank8` bounds the loop so callers can pass either a full-rank or a -// (rank-1)-length shape. +// Word d of the meta view bound at in1. +fn meta1(base : u32, d : u32) -> u32 { return bitcast(in1[base + d]); } + +// Row-major decode of a dispatch-global thread id `i` against a shape held in +// in1 from `base`, into a fixed-size local array. `rank8` bounds the loop so +// callers can pass either a full-rank or a (rank-1)-length shape. fn decode_idx(i : u32, base : u32, rank8 : u32, - out_idx : ptr>) { + out_idx : ptr>) { var rem = i; for (var d : i32 = i32(rank8) - 1; d >= 0; d = d - 1) { - let dim = bitcast(B[base + u32(d)]); + let dim = meta1(base, u32(d)); (*out_idx)[u32(d)] = rem % dim; rem = rem / dim; } } -// Row-major strides of a contiguous tensor whose shape is held in B at word -// offset `base`, length `rank` — used to address `A`, which pad_/fold_'s GPU +// Row-major strides of a contiguous tensor whose shape is held in in1 from +// `base`, length `rank` — used to address the source, which pad_/fold_'s GPU // dispatch (array.h's gpu_pad_/gpu_fold_) requires to be contiguous. fn a_strides_from_shape(base : u32, rank : u32, a_shape : ptr>, out_strides : ptr>) { var acc : u32 = 1u; for (var d : i32 = i32(rank) - 1; d >= 0; d = d - 1) { - let dim = bitcast(B[base + u32(d)]); + let dim = meta1(base, u32(d)); (*a_shape)[u32(d)] = dim; (*out_strides)[u32(d)] = acc; acc = acc * dim; } } -// B layout at word offset p.b_off: [out_shape(rank), a_shape(rank)] — webgpu.h -// reserves a fresh ring slot per call (see meta_reserve_slot_), so two pad/ -// fold calls batched into the same unflushed pass never share one offset. +// ---- im2col's pad/fold (M11) +// +// Both dispatch one invocation per OUTPUT element and gather from the source, +// rather than CUDA's scatter+atomicAdd: WGSL has no float atomicAdd. A gather +// needs no pre-zeroed output (every cell is written exactly once, by exactly +// one invocation) and no atomics — the tradeoff is fold's small bounded loop +// over the window indices that could cover a given output cell, in place of +// one atomicAdd per source element. + +// Views: a, meta [out_shape(rank), a_shape(rank)] -> out. +// Params: webgpu.h's pad_params (n = the output's element count). +struct PadArgs { v : Views, n : u32, rank : u32, axis : u32, before : u32 }; +@group(0) @binding(3) var pd : PadArgs; + @compute @workgroup_size(256, 1, 1) fn pad(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let rank = p._pad0; - let axis = p._pad1; - let before = i32(p._pad2); + if (i >= pd.n) { return; } + let rank = pd.rank; + let axis = pd.axis; + let before = i32(pd.before); var out_idx : array; - decode_idx(i, p.b_off, rank, &out_idx); + decode_idx(i, pd.v.v1, rank, &out_idx); var a_shape : array; var a_strides : array; - a_strides_from_shape(p.b_off + rank, rank, &a_shape, &a_strides); + a_strides_from_shape(pd.v.v1 + rank, rank, &a_shape, &a_strides); var src : u32 = 0u; var in_bounds = true; @@ -489,32 +516,36 @@ fn pad(@builtin(global_invocation_id) gid : vec3) { c = c - before; if (c < 0 || c >= i32(a_shape[d])) { in_bounds = false; } } - // Clamped even out of range: A[src] is still read below (WGSL's select - // evaluates both operands), so src must stay in bounds regardless of - // in_bounds — only the select, not the address, decides the result. + // Clamped even out of range: the source is still read below (WGSL's + // select evaluates both operands), so src must stay in bounds regardless + // of in_bounds — only the select, not the address, decides the result. let cc = clamp(c, 0, i32(a_shape[d]) - 1); src = src + u32(cc) * a_strides[d]; } - C[p.c_off + i] = select(0.0, A[p.a_off + src], in_bounds); + out2[pd.v.v2 + i] = select(0.0, in0[pd.v.v0 + src], in_bounds); } -// B layout at word offset p.b_off: [out_shape(rank-1), a_shape(rank)] — same -// per-call ring slot as pad above. a's last dim is the sliding window (size -// a_shape[rank-1]); a's `axis` dim (size a_shape[axis]) is the window count. +// Views: a, meta [out_shape(rank-1), a_shape(rank)] -> out. a's last dim is +// the sliding window (size a_shape[rank-1]); a's `axis` dim (size +// a_shape[axis]) is the window count. +// Params: webgpu.h's fold_params (n = the output's element count). +struct FoldArgs { v : Views, n : u32, rank : u32, axis : u32, step : u32 }; +@group(0) @binding(3) var fd : FoldArgs; + @compute @workgroup_size(256, 1, 1) fn fold(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let rank = p._pad0; - let axis = p._pad1; - let step = i32(p._pad2); + if (i >= fd.n) { return; } + let rank = fd.rank; + let axis = fd.axis; + let step = i32(fd.step); let out_rank = rank - 1u; var out_idx : array; - decode_idx(i, p.b_off, out_rank, &out_idx); + decode_idx(i, fd.v.v1, out_rank, &out_idx); var a_shape : array; var a_strides : array; - a_strides_from_shape(p.b_off + out_rank, rank, &a_shape, &a_strides); + a_strides_from_shape(fd.v.v1 + out_rank, rank, &a_shape, &a_strides); let win : i32 = i32(a_shape[rank - 1u]); let nwin : i32 = i32(a_shape[axis]); @@ -535,116 +566,187 @@ fn fold(@builtin(global_invocation_id) gid : vec3) { src = src + coord * a_strides[d]; } src = src + u32(k) * a_strides[rank - 1u]; - sum = sum + A[p.a_off + src]; + sum = sum + in0[fd.v.v0 + src]; } - C[p.c_off + i] = sum; + out2[fd.v.v2 + i] = sum; } // sum_to: sum `a` down to a smaller broadcast-target shape (the dual of // broadcast_to that every arithmetic op's backward uses to un-broadcast a // gradient). Gather, not scatter: one invocation per OUTPUT element sums -// every `a` element that broadcasts onto it, so -- like pad/fold above -- -// no write conflict and no atomics. B layout at word offset p.b_off: -// [a_shape(rank), a_strides(rank), acc(rank)] -- acc is a's shape -// broadcast-aligned against the output's own strides (array.h's -// broadcast_strides(target, out.strides(), a.shape())), 0 on a reduced -// axis. p._pad0 = rank, p._pad1 = reduced_n (product of a_shape over -// exactly the zero-acc axes; 1 if there are none). +// every `a` element that broadcasts onto it, so -- like pad/fold above -- no +// write conflict and no atomics. +// Views: a, meta [a_shape(rank), a_strides(rank), acc(rank)] -> out. acc is +// a's shape broadcast-aligned against the output's own strides (array.h's +// broadcast_strides(target, out.strides(), a.shape())), 0 on a reduced axis. +// Params: webgpu.h's sum_to_params (n = the output's element count; +// reduced_n = the product of a_shape over exactly the zero-acc axes, 1 if +// there are none). +struct SumToArgs { v : Views, n : u32, rank : u32, reduced_n : u32 }; +@group(0) @binding(3) var st : SumToArgs; + @compute @workgroup_size(256, 1, 1) fn sum_to(@builtin(global_invocation_id) gid : vec3) { let t = gid.x; - if (t >= p.M) { return; } - let rank = p._pad0; - let reduced_n = p._pad1; + if (t >= st.n) { return; } + let rank = st.rank; + let m = st.v.v1; var base : u32 = 0u; var red_axis : array; var red_count : u32 = 0u; for (var d : u32 = 0u; d < rank; d = d + 1u) { - let a_shape_d = bitcast(B[p.b_off + d]); - let acc_d = bitcast(B[p.b_off + 2u * rank + d]); + let a_shape_d = meta1(m, d); + let acc_d = meta1(m, 2u * rank + d); if (acc_d != 0u) { - let a_strides_d = bitcast(B[p.b_off + rank + d]); let idx = (t / acc_d) % a_shape_d; - base = base + idx * a_strides_d; + base = base + idx * meta1(m, rank + d); } else { red_axis[red_count] = d; red_count = red_count + 1u; } } var sum : f32 = 0.0; - for (var r : u32 = 0u; r < reduced_n; r = r + 1u) { + for (var r : u32 = 0u; r < st.reduced_n; r = r + 1u) { var rem = r; var off = base; for (var k : i32 = i32(red_count) - 1; k >= 0; k = k - 1) { let d = red_axis[u32(k)]; - let dim = bitcast(B[p.b_off + d]); + let dim = meta1(m, d); let coord = rem % dim; rem = rem / dim; - let stride = bitcast(B[p.b_off + rank + d]); - off = off + coord * stride; + off = off + coord * meta1(m, rank + d); } - sum = sum + A[p.a_off + off]; + sum = sum + in0[st.v.v0 + off]; } - C[p.c_off + t] = sum; + out2[st.v.v2 + t] = sum; } // concat_part: writes `a` (contiguous, one part of an N-ary Tensor.concat) -// into `out` at `p._pad1` (already before*out_strides[axis], a flat -// element offset) along one axis. One invocation per SOURCE element (this -// part's own count, p.M) -- unlike pad above, there is no zero border and -// no bounds check: concat's parts exhaustively and disjointly cover `out`, -// so this is a plain scatter that never collides across the separate -// per-part dispatches building up one `out`. B layout at word offset -// p.b_off: [a_shape(rank), out_strides(rank)] -- concat's own meta -// convention (mirrors cuda.h's upload_pad_fold_meta_), different from -// pad/fold's [out_shape, a_shape] above since this dispatch walks the -// SOURCE part's own shape, not the output's. p._pad0 = rank. +// into `out` along one axis, `shift` elements in (already before * +// out_strides[axis]). One invocation per SOURCE element -- unlike pad above, +// there is no zero border and no bounds check: concat's parts exhaustively +// and disjointly cover `out`, so this is a plain scatter that never collides +// across the separate per-part dispatches building up one `out`. +// Views: a, meta [a_shape(rank), out_strides(rank)] -> out -- walks the +// SOURCE part's own shape, not the output's, unlike pad/fold's meta. +// Params: webgpu.h's concat_params (n = this part's element count). +struct ConcatArgs { v : Views, n : u32, rank : u32, shift : u32 }; +@group(0) @binding(3) var cp : ConcatArgs; + @compute @workgroup_size(256, 1, 1) fn concat_part(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let rank = p._pad0; - let shift = p._pad1; + if (i >= cp.n) { return; } + let rank = cp.rank; var rem = i; var dst : u32 = 0u; for (var d : i32 = i32(rank) - 1; d >= 0; d = d - 1) { - let dim = bitcast(B[p.b_off + u32(d)]); + let dim = meta1(cp.v.v1, u32(d)); + let coord = rem % dim; + rem = rem / dim; + dst = dst + coord * meta1(cp.v.v1, rank + u32(d)); + } + out2[cp.v.v2 + dst + cp.shift] = in0[cp.v.v0 + i]; +} + +// ---- N-D broadcast binary (any rank) and N-D broadcast ternary select +// (Tensor.where's masking) -- the WGSL counterparts of tl_b*_nd/tl_where_nd +// in kernels/tensorlib_cuda.cu. Same flat-index decode as pad/fold above, +// against strides the host supplies (broadcast_strides(), 0 on a broadcast +// axis) rather than derived from a shape. Their meta sits at in2 and in3, +// after the tensor operands, so they decode it inline rather than through +// meta1 (WGSL has no templates to parameterize over which binding to read). + +// Views: a, b, meta [out_shape(rank), a_strides(rank), b_strides(rank)] -> +// out. Params: webgpu.h's bcast_nd_params; OP as ew_binary's. +struct BcastNdArgs { v : Views, n : u32, rank : u32, scale : f32, offset : f32 }; +@group(0) @binding(4) var bn : BcastNdArgs; + +@compute @workgroup_size(256, 1, 1) +fn ew_bcast_nd(@builtin(global_invocation_id) gid : vec3) { + let i = gid.x; + if (i >= bn.n) { return; } + let rank = bn.rank; + let base = bn.v.v2; + var rem = i; + var a_off : u32 = 0u; + var b_off : u32 = 0u; + for (var d : i32 = i32(rank) - 1; d >= 0; d = d - 1) { + let dim = bitcast(in2[base + u32(d)]); + let coord = rem % dim; + rem = rem / dim; + a_off = a_off + coord * bitcast(in2[base + rank + u32(d)]); + b_off = b_off + coord * bitcast(in2[base + 2u * rank + u32(d)]); + } + let av = in0[bn.v.v0 + a_off]; + let bv = in1[bn.v.v1 + b_off]; + out3[bn.v.v3 + i] = fma(binary_op(OP, av, bv), bn.scale, bn.offset); +} + +// Views: cond, a, b, meta [out_shape(rank), cond_strides(rank), +// a_strides(rank), b_strides(rank)] -> out. Params: webgpu.h's +// where_nd_params. +struct WhereNdArgs { v : Views, n : u32, rank : u32 }; +@group(0) @binding(5) var wn : WhereNdArgs; + +@compute @workgroup_size(256, 1, 1) +fn where_nd(@builtin(global_invocation_id) gid : vec3) { + let i = gid.x; + if (i >= wn.n) { return; } + let rank = wn.rank; + let base = wn.v.v3; + var rem = i; + var c_off : u32 = 0u; + var a_off : u32 = 0u; + var b_off : u32 = 0u; + for (var d : i32 = i32(rank) - 1; d >= 0; d = d - 1) { + let dim = bitcast(in3[base + u32(d)]); let coord = rem % dim; rem = rem / dim; - let os = bitcast(B[p.b_off + rank + u32(d)]); - dst = dst + coord * os; + c_off = c_off + coord * bitcast(in3[base + rank + u32(d)]); + a_off = a_off + coord * bitcast(in3[base + 2u * rank + u32(d)]); + b_off = b_off + coord * bitcast(in3[base + 3u * rank + u32(d)]); } - C[p.c_off + dst + shift] = A[p.a_off + i]; + let cv = in0[wn.v.v0 + c_off]; + let av = in1[wn.v.v1 + a_off]; + let bv = in2[wn.v.v2 + b_off]; + out4[wn.v.v4 + i] = select(bv, av, cv != 0.0); } -// RoPE (rotary position embedding), half-split (GPT-NeoX / HF-llama) -// convention -- mirrors cuda.h's/metal.h's own rope. A is [rows, D] -// contiguous (rows = H*T: a [H,T,D] tensor flattened, or [H,D] with -// T=1); row r's head-dim vector sits at position `pos + (r % T)`; pairs -// (j, j+D/2) rotate by angle = position * base^(-2j/D). Dispatched flat -// over rows*(D/2), one invocation per (r, j) pair. p.N = T, p.K = D, -// p._pad0 = pos, p._pad1 = D/2, p.scale = base (repurposed -- rope has -// no affine epilogue to compose with, same idea as clamp_'s lo/hi above). +// ---- RoPE (rotary position embedding), half-split (GPT-NeoX / HF-llama) +// convention -- mirrors cuda.h's/metal.h's own rope. x is [rows, D] +// contiguous (rows = H*T: a [H,T,D] tensor flattened, or [H,D] with T=1); +// row r's head-dim vector sits at position `pos + (r % T)`; pairs (j, j+D/2) +// rotate by angle = position * base^(-2j/D). Dispatched flat over +// rows*(D/2), one invocation per (r, j) pair. +// Views: x -> out. Params: webgpu.h's rope_params (n = rows * D/2). +struct RopeArgs { + v : Views, + n : u32, t : u32, d : u32, pos : u32, half : u32, + base : f32, +}; +@group(0) @binding(2) var rp : RopeArgs; + @compute @workgroup_size(256, 1, 1) fn rope(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let half = p._pad1; + if (i >= rp.n) { return; } + let half = rp.half; let r = i / half; let j = i % half; var t : u32 = 0u; - if (p.N > 0u) { t = r % p.N; } - let position = f32(p._pad0 + t); - let theta = pow(p.scale, -2.0 * f32(j) / f32(p.K)); + if (rp.t > 0u) { t = r % rp.t; } + let position = f32(rp.pos + t); + let theta = pow(rp.base, -2.0 * f32(j) / f32(rp.d)); let ang = position * theta; let c = cos(ang); let s = sin(ang); - let bi = r * p.K; - let x0 = A[p.a_off + bi + j]; - let x1 = A[p.a_off + bi + j + half]; - C[p.c_off + bi + j] = x0 * c - x1 * s; - C[p.c_off + bi + j + half] = x0 * s + x1 * c; + let bi = r * rp.d; + let x0 = in0[rp.v.v0 + bi + j]; + let x1 = in0[rp.v.v0 + bi + j + half]; + out1[rp.v.v1 + bi + j] = x0 * c - x1 * s; + out1[rp.v.v1 + bi + j + half] = x0 * s + x1 * c; } // ---- Embedding-table lookup (index_select/index_add) and pooling-style @@ -659,107 +761,55 @@ fn rope(@builtin(global_invocation_id) gid : vec3) { // per OUTPUT element, summing over every source row whose index matches it, // in place of scattering into a pre-zeroed buffer. -// A = a (table), B = idx, C = out. p._pad0 = row_size. +// Views: a (table), idx -> out. Params: gather_params. +struct GatherArgs { v : Views, row_size : u32, n : u32 }; +@group(0) @binding(3) var gs : GatherArgs; + @compute @workgroup_size(256, 1, 1) fn index_select(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let row_size = p._pad0; - let row = i / row_size; - let col = i % row_size; - let src_row = u32(B[p.b_off + row] + 0.5); - C[p.c_off + i] = A[p.a_off + src_row * row_size + col]; + if (i >= gs.n) { return; } + let row = i / gs.row_size; + let col = i % gs.row_size; + let src_row = u32(in1[gs.v.v1 + row] + 0.5); + out2[gs.v.v2 + i] = in0[gs.v.v0 + src_row * gs.row_size + col]; } -// A = idx, B = values, C = out. p._pad0 = row_size, p._pad1 = k (source rows). +// Views: idx, values -> out. Params: webgpu.h's index_add_params (n = the +// output's element count, k = the source rows to scan). +struct IndexAddArgs { v : Views, n : u32, row_size : u32, k : u32 }; +@group(0) @binding(3) var ia : IndexAddArgs; + @compute @workgroup_size(256, 1, 1) fn index_add(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let row_size = p._pad0; - let k_count = p._pad1; - let row = i / row_size; - let col = i % row_size; + if (i >= ia.n) { return; } + let row = i / ia.row_size; + let col = i % ia.row_size; var sum : f32 = 0.0; - for (var k : u32 = 0u; k < k_count; k = k + 1u) { - let idx_row = u32(A[p.a_off + k] + 0.5); + for (var k : u32 = 0u; k < ia.k; k = k + 1u) { + let idx_row = u32(in0[ia.v.v0 + k] + 0.5); if (idx_row == row) { - sum = sum + B[p.b_off + k * row_size + col]; + sum = sum + in1[ia.v.v1 + k * ia.row_size + col]; } } - C[p.c_off + i] = sum; + out2[ia.v.v2 + i] = sum; } -// A = idx, B = values, C = out. p._pad0 = size (the new trailing axis). // out[pos, k] = values[pos] where idx[pos] == k, else 0 -- gather, so no // zeroing needed, unlike CUDA's pre-zeroed scatter. -@compute @workgroup_size(256, 1, 1) -fn scatter_axis(@builtin(global_invocation_id) gid : vec3) { - let i = gid.x; - if (i >= p.M) { return; } - let size = p._pad0; - let pos = i / size; - let k = i % size; - let dst_k = u32(A[p.a_off + pos] + 0.5); - C[p.c_off + i] = select(0.0, B[p.b_off + pos], dst_k == k); -} +// Views: idx, values -> out. Params: webgpu.h's scatter_axis_params (n = the +// output's element count, size = the new trailing axis). +struct ScatterAxisArgs { v : Views, n : u32, size : u32 }; +@group(0) @binding(3) var sa : ScatterAxisArgs; -// ---- N-D broadcast binary (any rank) and N-D broadcast ternary select -// (Tensor.where's masking) -- the WGSL counterparts of tl_b*_nd/tl_where_nd -// in kernels/tensorlib_cuda.cu. Same flat-index decode as pad/fold above, -// against strides the host supplies (broadcast_strides(), 0 on a broadcast -// axis) rather than derived from a shape -- inlined per kernel rather than -// shared via decode_idx, since that helper always reads from B and here the -// meta buffer is D or E (WGSL has no templates to parameterize over which -// storage binding to read). - -// A = a, B = b, D = meta [out_shape(rank), a_strides(rank), b_strides(rank)], -// E unused. p._pad0 = rank, p._pad3 = meta's word offset into D. @compute @workgroup_size(256, 1, 1) -fn ew_bcast_nd(@builtin(global_invocation_id) gid : vec3) { - let i = gid.x; - if (i >= p.M) { return; } - let rank = p._pad0; - let base = p._pad3; - var rem = i; - var a_off : u32 = 0u; - var b_off : u32 = 0u; - for (var d : i32 = i32(rank) - 1; d >= 0; d = d - 1) { - let dim = bitcast(D[base + u32(d)]); - let coord = rem % dim; - rem = rem / dim; - a_off = a_off + coord * bitcast(D[base + rank + u32(d)]); - b_off = b_off + coord * bitcast(D[base + 2u * rank + u32(d)]); - } - let av = A[p.a_off + a_off]; - let bv = B[p.b_off + b_off]; - C[p.c_off + i] = fma(binary_op(p.op, av, bv), p.scale, p.offset); -} - -// A = cond, B = a, D = b, E = meta [out_shape(rank), cond_strides(rank), -// a_strides(rank), b_strides(rank)]. p._pad0 = rank, p._pad3 = b's element -// offset into D, p._pad4 = meta's word offset into E. -@compute @workgroup_size(256, 1, 1) -fn where_nd(@builtin(global_invocation_id) gid : vec3) { +fn scatter_axis(@builtin(global_invocation_id) gid : vec3) { let i = gid.x; - if (i >= p.M) { return; } - let rank = p._pad0; - let base = p._pad4; - var rem = i; - var c_off : u32 = 0u; - var a_off : u32 = 0u; - var b_off : u32 = 0u; - for (var d : i32 = i32(rank) - 1; d >= 0; d = d - 1) { - let dim = bitcast(E[base + u32(d)]); - let coord = rem % dim; - rem = rem / dim; - c_off = c_off + coord * bitcast(E[base + rank + u32(d)]); - a_off = a_off + coord * bitcast(E[base + 2u * rank + u32(d)]); - b_off = b_off + coord * bitcast(E[base + 3u * rank + u32(d)]); - } - let cv = A[p.a_off + c_off]; - let av = B[p.b_off + a_off]; - let bv = D[p._pad3 + b_off]; - C[p.c_off + i] = select(bv, av, cv != 0.0); + if (i >= sa.n) { return; } + let pos = i / sa.size; + let k = i % sa.size; + let dst_k = u32(in0[sa.v.v0 + pos] + 0.5); + out2[sa.v.v2 + i] = select(0.0, in1[sa.v.v1 + pos], dst_k == k); } )WGSL"