oidn-web 0.1.0

package/lib/WGPUComputePass.js

@@ -0,0 +1,220 @@
+function isStorageParamsEqual(params, other) {
+    return params.channels === other.channels;
+}
+export const WORKGROUP_SIZE = 8;
+export class WGPUComputePass {
+    _label;
+    _device;
+    _outputBuffers = {};
+    _pipeline;
+    _bindGroups = [];
+    _needsUpdatePipeline = true;
+    _inputs = [];
+    _outputs = [];
+    _uniforms = [];
+    _uniformBuffers = {};
+    _width = 10;
+    _height = 10;
+    _execWidth;
+    _execHeight;
+    _csCode = '';
+    _csMain;
+    _csDefine;
+    constructor(label, device, opts) {
+        this._label = label;
+        this._device = device;
+        this._csMain = opts.csMain;
+        this._csDefine = opts.csDefine;
+        this._inputs = opts.inputs;
+        this._outputs = opts.outputs;
+        this._uniforms = opts.uniforms;
+        opts.uniforms.forEach((uniform) => {
+            this._uniformBuffers[uniform.label] = device.createBuffer({
+                label: this._label,
+                size: uniform.data.byteLength,
+                usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
+            });
+            this._device.queue.writeBuffer(this._uniformBuffers[uniform.label], 0, uniform.data);
+        });
+    }
+    setSize(width, height) {
+        width = Math.ceil(width);
+        height = Math.ceil(height);
+        const sizeChanged = width !== this._width || height !== this._height;
+        this._width = width;
+        this._height = height;
+        if (sizeChanged) {
+            this._resizeOutputBuffers();
+            this._needsUpdatePipeline = true;
+        }
+    }
+    setExecuteSize(width, height) {
+        width = Math.ceil(width);
+        height = Math.ceil(height);
+        this._execWidth = width;
+        this._execHeight = height;
+    }
+    setOutputParams(outputParams) {
+        this._updateOutputBuffers(outputParams);
+        this._needsUpdatePipeline = true;
+    }
+    setUniform(label, data) {
+        const buffer = this._uniformBuffers[label];
+        this._device.queue.writeBuffer(buffer, 0, data);
+    }
+    getOutputBuffer(name) {
+        return this._outputBuffers[name].buffer;
+    }
+    dispose() {
+        Object.keys(this._uniformBuffers).forEach((key) => {
+            this._uniformBuffers[key].destroy();
+        });
+        Object.keys(this._outputBuffers).forEach((key) => {
+            this._outputBuffers[key].texture.destroy();
+        });
+    }
+    _createBuffer(params) {
+        // const byteLength = this._width * this._height * params.channels * 4;
+        // Buffer data needs to be aligned with 8, 16
+        const byteLength = this._width * this._height * 4 * 4;
+        return this._device.createBuffer({
+            label: this._label,
+            // webgpu needs buffer at least 80 bytes.
+            size: Math.max(byteLength, 80),
+            usage: GPUBufferUsage.STORAGE |
+                GPUBufferUsage.COPY_DST |
+                GPUBufferUsage.COPY_SRC
+        });
+    }
+    _resizeOutputBuffers() {
+        const outputBuffers = this._outputBuffers;
+        for (const key in outputBuffers) {
+            const { buffer, params } = outputBuffers[key];
+            buffer.destroy();
+            outputBuffers[key].buffer = this._createBuffer(params);
+        }
+    }
+    _updateOutputBuffers(outputParams) {
+        const outputBuffers = this._outputBuffers;
+        for (const key in outputParams) {
+            const params = outputParams[key];
+            if (!isStorageParamsEqual(params, outputBuffers[key]?.params || {})) {
+                outputBuffers[key]?.buffer.destroy();
+                const buffer = this._createBuffer(params);
+                outputBuffers[key] = {
+                    buffer,
+                    params
+                };
+            }
+        }
+    }
+    _updatePipeline(inputParams) {
+        if (!this._needsUpdatePipeline) {
+            return;
+        }
+        this._needsUpdatePipeline = false;
+        const device = this._device;
+        const csCode = this._getFullCs(inputParams);
+        if (csCode === this._csCode) {
+            return;
+        }
+        this._csCode = csCode;
+        this._pipeline = device.createComputePipeline({
+            label: this._label,
+            layout: 'auto',
+            compute: {
+                module: device.createShaderModule({
+                    label: this._label,
+                    code: csCode
+                }),
+                entryPoint: 'main'
+            }
+        });
+        this._updateBindGroups();
+    }
+    _getFullCs(inputParams) {
+        const inputs = this._inputs;
+        const hasInputs = inputs.length > 0;
+        const cs = `
+${inputs
+            .sort()
+            .map((bufferName, idx) => 
+        // TODO more channels option.
+        `@group(0) @binding(${idx}) var<storage, read> in_${bufferName}: array<vec${inputParams[bufferName].channels}f>;`)
+            .join('\n')}
+${this._uniforms
+            .map((uniform, idx) => `@group(${hasInputs ? 1 : 0}) @binding(${idx}) var<uniform> ${uniform.label}: ${uniform.type};`)
+            .join('\n')}
+
+${this._outputs
+            .map((name, idx) => `@group(${hasInputs ? 2 : 1}) @binding(${idx}) var<storage, read_write> out_${name}: array<vec${this._outputBuffers[name].params.channels}f>;`)
+            .join('\n')}
+${this._csDefine ?? ''}
+@compute @workgroup_size(${WORKGROUP_SIZE}, ${WORKGROUP_SIZE}, 1)
+fn main(@builtin(global_invocation_id) globalId: vec3u) {
+${this._csMain}
+}
+`;
+        return cs;
+    }
+    _updateBindGroups() {
+        const bindGroups = [];
+        const device = this._device;
+        //TODO
+        const uniformBindGroupIndex = this._inputs.length > 0 ? 1 : 0;
+        if (this._uniforms.length > 0) {
+            bindGroups[uniformBindGroupIndex] = device.createBindGroup({
+                label: this._label,
+                layout: this._pipeline.getBindGroupLayout(uniformBindGroupIndex),
+                entries: this._uniforms.map((uniform, idx) => ({
+                    binding: idx,
+                    resource: {
+                        buffer: this._uniformBuffers[uniform.label]
+                    }
+                }))
+            });
+        }
+        this._bindGroups = bindGroups;
+    }
+    createPass(commandEncoder, inputBuffers) {
+        this._updatePipeline(inputBuffers);
+        const hasInputs = this._inputs.length > 0;
+        // TODO createBindGroup every time?
+        if (hasInputs) {
+            this._bindGroups[0] = this._device.createBindGroup({
+                label: this._label,
+                layout: this._pipeline.getBindGroupLayout(0),
+                entries: this._inputs.map((bufferName, idx) => ({
+                    binding: idx,
+                    // TODO
+                    resource: {
+                        buffer: inputBuffers[bufferName].buffer
+                    }
+                }))
+            });
+        }
+        // Outputs
+        this._bindGroups[hasInputs ? 2 : 1] = this._device.createBindGroup({
+            label: this._label,
+            layout: this._pipeline.getBindGroupLayout(hasInputs ? 2 : 1),
+            entries: this._outputs.map((bufferName, idx) => ({
+                binding: idx,
+                resource: {
+                    buffer: this._outputBuffers[bufferName].buffer
+                }
+            }))
+        });
+        // Begin the render pass
+        const computePass = commandEncoder.beginComputePass();
+        // Draw a full quad
+        computePass.setPipeline(this._pipeline);
+        // Bind groups
+        this._bindGroups.forEach((bindGroup, idx) => {
+            computePass.setBindGroup(idx, bindGroup);
+        });
+        computePass.dispatchWorkgroups(Math.ceil((this._execWidth ?? this._width) / WORKGROUP_SIZE), Math.ceil((this._execHeight ?? this._height) / WORKGROUP_SIZE), 1);
+        // End the render pass
+        computePass.end();
+    }
+}
+//# sourceMappingURL=WGPUComputePass.js.map

