custom shader webgpu

additional/dom-webgpu.js

@@ -413,16 +413,18 @@ const hdrBG = {
 };
 
 // Group 1 changes between dispatches because `out` and `prev` swap on the
-// ping-pong. Cache by current out-is-ping bool + pipeline (each pipeline
-// has its own bgl1 instance even if the layout entries are identical).
-function getGroup1BG(pipe) {
-    const key = `${pipe.bgl1.label || ""}/${state.outIsPing ? 1 : 0}/${state.width}x${state.height}`;
+// ping-pong. Cached by current ping-pong direction and texture size; the
+// stored bind group is reusable across all pipelines that share a
+// layout-compatible bgl1 (all standard pipelines and custom shaders do,
+// since they declare identical group-1 entries per the contract).
+function getGroup1BG(bgl1) {
+    const key = `g1/${state.outIsPing ? 1 : 0}/${state.width}x${state.height}`;
     let bg = state.bindGroupCache.get(key);
     if (bg) return bg;
     const outView  = state.outIsPing ? state.pingView : state.pongView;
     const prevView = state.outIsPing ? state.pongView : state.pingView;
     bg = device.createBindGroup({
-        layout: pipe.bgl1,
+        layout: bgl1,
         entries: [
             { binding: 0, resource: outView },
             { binding: 1, resource: prevView },
@@ -662,7 +664,7 @@ function dispatchStandard(pipe, hdrBindGroup, headerPtr, gx, gy, itemsHandle, gr
     const off = writeHeader(headerPtr);
     state.pass.setPipeline(pipe.pipeline);
     state.pass.setBindGroup(0, hdrBindGroup, [off]);
-    state.pass.setBindGroup(1, getGroup1BG(pipe));
+    state.pass.setBindGroup(1, getGroup1BG(pipe.bgl1));
     state.pass.setBindGroup(2, getGroup2BG(pipe, itemsHandle));
     if (group3) state.pass.setBindGroup(3, group3);
     state.pass.dispatchWorkgroups(gx, gy, 1);
@@ -686,6 +688,136 @@ env.wgpuDispatchText = (itemsHandle, headerPtr, gx, gy, atlasHandle, sampHandle)
     dispatchStandard(pipeText, hdrBG.text, headerPtr, gx, gy, itemsHandle, g3);
 };
 
+// ─── custom user-authored shaders ─────────────────────────────────────
+//
+// Bind-group contract (mirrors :WebGPUComputeShader.cppm):
+//   group 0 binding 0 — uniform UIDispatchHeader (dynamic offset, 48b)
+//   group 1 binding 0 — texture_storage_2d<rgba8unorm, write> out
+//   group 1 binding 1 — texture_2d<f32>                     prev
+//   group 2+          — user-declared (UICustomBinding entries)
+//
+// Each UICustomBinding entry on the wasm side is 8 bytes:
+//   u8 group, u8 binding, u8 kind, u8 pad, u32 pushOffset
+// kind: 0 = read-only-storage SSBO, 1 = sampled tex 2d, 2 = filtering sampler.
+
+const customPipelines = new Map();   // handle → { pipeline, bgls, hdrBG, byGroup }
+
+env.wgpuLoadCustomShader = (wgslPtr, wgslLen, bindingsPtr, bindingsCount) => {
+    const wgsl = new TextDecoder().decode(memU8().subarray(wgslPtr, wgslPtr + wgslLen));
+    const bindings = [];
+    const dv = new DataView(memU8().buffer, bindingsPtr, bindingsCount * 8);
+    for (let i = 0; i < bindingsCount; i++) {
+        bindings.push({
+            group: dv.getUint8(i*8 + 0),
+            binding: dv.getUint8(i*8 + 1),
+            kind: dv.getUint8(i*8 + 2),
+            pushOffset: dv.getUint32(i*8 + 4, true),
+        });
+    }
+
+    // Group bindings by @group(N) for layout creation.
+    const byGroup = new Map();
+    for (const b of bindings) {
+        if (b.group < 2) {
+            console.error(`[crafter-wgpu] custom shader: @group(${b.group}) reserved; use groups >= 2`);
+            return 0;
+        }
+        if (!byGroup.has(b.group)) byGroup.set(b.group, []);
+        byGroup.get(b.group).push(b);
+    }
+
+    // Group 0 = header uniform, Group 1 = ping-pong out+prev — always injected.
+    const bgls = [
+        device.createBindGroupLayout({ entries: [
+            { binding: 0, visibility: GPUShaderStage.COMPUTE,
+              buffer: { type: "uniform", hasDynamicOffset: true, minBindingSize: 48 } },
+        ]}),
+        device.createBindGroupLayout({ entries: [
+            { binding: 0, visibility: GPUShaderStage.COMPUTE,
+              storageTexture: { format: "rgba8unorm", access: "write-only", viewDimension: "2d" } },
+            { binding: 1, visibility: GPUShaderStage.COMPUTE,
+              texture: { sampleType: "float", viewDimension: "2d" } },
+        ]}),
+    ];
+    // Sorted custom groups. Pad any gaps with empty bgls (WebGPU pipeline
+    // layouts require a contiguous array of GPUBindGroupLayout per group
+    // index up to the highest used).
+    const sortedGroups = [...byGroup.keys()].sort((a, b) => a - b);
+    const highest = sortedGroups.length ? sortedGroups[sortedGroups.length - 1] : 1;
+    for (let g = 2; g <= highest; g++) {
+        if (byGroup.has(g)) {
+            const entries = byGroup.get(g).map(b => {
+                const e = { binding: b.binding, visibility: GPUShaderStage.COMPUTE };
+                if      (b.kind === 0) e.buffer = { type: "read-only-storage" };
+                else if (b.kind === 1) e.texture = { sampleType: "float", viewDimension: "2d" };
+                else if (b.kind === 2) e.sampler = { type: "filtering" };
+                return e;
+            });
+            bgls.push(device.createBindGroupLayout({ entries }));
+        } else {
+            bgls.push(device.createBindGroupLayout({ entries: [] }));
+        }
+    }
+
+    let pipeline;
+    try {
+        const mod = device.createShaderModule({ code: wgsl });
+        const layout = device.createPipelineLayout({ bindGroupLayouts: bgls });
+        pipeline = device.createComputePipeline({ layout, compute: { module: mod, entryPoint: "main" } });
+    } catch (e) {
+        console.error("[crafter-wgpu] custom shader compile failed:", e);
+        return 0;
+    }
+
+    const hdrBG = device.createBindGroup({
+        layout: bgls[0],
+        entries: [{ binding: 0, resource: { buffer: state.headerRing, offset: 0, size: 48 } }],
+    });
+
+    const handle = newHandle();
+    customPipelines.set(handle, { pipeline, bgls, hdrBG, byGroup, sortedGroups });
+    return handle;
+};
+
+env.wgpuDispatchCustom = (pipelineHandle, pushPtr, pushBytes, handlesPtr, handlesCount,
+                          gx, gy, gz) => {
+    state.dispatchCustomCount = (state.dispatchCustomCount || 0) + 1;
+    if (!state.pass) return;
+    const pipe = customPipelines.get(pipelineHandle);
+    if (!pipe) {
+        console.error("[crafter-wgpu] wgpuDispatchCustom: unknown pipeline", pipelineHandle);
+        return;
+    }
+
+    // Write header (first 48 bytes of push).
+    const off = writeHeader(pushPtr);
+
+    state.pass.setPipeline(pipe.pipeline);
+    state.pass.setBindGroup(0, pipe.hdrBG, [off]);
+    state.pass.setBindGroup(1, getGroup1BG(pipe.bgls[1]));
+
+    // Walk bindings in declaration order and assemble bind groups.
+    // handles[] from wasm is in the SAME order as customBindings, so we
+    // pick up indices by walking byGroup in the same sorted order.
+    const handles = new Uint32Array(memU8().buffer, handlesPtr, handlesCount);
+    let handleIdx = 0;
+    for (const g of pipe.sortedGroups) {
+        const entries = pipe.byGroup.get(g).map(b => {
+            const h = handles[handleIdx++];
+            let resource;
+            if      (b.kind === 0) resource = { buffer: buffers.get(h) };
+            else if (b.kind === 1) resource = textureViews.get(h);
+            else if (b.kind === 2) resource = samplers.get(h);
+            return { binding: b.binding, resource };
+        });
+        const bg = device.createBindGroup({ layout: pipe.bgls[g], entries });
+        state.pass.setBindGroup(g, bg);
+    }
+
+    state.pass.dispatchWorkgroups(gx, gy, gz);
+    state.outIsPing = !state.outIsPing;
+};
+
 // Debug accessor for browser-console diagnostics.
 window.crafter_wgpu_state = state;
 window.crafter_wgpu_device = device;
@@ -708,6 +840,8 @@ window.crafter_wgpu_debug = () => ({
     lastWriteSize: state.lastWriteSize,
 });
 
+window.crafter_wgpu_bufferKeys = () => [...buffers.keys()];
+
 // Read back the first QuadItem from a registered buffer to verify the
 // GPU sees what the CPU wrote.
 window.crafter_wgpu_readBuffer = async (handle, byteSize = 64) => {