package/lib/WGPUComputePass.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"WGPUComputePass.js","sourceRoot":"","sources":["../src/WGPUComputePass.ts"],"names":[],"mappings":"AAAA,SAAS,oBAAoB,CAC3B,MAA6B,EAC7B,KAA4B;IAE5B,OAAO,MAAM,CAAC,QAAQ,KAAK,KAAK,CAAC,QAAQ,CAAC;AAC5C,CAAC;AAED,MAAM,CAAC,MAAM,cAAc,GAAG,CAAC,CAAC;AAgBhC,MAAM,OAAO,eAAe;IAClB,MAAM,CAAC;IAEP,OAAO,CAAC;IACR,cAAc,GAMlB,EAAE,CAAC;IAEC,SAAS,CAAsB;IAC/B,WAAW,GAAmB,EAAE,CAAC;IACjC,oBAAoB,GAAG,IAAI,CAAC;IAE5B,OAAO,GAAa,EAAE,CAAC;IACvB,QAAQ,GAAa,EAAE,CAAC;IACxB,SAAS,GAAc,EAAE,CAAC;IAC1B,eAAe,GAA8B,EAAE,CAAC;IAEhD,MAAM,GAAG,EAAE,CAAC;IACZ,OAAO,GAAG,EAAE,CAAC;IAEb,UAAU,CAAU;IACpB,WAAW,CAAU;IAErB,OAAO,GAAG,EAAE,CAAC;IACb,OAAO,CAAC;IACR,SAAS,CAAC;IAElB,YACE,KAAa,EACb,MAAiB,EACjB,IAMC;QAED,IAAI,CAAC,MAAM,GAAG,KAAK,CAAC;QACpB,IAAI,CAAC,OAAO,GAAG,MAAM,CAAC;QACtB,IAAI,CAAC,OAAO,GAAG,IAAI,CAAC,MAAM,CAAC;QAC3B,IAAI,CAAC,SAAS,GAAG,IAAI,CAAC,QAAQ,CAAC;QAC/B,IAAI,CAAC,OAAO,GAAG,IAAI,CAAC,MAAM,CAAC;QAC3B,IAAI,CAAC,QAAQ,GAAG,IAAI,CAAC,OAAO,CAAC;QAC7B,IAAI,CAAC,SAAS,GAAG,IAAI,CAAC,QAAQ,CAAC;QAE/B,IAAI,CAAC,QAAQ,CAAC,OAAO,CAAC,CAAC,OAAO,EAAE,EAAE;YAChC,IAAI,CAAC,eAAe,CAAC,OAAO,CAAC,KAAK,CAAC,GAAG,MAAM,CAAC,YAAY,CAAC;gBACxD,KAAK,EAAE,IAAI,CAAC,MAAM;gBAClB,IAAI,EAAE,OAAO,CAAC,IAAI,CAAC,UAAU;gBAC7B,KAAK,EAAE,cAAc,CAAC,OAAO,GAAG,cAAc,CAAC,QAAQ;aACxD,CAAC,CAAC;YACH,IAAI,CAAC,OAAO,CAAC,KAAK,CAAC,WAAW,CAC5B,IAAI,CAAC,eAAe,CAAC,OAAO,CAAC,KAAK,CAAC,EACnC,CAAC,EACD,OAAO,CAAC,IAAI,CACb,CAAC;QACJ,CAAC,CAAC,CAAC;IACL,CAAC;IAED,OAAO,CAAC,KAAa,EAAE,MAAc;QACnC,KAAK,GAAG,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;QACzB,MAAM,GAAG,IAAI,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;QAC3B,MAAM,WAAW,GAAG,KAAK,KAAK,IAAI,CAAC,MAAM,IAAI,MAAM,KAAK,IAAI,CAAC,OAAO,CAAC;QACrE,IAAI,CAAC,MAAM,GAAG,KAAK,CAAC;QACpB,IAAI,CAAC,OAAO,GAAG,MAAM,CAAC;QACtB,IAAI,WAAW,EAAE,CAAC;YAChB,IAAI,CAAC,oBAAoB,EAAE,CAAC;YAC5B,IAAI,CAAC,oBAAoB,GAAG,IAAI,CAAC;QACnC,CAAC;IACH,CAAC;IAED,cAAc,CAAC,KAAa,EAAE,MAAc;QAC1C,KAAK,GAAG,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;QACzB,MAAM,GAAG,IAAI,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;QAC3B,IAAI,CAAC,UAAU,GAAG,KAAK,CAAC;QACxB,IAAI,CAAC,WAAW,GAAG,MAAM,CAAC;IAC5B,CAAC;IAED,eAAe,CAAC,YAA8C;QAC5D,IAAI,CAAC,oBAAoB,CAAC,YAAY,CAAC,CAAC;QACxC,IAAI,CAAC,oBAAoB,GAAG,IAAI,CAAC;IACnC,CAAC;IAED,UAAU,CAAC,KAAa,EAAE,IAA6C;QACrE,MAAM,MAAM,GAAG,IAAI,CAAC,eAAe,CAAC,KAAK,CAAC,CAAC;QAC3C,IAAI,CAAC,OAAO,CAAC,KAAK,CAAC,WAAW,CAAC,MAAM,EAAE,CAAC,EAAE,IAAI,CAAC,CAAC;IAClD,CAAC;IAED,eAAe,CAAC,IAAO;QACrB,OAAO,IAAI,CAAC,cAAc,CAAC,IAAI,CAAC,CAAC,MAAM,CAAC;IAC1C,CAAC;IAED,OAAO;QACL,MAAM,CAAC,IAAI,CAAC,IAAI,CAAC,eAAe,CAAC,CAAC,OAAO,CAAC,CAAC,GAAG,EAAE,EAAE;YAC/C,IAAI,CAAC,eAAuB,CAAC,GAAG,CAAC,CAAC,OAAO,EAAE,CAAC;QAC/C,CAAC,CAAC,CAAC;QACH,MAAM,CAAC,IAAI,CAAC,IAAI,CAAC,cAAc,CAAC,CAAC,OAAO,CAAC,CAAC,GAAG,EAAE,EAAE;YAC9C,IAAI,CAAC,cAAsB,CAAC,GAAG,CAAC,CAAC,OAAO,CAAC,OAAO,EAAE,CAAC;QACtD,CAAC,CAAC,CAAC;IACL,CAAC;IAEO,aAAa,CAAC,MAA6B;QACjD,uEAAuE;QACvE,6CAA6C;QAC7C,MAAM,UAAU,GAAG,IAAI,CAAC,MAAM,GAAG,IAAI,CAAC,OAAO,GAAG,CAAC,GAAG,CAAC,CAAC;QACtD,OAAO,IAAI,CAAC,OAAO,CAAC,YAAY,CAAC;YAC/B,KAAK,EAAE,IAAI,CAAC,MAAM;YAClB,yCAAyC;YACzC,IAAI,EAAE,IAAI,CAAC,GAAG,CAAC,UAAU,EAAE,EAAE,CAAC;YAC9B,KAAK,EACH,cAAc,CAAC,OAAO;gBACtB,cAAc,CAAC,QAAQ;gBACvB,cAAc,CAAC,QAAQ;SAC1B,CAAC,CAAC;IACL,CAAC;IAEO,oBAAoB;QAC1B,MAAM,aAAa,GAAG,IAAI,CAAC,cAAc,CAAC;QAC1C,KAAK,MAAM,GAAG,IAAI,aAAa,EAAE,CAAC;YAChC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,GAAG,aAAa,CAAC,GAAG,CAAC,CAAC;YAC9C,MAAM,CAAC,OAAO,EAAE,CAAC;YACjB,aAAa,CAAC,GAAG,CAAC,CAAC,MAAM,GAAG,IAAI,CAAC,aAAa,CAAC,MAAM,CAAC,CAAC;QACzD,CAAC;IACH,CAAC;IAEO,oBAAoB,CAC1B,YAAmD;QAEnD,MAAM,aAAa,GAAG,IAAI,CAAC,cAAc,CAAC;QAC1C,KAAK,MAAM,GAAG,IAAI,YAAY,EAAE,CAAC;YAC/B,MAAM,MAAM,GAAG,YAAY,CAAC,GAAG,CAAC,CAAC;YACjC,IACE,CAAC,oBAAoB,CACnB,MAAM,EACN,aAAa,CAAC,GAAG,CAAC,EAAE,MAAM,IAAK,EAA4B,CAC5D,EACD,CAAC;gBACD,aAAa,CAAC,GAAG,CAAC,EAAE,MAAM,CAAC,OAAO,EAAE,CAAC;gBACrC,MAAM,MAAM,GAAG,IAAI,CAAC,aAAa,CAAC,MAAM,CAAC,CAAC;gBAC1C,aAAa,CAAC,GAAG,CAAC,GAAG;oBACnB,MAAM;oBACN,MAAM;iBACP,CAAC;YACJ,CAAC;QACH,CAAC;IACH,CAAC;IAEO,eAAe,CAAC,WAAiD;QACvE,IAAI,CAAC,IAAI,CAAC,oBAAoB,EAAE,CAAC;YAC/B,OAAO;QACT,CAAC;QACD,IAAI,CAAC,oBAAoB,GAAG,KAAK,CAAC;QAClC,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,CAAC;QAC5B,MAAM,MAAM,GAAG,IAAI,CAAC,UAAU,CAAC,WAAW,CAAC,CAAC;QAC5C,IAAI,MAAM,KAAK,IAAI,CAAC,OAAO,EAAE,CAAC;YAC5B,OAAO;QACT,CAAC;QACD,IAAI,CAAC,OAAO,GAAG,MAAM,CAAC;QACtB,IAAI,CAAC,SAAS,GAAG,MAAM,CAAC,qBAAqB,CAAC;YAC5C,KAAK,EAAE,IAAI,CAAC,MAAM;YAClB,MAAM,EAAE,MAAM;YACd,OAAO,EAAE;gBACP,MAAM,EAAE,MAAM,CAAC,kBAAkB,CAAC;oBAChC,KAAK,EAAE,IAAI,CAAC,MAAM;oBAClB,IAAI,EAAE,MAAM;iBACb,CAAC;gBACF,UAAU,EAAE,MAAM;aACnB;SACF,CAAC,CAAC;QACH,IAAI,CAAC,iBAAiB,EAAE,CAAC;IAC3B,CAAC;IAEO,UAAU,CAAC,WAAiD;QAClE,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,CAAC;QAC5B,MAAM,SAAS,GAAG,MAAM,CAAC,MAAM,GAAG,CAAC,CAAC;QACpC,MAAM,EAAE,GAAG;EACb,MAAM;aACL,IAAI,EAAE;aACN,GAAG,CACF,CAAC,UAAU,EAAE,GAAG,EAAE,EAAE;QAClB,6BAA6B;QAC7B,sBAAsB,GAAG,2BAA2B,UAAU,cAAc,WAAW,CAAC,UAAU,CAAC,CAAC,QAAQ,KAAK,CACpH;aACA,IAAI,CAAC,IAAI,CAAC;EACX,IAAI,CAAC,SAAS;aACb,GAAG,CACF,CAAC,OAAO,EAAE,GAAG,EAAE,EAAE,CACf,UAAU,SAAS,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,cAAc,GAAG,kBAC1C,OAAO,CAAC,KACV,KAAK,OAAO,CAAC,IAAI,GAAG,CACvB;aACA,IAAI,CAAC,IAAI,CAAC;;EAEX,IAAI,CAAC,QAAQ;aACZ,GAAG,CACF,CAAC,IAAI,EAAE,GAAG,EAAE,EAAE,CACZ,UACE,SAAS,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAClB,cAAc,GAAG,kCAAkC,IAAI,cACrD,IAAI,CAAC,cAAc,CAAC,IAAI,CAAC,CAAC,MAAM,CAAC,QACnC,KAAK,CACR;aACA,IAAI,CAAC,IAAI,CAAC;EACX,IAAI,CAAC,SAAS,IAAI,EAAE;2BACK,cAAc,KAAK,cAAc;;EAE1D,IAAI,CAAC,OAAO;;CAEb,CAAC;QAEE,OAAO,EAAE,CAAC;IACZ,CAAC;IAEO,iBAAiB;QACvB,MAAM,UAAU,GAAmB,EAAE,CAAC;QACtC,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,CAAC;QAE5B,MAAM;QACN,MAAM,qBAAqB,GAAG,IAAI,CAAC,OAAO,CAAC,MAAM,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QAC9D,IAAI,IAAI,CAAC,SAAS,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;YAC9B,UAAU,CAAC,qBAAqB,CAAC,GAAG,MAAM,CAAC,eAAe,CAAC;gBACzD,KAAK,EAAE,IAAI,CAAC,MAAM;gBAClB,MAAM,EAAE,IAAI,CAAC,SAAS,CAAC,kBAAkB,CAAC,qBAAqB,CAAC;gBAChE,OAAO,EAAE,IAAI,CAAC,SAAS,CAAC,GAAG,CACzB,CAAC,OAAO,EAAE,GAAG,EAAE,EAAE,CACf,CAAC;oBACC,OAAO,EAAE,GAAG;oBACZ,QAAQ,EAAE;wBACR,MAAM,EAAE,IAAI,CAAC,eAAe,CAAC,OAAO,CAAC,KAAK,CAAC;qBAC5C;iBACoB,CAAA,CAC1B;aACF,CAAC,CAAC;QACL,CAAC;QAED,IAAI,CAAC,WAAW,GAAG,UAAU,CAAC;IAChC,CAAC;IAED,UAAU,CACR,cAAiC,EACjC,YAA6C;QAE7C,IAAI,CAAC,eAAe,CAAC,YAAY,CAAC,CAAC;QAEnC,MAAM,SAAS,GAAG,IAAI,CAAC,OAAO,CAAC,MAAM,GAAG,CAAC,CAAC;QAC1C,mCAAmC;QACnC,IAAI,SAAS,EAAE,CAAC;YACd,IAAI,CAAC,WAAW,CAAC,CAAC,CAAC,GAAG,IAAI,CAAC,OAAO,CAAC,eAAe,CAAC;gBACjD,KAAK,EAAE,IAAI,CAAC,MAAM;gBAClB,MAAM,EAAE,IAAI,CAAC,SAAS,CAAC,kBAAkB,CAAC,CAAC,CAAC;gBAC5C,OAAO,EAAE,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,CAAC,UAAU,EAAE,GAAG,EAAE,EAAE,CAAC,CAAC;oBAC9C,OAAO,EAAE,GAAG;oBACZ,OAAO;oBACP,QAAQ,EAAE;wBACR,MAAM,EAAE,YAAY,CAAC,UAAe,CAAC,CAAC,MAAM;qBAC7C;iBACF,CAAC,CAAC;aACJ,CAAC,CAAC;QACL,CAAC;QAED,UAAU;QACV,IAAI,CAAC,WAAW,CAAC,SAAS,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,GAAG,IAAI,CAAC,OAAO,CAAC,eAAe,CAAC;YACjE,KAAK,EAAE,IAAI,CAAC,MAAM;YAClB,MAAM,EAAE,IAAI,CAAC,SAAS,CAAC,kBAAkB,CAAC,SAAS,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;YAC5D,OAAO,EAAE,IAAI,CAAC,QAAQ,CAAC,GAAG,CAAC,CAAC,UAAU,EAAE,GAAG,EAAE,EAAE,CAAC,CAAC;gBAC/C,OAAO,EAAE,GAAG;gBACZ,QAAQ,EAAE;oBACR,MAAM,EAAE,IAAI,CAAC,cAAc,CAAC,UAAU,CAAC,CAAC,MAAM;iBAC/C;aACF,CAAC,CAAC;SACJ,CAAC,CAAC;QAEH,wBAAwB;QACxB,MAAM,WAAW,GAAG,cAAc,CAAC,gBAAgB,EAAE,CAAC;QAEtD,mBAAmB;QACnB,WAAW,CAAC,WAAW,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;QACxC,cAAc;QACd,IAAI,CAAC,WAAW,CAAC,OAAO,CAAC,CAAC,SAAS,EAAE,GAAG,EAAE,EAAE;YAC1C,WAAW,CAAC,YAAY,CAAC,GAAG,EAAE,SAAS,CAAC,CAAC;QAC3C,CAAC,CAAC,CAAC;QACH,WAAW,CAAC,kBAAkB,CAC5B,IAAI,CAAC,IAAI,CAAC,CAAC,IAAI,CAAC,UAAU,IAAI,IAAI,CAAC,MAAM,CAAC,GAAG,cAAc,CAAC,EAC5D,IAAI,CAAC,IAAI,CAAC,CAAC,IAAI,CAAC,WAAW,IAAI,IAAI,CAAC,OAAO,CAAC,GAAG,cAAc,CAAC,EAC9D,CAAC,CACF,CAAC;QACF,sBAAsB;QACtB,WAAW,CAAC,GAAG,EAAE,CAAC;IACpB,CAAC;CACF"}

package/lib/WGPUFullQuadPass.d.ts

@@ -0,0 +1,51 @@
+/// <reference types="@webgpu/types/dist" />
+export declare const fullScreenQuadVertexShaderWGSL = "\n@vertex\nfn main(\n  @builtin(vertex_index) vertexIndex: u32\n) -> @builtin(position) vec4f {\n  const pos = array(\n    vec2(-1.0, -1.0), vec2(1.0, -1.0), vec2(-1.0, 1.0),\n    vec2(-1.0, 1.0), vec2(1.0, -1.0), vec2(1.0, 1.0),\n  );\n\n  return vec4<f32>(pos[vertexIndex], 0.0, 1.0);\n}\n";
+export interface Uniform {
+    label: string;
+    type: string;
+    data: Float32Array | Int32Array | Uint32Array;
+}
+export interface WGPUFullQuadPassOutput {
+    format: GPUTextureFormat;
+}
+export declare class WGPUFullQuadPass<I extends string, O extends string> {
+    private _label;
+    private _device;
+    private _outputTextures;
+    private _pipeline;
+    private _bindGroups;
+    private _needsUpdatePipeline;
+    /**
+     * When render to canvas
+     */
+    private _renderToScreen?;
+    private _inputs;
+    private _outputs;
+    private _uniforms;
+    private _uniformBuffers;
+    private _width;
+    private _height;
+    private _fsCode;
+    private _fsMain;
+    private _fsDefine;
+    constructor(label: string, device: GPUDevice, opts: {
+        inputs: I[];
+        outputs: O[];
+        fsMain: string;
+        fsDefine?: string;
+        uniforms: Uniform[];
+    });
+    setSize(width: number, height: number): void;
+    setOutputParams(outputParams: Record<O, WGPUFullQuadPassOutput>): void;
+    setRenderToScreen(screenTexture: GPUTexture, presentationFormat: GPUTextureFormat): void;
+    setUniform(label: string, data: Float32Array | Int32Array | Uint32Array): void;
+    getOutputTexture(name: O): GPUTexture;
+    dispose(): void;
+    private _createTexture;
+    private _resizeOutputTextures;
+    private _updateOutputTextures;
+    private _updatePipeline;
+    private _getFullFs;
+    private _updateBindGroups;
+    createPass(commandEncoder: GPUCommandEncoder, inputTextures: Record<I, GPUTexture>): void;
+}

package/lib/WGPUFullQuadPass.js

@@ -0,0 +1,261 @@
+export const fullScreenQuadVertexShaderWGSL = /*wgsl */ `
+@vertex
+fn main(
+  @builtin(vertex_index) vertexIndex: u32
+) -> @builtin(position) vec4f {
+  const pos = array(
+    vec2(-1.0, -1.0), vec2(1.0, -1.0), vec2(-1.0, 1.0),
+    vec2(-1.0, 1.0), vec2(1.0, -1.0), vec2(1.0, 1.0),
+  );
+
+  return vec4<f32>(pos[vertexIndex], 0.0, 1.0);
+}
+`;
+function isTextureParamsEqual(params, other) {
+    return params.format === other.format;
+}
+export class WGPUFullQuadPass {
+    _label;
+    _device;
+    _outputTextures = {};
+    _pipeline;
+    _bindGroups = [];
+    _needsUpdatePipeline = true;
+    /**
+     * When render to canvas
+     */
+    _renderToScreen;
+    _inputs = [];
+    _outputs = [];
+    _uniforms = [];
+    _uniformBuffers = {};
+    _width = 10;
+    _height = 10;
+    _fsCode = '';
+    _fsMain;
+    _fsDefine;
+    constructor(label, device, opts) {
+        this._label = label;
+        this._device = device;
+        this._fsMain = opts.fsMain;
+        this._fsDefine = opts.fsDefine;
+        this._inputs = opts.inputs;
+        this._outputs = opts.outputs;
+        this._uniforms = opts.uniforms;
+        opts.uniforms.forEach((uniform) => {
+            this._uniformBuffers[uniform.label] = device.createBuffer({
+                label: this._label,
+                size: uniform.data.byteLength,
+                usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
+            });
+            this._device.queue.writeBuffer(this._uniformBuffers[uniform.label], 0, uniform.data);
+        });
+    }
+    setSize(width, height) {
+        width = Math.ceil(width);
+        height = Math.ceil(height);
+        const sizeChanged = width !== this._width || height !== this._height;
+        this._width = width;
+        this._height = height;
+        if (sizeChanged) {
+            this._resizeOutputTextures();
+            this._needsUpdatePipeline = true;
+        }
+    }
+    setOutputParams(outputParams) {
+        this._renderToScreen = undefined;
+        this._updateOutputTextures(outputParams);
+        this._needsUpdatePipeline = true;
+    }
+    setRenderToScreen(screenTexture, presentationFormat) {
+        this._renderToScreen = {
+            screenTexture,
+            presentationFormat
+        };
+    }
+    setUniform(label, data) {
+        const buffer = this._uniformBuffers[label];
+        this._device.queue.writeBuffer(buffer, 0, data);
+    }
+    getOutputTexture(name) {
+        return this._outputTextures[name].texture;
+    }
+    dispose() {
+        Object.keys(this._uniformBuffers).forEach((key) => {
+            this._uniformBuffers[key].destroy();
+        });
+        Object.keys(this._outputTextures).forEach((key) => {
+            this._outputTextures[key].texture.destroy();
+        });
+    }
+    _createTexture(params) {
+        return this._device.createTexture({
+            label: this._label,
+            size: {
+                width: this._width,
+                height: this._height,
+                depthOrArrayLayers: 1
+            },
+            format: params.format,
+            usage: GPUTextureUsage.RENDER_ATTACHMENT | GPUTextureUsage.TEXTURE_BINDING
+        });
+    }
+    _resizeOutputTextures() {
+        const outputTextures = this._outputTextures;
+        for (const key in outputTextures) {
+            const { texture, params } = outputTextures[key];
+            texture.destroy();
+            outputTextures[key].texture = this._createTexture(params);
+        }
+    }
+    _updateOutputTextures(outputParams) {
+        const outputTextures = this._outputTextures;
+        for (const key in outputParams) {
+            const params = outputParams[key];
+            if (!isTextureParamsEqual(params, outputTextures[key]?.params || {})) {
+                outputTextures[key]?.texture.destroy();
+                const texture = this._createTexture(params);
+                outputTextures[key] = {
+                    texture,
+                    params
+                };
+            }
+        }
+    }
+    _updatePipeline() {
+        if (!this._needsUpdatePipeline) {
+            return;
+        }
+        this._needsUpdatePipeline = false;
+        const device = this._device;
+        const fsCode = this._getFullFs();
+        if (fsCode === this._fsCode) {
+            return;
+        }
+        const { screenTexture, presentationFormat } = this._renderToScreen || {};
+        this._fsCode = fsCode;
+        this._pipeline = device.createRenderPipeline({
+            label: this._label,
+            layout: 'auto',
+            vertex: {
+                module: device.createShaderModule({
+                    label: this._label,
+                    code: fullScreenQuadVertexShaderWGSL
+                }),
+                entryPoint: 'main'
+            },
+            fragment: {
+                module: device.createShaderModule({
+                    label: this._label,
+                    code: fsCode
+                }),
+                entryPoint: 'main',
+                targets: screenTexture
+                    ? [
+                        {
+                            format: presentationFormat
+                        }
+                    ]
+                    : this._outputs.map((key) => ({
+                        format: this._outputTextures[key].params.format
+                    }))
+            },
+            primitive: {
+                topology: 'triangle-list'
+            }
+        });
+        this._updateBindGroups();
+    }
+    _getFullFs() {
+        const inputs = this._inputs;
+        const hasInputs = inputs.length > 0;
+        const fs = `
+${inputs
+            .sort()
+            .map((textureName, idx) => `@group(0) @binding(${idx}) var ${textureName}: texture_2d<f32>;`)
+            .join('\n')}
+${this._uniforms
+            .map((uniform, idx) => `@group(${hasInputs ? 1 : 0}) @binding(${idx}) var<uniform> ${uniform.label}: ${uniform.type};`)
+            .join('\n')}
+
+struct FSOutput {
+${this._outputs
+            .map((name, idx) => `@location(${idx}) ${name}: vec4f,`)
+            .join('\n')}
+}
+${this._fsDefine ?? ''}
+@fragment
+fn main(
+  @builtin(position) coord: vec4f
+) -> FSOutput {
+  var uv = vec2i(floor(coord.xy));
+  var output: FSOutput;
+${this._fsMain}
+  return output;
+}
+`;
+        return fs;
+    }
+    _updateBindGroups() {
+        const bindGroups = [];
+        const device = this._device;
+        //TODO
+        const uniformBindGroupIndex = this._inputs.length > 0 ? 1 : 0;
+        if (this._uniforms.length > 0) {
+            bindGroups[uniformBindGroupIndex] = device.createBindGroup({
+                label: this._label,
+                layout: this._pipeline.getBindGroupLayout(uniformBindGroupIndex),
+                entries: this._uniforms.map((uniform, idx) => ({
+                    binding: idx,
+                    resource: {
+                        buffer: this._uniformBuffers[uniform.label]
+                    }
+                }))
+            });
+        }
+        this._bindGroups = bindGroups;
+    }
+    createPass(commandEncoder, inputTextures) {
+        this._updatePipeline();
+        // TODO createBindGrou every time?
+        if (this._inputs.length > 0) {
+            this._bindGroups[0] = this._device.createBindGroup({
+                label: this._label,
+                layout: this._pipeline.getBindGroupLayout(0),
+                entries: this._inputs.map((textureName, idx) => ({
+                    binding: idx,
+                    // TODO
+                    resource: inputTextures[textureName].createView()
+                }))
+            });
+        }
+        // Begin the render pass
+        const renderPass = commandEncoder.beginRenderPass({
+            colorAttachments: this._renderToScreen
+                ? [
+                    {
+                        view: this._renderToScreen.screenTexture.createView(),
+                        clearValue: { r: 0, g: 0, b: 0, a: 0 },
+                        storeOp: 'store',
+                        loadOp: 'clear'
+                    }
+                ]
+                : this._outputs.map((textureName) => ({
+                    view: this._outputTextures[textureName].texture.createView(),
+                    clearValue: { r: 0, g: 0, b: 0, a: 0 },
+                    loadOp: 'clear',
+                    storeOp: 'store'
+                }))
+        });
+        // Draw a full quad
+        renderPass.setPipeline(this._pipeline);
+        // Bind groups
+        this._bindGroups.forEach((bindGroup, idx) => {
+            renderPass.setBindGroup(idx, bindGroup);
+        });
+        renderPass.draw(6, 1, 0, 0);
+        // End the render pass
+        renderPass.end();
+    }
+}
+//# sourceMappingURL=WGPUFullQuadPass.js.map

package/lib/WGPUFullQuadPass.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"WGPUFullQuadPass.js","sourceRoot":"","sources":["../src/WGPUFullQuadPass.ts"],"names":[],"mappings":"AAAA,MAAM,CAAC,MAAM,8BAA8B,GAAG,SAAS,CAAC;;;;;;;;;;;;CAYvD,CAAC;AAEF,SAAS,oBAAoB,CAC3B,MAA8B,EAC9B,KAA6B;IAE7B,OAAO,MAAM,CAAC,MAAM,KAAK,KAAK,CAAC,MAAM,CAAC;AACxC,CAAC;AAWD,MAAM,OAAO,gBAAgB;IACnB,MAAM,CAAC;IAEP,OAAO,CAAC;IACR,eAAe,GAMnB,EAAE,CAAC;IAEC,SAAS,CAAqB;IAC9B,WAAW,GAAmB,EAAE,CAAC;IACjC,oBAAoB,GAAG,IAAI,CAAC;IACpC;;OAEG;IACK,eAAe,CAGrB;IACM,OAAO,GAAa,EAAE,CAAC;IACvB,QAAQ,GAAa,EAAE,CAAC;IACxB,SAAS,GAAc,EAAE,CAAC;IAC1B,eAAe,GAA8B,EAAE,CAAC;IAEhD,MAAM,GAAG,EAAE,CAAC;IACZ,OAAO,GAAG,EAAE,CAAC;IAEb,OAAO,GAAG,EAAE,CAAC;IACb,OAAO,CAAC;IACR,SAAS,CAAC;IAElB,YACE,KAAa,EACb,MAAiB,EACjB,IAMC;QAED,IAAI,CAAC,MAAM,GAAG,KAAK,CAAC;QACpB,IAAI,CAAC,OAAO,GAAG,MAAM,CAAC;QACtB,IAAI,CAAC,OAAO,GAAG,IAAI,CAAC,MAAM,CAAC;QAC3B,IAAI,CAAC,SAAS,GAAG,IAAI,CAAC,QAAQ,CAAC;QAC/B,IAAI,CAAC,OAAO,GAAG,IAAI,CAAC,MAAM,CAAC;QAC3B,IAAI,CAAC,QAAQ,GAAG,IAAI,CAAC,OAAO,CAAC;QAC7B,IAAI,CAAC,SAAS,GAAG,IAAI,CAAC,QAAQ,CAAC;QAE/B,IAAI,CAAC,QAAQ,CAAC,OAAO,CAAC,CAAC,OAAO,EAAE,EAAE;YAChC,IAAI,CAAC,eAAe,CAAC,OAAO,CAAC,KAAK,CAAC,GAAG,MAAM,CAAC,YAAY,CAAC;gBACxD,KAAK,EAAE,IAAI,CAAC,MAAM;gBAClB,IAAI,EAAE,OAAO,CAAC,IAAI,CAAC,UAAU;gBAC7B,KAAK,EAAE,cAAc,CAAC,OAAO,GAAG,cAAc,CAAC,QAAQ;aACxD,CAAC,CAAC;YACH,IAAI,CAAC,OAAO,CAAC,KAAK,CAAC,WAAW,CAC5B,IAAI,CAAC,eAAe,CAAC,OAAO,CAAC,KAAK,CAAC,EACnC,CAAC,EACD,OAAO,CAAC,IAAI,CACb,CAAC;QACJ,CAAC,CAAC,CAAC;IACL,CAAC;IAED,OAAO,CAAC,KAAa,EAAE,MAAc;QACnC,KAAK,GAAG,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;QACzB,MAAM,GAAG,IAAI,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;QAC3B,MAAM,WAAW,GAAG,KAAK,KAAK,IAAI,CAAC,MAAM,IAAI,MAAM,KAAK,IAAI,CAAC,OAAO,CAAC;QACrE,IAAI,CAAC,MAAM,GAAG,KAAK,CAAC;QACpB,IAAI,CAAC,OAAO,GAAG,MAAM,CAAC;QACtB,IAAI,WAAW,EAAE,CAAC;YAChB,IAAI,CAAC,qBAAqB,EAAE,CAAC;YAC7B,IAAI,CAAC,oBAAoB,GAAG,IAAI,CAAC;QACnC,CAAC;IACH,CAAC;IAED,eAAe,CAAC,YAA+C;QAC7D,IAAI,CAAC,eAAe,GAAG,SAAS,CAAC;QACjC,IAAI,CAAC,qBAAqB,CAAC,YAAY,CAAC,CAAC;QACzC,IAAI,CAAC,oBAAoB,GAAG,IAAI,CAAC;IACnC,CAAC;IAED,iBAAiB,CACf,aAAyB,EACzB,kBAAoC;QAEpC,IAAI,CAAC,eAAe,GAAG;YACrB,aAAa;YACb,kBAAkB;SACnB,CAAC;IACJ,CAAC;IAED,UAAU,CAAC,KAAa,EAAE,IAA6C;QACrE,MAAM,MAAM,GAAG,IAAI,CAAC,eAAe,CAAC,KAAK,CAAC,CAAC;QAC3C,IAAI,CAAC,OAAO,CAAC,KAAK,CAAC,WAAW,CAAC,MAAM,EAAE,CAAC,EAAE,IAAI,CAAC,CAAC;IAClD,CAAC;IAED,gBAAgB,CAAC,IAAO;QACtB,OAAO,IAAI,CAAC,eAAe,CAAC,IAAI,CAAC,CAAC,OAAO,CAAC;IAC5C,CAAC;IAED,OAAO;QACL,MAAM,CAAC,IAAI,CAAC,IAAI,CAAC,eAAe,CAAC,CAAC,OAAO,CAAC,CAAC,GAAG,EAAE,EAAE;YAC/C,IAAI,CAAC,eAAuB,CAAC,GAAG,CAAC,CAAC,OAAO,EAAE,CAAC;QAC/C,CAAC,CAAC,CAAC;QACH,MAAM,CAAC,IAAI,CAAC,IAAI,CAAC,eAAe,CAAC,CAAC,OAAO,CAAC,CAAC,GAAG,EAAE,EAAE;YAC/C,IAAI,CAAC,eAAuB,CAAC,GAAG,CAAC,CAAC,OAAO,CAAC,OAAO,EAAE,CAAC;QACvD,CAAC,CAAC,CAAC;IACL,CAAC;IAEO,cAAc,CAAC,MAA8B;QACnD,OAAO,IAAI,CAAC,OAAO,CAAC,aAAa,CAAC;YAChC,KAAK,EAAE,IAAI,CAAC,MAAM;YAClB,IAAI,EAAE;gBACJ,KAAK,EAAE,IAAI,CAAC,MAAM;gBAClB,MAAM,EAAE,IAAI,CAAC,OAAO;gBACpB,kBAAkB,EAAE,CAAC;aACtB;YACD,MAAM,EAAE,MAAM,CAAC,MAAM;YACrB,KAAK,EAAE,eAAe,CAAC,iBAAiB,GAAG,eAAe,CAAC,eAAe;SAC3E,CAAC,CAAC;IACL,CAAC;IAEO,qBAAqB;QAC3B,MAAM,cAAc,GAAG,IAAI,CAAC,eAAe,CAAC;QAC5C,KAAK,MAAM,GAAG,IAAI,cAAc,EAAE,CAAC;YACjC,MAAM,EAAE,OAAO,EAAE,MAAM,EAAE,GAAG,cAAc,CAAC,GAAG,CAAC,CAAC;YAChD,OAAO,CAAC,OAAO,EAAE,CAAC;YAClB,cAAc,CAAC,GAAG,CAAC,CAAC,OAAO,GAAG,IAAI,CAAC,cAAc,CAAC,MAAM,CAAC,CAAC;QAC5D,CAAC;IACH,CAAC;IAEO,qBAAqB,CAC3B,YAAoD;QAEpD,MAAM,cAAc,GAAG,IAAI,CAAC,eAAe,CAAC;QAC5C,KAAK,MAAM,GAAG,IAAI,YAAY,EAAE,CAAC;YAC/B,MAAM,MAAM,GAAG,YAAY,CAAC,GAAG,CAAC,CAAC;YACjC,IACE,CAAC,oBAAoB,CACnB,MAAM,EACN,cAAc,CAAC,GAAG,CAAC,EAAE,MAAM,IAAK,EAA6B,CAC9D,EACD,CAAC;gBACD,cAAc,CAAC,GAAG,CAAC,EAAE,OAAO,CAAC,OAAO,EAAE,CAAC;gBACvC,MAAM,OAAO,GAAG,IAAI,CAAC,cAAc,CAAC,MAAM,CAAC,CAAC;gBAC5C,cAAc,CAAC,GAAG,CAAC,GAAG;oBACpB,OAAO;oBACP,MAAM;iBACP,CAAC;YACJ,CAAC;QACH,CAAC;IACH,CAAC;IAEO,eAAe;QACrB,IAAI,CAAC,IAAI,CAAC,oBAAoB,EAAE,CAAC;YAC/B,OAAO;QACT,CAAC;QACD,IAAI,CAAC,oBAAoB,GAAG,KAAK,CAAC;QAClC,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,CAAC;QAC5B,MAAM,MAAM,GAAG,IAAI,CAAC,UAAU,EAAE,CAAC;QACjC,IAAI,MAAM,KAAK,IAAI,CAAC,OAAO,EAAE,CAAC;YAC5B,OAAO;QACT,CAAC;QACD,MAAM,EAAE,aAAa,EAAE,kBAAkB,EAAE,GAAG,IAAI,CAAC,eAAe,IAAI,EAAE,CAAC;QACzE,IAAI,CAAC,OAAO,GAAG,MAAM,CAAC;QACtB,IAAI,CAAC,SAAS,GAAG,MAAM,CAAC,oBAAoB,CAAC;YAC3C,KAAK,EAAE,IAAI,CAAC,MAAM;YAClB,MAAM,EAAE,MAAM;YACd,MAAM,EAAE;gBACN,MAAM,EAAE,MAAM,CAAC,kBAAkB,CAAC;oBAChC,KAAK,EAAE,IAAI,CAAC,MAAM;oBAClB,IAAI,EAAE,8BAA8B;iBACrC,CAAC;gBACF,UAAU,EAAE,MAAM;aACnB;YACD,QAAQ,EAAE;gBACR,MAAM,EAAE,MAAM,CAAC,kBAAkB,CAAC;oBAChC,KAAK,EAAE,IAAI,CAAC,MAAM;oBAClB,IAAI,EAAE,MAAM;iBACb,CAAC;gBACF,UAAU,EAAE,MAAM;gBAClB,OAAO,EAAE,aAAa;oBACpB,CAAC,CAAC;wBACE;4BACE,MAAM,EAAE,kBAAmB;yBAC5B;qBACF;oBACH,CAAC,CAAC,IAAI,CAAC,QAAQ,CAAC,GAAG,CAAC,CAAC,GAAG,EAAE,EAAE,CAAC,CAAC;wBAC1B,MAAM,EAAE,IAAI,CAAC,eAAe,CAAC,GAAG,CAAC,CAAC,MAAM,CAAC,MAAM;qBAChD,CAAC,CAAC;aACR;YACD,SAAS,EAAE;gBACT,QAAQ,EAAE,eAAe;aAC1B;SACF,CAAC,CAAC;QACH,IAAI,CAAC,iBAAiB,EAAE,CAAC;IAC3B,CAAC;IAEO,UAAU;QAChB,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,CAAC;QAC5B,MAAM,SAAS,GAAG,MAAM,CAAC,MAAM,GAAG,CAAC,CAAC;QACpC,MAAM,EAAE,GAAG;EACb,MAAM;aACL,IAAI,EAAE;aACN,GAAG,CACF,CAAC,WAAW,EAAE,GAAG,EAAE,EAAE,CACnB,sBAAsB,GAAG,SAAS,WAAW,oBAAoB,CACpE;aACA,IAAI,CAAC,IAAI,CAAC;EACX,IAAI,CAAC,SAAS;aACb,GAAG,CACF,CAAC,OAAO,EAAE,GAAG,EAAE,EAAE,CACf,UAAU,SAAS,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,cAAc,GAAG,kBAC1C,OAAO,CAAC,KACV,KAAK,OAAO,CAAC,IAAI,GAAG,CACvB;aACA,IAAI,CAAC,IAAI,CAAC;;;EAGX,IAAI,CAAC,QAAQ;aACZ,GAAG,CAAC,CAAC,IAAI,EAAE,GAAG,EAAE,EAAE,CAAC,aAAa,GAAG,KAAK,IAAI,UAAU,CAAC;aACvD,IAAI,CAAC,IAAI,CAAC;;EAEX,IAAI,CAAC,SAAS,IAAI,EAAE;;;;;;;EAOpB,IAAI,CAAC,OAAO;;;CAGb,CAAC;QAEE,OAAO,EAAE,CAAC;IACZ,CAAC;IAEO,iBAAiB;QACvB,MAAM,UAAU,GAAmB,EAAE,CAAC;QACtC,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,CAAC;QAE5B,MAAM;QACN,MAAM,qBAAqB,GAAG,IAAI,CAAC,OAAO,CAAC,MAAM,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;QAC9D,IAAI,IAAI,CAAC,SAAS,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;YAC9B,UAAU,CAAC,qBAAqB,CAAC,GAAG,MAAM,CAAC,eAAe,CAAC;gBACzD,KAAK,EAAE,IAAI,CAAC,MAAM;gBAClB,MAAM,EAAE,IAAI,CAAC,SAAS,CAAC,kBAAkB,CAAC,qBAAqB,CAAC;gBAChE,OAAO,EAAE,IAAI,CAAC,SAAS,CAAC,GAAG,CACzB,CAAC,OAAO,EAAE,GAAG,EAAE,EAAE,CACf,CAAC;oBACC,OAAO,EAAE,GAAG;oBACZ,QAAQ,EAAE;wBACR,MAAM,EAAE,IAAI,CAAC,eAAe,CAAC,OAAO,CAAC,KAAK,CAAC;qBAC5C;iBAC0B,CAAC,CACjC;aACF,CAAC,CAAC;QACL,CAAC;QAED,IAAI,CAAC,WAAW,GAAG,UAAU,CAAC;IAChC,CAAC;IAED,UAAU,CACR,cAAiC,EACjC,aAAoC;QAEpC,IAAI,CAAC,eAAe,EAAE,CAAC;QAEvB,kCAAkC;QAClC,IAAI,IAAI,CAAC,OAAO,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC;YAC5B,IAAI,CAAC,WAAW,CAAC,CAAC,CAAC,GAAG,IAAI,CAAC,OAAO,CAAC,eAAe,CAAC;gBACjD,KAAK,EAAE,IAAI,CAAC,MAAM;gBAClB,MAAM,EAAE,IAAI,CAAC,SAAS,CAAC,kBAAkB,CAAC,CAAC,CAAC;gBAC5C,OAAO,EAAE,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,CAAC,WAAW,EAAE,GAAG,EAAE,EAAE,CAAC,CAAC;oBAC/C,OAAO,EAAE,GAAG;oBACZ,OAAO;oBACP,QAAQ,EAAE,aAAa,CAAC,WAAgB,CAAC,CAAC,UAAU,EAAE;iBACvD,CAAC,CAAC;aACJ,CAAC,CAAC;QACL,CAAC;QACD,wBAAwB;QACxB,MAAM,UAAU,GAAG,cAAc,CAAC,eAAe,CAAC;YAChD,gBAAgB,EAAE,IAAI,CAAC,eAAe;gBACpC,CAAC,CAAC;oBACE;wBACE,IAAI,EAAE,IAAI,CAAC,eAAe,CAAC,aAAa,CAAC,UAAU,EAAE;wBACrD,UAAU,EAAE,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE;wBACtC,OAAO,EAAE,OAAqB;wBAC9B,MAAM,EAAE,OAAoB;qBAC7B;iBACF;gBACH,CAAC,CAAC,IAAI,CAAC,QAAQ,CAAC,GAAG,CACf,CAAC,WAAW,EAAE,EAAE,CACd,CAAC;oBACC,IAAI,EAAE,IAAI,CAAC,eAAe,CAAC,WAAW,CAAC,CAAC,OAAO,CAAC,UAAU,EAAE;oBAC5D,UAAU,EAAE,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE;oBACtC,MAAM,EAAE,OAAoB;oBAC5B,OAAO,EAAE,OAAqB;iBACQ,CAAC,CAC5C;SACN,CAAC,CAAC;QAEH,mBAAmB;QACnB,UAAU,CAAC,WAAW,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;QACvC,cAAc;QACd,IAAI,CAAC,WAAW,CAAC,OAAO,CAAC,CAAC,SAAS,EAAE,GAAG,EAAE,EAAE;YAC1C,UAAU,CAAC,YAAY,CAAC,GAAG,EAAE,SAAS,CAAC,CAAC;QAC1C,CAAC,CAAC,CAAC;QACH,UAAU,CAAC,IAAI,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC;QAC5B,sBAAsB;QACtB,UAAU,CAAC,GAAG,EAAE,CAAC;IACnB,CAAC;CACF"}

package/lib/backend.d.ts

@@ -0,0 +1,5 @@
+/// <reference types="@webgpu/types/dist" />
+import { WebGPUBackend } from '@tensorflow/tfjs-backend-webgpu/dist/base';
+import './kernels';
+export declare function initWebGPUBackend(): Promise<WebGPUBackend | undefined>;
+export declare function initWebGPUBackendWithDevice(device: GPUDevice, adapter: GPUAdapterInfo): Promise<WebGPUBackend>;

package/lib/backend.js

@@ -0,0 +1,41 @@
+import { WebGPUBackend } from '@tensorflow/tfjs-backend-webgpu/dist/base';
+import { ENGINE } from '@tensorflow/tfjs-core/dist/engine';
+import './kernels';
+export async function initWebGPUBackend() {
+    try {
+        const gpuDescriptor = {
+            powerPreference: 'high-performance'
+        };
+        const adapter = (await navigator.gpu.requestAdapter(gpuDescriptor));
+        const deviceDescriptor = {};
+        const requiredFeatures = [];
+        if (adapter.features.has('timestamp-query')) {
+            requiredFeatures.push('timestamp-query');
+        }
+        if (adapter.features.has('bgra8unorm-storage')) {
+            requiredFeatures.push(['bgra8unorm-storage']);
+        }
+        deviceDescriptor.requiredFeatures =
+            requiredFeatures;
+        const adapterLimits = adapter.limits;
+        deviceDescriptor.requiredLimits = {
+            maxComputeWorkgroupStorageSize: adapterLimits.maxComputeWorkgroupStorageSize,
+            maxComputeWorkgroupsPerDimension: adapterLimits.maxComputeWorkgroupsPerDimension,
+            maxStorageBufferBindingSize: adapterLimits.maxStorageBufferBindingSize,
+            maxBufferSize: adapterLimits.maxBufferSize,
+            maxComputeWorkgroupSizeX: adapterLimits.maxComputeWorkgroupSizeX,
+            maxComputeInvocationsPerWorkgroup: adapterLimits.maxComputeInvocationsPerWorkgroup
+        };
+        const device = await adapter.requestDevice(deviceDescriptor);
+        const adapterInfo = await adapter.requestAdapterInfo();
+        return initWebGPUBackendWithDevice(device, adapterInfo);
+    }
+    catch (e) { }
+}
+export async function initWebGPUBackendWithDevice(device, adapter) {
+    const backend = new WebGPUBackend(device, adapter);
+    ENGINE.registerBackend('webgpu-oidn', () => backend);
+    await ENGINE.setBackend('webgpu-oidn');
+    return backend;
+}
+//# sourceMappingURL=backend.js.map

package/lib/backend.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"backend.js","sourceRoot":"","sources":["../src/backend.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,aAAa,EAAE,MAAM,2CAA2C,CAAC;AAC1E,OAAO,EAAE,MAAM,EAAE,MAAM,mCAAmC,CAAC;AAE3D,OAAO,WAAW,CAAC;AAEnB,MAAM,CAAC,KAAK,UAAU,iBAAiB;IACrC,IAAI,CAAC;QACH,MAAM,aAAa,GAA6B;YAC9C,eAAe,EAAE,kBAAkB;SACpC,CAAC;QAEF,MAAM,OAAO,GAAG,CAAC,MAAM,SAAS,CAAC,GAAG,CAAC,cAAc,CAAC,aAAa,CAAC,CAAE,CAAC;QACrE,MAAM,gBAAgB,GAAwB,EAAE,CAAC;QAEjD,MAAM,gBAAgB,GAAG,EAAE,CAAC;QAC5B,IAAI,OAAO,CAAC,QAAQ,CAAC,GAAG,CAAC,iBAAiB,CAAC,EAAE,CAAC;YAC5C,gBAAgB,CAAC,IAAI,CAAC,iBAAiB,CAAC,CAAC;QAC3C,CAAC;QACD,IAAI,OAAO,CAAC,QAAQ,CAAC,GAAG,CAAC,oBAAoB,CAAC,EAAE,CAAC;YAC/C,gBAAgB,CAAC,IAAI,CAAC,CAAC,oBAAoB,CAAC,CAAC,CAAC;QAChD,CAAC;QACD,gBAAgB,CAAC,gBAAgB;YAC/B,gBAA4C,CAAC;QAE/C,MAAM,aAAa,GAAG,OAAO,CAAC,MAAM,CAAC;QACrC,gBAAgB,CAAC,cAAc,GAAG;YAChC,8BAA8B,EAC5B,aAAa,CAAC,8BAA8B;YAC9C,gCAAgC,EAC9B,aAAa,CAAC,gCAAgC;YAChD,2BAA2B,EAAE,aAAa,CAAC,2BAA2B;YACtE,aAAa,EAAE,aAAa,CAAC,aAAa;YAC1C,wBAAwB,EAAE,aAAa,CAAC,wBAAwB;YAChE,iCAAiC,EAC/B,aAAa,CAAC,iCAAiC;SAClD,CAAC;QACF,MAAM,MAAM,GAAG,MAAM,OAAO,CAAC,aAAa,CAAC,gBAAgB,CAAC,CAAC;QAC7D,MAAM,WAAW,GAAG,MAAM,OAAO,CAAC,kBAAkB,EAAE,CAAC;QAEvD,OAAO,2BAA2B,CAAC,MAAM,EAAE,WAAW,CAAC,CAAC;IAC1D,CAAC;IAAC,OAAO,CAAC,EAAE,CAAC,CAAA,CAAC;AAChB,CAAC;AAED,MAAM,CAAC,KAAK,UAAU,2BAA2B,CAC/C,MAAiB,EACjB,OAAuB;IAEvB,MAAM,OAAO,GAAG,IAAI,aAAa,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;IACnD,MAAM,CAAC,eAAe,CAAC,aAAa,EAAE,GAAG,EAAE,CAAC,OAAO,CAAC,CAAC;IACrD,MAAM,MAAM,CAAC,UAAU,CAAC,aAAa,CAAC,CAAC;IAEvC,OAAO,OAAO,CAAC;AACjB,CAAC"}

package/lib/hdr.d.ts

@@ -0,0 +1,31 @@
+/**
+ * Convert 4 byte uint8 buffer to 3 channels float data
+ * @param rgbe input uint8 buffer
+ * @param float output float data
+ */
+declare function rgbe2float(rgbe: Uint8Array, float: Float32Array): void;
+/**
+ * Convert 3 channels float data to 4 byte uint8 buffer
+ * @param float input float data
+ * @param rgbe output uint8 buffer
+ */
+declare function float2rgbe(float: Float32Array, rgbe: Uint8Array): void;
+/**
+ * Write float data to RGBE(.hdr) file buffer
+ * @param w image width
+ * @param h image height
+ * @param data float data, RGB 3 channels.
+ * @returns file buffer
+ */
+declare function write_hdr(w: number, h: number, data: Float32Array): Uint8Array;
+/**
+ * Read a RGBE(.hdr) file from buffer
+ * @param uint8 RGBE(.hdr) file buffer
+ * @returns Failure reason or resolved data.
+ */
+declare function read_hdr(uint8: Uint8Array): string | {
+    rgbFloat: Float32Array;
+    width: number;
+    height: number;
+};
+export { read_hdr as readHDR, write_hdr as writeHDR, float2rgbe, rgbe2float };