oidn-webgpu 0.0.0-stage → 0.6.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/CHANGELOG.md +182 -0
- package/LICENSE +21 -0
- package/README.md +89 -2
- package/dist/oidn.js +6140 -0
- package/dist/oidn.umd.cjs +1513 -0
- package/lib/UNet.d.ts +161 -0
- package/lib/UNet.js +536 -0
- package/lib/UNet.js.map +1 -0
- package/lib/WGPUComputePass.d.ts +62 -0
- package/lib/WGPUComputePass.js +278 -0
- package/lib/WGPUComputePass.js.map +1 -0
- package/lib/WGPUFullQuadPass.d.ts +50 -0
- package/lib/WGPUFullQuadPass.js +261 -0
- package/lib/WGPUFullQuadPass.js.map +1 -0
- package/lib/backend.d.ts +1 -0
- package/lib/backend.js +39 -0
- package/lib/backend.js.map +1 -0
- package/lib/finalRgbShader.d.ts +13 -0
- package/lib/finalRgbShader.js +160 -0
- package/lib/finalRgbShader.js.map +1 -0
- package/lib/gpuTypes.d.ts +158 -0
- package/lib/gpuTypes.js +27 -0
- package/lib/gpuTypes.js.map +1 -0
- package/lib/graphOptimizer.d.ts +54 -0
- package/lib/graphOptimizer.js +216 -0
- package/lib/graphOptimizer.js.map +1 -0
- package/lib/hdr.d.ts +31 -0
- package/lib/hdr.js +340 -0
- package/lib/hdr.js.map +1 -0
- package/lib/hdrTransfer.d.ts +22 -0
- package/lib/hdrTransfer.js +128 -0
- package/lib/hdrTransfer.js.map +1 -0
- package/lib/main.d.ts +57 -0
- package/lib/main.js +24 -0
- package/lib/main.js.map +1 -0
- package/lib/modelSpec.d.ts +80 -0
- package/lib/modelSpec.js +270 -0
- package/lib/modelSpec.js.map +1 -0
- package/lib/native/denoise.d.ts +25 -0
- package/lib/native/denoise.js +223 -0
- package/lib/native/denoise.js.map +1 -0
- package/lib/native/program.d.ts +53 -0
- package/lib/native/program.js +310 -0
- package/lib/native/program.js.map +1 -0
- package/lib/native/recordingDevice.d.ts +125 -0
- package/lib/native/recordingDevice.js +247 -0
- package/lib/native/recordingDevice.js.map +1 -0
- package/lib/nativeUNet.d.ts +25 -0
- package/lib/nativeUNet.js +184 -0
- package/lib/nativeUNet.js.map +1 -0
- package/lib/nativeUNetCore.d.ts +202 -0
- package/lib/nativeUNetCore.js +2945 -0
- package/lib/nativeUNetCore.js.map +1 -0
- package/lib/process.d.ts +37 -0
- package/lib/process.js +328 -0
- package/lib/process.js.map +1 -0
- package/lib/resourceTracker.d.ts +26 -0
- package/lib/resourceTracker.js +65 -0
- package/lib/resourceTracker.js.map +1 -0
- package/lib/sanitize.d.ts +15 -0
- package/lib/sanitize.js +28 -0
- package/lib/sanitize.js.map +1 -0
- package/lib/tileScheduler.d.ts +61 -0
- package/lib/tileScheduler.js +199 -0
- package/lib/tileScheduler.js.map +1 -0
- package/lib/tza.d.ts +15 -0
- package/lib/tza.js +114 -0
- package/lib/tza.js.map +1 -0
- package/lib/webnnUNet.d.ts +52 -0
- package/lib/webnnUNet.js +535 -0
- package/lib/webnnUNet.js.map +1 -0
- package/package.json +43 -4
- package/src/UNet.ts +861 -0
- package/src/WGPUComputePass.ts +382 -0
- package/src/WGPUFullQuadPass.ts +348 -0
- package/src/backend.ts +42 -0
- package/src/finalRgbShader.ts +186 -0
- package/src/gpuTypes.ts +179 -0
- package/src/graphOptimizer.ts +301 -0
- package/src/hdr.ts +398 -0
- package/src/hdrTransfer.ts +169 -0
- package/src/main.ts +126 -0
- package/src/modelSpec.ts +414 -0
- package/src/native/core-profile.json +40 -0
- package/src/native/denoise.ts +268 -0
- package/src/native/program.ts +326 -0
- package/src/native/recordingDevice.ts +341 -0
- package/src/nativeUNet.ts +269 -0
- package/src/nativeUNetCore.ts +3735 -0
- package/src/process.ts +368 -0
- package/src/resourceTracker.ts +94 -0
- package/src/sanitize.ts +29 -0
- package/src/tileScheduler.ts +330 -0
- package/src/tza.ts +136 -0
- package/src/webnnUNet.ts +812 -0
|
@@ -0,0 +1,1513 @@
|
|
|
1
|
+
(function(X,ct){typeof exports=="object"&&typeof module<"u"?ct(exports):typeof define=="function"&&define.amd?define(["exports"],ct):(X=typeof globalThis<"u"?globalThis:X||self,ct(X.oidn={}))})(this,function(X){"use strict";var sa=Object.defineProperty;var ca=(X,ct,Dt)=>ct in X?sa(X,ct,{enumerable:!0,configurable:!0,writable:!0,value:Dt}):X[ct]=Dt;var y=(X,ct,Dt)=>ca(X,typeof ct!="symbol"?ct+"":ct,Dt);class ct{constructor(){y(this,"dims",[]);y(this,"paddedDims",[]);y(this,"layout","x");y(this,"dataType","Float32")}getByteSize(){let t=1;for(const e of this.paddedDims)t*=e;return this.dataType==="Float32"?t*=4:this.dataType==="Float16"&&(t*=2),t}}class Dt{constructor(t,e){this.desc=t,this.data=e}}class on{constructor(t){y(this,"offset",0);this._view=t}read(t){const e=this._view,i=this.offset;switch(this.offset+=t,t){case 1:return e.getUint8(i);case 2:return e.getUint16(i,!0);case 4:return e.getUint32(i,!0);case 8:return Number(e.getBigUint64(i,!0));default:throw new Error("unsupported read size")}}}function bi(r){const t=new Uint8Array(r),e=new on(new DataView(r));if(e.read(2)!==16855)throw new Error("invalid or corrupted weights blob");const n=e.read(1);if(e.read(1),n!==2)throw new Error("unsupported weights blob version");const o=e.read(8);e.offset=o;const a=e.read(4),s=new Map;for(let u=0;u<a;++u){const l=new ct,p=e.read(2),c=new TextDecoder().decode(t.subarray(e.offset,e.offset+p));e.offset+=p;const d=e.read(1);for(let m=0;m<d;++m)l.dims.push(e.read(4));l.paddedDims=[...l.dims],new TextDecoder().decode(t.subarray(e.offset,e.offset+d))==="oihw"&&(l.layout="oihw"),e.offset+=d;const g=String.fromCharCode(e.read(1));if(g==="f")l.dataType="Float32";else if(g==="h")l.dataType="Float16";else throw new Error("invalid tensor data type");const f=e.read(8),_=t.slice(f,f+l.getByteSize());s.set(c,new Dt(l,_))}return s}function an(r,t){return r.channels===t.channels}const me=8;class $e{constructor(t,e,i){y(this,"autoUpdateOutputBuffer",!0);y(this,"_label");y(this,"_device");y(this,"_outputBuffers",{});y(this,"_pipeline");y(this,"_bindGroups",[]);y(this,"_needsUpdatePipeline",!0);y(this,"_needsResizeBuffer",!0);y(this,"_inputs",[]);y(this,"_outputs",[]);y(this,"_uniforms",[]);y(this,"_uniformBuffers",{});y(this,"_width",10);y(this,"_height",10);y(this,"_execWidth");y(this,"_execHeight");y(this,"_csCode","");y(this,"_csMain");y(this,"_csDefine");y(this,"_groupOffsets",{inputs:0,uniforms:1,outputs:2});this._label=t,this._device=e,this._csMain=i.csMain,this._csDefine=i.csDefine,this._inputs=i.inputs,this._outputs=i.outputs,this._uniforms=i.uniforms,this.autoUpdateOutputBuffer=i.autoUpdateOutputBuffer??!0,i.uniforms.forEach(n=>{this._uniformBuffers[n.label]=e.createBuffer({label:this._label,size:n.data.byteLength,usage:GPUBufferUsage.UNIFORM|GPUBufferUsage.COPY_DST}),this._device.queue.writeBuffer(this._uniformBuffers[n.label],0,n.data)})}setCSCode({csDefine:t,csMain:e}){this._csDefine=t,this._csMain=e,this._needsUpdatePipeline=!0}setSize(t,e){t=Math.ceil(t),e=Math.ceil(e);const i=t!==this._width||e!==this._height;this._width=t,this._height=e,i&&(this._needsResizeBuffer=!0,this._needsUpdatePipeline=!0)}setExecuteSize(t,e){t=Math.ceil(t),e=Math.ceil(e),this._execWidth=t,this._execHeight=e}setOutputParams(t){this.autoUpdateOutputBuffer&&this._updateOutputBuffers(t),this._needsUpdatePipeline=!0}setOutputBuffers(t){this._outputBuffers=Object.keys(t).reduce((e,i)=>(e[i]={buffer:t[i],params:{channels:4}},e),{})}setUniform(t,e){const i=this._uniformBuffers[t];this._device.queue.writeBuffer(i,0,e)}getOutput(t){return this._needsResizeBuffer&&this.autoUpdateOutputBuffer&&(this._resizeOutputBuffers(),this._needsResizeBuffer=!1),this._outputBuffers[t].buffer}dispose(t=!0){Object.keys(this._uniformBuffers).forEach(e=>{this._uniformBuffers[e].destroy()}),t&&Object.keys(this._outputBuffers).forEach(e=>{this._outputBuffers[e].buffer.destroy()})}_createBuffer(t){const e=this._width*this._height*4*4;return this._device.createBuffer({label:this._label,size:Math.max(e,80),usage:GPUBufferUsage.STORAGE|GPUBufferUsage.COPY_DST|GPUBufferUsage.COPY_SRC})}_resizeOutputBuffers(){const t=this._outputBuffers;for(const e in t){const{buffer:i,params:n}=t[e];i.destroy(),t[e].buffer=this._createBuffer(n)}}_updateOutputBuffers(t){var i,n;const e=this._outputBuffers;for(const o in t){const a=t[o];if(!an(a,((i=e[o])==null?void 0:i.params)||{})){(n=e[o])==null||n.buffer.destroy();const s=this._createBuffer(a);e[o]={buffer:s,params:a}}}}_updatePipeline(t,e){if(!this._needsUpdatePipeline)return;this._needsUpdatePipeline=!1;const i=this._device,n=this._getFullCs(t,e);n!==this._csCode&&(this._csCode=n,this._pipeline=i.createComputePipeline({label:this._label,layout:"auto",compute:{module:i.createShaderModule({label:this._label,code:n}),entryPoint:"main"}}),this._updateBindGroups())}_getFullCs(t,e){const i=this._inputs,n=this._uniforms;let o=0;const a=this._groupOffsets={inputs:0,uniforms:0,outputs:0};return i.length>0&&o++,n.length>0&&(a.uniforms=o,o++),a.outputs=o,`
|
|
2
|
+
${i.sort().map((u,l)=>{const p=`@group(${a.inputs}) @binding(${l}) `,c=`in_${u}`;return e[u]==="texture"?`${p} var ${c}: texture_2d<f32>;`:`${p} var<storage, read> ${c}: array<vec${t[u].channels}f>;`}).join(`
|
|
3
|
+
`)}
|
|
4
|
+
${this._uniforms.map((u,l)=>`@group(${a.uniforms}) @binding(${l}) var<uniform> ${u.label}: ${u.type};`).join(`
|
|
5
|
+
`)}
|
|
6
|
+
|
|
7
|
+
${this._outputs.map((u,l)=>`@group(${a.outputs}) @binding(${l}) var<storage, read_write> out_${u}: array<vec${this._outputBuffers[u].params.channels}f>;`).join(`
|
|
8
|
+
`)}
|
|
9
|
+
${this._csDefine??""}
|
|
10
|
+
@compute @workgroup_size(${me}, ${me}, 1)
|
|
11
|
+
fn main(@builtin(global_invocation_id) globalId: vec3u) {
|
|
12
|
+
${this._csMain}
|
|
13
|
+
}
|
|
14
|
+
`}_updateBindGroups(){const t=[],e=this._device,i=this._groupOffsets;this._uniforms.length>0&&(t[i.uniforms]=e.createBindGroup({label:this._label,layout:this._pipeline.getBindGroupLayout(i.uniforms),entries:this._uniforms.map((n,o)=>({binding:o,resource:{buffer:this._uniformBuffers[n.label]}}))})),this._bindGroups=t}createPass(t,e){this._needsResizeBuffer&&this.autoUpdateOutputBuffer&&(this._resizeOutputBuffers(),this._needsResizeBuffer=!1);const i=this._inputs.reduce((a,s)=>(a[s]=e[s].buffer?"buffer":"texture",a),{});this._updatePipeline(e,i);const n=this._groupOffsets;this._inputs.length>0&&(this._bindGroups[n.inputs]=this._device.createBindGroup({label:this._label,layout:this._pipeline.getBindGroupLayout(n.inputs),entries:this._inputs.map((a,s)=>({binding:s,resource:e[a].buffer?{buffer:e[a].buffer}:e[a].texture.createView()}))})),this._bindGroups[n.outputs]=this._device.createBindGroup({label:this._label,layout:this._pipeline.getBindGroupLayout(n.outputs),entries:this._outputs.map((a,s)=>({binding:s,resource:{buffer:this._outputBuffers[a].buffer}}))});const o=t.beginComputePass();o.setPipeline(this._pipeline),this._bindGroups.forEach((a,s)=>{o.setBindGroup(s,a)}),o.dispatchWorkgroups(Math.ceil((this._execWidth??this._width)/me),Math.ceil((this._execHeight??this._height)/me),1),o.end()}}const ie=34028234663852886e22;function re(r,t,e){return Number.isNaN(r)?0:Math.min(Math.max(r,t),e)}const ki=`
|
|
15
|
+
const FLT_MAX = 3.4028234663852886e38;
|
|
16
|
+
fn sanitize(v: vec3f, lo: f32, hi: f32) -> vec3f {
|
|
17
|
+
let bits = bitcast<vec3u>(v);
|
|
18
|
+
let magnitude = bits & vec3u(0x7fffffffu);
|
|
19
|
+
let negative = (bits & vec3u(0x80000000u)) != vec3u(0u);
|
|
20
|
+
let infinite = select(vec3f(hi), vec3f(lo), negative);
|
|
21
|
+
let clamped = select(clamp(v, vec3f(lo), vec3f(hi)), infinite, magnitude == vec3u(0x7f800000u));
|
|
22
|
+
return select(clamped, vec3f(0.0), magnitude > vec3u(0x7f800000u));
|
|
23
|
+
}
|
|
24
|
+
`,Re=1412.83765,Xe=1.64593172,Ye=.431384981,Fe=-.00294139609,Ge=.192653254,He=.00626026094,Ve=.998620152,Ti=15794576e-13,Bi=.0322087631,Ii=.00223151711,Pi=.370974749,Ci=65504,Si=Oi(Ci),Mi=1/Si,Ei=Si,ye=Math.log(Ci+1),Ai=1/ye;function Oi(r){return r<=Ti?Re*r:r<=Bi?Xe*Math.pow(r,Ye)+Fe:Ge*Math.log(r+He)+Ve}function un(r){return r<=Ii?r/Re:r<=Pi?Math.pow((r-Fe)/Xe,1/Ye):Math.exp((r-Ve)/Ge)-He}function sn(r,t){return t==="log"?Math.log(r+1)*Ai:Oi(r)*Mi}function cn(r,t){return t==="log"?Math.exp(r*ye)-1:un(r*Ei)}function ln({data:r,channels:t,inputScale:e,transfer:i="pu"}){const n=new Float32Array(r);for(let o=0;o<n.length;o+=t)for(let a=0;a<3;a++)n[o+a]=sn(re(n[o+a]*e,0,ie),i);return n}function pn({data:r,channels:t,inputScale:e,transfer:i="pu"}){const n=new Float32Array(r),o=1/e;for(let a=0;a<n.length;a+=t)for(let s=0;s<3;s++)n[a+s]=cn(re(n[a+s],0,ie),i)*o;return n}function dn({data:r,channels:t}){let e=0;for(let a=0;a<r.length;a+=t){const s=re(r[a],0,ie),u=re(r[a+1],0,ie),l=re(r[a+2],0,ie),p=.212671*s+.71516*u+.072169*l;e+=Math.log2(p+1e-4)}const i=r.length/t,n=e/i;return .18/Math.pow(2,n)}const Ni=`
|
|
25
|
+
const a = ${Re};
|
|
26
|
+
const b = ${Xe};
|
|
27
|
+
const c = ${Ye};
|
|
28
|
+
const d = ${Fe};
|
|
29
|
+
const e = ${Ge};
|
|
30
|
+
const f = ${He};
|
|
31
|
+
const g = ${Ve};
|
|
32
|
+
const y0 =${Ti};
|
|
33
|
+
const y1 =${Bi};
|
|
34
|
+
const x0 =${Ii};
|
|
35
|
+
const x1 =${Pi};
|
|
36
|
+
|
|
37
|
+
const normScale = ${Mi};
|
|
38
|
+
const rcpNormScale = ${Ei};
|
|
39
|
+
`;function hn(r){return`
|
|
40
|
+
${Ni}
|
|
41
|
+
const logXMax = ${ye};
|
|
42
|
+
const logNormScale = ${Ai};
|
|
43
|
+
${r==="log"?"fn HDRForward(y: f32) -> f32 { return log(y + 1.0) * logNormScale; }":`fn HDRForward(y: f32) -> f32 {
|
|
44
|
+
if (y <= y0) { return a * y * normScale; }
|
|
45
|
+
else if (y <= y1) { return (b * pow(y, c) + d) * normScale; }
|
|
46
|
+
else { return (e * log(y + f) + g) * normScale; }
|
|
47
|
+
}`}`}function fn(r){return`
|
|
48
|
+
${Ni}
|
|
49
|
+
const logXMax = ${ye};
|
|
50
|
+
${r==="log"?"fn HDRInverse(x: f32) -> f32 { return exp(x * logXMax) - 1.0; }":`fn HDRInverse(x: f32) -> f32 {
|
|
51
|
+
let y = x * rcpNormScale;
|
|
52
|
+
if (y <= x0) { return y / a; }
|
|
53
|
+
else if (y <= x1) { return pow((y - d) / b, 1 / c); }
|
|
54
|
+
else { return exp((y - g) / e) - f; }
|
|
55
|
+
}`}
|
|
56
|
+
`}class je{constructor(t,e,i,n){this.x=t,this.y=e,this.width=i,this.height=n}}class gn{constructor(t,e,i="pu"){y(this,"_inputPassAux");y(this,"_inputPassColor");y(this,"_outputPass");y(this,"_copyPass");y(this,"_isInputTexture");this._device=t,this._isHDR=e,this._hdrTransfer=i;const n=[{label:"inputScale",type:"f32",data:new Float32Array([1])},{label:"inputSize",type:"vec2i",data:new Int32Array(2)},{label:"outputSize",type:"vec2i",data:new Int32Array(2)},{label:"inputOffset",type:"vec2i",data:new Int32Array(2)}];this._inputPassAux=new $e("inputPassAux",this._device,{inputs:["color","albedo","normal"],outputs:["color","albedo","normal"],uniforms:n,csDefine:"",csMain:""}),this._inputPassColor=new $e("inputPassColor",this._device,{inputs:["color"],outputs:["color"],uniforms:n,csDefine:"",csMain:""}),this._outputPass=new $e("outputPass",this._device,{inputs:["color","raw"],outputs:["color"],uniforms:[{label:"inputScale",type:"f32",data:new Float32Array([1])},{label:"inputSize",type:"vec2i",data:new Int32Array(2)},{label:"outputSize",type:"vec2i",data:new Int32Array(2)},{label:"imageSize",type:"vec2i",data:new Int32Array(2)},{label:"inputOffset",type:"vec2i",data:new Int32Array(2)},{label:"outputOffset",type:"vec2i",data:new Int32Array(2)}],csDefine:"",csMain:""}),this._copyPass=new $e("copyPass",this._device,{inputs:["color"],outputs:["color"],autoUpdateOutputBuffer:!1,uniforms:[{label:"size",type:"vec2i",data:new Int32Array(2)}],csMain:`
|
|
57
|
+
let outIdx = i32(globalId.x + globalId.y * u32(size.x));
|
|
58
|
+
out_color[outIdx] = textureLoad(in_color, globalId.xy, 0);
|
|
59
|
+
`}),this._inputPassAux.setOutputParams({color:{channels:3},albedo:{channels:3},normal:{channels:3}}),this._inputPassColor.setOutputParams({color:{channels:3}}),this._outputPass.setOutputParams({color:{channels:4}})}_updatePasses(t,e=!1){if(this._isInputTexture!=null&&this._isInputTexture===t)return;this._isInputTexture=t;const i=this._isHDR,n=hn(this._hdrTransfer)+ki;function o(s){return t?`textureLoad(in_${s}, vec2u(inputPosition), 0)`:`in_${s}[inIdx]`}const a=`
|
|
60
|
+
let x = i32(globalId.x);
|
|
61
|
+
let y = i32(globalId.y);
|
|
62
|
+
let inputPosition = clamp(
|
|
63
|
+
vec2i(x, y) + inputOffset,
|
|
64
|
+
vec2i(0),
|
|
65
|
+
inputSize - vec2i(1)
|
|
66
|
+
);
|
|
67
|
+
let inIdx = inputPosition.y * inputSize.x + inputPosition.x;
|
|
68
|
+
let col = ${o("color")};
|
|
69
|
+
|
|
70
|
+
let outIdx = y * outputSize.x + x;
|
|
71
|
+
|
|
72
|
+
if (${e}) {
|
|
73
|
+
// Denoise the inversed alpha. Or the anti aliased edge will be too dark after denoised
|
|
74
|
+
out_color[outIdx] = sanitize(vec3f(1.0 - col.a), 0.0, 1.0);
|
|
75
|
+
}
|
|
76
|
+
else if (${i}) {
|
|
77
|
+
let scaled = sanitize(col.rgb * inputScale, 0.0, FLT_MAX);
|
|
78
|
+
out_color[outIdx] = vec3f(HDRForward(scaled.r), HDRForward(scaled.g), HDRForward(scaled.b));
|
|
79
|
+
}
|
|
80
|
+
else {
|
|
81
|
+
out_color[outIdx] = sanitize(col.rgb, 0.0, 1.0);
|
|
82
|
+
}
|
|
83
|
+
`;this._inputPassAux.setCSCode({csDefine:n,csMain:`
|
|
84
|
+
${a}
|
|
85
|
+
let alb = ${o("albedo")};
|
|
86
|
+
let nor = ${o("normal")};
|
|
87
|
+
// Normals arrive [0, 1]-encoded, so OIDN's [-1, 1] clamp becomes [0, 1].
|
|
88
|
+
out_normal[outIdx] = sanitize(nor.rgb, 0.0, 1.0);
|
|
89
|
+
out_albedo[outIdx] = sanitize(alb.rgb, 0.0, 1.0);
|
|
90
|
+
`}),this._inputPassColor.setCSCode({csDefine:n,csMain:`
|
|
91
|
+
${a}
|
|
92
|
+
`}),this._outputPass.setCSCode({csDefine:fn(this._hdrTransfer)+ki,csMain:`
|
|
93
|
+
let x = i32(globalId.x);
|
|
94
|
+
let y = i32(globalId.y);
|
|
95
|
+
if (x >= outputSize.x || y >= outputSize.y) {
|
|
96
|
+
return;
|
|
97
|
+
}
|
|
98
|
+
let inIdx = (y + inputOffset.y) * inputSize.x + x + inputOffset.x;
|
|
99
|
+
let outIdx = (y + outputOffset.y) * imageSize.x + x + outputOffset.x;
|
|
100
|
+
let col = in_color[inIdx];
|
|
101
|
+
let raw = ${t?"textureLoad(in_raw, globalId.xy + vec2u(outputOffset), 0)":"in_raw[outIdx]"};
|
|
102
|
+
|
|
103
|
+
// The network output may contain negative values or NaNs.
|
|
104
|
+
let denoised = sanitize(col.rgb, 0.0, ${i&&!e?"FLT_MAX":"1.0"});
|
|
105
|
+
if (${e}) {
|
|
106
|
+
out_color[outIdx] = vec4f(raw.rgb, 1.0 - denoised.r);
|
|
107
|
+
}
|
|
108
|
+
else if (${i}) {
|
|
109
|
+
out_color[outIdx] = vec4f(
|
|
110
|
+
vec3f(HDRInverse(denoised.r), HDRInverse(denoised.g), HDRInverse(denoised.b)) / inputScale,
|
|
111
|
+
// Pick the alpha
|
|
112
|
+
raw.a
|
|
113
|
+
);
|
|
114
|
+
}
|
|
115
|
+
else {
|
|
116
|
+
out_color[outIdx] = vec4f(denoised, raw.a);
|
|
117
|
+
}
|
|
118
|
+
`})}setImageSize(t,e){this._inputPassAux.setUniform("inputSize",new Int32Array([t,e])),this._inputPassColor.setUniform("inputSize",new Int32Array([t,e])),this._outputPass.setUniform("imageSize",new Int32Array([t,e])),this._outputPass.setSize(t,e),this._copyPass.setSize(t,e),this._copyPass.setUniform("size",new Int32Array([t,e]))}setInputTile(t){const e=new Int32Array([t.width,t.height]);[this._inputPassAux,this._inputPassColor].forEach(i=>{i.setUniform("inputOffset",new Int32Array([t.x,t.y])),i.setUniform("outputSize",e),i.setSize(e[0],e[1])}),this._outputPass.setUniform("inputSize",e)}setOutputTile(t,e){const i=this._outputPass,n=new Int32Array([t.width,t.height]),o=t.x-e.x,a=t.y-e.y;i.setUniform("outputSize",n),i.setUniform("inputOffset",new Int32Array([o,a])),i.setUniform("outputOffset",new Int32Array([t.x,t.y])),i.setExecuteSize(n[0],n[1])}forward(t,e,i,n){const o=t instanceof GPUTexture;this._updatePasses(o,n);const a=this._inputPassAux,s=this._inputPassColor,u=this._device.createCommandEncoder();function l(p){return p instanceof GPUTexture?{texture:p,channels:4}:{buffer:p,channels:4}}return e&&i?a.createPass(u,{color:l(t),albedo:l(e),normal:l(i)}):s.createPass(u,{color:l(t)}),this._device.queue.submit([u.finish()]),e&&i?{color:a.getOutput("color"),albedo:a.getOutput("albedo"),normal:a.getOutput("normal")}:{color:s.getOutput("color")}}inverse(t,e){const n=this._device.createCommandEncoder(),o=this._outputPass;return o.createPass(n,{color:{buffer:t,channels:4},raw:e instanceof GPUBuffer?{buffer:e,channels:4}:{texture:e,channels:4}}),this._device.queue.submit([n.finish()]),o.getOutput("color")}copyInputDataToOutput(t){const e=this._device.createCommandEncoder(),n=this._outputPass.getOutput("color"),o=this._copyPass;t instanceof GPUTexture?(o.setOutputBuffers({color:n}),o.createPass(e,{color:{texture:t,channels:4}})):e.copyBufferToBuffer(t,0,n,0,n.size),this._device.queue.submit([e.finish()])}dispose(){this._outputPass.dispose(),this._inputPassAux.dispose(),this._inputPassColor.dispose(),this._copyPass.dispose(!1)}}const mn=256,$n=432,yn=16,_n=16,J=16;function Rt(r,t){return Math.ceil(r/t)*t}function zi(r,t){return Math.floor(r/t)*t}function Ui(r,t,e){const i=Math.round(e*r/t),n=Math.round((e+1)*r/t);return{start:i,end:n}}function _e(r,t,e,i=0){const n=t-r,o=Math.max(i,Rt(n,J));let a=r-Math.floor((o-n)/2);return a=we(a,0,Math.max(0,e-o)),{start:a,end:a+o}}function Ke(r,t,e,i){if(!Number.isInteger(r)||r<=0||!Number.isInteger(t)||t<=0)throw new Error("Tile grid dimensions must be positive integers");if(!Number.isFinite(e)||e<=0)throw new Error("Maximum tile size must be positive");if(!Number.isFinite(i)||i<0)throw new Error("Tile overlap must be non-negative");const n=Math.max(J,zi(e,J)),o=Rt(i,J),a=Math.max(1,Math.ceil(r/n)),s=Math.max(1,Math.ceil(t/n)),u=[];let l=0,p=0;for(let g=0;g<s;g++){const f=Ui(t,s,g);for(let _=0;_<a;_++){const m=Ui(r,a,_),v={start:_===0?m.start:Math.max(0,m.start-o),end:_===a-1?m.end:Math.min(r,m.end+o)},w={start:g===0?f.start:Math.max(0,f.start-o),end:g===s-1?f.end:Math.min(t,f.end+o)},T=_e(v.start,v.end,r),$=_e(w.start,w.end,t),b={x:m.start,y:f.start,width:m.end-m.start,height:f.end-f.start},x={x:T.start,y:$.start,width:T.end-T.start,height:$.end-$.start};l=Math.max(l,b.width),p=Math.max(p,b.height),u.push({column:_,row:g,input:x,output:b})}}const c=()=>new Set(u.map(({input:g})=>`${g.width}x${g.height}`));if(c().size>2){const g=Math.max(...u.map(({input:v})=>v.width)),f=Math.max(...u.map(({input:v})=>v.height)),_=u.reduce((v,{input:w})=>v+g*w.height,0),m=u.reduce((v,{input:w})=>v+w.width*f,0);if(_<=m)for(const v of u){const w=_e(v.input.x,v.input.x+v.input.width,r,g);v.input.x=w.start,v.input.width=w.end-w.start}else for(const v of u){const w=_e(v.input.y,v.input.y+v.input.height,t,f);v.input.y=w.start,v.input.height=w.end-w.start}}const d=c(),h=u.reduce((g,{input:f})=>g+f.width*f.height,0);return{columns:a,rows:s,overlap:o,maxOutputWidth:l,maxOutputHeight:p,inputPixelCount:h,inputShapeCount:d.size,tiles:u}}function we(r,t,e){return Math.min(Math.max(r,t),e)}function wn(r,t){const e=[...r].sort((i,n)=>i-n);return e[Math.ceil(e.length*t)-1]}class vn{constructor(t,e=!0){y(this,"enabled");y(this,"maxTileSize");y(this,"minTileSize");y(this,"targetTileTimeMs");y(this,"_tileSize");y(this,"_adjustmentStep");y(this,"_smoothedTileTimeMs");y(this,"_completeExecutionsSinceChange",0);const i=typeof e=="object"?e:{};this.enabled=e!==!1,this.maxTileSize=Math.max(J,zi(t,J)),this.minTileSize=we(Rt(i.minTileSize??mn,J),J,this.maxTileSize),this.targetTileTimeMs=Math.max(1,i.targetTileTimeMs??yn),this._adjustmentStep=Math.max(J,Rt(i.adjustmentStep??_n,J)),this._tileSize=this.enabled?we(Rt(i.initialTileSize??$n,J),this.minTileSize,this.maxTileSize):this.maxTileSize}get tileSize(){return this._tileSize}observe(t){if(!this.enabled||t.length===0)return!1;const e=t.filter(a=>Number.isFinite(a)&&a>=0);if(e.length===0)return!1;const i=e.length>=3?e.slice(1):e,n=wn(i,.75);if(this._smoothedTileTimeMs=this._smoothedTileTimeMs===void 0?n:this._smoothedTileTimeMs*.65+n*.35,this._completeExecutionsSinceChange++,this._completeExecutionsSinceChange<2)return!1;let o=this._tileSize;return this._smoothedTileTimeMs>this.targetTileTimeMs*1.25?o-=this._adjustmentStep:this._smoothedTileTimeMs<this.targetTileTimeMs*.65&&(o+=this._adjustmentStep),o=we(Rt(o,J),this.minTileSize,this.maxTileSize),o===this._tileSize?!1:(this._tileSize=o,this._smoothedTileTimeMs=void 0,this._completeExecutionsSinceChange=0,!0)}}function S(r,t){return{op:"conv2d",id:r,input:t,weight:`${r}.weight`,bias:`${r}.bias`,activation:"relu",padding:"same"}}function Bt(r,t){return{op:"maxPool2d",id:r,input:t,size:2,stride:2,padding:"same"}}function It(r,t){return{op:"upsample2d",id:r,input:t,scale:2,mode:"nearest"}}function Pt(r,t,e){return{op:"concat",id:r,inputs:[t,e],axis:"channels"}}const Wi={schemaVersion:1,id:"oidn-unet-small-v1",family:"oidn-unet-small",input:"input",output:"dec_conv0",receptiveField:174,nodes:[S("enc_conv0","input"),S("enc_conv1","enc_conv0"),Bt("pool1","enc_conv1"),S("enc_conv2","pool1"),Bt("pool2","enc_conv2"),S("enc_conv3","pool2"),Bt("pool3","enc_conv3"),S("enc_conv4","pool3"),Bt("pool4","enc_conv4"),S("enc_conv5a","pool4"),S("enc_conv5b","enc_conv5a"),It("up4","enc_conv5b"),Pt("concat4","up4","pool3"),S("dec_conv4a","concat4"),S("dec_conv4b","dec_conv4a"),It("up3","dec_conv4b"),Pt("concat3","up3","pool2"),S("dec_conv3a","concat3"),S("dec_conv3b","dec_conv3a"),It("up2","dec_conv3b"),Pt("concat2","up2","pool1"),S("dec_conv2a","concat2"),S("dec_conv2b","dec_conv2a"),It("up1","dec_conv2b"),Pt("concat1","up1","input"),S("dec_conv1a","concat1"),S("dec_conv1b","dec_conv1a"),S("dec_conv0","dec_conv1b")]},Li={schemaVersion:1,id:"oidn-unet-large-v1",family:"oidn-unet-large",input:"input",output:"dec_conv1c",receptiveField:202,nodes:[S("enc_conv1a","input"),S("enc_conv1b","enc_conv1a"),Bt("pool1","enc_conv1b"),S("enc_conv2a","pool1"),S("enc_conv2b","enc_conv2a"),Bt("pool2","enc_conv2b"),S("enc_conv3a","pool2"),S("enc_conv3b","enc_conv3a"),Bt("pool3","enc_conv3b"),S("enc_conv4a","pool3"),S("enc_conv4b","enc_conv4a"),Bt("pool4","enc_conv4b"),S("enc_conv5a","pool4"),S("enc_conv5b","enc_conv5a"),It("up4","enc_conv5b"),Pt("concat4","up4","pool3"),S("dec_conv4a","concat4"),S("dec_conv4b","dec_conv4a"),It("up3","dec_conv4b"),Pt("concat3","up3","pool2"),S("dec_conv3a","concat3"),S("dec_conv3b","dec_conv3a"),It("up2","dec_conv3b"),Pt("concat2","up2","pool1"),S("dec_conv2a","concat2"),S("dec_conv2b","dec_conv2a"),It("up1","dec_conv2b"),Pt("concat1","up1","input"),S("dec_conv1a","concat1"),S("dec_conv1b","dec_conv1a"),S("dec_conv1c","dec_conv1b")]},xn=[Wi,Li];function Di(r){const t=new Set;for(const e of r.nodes)e.op==="conv2d"&&(t.add(e.weight),t.add(e.bias));return t}function Ri(r){return r.desc.getByteSize()}function Xi(r){return[...r].sort().join(", ")}function qe(r,t=xn){const e=t.filter(i=>{const n=Di(i);return[...n].some(o=>!r.has(o))?!1:i.allowAdditionalTensors===!0||[...r.keys()].every(o=>n.has(o))});if(e.length===1)return e[0];throw e.length>1?new Error(`Ambiguous OIDN model topology: ${e.map(i=>i.id).join(", ")}`):new Error(`Unsupported OIDN model topology. TZA tensors: ${Xi([...r.keys()])}`)}function Yi(r,t,e){const i=r.get(t);if(!i)throw new Error(`Model ${e} is missing tensor ${t}`);if(i.data.byteLength!==Ri(i))throw new Error(`Tensor ${t} has ${i.data.byteLength} bytes, expected ${Ri(i)}`);return i}function Fi(r,t=qe(r)){if(t.schemaVersion!==1)throw new Error(`Unsupported model descriptor schema ${t.schemaVersion}`);const e=Di(t);if(!t.allowAdditionalTensors){const c=[...r.keys()].filter(d=>!e.has(d));if(c.length>0)throw new Error(`Model ${t.id} has unexpected tensors: ${Xi(c)}`)}const i=new Map,n=new Map,o=new Map,a=new Set([t.input]);let s,u;const l=(c,d)=>{const h=i.get(c);if(h===void 0)throw new Error(`Model ${t.id} node ${d} reads unknown or forward value ${c}`);return h};for(const c of t.nodes){if(a.has(c.id))throw new Error(`Model ${t.id} produces duplicate value ${c.id}`);if(c.op==="conv2d"){const d=Yi(r,c.weight,t.id),h=Yi(r,c.bias,t.id),g=d.desc.dims;if(d.desc.layout!=="oihw"||g.length!==4)throw new Error(`Tensor ${c.weight} must use OIHW layout`);if(g[2]!==3||g[3]!==3)throw new Error(`Tensor ${c.weight} must use a 3x3 kernel`);if(h.desc.layout!=="x"||h.desc.dims.length!==1)throw new Error(`Tensor ${c.bias} must be a one-dimensional bias`);if(h.desc.dims[0]!==g[0])throw new Error(`Tensor ${c.bias} has ${h.desc.dims[0]} channels, expected ${g[0]}`);if(d.desc.dataType!==h.desc.dataType)throw new Error(`Weight and bias dtype differ for ${c.id}`);if(u&&u!==d.desc.dataType)throw new Error(`Mixed tensor dtypes are not supported by model ${t.id}`);u=d.desc.dataType,c.input===t.input&&s===void 0&&(s=g[1],i.set(t.input,s));const f=l(c.input,c.id);if(f!==g[1])throw new Error(`Tensor ${c.weight} expects ${g[1]} input channels, but ${c.input} provides ${f}`);i.set(c.id,g[0]),n.set(c.id,{weight:d,bias:h,inputChannels:g[1],outputChannels:g[0],kernelHeight:g[2],kernelWidth:g[3]}),o.set(c.id,{inputChannels:g[1],outputChannels:g[0]})}else if(c.op==="concat"){if(c.inputs.length<2)throw new Error(`Concat ${c.id} requires at least two inputs`);const d=c.inputs.reduce((h,g)=>h+l(g,c.id),0);i.set(c.id,d)}else i.set(c.id,l(c.input,c.id));a.add(c.id)}if(s===void 0||u===void 0)throw new Error(`Model ${t.id} has no convolution reading its input`);const p=i.get(t.output);if(p===void 0)throw new Error(`Model ${t.id} output ${t.output} is not produced`);if(p!==3)throw new Error(`Model ${t.id} must produce 3 channels, got ${p}`);return{spec:t,inputChannels:s,outputChannels:p,tensorDataType:u,channelsByValue:i,convChannels:o,convTensors:n}}function bn(r){return r.op==="concat"?r.inputs:[r.input]}function kn(r){const t=new Map;for(const e of r.nodes)for(const i of bn(e)){const n=t.get(i)??[];n.push(e),t.set(i,n)}return t}function Ze(r,t,e){const i=r.get(t);if(!(i===void 0||i.length!==1||i[0].op!==e))return i[0]}function ve(r,t={}){var p;const e=r.spec,i=kn(e),n=new Map(e.nodes.map(c=>[c.id,c])),o=new Set,a=new Map;let s=0,u=0;if(t.fuseConvPool!==!1)for(const c of e.nodes){if(c.op!=="conv2d"||c.activation!=="relu")continue;const d=Ze(i,c.id,"maxPool2d");!d||d.op!=="maxPool2d"||d.size!==2||d.stride!==2||(o.add(c.id),a.set(d.id,{op:"fusedConvReluMaxPool2d",id:d.id,input:c.input,conv:c,pool:d}),s++)}if(t.fuseUpsampleConcatConv!==!1)for(const c of e.nodes){if(c.op!=="conv2d")continue;const d=n.get(c.input);if((d==null?void 0:d.op)!=="concat"||d.inputs.length!==2||((p=Ze(i,d.id,"conv2d"))==null?void 0:p.id)!==c.id)continue;const h=r.channelsByValue.get(d.inputs[0]);if(h===void 0||h%4!==0)continue;const g=d.inputs.map(_=>{var v;const m=n.get(_);return(m==null?void 0:m.op)==="upsample2d"&&m.scale===2&&m.mode==="nearest"&&((v=Ze(i,m.id,"concat"))==null?void 0:v.id)===d.id?{value:m.input,upsample:m}:{value:_}});if(g.filter(_=>_.upsample).length===1){o.add(d.id);for(const _ of d.inputs){const m=n.get(_);(m==null?void 0:m.op)==="upsample2d"&&o.add(m.id)}a.set(c.id,{op:"fusedUpsampleConcatConv2d",id:c.id,inputs:g,conv:c}),u++}}const l=[];for(const c of e.nodes){const d=a.get(c.id);d?l.push(d):o.has(c.id)||l.push(c)}return{spec:e,nodes:l,fusions:{convPool:s,upsampleConcatConv:u}}}function Tn(r){return r.op==="concat"?r.inputs:r.op==="fusedUpsampleConcatConv2d"?r.inputs.map(t=>t.value):[r.input]}function Gi(r,t){return r.width===t.width&&r.height===t.height}function Hi(r,t,e,i){if(!Number.isInteger(t)||t<=0||!Number.isInteger(e)||e<=0)throw new Error(`Invalid model input size ${t}x${e}`);const n=ve(r,i),o={width:t,height:e,channels:r.inputChannels},a=new Map([[r.spec.input,o]]),s=[],u=(c,d)=>{const h=a.get(c);if(!h)throw new Error(`Planned node ${d} reads missing value ${c}`);return h};for(const c of n.nodes){let d;if(c.op==="conv2d"){const h=u(c.input,c.id);d={width:h.width,height:h.height,channels:r.convChannels.get(c.id).outputChannels}}else if(c.op==="maxPool2d"){const h=u(c.input,c.id);d={width:Math.ceil(h.width/2),height:Math.ceil(h.height/2),channels:h.channels}}else if(c.op==="upsample2d"){const h=u(c.input,c.id);d={width:h.width*2,height:h.height*2,channels:h.channels}}else if(c.op==="concat"){const h=c.inputs.map(g=>u(g,c.id));if(h.some(g=>!Gi(g,h[0])))throw new Error(`Concat ${c.id} has mismatched spatial shapes`);d={width:h[0].width,height:h[0].height,channels:h.reduce((g,f)=>g+f.channels,0)}}else if(c.op==="fusedConvReluMaxPool2d"){const h=u(c.input,c.id);d={width:Math.ceil(h.width/2),height:Math.ceil(h.height/2),channels:r.convChannels.get(c.conv.id).outputChannels}}else{const h=c.inputs.map(g=>{const f=u(g.value,c.id);return g.upsample?{...f,width:f.width*2,height:f.height*2}:f});if(h.some(g=>!Gi(g,h[0])))throw new Error(`Fused decoder ${c.id} has mismatched spatial shapes`);d={width:h[0].width,height:h[0].height,channels:r.convChannels.get(c.conv.id).outputChannels}}a.set(c.id,d),s.push(d)}const l=new Map;n.nodes.forEach((c,d)=>{for(const h of Tn(c))l.set(h,d)}),l.set(r.spec.output,n.nodes.length);const p=n.nodes.map((c,d)=>({node:c,outputShape:s[d],lastUse:l.get(c.id)??d}));return{...n,inputShape:o,valueShapes:a,plannedNodes:p}}const ot={MAP_READ:1,MAP_WRITE:2,COPY_SRC:4,COPY_DST:8,INDEX:16,VERTEX:32,UNIFORM:64,STORAGE:128,INDIRECT:256,QUERY_RESOLVE:512},Vi={READ:1,WRITE:2},yt=8,Et=yt+2,ji=3*3;function Ki(r,t,e=!1){const i=r==="fp16"?2:4,n=Et*Et*t,o=e?ji*t*4:0;return(n+o)*4*i}function Bn(r){return r==="fp16"?"vec4<f16>":"vec4<f32>"}function In(r){return r==="fp16"?`enable f16;
|
|
119
|
+
`:""}function Pn(r,t,e){const i=n=>`${e}[weightBase + ${n}u]`;return r==="fp16"?`
|
|
120
|
+
let inputValue = vec4<f16>(${t});
|
|
121
|
+
var partial = vec4<f16>(0.0h);
|
|
122
|
+
partial = fma(${i(0)}, vec4<f16>(inputValue.x), partial);
|
|
123
|
+
partial = fma(${i(1)}, vec4<f16>(inputValue.y), partial);
|
|
124
|
+
partial = fma(${i(2)}, vec4<f16>(inputValue.z), partial);
|
|
125
|
+
partial = fma(${i(3)}, vec4<f16>(inputValue.w), partial);
|
|
126
|
+
acc += vec4<f32>(partial);
|
|
127
|
+
`:`
|
|
128
|
+
let inputValue = vec4<f32>(${t});
|
|
129
|
+
acc = fma(vec4<f32>(${i(0)}), vec4<f32>(inputValue.x), acc);
|
|
130
|
+
acc = fma(vec4<f32>(${i(1)}), vec4<f32>(inputValue.y), acc);
|
|
131
|
+
acc = fma(vec4<f32>(${i(2)}), vec4<f32>(inputValue.z), acc);
|
|
132
|
+
acc = fma(vec4<f32>(${i(3)}), vec4<f32>(inputValue.w), acc);
|
|
133
|
+
`}function Cn(r,t,e,i=!1){if(!Number.isInteger(e)||e<1)throw new Error(`Final RGB shader requires positive input blocks, got ${e}`);const n=Bn(r),o="vec4<f32>",a=Et*Et*e,s=ji*e*4,u=t==="relu"?"max(acc, vec4<f32>(0.0))":"acc",l=Pn(r,"inputTile[patchBase + inputBlock]",i?"weightTile":"weights");return`${In(r)}
|
|
134
|
+
struct Params {
|
|
135
|
+
inputWidth: u32,
|
|
136
|
+
inputHeight: u32,
|
|
137
|
+
outputWidth: u32,
|
|
138
|
+
outputHeight: u32,
|
|
139
|
+
inputBlocks: u32,
|
|
140
|
+
outputBlocks: u32,
|
|
141
|
+
}
|
|
142
|
+
|
|
143
|
+
@group(0) @binding(0) var<storage, read> inputData: array<${n}>;
|
|
144
|
+
@group(0) @binding(1) var<storage, read> weights: array<${n}>;
|
|
145
|
+
@group(0) @binding(2) var<storage, read> bias: array<vec4<f32>>;
|
|
146
|
+
@group(0) @binding(3) var<storage, read_write> outputData: array<${o}>;
|
|
147
|
+
@group(0) @binding(4) var<uniform> params: Params;
|
|
148
|
+
|
|
149
|
+
var<workgroup> inputTile: array<${n}, ${a}>;
|
|
150
|
+
${i?`var<workgroup> weightTile: array<${n}, ${s}>;`:""}
|
|
151
|
+
|
|
152
|
+
@compute @workgroup_size(${yt}, ${yt}, 1)
|
|
153
|
+
fn main(
|
|
154
|
+
@builtin(local_invocation_id) localId: vec3<u32>,
|
|
155
|
+
@builtin(global_invocation_id) gid: vec3<u32>,
|
|
156
|
+
@builtin(workgroup_id) workgroupId: vec3<u32>
|
|
157
|
+
) {
|
|
158
|
+
let localLinear = localId.y * ${yt}u + localId.x;
|
|
159
|
+
for (
|
|
160
|
+
var loadIndex = localLinear;
|
|
161
|
+
loadIndex < ${a}u;
|
|
162
|
+
loadIndex += ${yt*yt}u
|
|
163
|
+
) {
|
|
164
|
+
let tilePixel = loadIndex / ${e}u;
|
|
165
|
+
let inputBlock = loadIndex % ${e}u;
|
|
166
|
+
let tileX = tilePixel % ${Et}u;
|
|
167
|
+
let tileY = tilePixel / ${Et}u;
|
|
168
|
+
let inputX = i32(workgroupId.x * ${yt}u + tileX) - 1;
|
|
169
|
+
let inputY = i32(workgroupId.y * ${yt}u + tileY) - 1;
|
|
170
|
+
var value = ${n}(0.0);
|
|
171
|
+
if (
|
|
172
|
+
inputX >= 0 && inputX < i32(params.inputWidth) &&
|
|
173
|
+
inputY >= 0 && inputY < i32(params.inputHeight)
|
|
174
|
+
) {
|
|
175
|
+
let inputIndex =
|
|
176
|
+
(u32(inputY) * params.inputWidth + u32(inputX)) *
|
|
177
|
+
${e}u + inputBlock;
|
|
178
|
+
value = inputData[inputIndex];
|
|
179
|
+
}
|
|
180
|
+
inputTile[loadIndex] = value;
|
|
181
|
+
}
|
|
182
|
+
|
|
183
|
+
${i?` for (
|
|
184
|
+
var loadIndex = localLinear;
|
|
185
|
+
loadIndex < ${s}u;
|
|
186
|
+
loadIndex += ${yt*yt}u
|
|
187
|
+
) {
|
|
188
|
+
weightTile[loadIndex] = weights[loadIndex];
|
|
189
|
+
}
|
|
190
|
+
|
|
191
|
+
`:""} // Out-of-range invocations must reach this barrier before returning.
|
|
192
|
+
workgroupBarrier();
|
|
193
|
+
|
|
194
|
+
let outputInBounds =
|
|
195
|
+
gid.x < params.outputWidth &&
|
|
196
|
+
gid.y < params.outputHeight &&
|
|
197
|
+
gid.z < params.outputBlocks;
|
|
198
|
+
if (!outputInBounds) {
|
|
199
|
+
return;
|
|
200
|
+
}
|
|
201
|
+
|
|
202
|
+
var acc = bias[gid.z];
|
|
203
|
+
for (var ky = 0u; ky < 3u; ky++) {
|
|
204
|
+
let inputY = i32(gid.y) + i32(ky) - 1;
|
|
205
|
+
if (inputY < 0 || inputY >= i32(params.inputHeight)) {
|
|
206
|
+
continue;
|
|
207
|
+
}
|
|
208
|
+
for (var kx = 0u; kx < 3u; kx++) {
|
|
209
|
+
let inputX = i32(gid.x) + i32(kx) - 1;
|
|
210
|
+
if (inputX < 0 || inputX >= i32(params.inputWidth)) {
|
|
211
|
+
continue;
|
|
212
|
+
}
|
|
213
|
+
let patchBase =
|
|
214
|
+
((localId.y + ky) * ${Et}u + localId.x + kx) *
|
|
215
|
+
${e}u;
|
|
216
|
+
for (var inputBlock = 0u; inputBlock < ${e}u; inputBlock++) {
|
|
217
|
+
let weightBase =
|
|
218
|
+
((ky * 3u + kx) * ${e}u + inputBlock) * 4u;
|
|
219
|
+
${l}
|
|
220
|
+
}
|
|
221
|
+
}
|
|
222
|
+
}
|
|
223
|
+
|
|
224
|
+
let outputIndex =
|
|
225
|
+
(gid.y * params.outputWidth + gid.x) * 1u + gid.z;
|
|
226
|
+
outputData[outputIndex] = ${u};
|
|
227
|
+
}
|
|
228
|
+
`}class qi{constructor(){y(this,"_stats",new Map)}track(t,e){let i=this._stats.get(t);return i||(i={created:0,destroyed:0,live:0,peakLive:0,resources:new Set},this._stats.set(t,i)),i.resources.has(e)||(i.resources.add(e),i.created++,i.live++,i.peakLive=Math.max(i.peakLive,i.live)),e}release(t,e,i){if(!e)return!1;const n=this._stats.get(t);if(!(n!=null&&n.resources.delete(e)))return!1;try{i()}finally{n.destroyed++,n.live--}return!0}snapshot(t=0){let e=0,i=0,n=0,o=0;const a={};for(const[s,u]of this._stats){const l={created:u.created,destroyed:u.destroyed,live:u.live,peakLive:u.peakLive};a[s]=l,e+=l.live,i+=l.created,n+=l.destroyed,o+=l.peakLive}return{live:e,created:i,destroyed:n,peakLive:o,pending:t,byKind:a}}}const Zi=new WeakMap;function Sn(r){let t=Zi.get(r);return t||(t=new Map,Zi.set(r,t)),t}const tt=8,Qi=8,j=8,pt=j+2;function dt(r){const[t,e]=r.workgroupSize;return{workgroupX:t,workgroupY:e,rowsPerThread:r.rowsPerThread,tileM:e*r.rowsPerThread,tileNBlocks:t,tileKBlocks:Qi}}function bt(r){const t=dt(r),e=r.sharedLayout==="padded"||r.sharedLayout==="padded-input",i=r.sharedLayout==="padded"||r.sharedLayout==="padded-weights";return{input:t.tileM*t.tileKBlocks+(e?t.workgroupY:0),weights:t.tileKBlocks*t.tileNBlocks*(i?5:4)}}function Xt(r,t){return t.sharedLayout==="padded"||t.sharedLayout==="padded-input"?`(${r}) + (${r}) / ${t.rowsPerThread*Qi}u`:r}function Qe(r,t){return t.sharedLayout==="padded"||t.sharedLayout==="padded-weights"?`(${r}) + (${r}) / 4u`:r}function Je(r,t){return Math.ceil(r/t)*t}function M(r){return Math.ceil(r/4)}function Mn(r,t){return r.width*r.height*M(r.channels)*4*t}let ne;function Ji(r){const t=r&32768?-1:1,e=r>>>10&31,i=r&1023;return e===0?t*i*2**-24:e===31?i===0?t*(1/0):NaN:t*(1+i/1024)*2**(e-15)}const tr=new DataView(new ArrayBuffer(4));function er(r){tr.setFloat32(0,r);const t=tr.getUint32(0),e=t>>>16&32768,i=t>>>23&255,n=t&8388607;if(i===255)return e|31744|(n?512:0);const o=i-127+15;if(o>=31)return e|31744;if(o<=0){if(o<-10)return e;const u=n|8388608,l=14-o,p=u>>>l,c=u&(1<<l)-1,d=1<<l-1,h=c>d||c===d&&p&1?p+1:p;return e|h}const a=o<<10|n>>>13,s=n&8191;return e|(s>4096||s===4096&&a&1?a+1:a)}function En(){if(!ne){ne=new Float32Array(65536);for(let r=0;r<ne.length;r++)ne[r]=Ji(r)}return ne}function xe(r){if(r.desc.dataType==="Float32")return new Float32Array(r.data.buffer,r.data.byteOffset,r.data.byteLength/4);const t=new Uint16Array(r.data.buffer,r.data.byteOffset,r.data.byteLength/2),e=new Float32Array(t.length);if(t.length<4096)for(let i=0;i<t.length;i++)e[i]=Ji(t[i]);else{const i=En();for(let n=0;n<t.length;n++)e[n]=i[t[n]]}return e}function oe(r,t,e,i){const n=Je(e.byteLength,4),o=r.createBuffer({label:t,size:n,usage:i,mappedAtCreation:!0});return new Uint8Array(o.getMappedRange()).set(e),o.unmap(),o}function ir(r,t,e){const i=new Uint32Array(Je(e.length,4));return i.set(e),oe(r,t,new Uint8Array(i.buffer,i.byteOffset,i.byteLength),ot.UNIFORM)}function An(r,t,e,i,n){const o=M(e.inputChannels),a=M(e.outputChannels),s=a*e.kernelHeight*e.kernelWidth*o*4*4,u=i==="fp16"&&e.weight.desc.dataType==="Float16",l=i==="fp16"?new Uint16Array(s):void 0,p=l?void 0:new Float32Array(s),c=u?new Uint16Array(e.weight.data.buffer,e.weight.data.byteOffset,e.weight.data.byteLength/2):void 0,d=c?void 0:xe(e.weight);for(let f=0;f<a;f++)for(let _=0;_<e.kernelHeight;_++)for(let m=0;m<e.kernelWidth;m++)for(let v=0;v<o;v++)for(let w=0;w<4;w++){const T=f*4+w;for(let $=0;$<4;$++){const b=v*4+$,B=(n==="k-major"?(((_*e.kernelWidth+m)*o+v)*a+f)*16:(((f*e.kernelHeight+_)*e.kernelWidth+m)*o+v)*16)+$*4+w;if(T<e.outputChannels&&b<e.inputChannels){const k=((T*e.inputChannels+b)*e.kernelHeight+_)*e.kernelWidth+m;l?l[B]=c?c[k]:er(d[k]):p[B]=d[k]}}}const h=new Float32Array(a*4);h.set(xe(e.bias));const g=oe(r,`oidn/${t}/weights/${i}/${n}`,l?new Uint8Array(l.buffer):new Uint8Array(p.buffer),ot.STORAGE);try{return{weights:g,weightLayout:n,bias:oe(r,`oidn/${t}/bias`,new Uint8Array(h.buffer),ot.STORAGE)}}catch(f){throw g.destroy(),f}}function D(r){return r==="fp16"?"vec4<f16>":"vec4<f32>"}function ft(r){return r==="fp16"?`enable f16;
|
|
229
|
+
`:""}function ut(r,t){return t==="fp16"?`vec4<f16>(${r})`:r}function ht(r,t){return t==="relu"?`max(${r}, vec4<f32>(0.0))`:r}function ae(r,t,e){return e==="fp32"?`
|
|
230
|
+
let inputValue = vec4<f32>(${r});
|
|
231
|
+
let weightBase = ${t};
|
|
232
|
+
acc = fma(vec4<f32>(weights[weightBase]), vec4<f32>(inputValue.x), acc);
|
|
233
|
+
acc = fma(vec4<f32>(weights[weightBase + 1u]), vec4<f32>(inputValue.y), acc);
|
|
234
|
+
acc = fma(vec4<f32>(weights[weightBase + 2u]), vec4<f32>(inputValue.z), acc);
|
|
235
|
+
acc = fma(vec4<f32>(weights[weightBase + 3u]), vec4<f32>(inputValue.w), acc);
|
|
236
|
+
`:`
|
|
237
|
+
let inputValue = vec4<f16>(${r});
|
|
238
|
+
let weightBase = ${t};
|
|
239
|
+
var partial = vec4<f16>(0.0h);
|
|
240
|
+
partial = fma(weights[weightBase], vec4<f16>(inputValue.x), partial);
|
|
241
|
+
partial = fma(weights[weightBase + 1u], vec4<f16>(inputValue.y), partial);
|
|
242
|
+
partial = fma(weights[weightBase + 2u], vec4<f16>(inputValue.z), partial);
|
|
243
|
+
partial = fma(weights[weightBase + 3u], vec4<f16>(inputValue.w), partial);
|
|
244
|
+
acc += vec4<f32>(partial);
|
|
245
|
+
`}function On(r,t,e){const i=e==="fp16"?"vec4<f16>":"vec4<f32>",n=e==="fp16"?`var partial = vec4<f16>(0.0h);
|
|
246
|
+
partial = fma(subgroupBroadcast(weights[weightBase], 0u), ${i}(inputValue.x), partial);
|
|
247
|
+
partial = fma(subgroupBroadcast(weights[weightBase + 1u], 0u), ${i}(inputValue.y), partial);
|
|
248
|
+
partial = fma(subgroupBroadcast(weights[weightBase + 2u], 0u), ${i}(inputValue.z), partial);
|
|
249
|
+
partial = fma(subgroupBroadcast(weights[weightBase + 3u], 0u), ${i}(inputValue.w), partial);
|
|
250
|
+
acc += vec4<f32>(partial);`:`acc = fma(subgroupBroadcast(weights[weightBase], 0u), vec4<f32>(inputValue.x), acc);
|
|
251
|
+
acc = fma(subgroupBroadcast(weights[weightBase + 1u], 0u), vec4<f32>(inputValue.y), acc);
|
|
252
|
+
acc = fma(subgroupBroadcast(weights[weightBase + 2u], 0u), vec4<f32>(inputValue.z), acc);
|
|
253
|
+
acc = fma(subgroupBroadcast(weights[weightBase + 3u], 0u), vec4<f32>(inputValue.w), acc);`;return`
|
|
254
|
+
let inputValue = ${i}(${r});
|
|
255
|
+
let weightBase = ${t};
|
|
256
|
+
${n}
|
|
257
|
+
`}function Yt(r,t,e){return r==="fp16"&&(t.loadMode==="packed-all"||e&&t.loadMode==="packed-weights")}function At(r,t){return t?`bitcast<vec4<f16>>(${r})`:r}function ti(r,t){const{rowsPerThread:e,tileNBlocks:i,tileKBlocks:n}=dt(t);if(!t.unrollRows||t.accumulationOrder==="row-major")return Nn(r,t);const o=`(tileK * ${i}u + localId.x) * ${t.sharedLayout==="padded"||t.sharedLayout==="padded-weights"?5:4}u`,a=Array.from({length:e},(u,l)=>l),s=D(r);return`
|
|
258
|
+
${r==="fp16"?a.map(u=>`var partial${u} = vec4<f16>(0.0h);`).join(`
|
|
259
|
+
`):""}
|
|
260
|
+
for (var tileK = 0u; tileK < ${n}u; tileK++) {
|
|
261
|
+
let weightBase = ${o};
|
|
262
|
+
let w0 = weightTile[weightBase];
|
|
263
|
+
let w1 = weightTile[weightBase + 1u];
|
|
264
|
+
let w2 = weightTile[weightBase + 2u];
|
|
265
|
+
let w3 = weightTile[weightBase + 3u];
|
|
266
|
+
${a.map(u=>{const l=Xt(`(localId.y * ${e}u + ${u}u) * ${n}u + tileK`,t),p=r==="fp16"?`partial${u}`:`acc[${u}]`;return`let v${u} = inputTile[${l}];
|
|
267
|
+
`+["x","y","z","w"].map((c,d)=>`${p} = fma(w${d}, ${s}(v${u}.${c}), ${p});`).join(`
|
|
268
|
+
`)}).join(`
|
|
269
|
+
`)}
|
|
270
|
+
}
|
|
271
|
+
${r==="fp16"?a.map(u=>`acc[${u}] += vec4<f32>(partial${u});`).join(`
|
|
272
|
+
`):""}
|
|
273
|
+
`}function Nn(r,t){const{rowsPerThread:e,tileNBlocks:i,tileKBlocks:n}=dt(t),o=Xt(`tileSpatial * ${n}u + tileK`,t),a=`(tileK * ${i}u + localId.x) * ${t.sharedLayout==="padded"||t.sharedLayout==="padded-weights"?5:4}u`;if(t.accumulationOrder==="row-major"){const s=D(r),u=r==="fp16"?"partial":"acc[row]";return`
|
|
274
|
+
for (var row = 0u; row < ${e}u; row++) {
|
|
275
|
+
${r==="fp16"?"var partial = vec4<f16>(0.0h);":""}
|
|
276
|
+
let tileSpatial = localId.y * ${e}u + row;
|
|
277
|
+
for (var tileK = 0u; tileK < ${n}u; tileK++) {
|
|
278
|
+
let weightBase = ${a};
|
|
279
|
+
let inputValue = inputTile[${o}];
|
|
280
|
+
${u} = fma(weightTile[weightBase], ${s}(inputValue.x), ${u});
|
|
281
|
+
${u} = fma(weightTile[weightBase + 1u], ${s}(inputValue.y), ${u});
|
|
282
|
+
${u} = fma(weightTile[weightBase + 2u], ${s}(inputValue.z), ${u});
|
|
283
|
+
${u} = fma(weightTile[weightBase + 3u], ${s}(inputValue.w), ${u});
|
|
284
|
+
}
|
|
285
|
+
${r==="fp16"?"acc[row] += vec4<f32>(partial);":""}
|
|
286
|
+
}
|
|
287
|
+
`}return r==="fp16"?`
|
|
288
|
+
var partial: array<vec4<f16>, ${e}>;
|
|
289
|
+
for (var row = 0u; row < ${e}u; row++) {
|
|
290
|
+
partial[row] = vec4<f16>(0.0h);
|
|
291
|
+
}
|
|
292
|
+
for (var tileK = 0u; tileK < ${n}u; tileK++) {
|
|
293
|
+
let weightBase =
|
|
294
|
+
${a};
|
|
295
|
+
for (var row = 0u; row < ${e}u; row++) {
|
|
296
|
+
let tileSpatial =
|
|
297
|
+
localId.y * ${e}u + row;
|
|
298
|
+
let inputValue =
|
|
299
|
+
inputTile[${o}];
|
|
300
|
+
partial[row] = fma(
|
|
301
|
+
weightTile[weightBase],
|
|
302
|
+
vec4<f16>(inputValue.x),
|
|
303
|
+
partial[row]
|
|
304
|
+
);
|
|
305
|
+
partial[row] = fma(
|
|
306
|
+
weightTile[weightBase + 1u],
|
|
307
|
+
vec4<f16>(inputValue.y),
|
|
308
|
+
partial[row]
|
|
309
|
+
);
|
|
310
|
+
partial[row] = fma(
|
|
311
|
+
weightTile[weightBase + 2u],
|
|
312
|
+
vec4<f16>(inputValue.z),
|
|
313
|
+
partial[row]
|
|
314
|
+
);
|
|
315
|
+
partial[row] = fma(
|
|
316
|
+
weightTile[weightBase + 3u],
|
|
317
|
+
vec4<f16>(inputValue.w),
|
|
318
|
+
partial[row]
|
|
319
|
+
);
|
|
320
|
+
}
|
|
321
|
+
}
|
|
322
|
+
for (var row = 0u; row < ${e}u; row++) {
|
|
323
|
+
acc[row] += vec4<f32>(partial[row]);
|
|
324
|
+
}
|
|
325
|
+
`:`
|
|
326
|
+
for (var tileK = 0u; tileK < ${n}u; tileK++) {
|
|
327
|
+
let weightBase =
|
|
328
|
+
${a};
|
|
329
|
+
for (var row = 0u; row < ${e}u; row++) {
|
|
330
|
+
let tileSpatial =
|
|
331
|
+
localId.y * ${e}u + row;
|
|
332
|
+
let inputValue =
|
|
333
|
+
inputTile[${o}];
|
|
334
|
+
acc[row] = fma(
|
|
335
|
+
weightTile[weightBase],
|
|
336
|
+
vec4<f32>(inputValue.x),
|
|
337
|
+
acc[row]
|
|
338
|
+
);
|
|
339
|
+
acc[row] = fma(
|
|
340
|
+
weightTile[weightBase + 1u],
|
|
341
|
+
vec4<f32>(inputValue.y),
|
|
342
|
+
acc[row]
|
|
343
|
+
);
|
|
344
|
+
acc[row] = fma(
|
|
345
|
+
weightTile[weightBase + 2u],
|
|
346
|
+
vec4<f32>(inputValue.z),
|
|
347
|
+
acc[row]
|
|
348
|
+
);
|
|
349
|
+
acc[row] = fma(
|
|
350
|
+
weightTile[weightBase + 3u],
|
|
351
|
+
vec4<f32>(inputValue.w),
|
|
352
|
+
acc[row]
|
|
353
|
+
);
|
|
354
|
+
}
|
|
355
|
+
}
|
|
356
|
+
`}function zn(r,t,e,i,n){const o=D(r),a=D(r),s=D(t),u=ut(ht("acc",e),t);return`${ft(r)}
|
|
357
|
+
struct Params {
|
|
358
|
+
inputWidth: u32,
|
|
359
|
+
inputHeight: u32,
|
|
360
|
+
outputWidth: u32,
|
|
361
|
+
outputHeight: u32,
|
|
362
|
+
inputBlocks: u32,
|
|
363
|
+
outputBlocks: u32,
|
|
364
|
+
}
|
|
365
|
+
|
|
366
|
+
@group(0) @binding(0) var<storage, read> inputData: array<${o}>;
|
|
367
|
+
@group(0) @binding(1) var<storage, read> weights: array<${a}>;
|
|
368
|
+
@group(0) @binding(2) var<storage, read> bias: array<vec4<f32>>;
|
|
369
|
+
@group(0) @binding(3) var<storage, read_write> outputData: array<${s}>;
|
|
370
|
+
@group(0) @binding(4) var<uniform> params: Params;
|
|
371
|
+
|
|
372
|
+
@compute @workgroup_size(${tt}, ${tt}, 1)
|
|
373
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
374
|
+
if (gid.x >= params.outputWidth || gid.y >= params.outputHeight || gid.z >= ${n}u) {
|
|
375
|
+
return;
|
|
376
|
+
}
|
|
377
|
+
var acc = bias[gid.z];
|
|
378
|
+
for (var ky = 0u; ky < 3u; ky++) {
|
|
379
|
+
let inputY = i32(gid.y) + i32(ky) - 1;
|
|
380
|
+
if (inputY < 0 || inputY >= i32(params.inputHeight)) { continue; }
|
|
381
|
+
for (var kx = 0u; kx < 3u; kx++) {
|
|
382
|
+
let inputX = i32(gid.x) + i32(kx) - 1;
|
|
383
|
+
if (inputX < 0 || inputX >= i32(params.inputWidth)) { continue; }
|
|
384
|
+
let pixelBase = (u32(inputY) * params.inputWidth + u32(inputX)) * ${i}u;
|
|
385
|
+
for (var inputBlock = 0u; inputBlock < ${i}u; inputBlock++) {
|
|
386
|
+
${ae("inputData[pixelBase + inputBlock]",`((((gid.z * 3u + ky) * 3u + kx) * ${i}u + inputBlock) * 4u)`,r)}
|
|
387
|
+
}
|
|
388
|
+
}
|
|
389
|
+
}
|
|
390
|
+
let outputIndex = (gid.y * params.outputWidth + gid.x) * ${n}u + gid.z;
|
|
391
|
+
outputData[outputIndex] = ${u};
|
|
392
|
+
}
|
|
393
|
+
`}function Un(r,t,e,i,n){const o=D(r),a=D(r),s=D(t),u=ut(ht("acc",e),t);return`${ft(r)}
|
|
394
|
+
enable subgroups;
|
|
395
|
+
struct Params {
|
|
396
|
+
inputWidth: u32,
|
|
397
|
+
inputHeight: u32,
|
|
398
|
+
outputWidth: u32,
|
|
399
|
+
outputHeight: u32,
|
|
400
|
+
inputBlocks: u32,
|
|
401
|
+
outputBlocks: u32,
|
|
402
|
+
}
|
|
403
|
+
@group(0) @binding(0) var<storage, read> inputData: array<${o}>;
|
|
404
|
+
@group(0) @binding(1) var<storage, read> weights: array<${a}>;
|
|
405
|
+
@group(0) @binding(2) var<storage, read> bias: array<vec4<f32>>;
|
|
406
|
+
@group(0) @binding(3) var<storage, read_write> outputData: array<${s}>;
|
|
407
|
+
@group(0) @binding(4) var<uniform> params: Params;
|
|
408
|
+
|
|
409
|
+
@compute @workgroup_size(${tt}, ${tt}, 1)
|
|
410
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
411
|
+
let outputInBounds =
|
|
412
|
+
gid.x < params.outputWidth && gid.y < params.outputHeight;
|
|
413
|
+
var acc = bias[gid.z];
|
|
414
|
+
for (var ky = 0u; ky < 3u; ky++) {
|
|
415
|
+
let inputY = i32(gid.y) + i32(ky) - 1;
|
|
416
|
+
let clampedY = u32(clamp(inputY, 0, i32(params.inputHeight) - 1));
|
|
417
|
+
for (var kx = 0u; kx < 3u; kx++) {
|
|
418
|
+
let inputX = i32(gid.x) + i32(kx) - 1;
|
|
419
|
+
let clampedX = u32(clamp(inputX, 0, i32(params.inputWidth) - 1));
|
|
420
|
+
let inputInBounds =
|
|
421
|
+
outputInBounds && inputX >= 0 && inputY >= 0 &&
|
|
422
|
+
inputX < i32(params.inputWidth) && inputY < i32(params.inputHeight);
|
|
423
|
+
let pixelBase =
|
|
424
|
+
(clampedY * params.inputWidth + clampedX) * ${i}u;
|
|
425
|
+
for (var inputBlock = 0u; inputBlock < ${i}u; inputBlock++) {
|
|
426
|
+
${On(`select(${o}(0.0), inputData[pixelBase + inputBlock], inputInBounds)`,`((((gid.z * 3u + ky) * 3u + kx) * ${i}u + inputBlock) * 4u)`,r)}
|
|
427
|
+
}
|
|
428
|
+
}
|
|
429
|
+
}
|
|
430
|
+
if (outputInBounds) {
|
|
431
|
+
let outputIndex =
|
|
432
|
+
(gid.y * params.outputWidth + gid.x) * ${n}u + gid.z;
|
|
433
|
+
outputData[outputIndex] = ${u};
|
|
434
|
+
}
|
|
435
|
+
}
|
|
436
|
+
`}function Wn(r,t,e,i,n){const o=D(r),a=D(t),s=ut(ht("acc",e),t),u=pt*pt*i,l=j*j;return`${ft(r)}
|
|
437
|
+
struct Params {
|
|
438
|
+
inputWidth: u32,
|
|
439
|
+
inputHeight: u32,
|
|
440
|
+
outputWidth: u32,
|
|
441
|
+
outputHeight: u32,
|
|
442
|
+
inputBlocks: u32,
|
|
443
|
+
outputBlocks: u32,
|
|
444
|
+
}
|
|
445
|
+
@group(0) @binding(0) var<storage, read> inputData: array<${o}>;
|
|
446
|
+
@group(0) @binding(1) var<storage, read> weights: array<${o}>;
|
|
447
|
+
@group(0) @binding(2) var<storage, read> bias: array<vec4<f32>>;
|
|
448
|
+
@group(0) @binding(3) var<storage, read_write> outputData: array<${a}>;
|
|
449
|
+
@group(0) @binding(4) var<uniform> params: Params;
|
|
450
|
+
|
|
451
|
+
var<workgroup> inputPatch: array<${o}, ${u}>;
|
|
452
|
+
|
|
453
|
+
@compute @workgroup_size(${j}, ${j}, 1)
|
|
454
|
+
fn main(
|
|
455
|
+
@builtin(local_invocation_id) localId: vec3<u32>,
|
|
456
|
+
@builtin(workgroup_id) workgroupId: vec3<u32>
|
|
457
|
+
) {
|
|
458
|
+
let localLinear =
|
|
459
|
+
localId.y * ${j}u + localId.x;
|
|
460
|
+
for (
|
|
461
|
+
var loadIndex = localLinear;
|
|
462
|
+
loadIndex < ${u}u;
|
|
463
|
+
loadIndex += ${l}u
|
|
464
|
+
) {
|
|
465
|
+
let patchPixel = loadIndex / ${i}u;
|
|
466
|
+
let inputBlock = loadIndex % ${i}u;
|
|
467
|
+
let patchX = patchPixel % ${pt}u;
|
|
468
|
+
let patchY = patchPixel / ${pt}u;
|
|
469
|
+
let inputX =
|
|
470
|
+
i32(workgroupId.x * ${j}u + patchX) - 1;
|
|
471
|
+
let inputY =
|
|
472
|
+
i32(workgroupId.y * ${j}u + patchY) - 1;
|
|
473
|
+
var value = ${o}(0.0);
|
|
474
|
+
if (
|
|
475
|
+
inputX >= 0 && inputX < i32(params.inputWidth) &&
|
|
476
|
+
inputY >= 0 && inputY < i32(params.inputHeight)
|
|
477
|
+
) {
|
|
478
|
+
let inputIndex =
|
|
479
|
+
(u32(inputY) * params.inputWidth + u32(inputX)) *
|
|
480
|
+
${i}u + inputBlock;
|
|
481
|
+
value = inputData[inputIndex];
|
|
482
|
+
}
|
|
483
|
+
inputPatch[loadIndex] = value;
|
|
484
|
+
}
|
|
485
|
+
workgroupBarrier();
|
|
486
|
+
|
|
487
|
+
let outputX =
|
|
488
|
+
workgroupId.x * ${j}u + localId.x;
|
|
489
|
+
let outputY =
|
|
490
|
+
workgroupId.y * ${j}u + localId.y;
|
|
491
|
+
let outputBlock = workgroupId.z;
|
|
492
|
+
if (
|
|
493
|
+
outputX >= params.outputWidth || outputY >= params.outputHeight ||
|
|
494
|
+
outputBlock >= ${n}u
|
|
495
|
+
) {
|
|
496
|
+
return;
|
|
497
|
+
}
|
|
498
|
+
|
|
499
|
+
var acc = bias[outputBlock];
|
|
500
|
+
for (var ky = 0u; ky < 3u; ky++) {
|
|
501
|
+
for (var kx = 0u; kx < 3u; kx++) {
|
|
502
|
+
let patchBase =
|
|
503
|
+
((localId.y + ky) * ${pt}u + localId.x + kx) *
|
|
504
|
+
${i}u;
|
|
505
|
+
for (var inputBlock = 0u; inputBlock < ${i}u; inputBlock++) {
|
|
506
|
+
${ae("inputPatch[patchBase + inputBlock]",`((((outputBlock * 3u + ky) * 3u + kx) * ${i}u + inputBlock) * 4u)`,r)}
|
|
507
|
+
}
|
|
508
|
+
}
|
|
509
|
+
}
|
|
510
|
+
let outputIndex =
|
|
511
|
+
(outputY * params.outputWidth + outputX) * ${n}u + outputBlock;
|
|
512
|
+
outputData[outputIndex] = ${s};
|
|
513
|
+
}
|
|
514
|
+
`}function rr(r,t,e,i){const{workgroupX:n,workgroupY:o,tileM:a,tileKBlocks:s}=dt(e),u=t?"params.outputWidth":"params.inputWidth",l=t?"params.outputHeight":"params.inputHeight",p=e.addressMode==="base-offset",c=i?i.blocks.map((h,g)=>`cachedBase${g}[load] = (${g===i.upsampled?"y / 2u":"y"} * params.source${g}Width + ${g===i.upsampled?"x / 2u":"x"}) * ${h}u;`).join(`
|
|
515
|
+
`):`cachedBase0[load] = (y * params.inputWidth + x) * ${r}u;`,d=a*s/(n*o);return`
|
|
516
|
+
${p?`var cachedBase0: array<u32, ${d}>;
|
|
517
|
+
${i?`var cachedBase1: array<u32, ${d}>;`:""}`:`var cachedX: array<u32, ${d}>;
|
|
518
|
+
var cachedY: array<u32, ${d}>;`}
|
|
519
|
+
var cachedMask: array<u32, ${d}>;
|
|
520
|
+
for (var load = 0u; load < ${d}u; load++) {
|
|
521
|
+
let loadIndex = localLinear + load * ${n*o}u;
|
|
522
|
+
let spatial = workgroupId.x * ${a}u + loadIndex / ${s}u;
|
|
523
|
+
let x = spatial % params.outputWidth;
|
|
524
|
+
let y = spatial / params.outputWidth;
|
|
525
|
+
${p?c:`cachedX[load] = x;
|
|
526
|
+
cachedY[load] = y;`}
|
|
527
|
+
let columns =
|
|
528
|
+
select(0u, 0x049u, x > 0u && x - 1u < ${u}) |
|
|
529
|
+
select(0u, 0x092u, x < ${u}) |
|
|
530
|
+
select(0u, 0x124u, x + 1u < ${u});
|
|
531
|
+
let rows =
|
|
532
|
+
select(0u, 0x007u, y > 0u && y - 1u < ${l}) |
|
|
533
|
+
select(0u, 0x038u, y < ${l}) |
|
|
534
|
+
select(0u, 0x1c0u, y + 1u < ${l});
|
|
535
|
+
cachedMask[load] = select(0u, columns & rows, spatial < spatialCount)${p&&i?" | ((x & 1u) << 9u) | ((y & 1u) << 10u)":""};
|
|
536
|
+
}
|
|
537
|
+
var channelBlock = (localLinear % ${s}u) % ${r}u;
|
|
538
|
+
var filterPosition = (localLinear % ${s}u) / ${r}u;
|
|
539
|
+
`}function nr(r,t){const{tileKBlocks:e}=dt(t);return`
|
|
540
|
+
channelBlock += ${e%r}u;
|
|
541
|
+
let carry = channelBlock >= ${r}u;
|
|
542
|
+
channelBlock -= select(0u, ${r}u, carry);
|
|
543
|
+
filterPosition += ${Math.floor(e/r)}u + select(0u, 1u, carry);
|
|
544
|
+
`}function be(r,t,e){const{workgroupX:i,workgroupY:n,tileM:o,tileKBlocks:a}=dt(e);return`
|
|
545
|
+
let inputBlock = channelBlock;
|
|
546
|
+
let kernelY = filterPosition / 3u;
|
|
547
|
+
let kernelX = filterPosition % 3u;
|
|
548
|
+
for (var load = 0u; load < ${o*a/(i*n)}u; load++) {
|
|
549
|
+
let loadIndex = localLinear + load * ${i*n}u;
|
|
550
|
+
var value = ${r}(0.0);
|
|
551
|
+
if (filterPosition < 9u && (cachedMask[load] & (1u << filterPosition)) != 0u) {
|
|
552
|
+
${e.addressMode==="base-offset"?"":`let inputX = i32(cachedX[load]) + i32(kernelX) - 1;
|
|
553
|
+
let inputY = i32(cachedY[load]) + i32(kernelY) - 1;`}
|
|
554
|
+
${t}
|
|
555
|
+
}
|
|
556
|
+
inputTile[${Xt("loadIndex",e)}] = value;
|
|
557
|
+
}
|
|
558
|
+
`}function Ln(r,t,e,i,n,o){const{workgroupX:a,workgroupY:s,rowsPerThread:u,tileM:l,tileNBlocks:p,tileKBlocks:c}=dt(o),{addressMode:d,weightLayout:h}=o,g=D(r),f=D(t),_=Yt(r,o,!1),m=Yt(r,o,!0),v=_?"vec2<u32>":g,w=m?"vec2<u32>":g,T=ut(ht("acc",e),t),$=a*s,b=l*c,x=c*p*4;return`${ft(r)}
|
|
559
|
+
struct Params {
|
|
560
|
+
inputWidth: u32,
|
|
561
|
+
inputHeight: u32,
|
|
562
|
+
outputWidth: u32,
|
|
563
|
+
outputHeight: u32,
|
|
564
|
+
inputBlocks: u32,
|
|
565
|
+
outputBlocks: u32,
|
|
566
|
+
}
|
|
567
|
+
@group(0) @binding(0) var<storage, read> inputData: array<${v}>;
|
|
568
|
+
@group(0) @binding(1) var<storage, read> weights: array<${w}>;
|
|
569
|
+
@group(0) @binding(2) var<storage, read> bias: array<vec4<f32>>;
|
|
570
|
+
@group(0) @binding(3) var<storage, read_write> outputData: array<${f}>;
|
|
571
|
+
@group(0) @binding(4) var<uniform> params: Params;
|
|
572
|
+
|
|
573
|
+
var<workgroup> inputTile: array<${g}, ${bt(o).input}>;
|
|
574
|
+
var<workgroup> weightTile: array<${g}, ${bt(o).weights}>;
|
|
575
|
+
|
|
576
|
+
@compute @workgroup_size(${a}, ${s}, 1)
|
|
577
|
+
fn main(
|
|
578
|
+
@builtin(local_invocation_id) localId: vec3<u32>,
|
|
579
|
+
@builtin(workgroup_id) workgroupId: vec3<u32>
|
|
580
|
+
) {
|
|
581
|
+
let spatialBase =
|
|
582
|
+
workgroupId.x * ${l}u +
|
|
583
|
+
localId.y * ${u}u;
|
|
584
|
+
let outputBlock =
|
|
585
|
+
workgroupId.y * ${p}u + localId.x;
|
|
586
|
+
let spatialCount = params.outputWidth * params.outputHeight;
|
|
587
|
+
var acc: array<vec4<f32>, ${u}>;
|
|
588
|
+
if (outputBlock < ${n}u) {
|
|
589
|
+
for (var row = 0u; row < ${u}u; row++) {
|
|
590
|
+
acc[row] = bias[outputBlock];
|
|
591
|
+
}
|
|
592
|
+
}
|
|
593
|
+
|
|
594
|
+
let localLinear =
|
|
595
|
+
localId.y * ${a}u + localId.x;
|
|
596
|
+
${d!=="analytic"?rr(i,!1,o):""}
|
|
597
|
+
let totalK = ${i*9}u;
|
|
598
|
+
for (var kBase = 0u; kBase < totalK; kBase += ${c}u) {
|
|
599
|
+
${d!=="analytic"?be(g,`
|
|
600
|
+
${d==="base-offset"?`let offset = (i32(kernelY) - 1) * i32(params.inputWidth) + i32(kernelX) - 1;
|
|
601
|
+
let inputIndex = cachedBase0[load] + u32(offset * ${i}) + inputBlock;`:`let inputIndex = (u32(inputY) * params.inputWidth + u32(inputX)) *
|
|
602
|
+
${i}u + inputBlock;`}
|
|
603
|
+
value = ${At("inputData[inputIndex]",_)};
|
|
604
|
+
`,o):`
|
|
605
|
+
for (
|
|
606
|
+
var loadIndex = localLinear;
|
|
607
|
+
loadIndex < ${b}u;
|
|
608
|
+
loadIndex += ${$}u
|
|
609
|
+
) {
|
|
610
|
+
let tileSpatial = loadIndex / ${c}u;
|
|
611
|
+
let tileK = loadIndex % ${c}u;
|
|
612
|
+
let inputSpatialIndex =
|
|
613
|
+
workgroupId.x * ${l}u + tileSpatial;
|
|
614
|
+
let kIndex = kBase + tileK;
|
|
615
|
+
var value = ${g}(0.0);
|
|
616
|
+
if (inputSpatialIndex < spatialCount && kIndex < totalK) {
|
|
617
|
+
let outputY = inputSpatialIndex / params.outputWidth;
|
|
618
|
+
let outputX = inputSpatialIndex % params.outputWidth;
|
|
619
|
+
let inputBlock = kIndex % ${i}u;
|
|
620
|
+
let kernelIndex = kIndex / ${i}u;
|
|
621
|
+
let kernelY = kernelIndex / 3u;
|
|
622
|
+
let kernelX = kernelIndex % 3u;
|
|
623
|
+
let inputY = i32(outputY) + i32(kernelY) - 1;
|
|
624
|
+
let inputX = i32(outputX) + i32(kernelX) - 1;
|
|
625
|
+
if (
|
|
626
|
+
inputY >= 0 && inputY < i32(params.inputHeight) &&
|
|
627
|
+
inputX >= 0 && inputX < i32(params.inputWidth)
|
|
628
|
+
) {
|
|
629
|
+
let inputIndex =
|
|
630
|
+
(u32(inputY) * params.inputWidth + u32(inputX)) *
|
|
631
|
+
${i}u + inputBlock;
|
|
632
|
+
value = ${At("inputData[inputIndex]",_)};
|
|
633
|
+
}
|
|
634
|
+
}
|
|
635
|
+
inputTile[${Xt("loadIndex",o)}] = value;
|
|
636
|
+
}
|
|
637
|
+
`}
|
|
638
|
+
|
|
639
|
+
for (
|
|
640
|
+
var loadIndex = localLinear;
|
|
641
|
+
loadIndex < ${x}u;
|
|
642
|
+
loadIndex += ${$}u
|
|
643
|
+
) {
|
|
644
|
+
let tileK = loadIndex / ${p*4}u;
|
|
645
|
+
let outputRemainder = loadIndex % ${p*4}u;
|
|
646
|
+
let tileOutputBlock = outputRemainder / 4u;
|
|
647
|
+
let outputLane = outputRemainder % 4u;
|
|
648
|
+
let loadedOutputBlock =
|
|
649
|
+
workgroupId.y * ${p}u + tileOutputBlock;
|
|
650
|
+
let kIndex = kBase + tileK;
|
|
651
|
+
var value = ${g}(0.0);
|
|
652
|
+
if (loadedOutputBlock < ${n}u && kIndex < totalK) {
|
|
653
|
+
${h==="k-major"?`
|
|
654
|
+
let weightIndex = (kIndex * ${n}u + loadedOutputBlock) * 4u + outputLane;
|
|
655
|
+
`:d!=="analytic"?`
|
|
656
|
+
let weightIndex = (loadedOutputBlock * ${i*9}u + kIndex) * 4u + outputLane;
|
|
657
|
+
`:`
|
|
658
|
+
let inputBlock = kIndex % ${i}u;
|
|
659
|
+
let kernelIndex = kIndex / ${i}u;
|
|
660
|
+
let kernelY = kernelIndex / 3u;
|
|
661
|
+
let kernelX = kernelIndex % 3u;
|
|
662
|
+
let weightIndex =
|
|
663
|
+
((((loadedOutputBlock * 3u + kernelY) * 3u + kernelX) *
|
|
664
|
+
${i}u + inputBlock) * 4u + outputLane);
|
|
665
|
+
`}
|
|
666
|
+
value = ${At("weights[weightIndex]",m)};
|
|
667
|
+
}
|
|
668
|
+
weightTile[${Qe("loadIndex",o)}] = value;
|
|
669
|
+
}
|
|
670
|
+
|
|
671
|
+
workgroupBarrier();
|
|
672
|
+
${ti(r,o)}
|
|
673
|
+
workgroupBarrier();
|
|
674
|
+
${d!=="analytic"?nr(i,o):""}
|
|
675
|
+
}
|
|
676
|
+
|
|
677
|
+
if (outputBlock < ${n}u) {
|
|
678
|
+
for (var row = 0u; row < ${u}u; row++) {
|
|
679
|
+
let spatialIndex = spatialBase + row;
|
|
680
|
+
if (spatialIndex < spatialCount) {
|
|
681
|
+
let outputIndex = spatialIndex * ${n}u + outputBlock;
|
|
682
|
+
outputData[outputIndex] = ${T.replaceAll("acc","acc[row]")};
|
|
683
|
+
}
|
|
684
|
+
}
|
|
685
|
+
}
|
|
686
|
+
}
|
|
687
|
+
`}function Dn(r,t,e,i){const n=D(r),o=ht("acc",t);return`${ft(r)}
|
|
688
|
+
struct Params {
|
|
689
|
+
inputWidth: u32,
|
|
690
|
+
inputHeight: u32,
|
|
691
|
+
outputWidth: u32,
|
|
692
|
+
outputHeight: u32,
|
|
693
|
+
inputBlocks: u32,
|
|
694
|
+
outputBlocks: u32,
|
|
695
|
+
}
|
|
696
|
+
@group(0) @binding(0) var<storage, read> inputData: array<${n}>;
|
|
697
|
+
@group(0) @binding(1) var<storage, read> weights: array<${n}>;
|
|
698
|
+
@group(0) @binding(2) var<storage, read> bias: array<vec4<f32>>;
|
|
699
|
+
@group(0) @binding(3) var<storage, read_write> outputData: array<${n}>;
|
|
700
|
+
@group(0) @binding(4) var<uniform> params: Params;
|
|
701
|
+
|
|
702
|
+
@compute @workgroup_size(${tt}, ${tt}, 1)
|
|
703
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
704
|
+
if (gid.x >= params.outputWidth || gid.y >= params.outputHeight || gid.z >= ${i}u) {
|
|
705
|
+
return;
|
|
706
|
+
}
|
|
707
|
+
var pooled = vec4<f32>(-3.402823466e+38);
|
|
708
|
+
for (var py = 0u; py < 2u; py++) {
|
|
709
|
+
let centerY = gid.y * 2u + py;
|
|
710
|
+
if (centerY >= params.inputHeight) { continue; }
|
|
711
|
+
for (var px = 0u; px < 2u; px++) {
|
|
712
|
+
let centerX = gid.x * 2u + px;
|
|
713
|
+
if (centerX >= params.inputWidth) { continue; }
|
|
714
|
+
var acc = bias[gid.z];
|
|
715
|
+
for (var ky = 0u; ky < 3u; ky++) {
|
|
716
|
+
let inputY = i32(centerY) + i32(ky) - 1;
|
|
717
|
+
if (inputY < 0 || inputY >= i32(params.inputHeight)) { continue; }
|
|
718
|
+
for (var kx = 0u; kx < 3u; kx++) {
|
|
719
|
+
let inputX = i32(centerX) + i32(kx) - 1;
|
|
720
|
+
if (inputX < 0 || inputX >= i32(params.inputWidth)) { continue; }
|
|
721
|
+
let pixelBase = (u32(inputY) * params.inputWidth + u32(inputX)) * ${e}u;
|
|
722
|
+
for (var inputBlock = 0u; inputBlock < ${e}u; inputBlock++) {
|
|
723
|
+
${ae("inputData[pixelBase + inputBlock]",`((((gid.z * 3u + ky) * 3u + kx) * ${e}u + inputBlock) * 4u)`,r)}
|
|
724
|
+
}
|
|
725
|
+
}
|
|
726
|
+
}
|
|
727
|
+
pooled = max(pooled, ${o});
|
|
728
|
+
}
|
|
729
|
+
}
|
|
730
|
+
let outputIndex = (gid.y * params.outputWidth + gid.x) * ${i}u + gid.z;
|
|
731
|
+
outputData[outputIndex] = ${ut("pooled",r)};
|
|
732
|
+
}
|
|
733
|
+
`}function Rn(r,t,e){const i=D(r);return`${ft(r)}
|
|
734
|
+
struct Params {
|
|
735
|
+
inputWidth: u32,
|
|
736
|
+
inputHeight: u32,
|
|
737
|
+
outputWidth: u32,
|
|
738
|
+
outputHeight: u32,
|
|
739
|
+
outputBlocks: u32,
|
|
740
|
+
}
|
|
741
|
+
@group(0) @binding(0) var<storage, read> inputData: array<${i}>;
|
|
742
|
+
@group(0) @binding(1) var<storage, read_write> outputData: array<${i}>;
|
|
743
|
+
@group(0) @binding(2) var<uniform> params: Params;
|
|
744
|
+
|
|
745
|
+
@compute @workgroup_size(${tt}, ${tt}, 1)
|
|
746
|
+
fn main(
|
|
747
|
+
@builtin(global_invocation_id) gid: vec3<u32>${e?`,
|
|
748
|
+
@builtin(num_workgroups) groupCount: vec3<u32>`:""}
|
|
749
|
+
) {
|
|
750
|
+
${e?`let linearX = gid.x + gid.z * groupCount.x * ${tt}u;
|
|
751
|
+
let outputBlock = linearX % ${t}u;
|
|
752
|
+
let outputX = linearX / ${t}u;`:`let outputBlock = gid.z;
|
|
753
|
+
let outputX = gid.x;`}
|
|
754
|
+
if (
|
|
755
|
+
outputX >= params.outputWidth ||
|
|
756
|
+
gid.y >= params.outputHeight ||
|
|
757
|
+
outputBlock >= ${t}u
|
|
758
|
+
) {
|
|
759
|
+
return;
|
|
760
|
+
}
|
|
761
|
+
var pooled = vec4<f32>(-3.402823466e+38);
|
|
762
|
+
for (var py = 0u; py < 2u; py++) {
|
|
763
|
+
let inputY = gid.y * 2u + py;
|
|
764
|
+
if (inputY >= params.inputHeight) { continue; }
|
|
765
|
+
for (var px = 0u; px < 2u; px++) {
|
|
766
|
+
let inputX = outputX * 2u + px;
|
|
767
|
+
if (inputX >= params.inputWidth) { continue; }
|
|
768
|
+
let inputIndex =
|
|
769
|
+
(inputY * params.inputWidth + inputX) * ${t}u + outputBlock;
|
|
770
|
+
pooled = max(pooled, vec4<f32>(inputData[inputIndex]));
|
|
771
|
+
}
|
|
772
|
+
}
|
|
773
|
+
let outputIndex =
|
|
774
|
+
(gid.y * params.outputWidth + outputX) * ${t}u + outputBlock;
|
|
775
|
+
outputData[outputIndex] = ${ut("pooled",r)};
|
|
776
|
+
}
|
|
777
|
+
`}function Xn(r,t,e,i,n,o,a){const{workgroupX:s,workgroupY:u,rowsPerThread:l,tileM:p,tileNBlocks:c,tileKBlocks:d}=dt(a),{addressMode:h,weightLayout:g}=a,f=D(r),_=D(t),m=Yt(r,a,!1),v=Yt(r,a,!0),w=m?"vec2<u32>":f,T=v?"vec2<u32>":f,$=i[0]+i[1],b=ut(ht("acc[row]",e),t),x=s*u,B=p*d,k=d*c*4,I=(P,O)=>`
|
|
778
|
+
{
|
|
779
|
+
let sourceBlock = ${O};
|
|
780
|
+
${h==="base-offset"?`let dx = ${P===n?"(i32((cachedMask[load] >> 9u) & 1u) + i32(kernelX) - 1) >> 1u":"i32(kernelX) - 1"};
|
|
781
|
+
let dy = ${P===n?"(i32((cachedMask[load] >> 10u) & 1u) + i32(kernelY) - 1) >> 1u":"i32(kernelY) - 1"};
|
|
782
|
+
let offset = (dy * i32(params.source${P}Width) + dx) * ${i[P]};
|
|
783
|
+
let sourceIndex = cachedBase${P}[load] + u32(offset) + sourceBlock;`:`let sourceX = ${P===n?"u32(inputX) / 2u":"u32(inputX)"};
|
|
784
|
+
let sourceY = ${P===n?"u32(inputY) / 2u":"u32(inputY)"};
|
|
785
|
+
let sourceIndex =
|
|
786
|
+
(sourceY * params.source${P}Width + sourceX) *
|
|
787
|
+
${i[P]}u + sourceBlock;`}
|
|
788
|
+
value = ${At(`input${P}[sourceIndex]`,m)};
|
|
789
|
+
}`;return`${ft(r)}
|
|
790
|
+
struct Params {
|
|
791
|
+
outputWidth: u32,
|
|
792
|
+
outputHeight: u32,
|
|
793
|
+
outputBlocks: u32,
|
|
794
|
+
inputBlocks: u32,
|
|
795
|
+
source0Width: u32,
|
|
796
|
+
source0Height: u32,
|
|
797
|
+
source1Width: u32,
|
|
798
|
+
source1Height: u32,
|
|
799
|
+
}
|
|
800
|
+
@group(0) @binding(0) var<storage, read> input0: array<${w}>;
|
|
801
|
+
@group(0) @binding(1) var<storage, read> input1: array<${w}>;
|
|
802
|
+
@group(0) @binding(2) var<storage, read> weights: array<${T}>;
|
|
803
|
+
@group(0) @binding(3) var<storage, read> bias: array<vec4<f32>>;
|
|
804
|
+
@group(0) @binding(4) var<storage, read_write> outputData: array<${_}>;
|
|
805
|
+
@group(0) @binding(5) var<uniform> params: Params;
|
|
806
|
+
|
|
807
|
+
var<workgroup> inputTile: array<${f}, ${bt(a).input}>;
|
|
808
|
+
var<workgroup> weightTile: array<${f}, ${bt(a).weights}>;
|
|
809
|
+
|
|
810
|
+
@compute @workgroup_size(${s}, ${u}, 1)
|
|
811
|
+
fn main(
|
|
812
|
+
@builtin(local_invocation_id) localId: vec3<u32>,
|
|
813
|
+
@builtin(workgroup_id) workgroupId: vec3<u32>
|
|
814
|
+
) {
|
|
815
|
+
let spatialBase =
|
|
816
|
+
workgroupId.x * ${p}u +
|
|
817
|
+
localId.y * ${l}u;
|
|
818
|
+
let outputBlock =
|
|
819
|
+
workgroupId.y * ${c}u + localId.x;
|
|
820
|
+
let spatialCount = params.outputWidth * params.outputHeight;
|
|
821
|
+
var acc: array<vec4<f32>, ${l}>;
|
|
822
|
+
if (outputBlock < ${o}u) {
|
|
823
|
+
for (var row = 0u; row < ${l}u; row++) {
|
|
824
|
+
acc[row] = bias[outputBlock];
|
|
825
|
+
}
|
|
826
|
+
}
|
|
827
|
+
|
|
828
|
+
let localLinear =
|
|
829
|
+
localId.y * ${s}u + localId.x;
|
|
830
|
+
${h!=="analytic"?rr($,!0,a,{blocks:i,upsampled:n}):""}
|
|
831
|
+
let totalK = ${$*9}u;
|
|
832
|
+
for (var kBase = 0u; kBase < totalK; kBase += ${d}u) {
|
|
833
|
+
${h!=="analytic"?a.decoderLoad==="source-first"?`
|
|
834
|
+
if (channelBlock < ${i[0]}u) {
|
|
835
|
+
${be(f,I(0,"inputBlock"),a)}
|
|
836
|
+
} else {
|
|
837
|
+
${be(f,I(1,`inputBlock - ${i[0]}u`),a)}
|
|
838
|
+
}
|
|
839
|
+
`:be(f,`
|
|
840
|
+
if (inputBlock < ${i[0]}u) {
|
|
841
|
+
${I(0,"inputBlock")}
|
|
842
|
+
} else {
|
|
843
|
+
${I(1,`inputBlock - ${i[0]}u`)}
|
|
844
|
+
}
|
|
845
|
+
`,a):`
|
|
846
|
+
for (
|
|
847
|
+
var loadIndex = localLinear;
|
|
848
|
+
loadIndex < ${B}u;
|
|
849
|
+
loadIndex += ${x}u
|
|
850
|
+
) {
|
|
851
|
+
let tileSpatial = loadIndex / ${d}u;
|
|
852
|
+
let tileK = loadIndex % ${d}u;
|
|
853
|
+
let outputSpatialIndex =
|
|
854
|
+
workgroupId.x * ${p}u + tileSpatial;
|
|
855
|
+
let kIndex = kBase + tileK;
|
|
856
|
+
var value = ${f}(0.0);
|
|
857
|
+
if (outputSpatialIndex < spatialCount && kIndex < totalK) {
|
|
858
|
+
let outputY = outputSpatialIndex / params.outputWidth;
|
|
859
|
+
let outputX = outputSpatialIndex % params.outputWidth;
|
|
860
|
+
let inputBlock = kIndex % ${$}u;
|
|
861
|
+
let kernelIndex = kIndex / ${$}u;
|
|
862
|
+
let kernelY = kernelIndex / 3u;
|
|
863
|
+
let kernelX = kernelIndex % 3u;
|
|
864
|
+
let inputY = i32(outputY) + i32(kernelY) - 1;
|
|
865
|
+
let inputX = i32(outputX) + i32(kernelX) - 1;
|
|
866
|
+
if (
|
|
867
|
+
inputY >= 0 && inputY < i32(params.outputHeight) &&
|
|
868
|
+
inputX >= 0 && inputX < i32(params.outputWidth)
|
|
869
|
+
) {
|
|
870
|
+
if (inputBlock < ${i[0]}u) {
|
|
871
|
+
${I(0,"inputBlock")}
|
|
872
|
+
} else {
|
|
873
|
+
${I(1,`inputBlock - ${i[0]}u`)}
|
|
874
|
+
}
|
|
875
|
+
}
|
|
876
|
+
}
|
|
877
|
+
inputTile[${Xt("loadIndex",a)}] = value;
|
|
878
|
+
}
|
|
879
|
+
`}
|
|
880
|
+
|
|
881
|
+
for (
|
|
882
|
+
var loadIndex = localLinear;
|
|
883
|
+
loadIndex < ${k}u;
|
|
884
|
+
loadIndex += ${x}u
|
|
885
|
+
) {
|
|
886
|
+
let tileK = loadIndex / ${c*4}u;
|
|
887
|
+
let outputRemainder = loadIndex % ${c*4}u;
|
|
888
|
+
let tileOutputBlock = outputRemainder / 4u;
|
|
889
|
+
let outputLane = outputRemainder % 4u;
|
|
890
|
+
let loadedOutputBlock =
|
|
891
|
+
workgroupId.y * ${c}u + tileOutputBlock;
|
|
892
|
+
let kIndex = kBase + tileK;
|
|
893
|
+
var value = ${f}(0.0);
|
|
894
|
+
if (loadedOutputBlock < ${o}u && kIndex < totalK) {
|
|
895
|
+
${g==="k-major"?`
|
|
896
|
+
let weightIndex = (kIndex * ${o}u + loadedOutputBlock) * 4u + outputLane;
|
|
897
|
+
`:h!=="analytic"?`
|
|
898
|
+
let weightIndex = (loadedOutputBlock * ${$*9}u + kIndex) * 4u + outputLane;
|
|
899
|
+
`:`
|
|
900
|
+
let inputBlock = kIndex % ${$}u;
|
|
901
|
+
let kernelIndex = kIndex / ${$}u;
|
|
902
|
+
let kernelY = kernelIndex / 3u;
|
|
903
|
+
let kernelX = kernelIndex % 3u;
|
|
904
|
+
let weightIndex =
|
|
905
|
+
((((loadedOutputBlock * 3u + kernelY) * 3u + kernelX) *
|
|
906
|
+
${$}u + inputBlock) * 4u + outputLane);
|
|
907
|
+
`}
|
|
908
|
+
value = ${At("weights[weightIndex]",v)};
|
|
909
|
+
}
|
|
910
|
+
weightTile[${Qe("loadIndex",a)}] = value;
|
|
911
|
+
}
|
|
912
|
+
|
|
913
|
+
workgroupBarrier();
|
|
914
|
+
${ti(r,a)}
|
|
915
|
+
workgroupBarrier();
|
|
916
|
+
${h!=="analytic"?nr($,a):""}
|
|
917
|
+
}
|
|
918
|
+
|
|
919
|
+
if (outputBlock < ${o}u) {
|
|
920
|
+
for (var row = 0u; row < ${l}u; row++) {
|
|
921
|
+
let spatialIndex = spatialBase + row;
|
|
922
|
+
if (spatialIndex < spatialCount) {
|
|
923
|
+
let outputIndex = spatialIndex * ${o}u + outputBlock;
|
|
924
|
+
outputData[outputIndex] = ${b};
|
|
925
|
+
}
|
|
926
|
+
}
|
|
927
|
+
}
|
|
928
|
+
}
|
|
929
|
+
`}function Yn(r,t,e,i,n,o,a){const s=dt(a),{workgroupX:u,workgroupY:l,rowsPerThread:p,tileM:c,tileNBlocks:d,tileKBlocks:h}=s,g=u*l,f=c*h/g,_=D(r),m=Yt(r,a,!1),v=Yt(r,a,!0),w=i[n],T=1-n;i[T];const $=(x,B)=>{const k=i[x],I=B?4:9,P=B?`phase * ${4*w*o*4}u`:`${ue*w*o*4}u`,O=B?`(i32(tap / 2u) + i32(phaseY) - 1) * i32(params.source${x}Width) + i32(tap % 2u) + i32(phaseX) - 1`:`(i32(tap / 3u) - 1) * i32(params.source${x}Width) + i32(tap % 3u) - 1`;return`
|
|
930
|
+
for (var kBase = 0u; kBase < ${I*k}u; kBase += ${h}u) {
|
|
931
|
+
let inputK = kBase + localLinear % ${h}u;
|
|
932
|
+
let tap = inputK / ${k}u;
|
|
933
|
+
let channel = inputK % ${k}u;
|
|
934
|
+
for (var load = 0u; load < ${f}u; load++) {
|
|
935
|
+
let i = localLinear + load * ${g}u;
|
|
936
|
+
var value = ${_}(0.0);
|
|
937
|
+
if (tap < ${I}u && (mask${x}[load] & (1u << tap)) != 0u) {
|
|
938
|
+
let offset = ${O};
|
|
939
|
+
let index = u32(base${x}[load] + offset) * ${k}u + channel;
|
|
940
|
+
value = ${At(`input${x}[index]`,m)};
|
|
941
|
+
}
|
|
942
|
+
inputTile[${Xt("i",a)}] = value;
|
|
943
|
+
}
|
|
944
|
+
for (var i = localLinear; i < ${h*d*4}u; i += ${g}u) {
|
|
945
|
+
let k = kBase + i / ${d*4}u;
|
|
946
|
+
let outputBlock = workgroupId.y * ${d}u + (i % ${d*4}u) / 4u;
|
|
947
|
+
var value = ${_}(0.0);
|
|
948
|
+
if (k < ${I*k}u && outputBlock < ${o}u) {
|
|
949
|
+
let index = ${P} + (k * ${o}u + outputBlock) * 4u + i % 4u;
|
|
950
|
+
value = ${At("weights[index]",v)};
|
|
951
|
+
}
|
|
952
|
+
weightTile[${Qe("i",a)}] = value;
|
|
953
|
+
}
|
|
954
|
+
workgroupBarrier();
|
|
955
|
+
${ti(r,a)}
|
|
956
|
+
workgroupBarrier();
|
|
957
|
+
}`},b=(x,B)=>{const k=B?4:9,I=B?2:3;return`
|
|
958
|
+
let x${x} = i32(${B?"x":"x * 2u + phaseX"});
|
|
959
|
+
let y${x} = i32(${B?"y":"y * 2u + phaseY"});
|
|
960
|
+
base${x}[load] = y${x} * i32(params.source${x}Width) + x${x};
|
|
961
|
+
var bits${x} = 0u;
|
|
962
|
+
for (var tap = 0u; tap < ${k}u; tap++) {
|
|
963
|
+
let sx = x${x} + i32(tap % ${I}u) ${B?"+ i32(phaseX)":""} - 1;
|
|
964
|
+
let sy = y${x} + i32(tap / ${I}u) ${B?"+ i32(phaseY)":""} - 1;
|
|
965
|
+
if (valid && sx >= 0 && sy >= 0 && sx < i32(params.source${x}Width) && sy < i32(params.source${x}Height)) {
|
|
966
|
+
bits${x} |= 1u << tap;
|
|
967
|
+
}
|
|
968
|
+
}
|
|
969
|
+
mask${x}[load] = bits${x};`};return`${ft(r)}
|
|
970
|
+
struct Params {
|
|
971
|
+
outputWidth: u32, outputHeight: u32, outputBlocks: u32, inputBlocks: u32,
|
|
972
|
+
source0Width: u32, source0Height: u32, source1Width: u32, source1Height: u32,
|
|
973
|
+
}
|
|
974
|
+
@group(0) @binding(0) var<storage, read> input0: array<${m?"vec2<u32>":_}>;
|
|
975
|
+
@group(0) @binding(1) var<storage, read> input1: array<${m?"vec2<u32>":_}>;
|
|
976
|
+
@group(0) @binding(2) var<storage, read> weights: array<${v?"vec2<u32>":_}>;
|
|
977
|
+
@group(0) @binding(3) var<storage, read> bias: array<vec4<f32>>;
|
|
978
|
+
@group(0) @binding(4) var<storage, read_write> outputData: array<${D(t)}>;
|
|
979
|
+
@group(0) @binding(5) var<uniform> params: Params;
|
|
980
|
+
var<workgroup> inputTile: array<${_}, ${bt(a).input}>;
|
|
981
|
+
var<workgroup> weightTile: array<${_}, ${bt(a).weights}>;
|
|
982
|
+
|
|
983
|
+
@compute @workgroup_size(${u}, ${l}, 1)
|
|
984
|
+
fn main(@builtin(local_invocation_id) localId: vec3<u32>, @builtin(workgroup_id) workgroupId: vec3<u32>) {
|
|
985
|
+
let phase = workgroupId.z;
|
|
986
|
+
let phaseX = phase & 1u;
|
|
987
|
+
let phaseY = phase >> 1u;
|
|
988
|
+
let phaseWidth = (params.outputWidth + 1u) / 2u;
|
|
989
|
+
let localLinear = localId.y * ${u}u + localId.x;
|
|
990
|
+
let outputBlock = workgroupId.y * ${d}u + localId.x;
|
|
991
|
+
var acc: array<vec4<f32>, ${p}>;
|
|
992
|
+
if (outputBlock < ${o}u) {
|
|
993
|
+
for (var row = 0u; row < ${p}u; row++) { acc[row] = bias[outputBlock]; }
|
|
994
|
+
}
|
|
995
|
+
var base0: array<i32, ${f}>;
|
|
996
|
+
var base1: array<i32, ${f}>;
|
|
997
|
+
var mask0: array<u32, ${f}>;
|
|
998
|
+
var mask1: array<u32, ${f}>;
|
|
999
|
+
for (var load = 0u; load < ${f}u; load++) {
|
|
1000
|
+
let spatial = workgroupId.x * ${c}u + (localLinear + load * ${g}u) / ${h}u;
|
|
1001
|
+
let x = spatial % phaseWidth;
|
|
1002
|
+
let y = spatial / phaseWidth;
|
|
1003
|
+
let valid = x * 2u + phaseX < params.outputWidth && y * 2u + phaseY < params.outputHeight;
|
|
1004
|
+
${b(n,!0)}
|
|
1005
|
+
${b(T,!1)}
|
|
1006
|
+
}
|
|
1007
|
+
${$(n,!0)}
|
|
1008
|
+
${$(T,!1)}
|
|
1009
|
+
if (outputBlock < ${o}u) {
|
|
1010
|
+
for (var row = 0u; row < ${p}u; row++) {
|
|
1011
|
+
let spatial = workgroupId.x * ${c}u + localId.y * ${p}u + row;
|
|
1012
|
+
let x = (spatial % phaseWidth) * 2u + phaseX;
|
|
1013
|
+
let y = (spatial / phaseWidth) * 2u + phaseY;
|
|
1014
|
+
if (x < params.outputWidth && y < params.outputHeight) {
|
|
1015
|
+
outputData[(y * params.outputWidth + x) * ${o}u + outputBlock] = ${ut(ht("acc[row]",e),t)};
|
|
1016
|
+
}
|
|
1017
|
+
}
|
|
1018
|
+
}
|
|
1019
|
+
}
|
|
1020
|
+
`}function Fn(r,t,e,i,n,o){const a=D(r),s=D(t),u=i[0]+i[1],l=(c,d)=>{const h=c===n;return`
|
|
1021
|
+
{
|
|
1022
|
+
let sourceX = ${h?"u32(inputX) / 2u":"u32(inputX)"};
|
|
1023
|
+
let sourceY = ${h?"u32(inputY) / 2u":"u32(inputY)"};
|
|
1024
|
+
let sourcePixelBase = (sourceY * params.source${c}Width + sourceX) * ${i[c]}u;
|
|
1025
|
+
for (var sourceBlock = 0u; sourceBlock < ${i[c]}u; sourceBlock++) {
|
|
1026
|
+
let inputBlock = ${d}u + sourceBlock;
|
|
1027
|
+
${ae(`input${c}[sourcePixelBase + sourceBlock]`,`((((gid.z * 3u + ky) * 3u + kx) * ${u}u + inputBlock) * 4u)`,r)}
|
|
1028
|
+
}
|
|
1029
|
+
}
|
|
1030
|
+
`},p=ut(ht("acc",e),t);return`${ft(r)}
|
|
1031
|
+
struct Params {
|
|
1032
|
+
outputWidth: u32,
|
|
1033
|
+
outputHeight: u32,
|
|
1034
|
+
outputBlocks: u32,
|
|
1035
|
+
inputBlocks: u32,
|
|
1036
|
+
source0Width: u32,
|
|
1037
|
+
source0Height: u32,
|
|
1038
|
+
source1Width: u32,
|
|
1039
|
+
source1Height: u32,
|
|
1040
|
+
}
|
|
1041
|
+
@group(0) @binding(0) var<storage, read> input0: array<${a}>;
|
|
1042
|
+
@group(0) @binding(1) var<storage, read> input1: array<${a}>;
|
|
1043
|
+
@group(0) @binding(2) var<storage, read> weights: array<${a}>;
|
|
1044
|
+
@group(0) @binding(3) var<storage, read> bias: array<vec4<f32>>;
|
|
1045
|
+
@group(0) @binding(4) var<storage, read_write> outputData: array<${s}>;
|
|
1046
|
+
@group(0) @binding(5) var<uniform> params: Params;
|
|
1047
|
+
|
|
1048
|
+
@compute @workgroup_size(${tt}, ${tt}, 1)
|
|
1049
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
1050
|
+
if (gid.x >= params.outputWidth || gid.y >= params.outputHeight || gid.z >= ${o}u) {
|
|
1051
|
+
return;
|
|
1052
|
+
}
|
|
1053
|
+
var acc = bias[gid.z];
|
|
1054
|
+
for (var ky = 0u; ky < 3u; ky++) {
|
|
1055
|
+
let inputY = i32(gid.y) + i32(ky) - 1;
|
|
1056
|
+
if (inputY < 0 || inputY >= i32(params.outputHeight)) { continue; }
|
|
1057
|
+
for (var kx = 0u; kx < 3u; kx++) {
|
|
1058
|
+
let inputX = i32(gid.x) + i32(kx) - 1;
|
|
1059
|
+
if (inputX < 0 || inputX >= i32(params.outputWidth)) { continue; }
|
|
1060
|
+
${l(0,0)}
|
|
1061
|
+
${l(1,i[0])}
|
|
1062
|
+
}
|
|
1063
|
+
}
|
|
1064
|
+
let outputIndex = (gid.y * params.outputWidth + gid.x) * ${o}u + gid.z;
|
|
1065
|
+
outputData[outputIndex] = ${p};
|
|
1066
|
+
}
|
|
1067
|
+
`}function Gn(r,t,e,i,n,o){const a=D(r),s=D(t),u=i[0]+i[1],l=pt*pt*u,p=j*j,c=ut(ht("acc",e),t),d=(h,g)=>`
|
|
1068
|
+
let sourceBlock = ${g};
|
|
1069
|
+
let sourceIndex =
|
|
1070
|
+
(${h===n?"u32(inputY) / 2u":"u32(inputY)"} * params.source${h}Width + ${h===n?"u32(inputX) / 2u":"u32(inputX)"}) *
|
|
1071
|
+
${i[h]}u + sourceBlock;
|
|
1072
|
+
value = input${h}[sourceIndex];`;return`${ft(r)}
|
|
1073
|
+
struct Params {
|
|
1074
|
+
outputWidth: u32,
|
|
1075
|
+
outputHeight: u32,
|
|
1076
|
+
outputBlocks: u32,
|
|
1077
|
+
inputBlocks: u32,
|
|
1078
|
+
source0Width: u32,
|
|
1079
|
+
source0Height: u32,
|
|
1080
|
+
source1Width: u32,
|
|
1081
|
+
source1Height: u32,
|
|
1082
|
+
}
|
|
1083
|
+
@group(0) @binding(0) var<storage, read> input0: array<${a}>;
|
|
1084
|
+
@group(0) @binding(1) var<storage, read> input1: array<${a}>;
|
|
1085
|
+
@group(0) @binding(2) var<storage, read> weights: array<${a}>;
|
|
1086
|
+
@group(0) @binding(3) var<storage, read> bias: array<vec4<f32>>;
|
|
1087
|
+
@group(0) @binding(4) var<storage, read_write> outputData: array<${s}>;
|
|
1088
|
+
@group(0) @binding(5) var<uniform> params: Params;
|
|
1089
|
+
|
|
1090
|
+
var<workgroup> inputPatch: array<${a}, ${l}>;
|
|
1091
|
+
|
|
1092
|
+
@compute @workgroup_size(${j}, ${j}, 1)
|
|
1093
|
+
fn main(
|
|
1094
|
+
@builtin(local_invocation_id) localId: vec3<u32>,
|
|
1095
|
+
@builtin(workgroup_id) workgroupId: vec3<u32>
|
|
1096
|
+
) {
|
|
1097
|
+
let localLinear =
|
|
1098
|
+
localId.y * ${j}u + localId.x;
|
|
1099
|
+
for (
|
|
1100
|
+
var loadIndex = localLinear;
|
|
1101
|
+
loadIndex < ${l}u;
|
|
1102
|
+
loadIndex += ${p}u
|
|
1103
|
+
) {
|
|
1104
|
+
let patchPixel = loadIndex / ${u}u;
|
|
1105
|
+
let inputBlock = loadIndex % ${u}u;
|
|
1106
|
+
let patchX = patchPixel % ${pt}u;
|
|
1107
|
+
let patchY = patchPixel / ${pt}u;
|
|
1108
|
+
let inputX =
|
|
1109
|
+
i32(workgroupId.x * ${j}u + patchX) - 1;
|
|
1110
|
+
let inputY =
|
|
1111
|
+
i32(workgroupId.y * ${j}u + patchY) - 1;
|
|
1112
|
+
var value = ${a}(0.0);
|
|
1113
|
+
if (
|
|
1114
|
+
inputX >= 0 && inputX < i32(params.outputWidth) &&
|
|
1115
|
+
inputY >= 0 && inputY < i32(params.outputHeight)
|
|
1116
|
+
) {
|
|
1117
|
+
if (inputBlock < ${i[0]}u) {
|
|
1118
|
+
${d(0,"inputBlock")}
|
|
1119
|
+
} else {
|
|
1120
|
+
${d(1,`inputBlock - ${i[0]}u`)}
|
|
1121
|
+
}
|
|
1122
|
+
}
|
|
1123
|
+
inputPatch[loadIndex] = value;
|
|
1124
|
+
}
|
|
1125
|
+
workgroupBarrier();
|
|
1126
|
+
|
|
1127
|
+
let outputX =
|
|
1128
|
+
workgroupId.x * ${j}u + localId.x;
|
|
1129
|
+
let outputY =
|
|
1130
|
+
workgroupId.y * ${j}u + localId.y;
|
|
1131
|
+
let outputBlock = workgroupId.z;
|
|
1132
|
+
if (
|
|
1133
|
+
outputX >= params.outputWidth || outputY >= params.outputHeight ||
|
|
1134
|
+
outputBlock >= ${o}u
|
|
1135
|
+
) {
|
|
1136
|
+
return;
|
|
1137
|
+
}
|
|
1138
|
+
|
|
1139
|
+
var acc = bias[outputBlock];
|
|
1140
|
+
for (var ky = 0u; ky < 3u; ky++) {
|
|
1141
|
+
for (var kx = 0u; kx < 3u; kx++) {
|
|
1142
|
+
let patchBase =
|
|
1143
|
+
((localId.y + ky) * ${pt}u + localId.x + kx) *
|
|
1144
|
+
${u}u;
|
|
1145
|
+
for (var inputBlock = 0u; inputBlock < ${u}u; inputBlock++) {
|
|
1146
|
+
${ae("inputPatch[patchBase + inputBlock]",`((((outputBlock * 3u + ky) * 3u + kx) * ${u}u + inputBlock) * 4u)`,r)}
|
|
1147
|
+
}
|
|
1148
|
+
}
|
|
1149
|
+
}
|
|
1150
|
+
let outputIndex =
|
|
1151
|
+
(outputY * params.outputWidth + outputX) * ${o}u + outputBlock;
|
|
1152
|
+
outputData[outputIndex] = ${c};
|
|
1153
|
+
}
|
|
1154
|
+
`}function or(r,t){const e=D(r),i=Array.from({length:t},(s,u)=>`@group(0) @binding(${u}) var<storage, read> input${u}: array<vec4<f32>>;`).join(`
|
|
1155
|
+
`),n=Array.from({length:t},(s,u)=>{const l=u*3;return`if (channel < ${l+3}u) { return input${u}[pixel][channel - ${l}u]; }`}).join(`
|
|
1156
|
+
`),o=t,a=t+1;return`${ft(r)}
|
|
1157
|
+
struct Params {
|
|
1158
|
+
width: u32,
|
|
1159
|
+
height: u32,
|
|
1160
|
+
outputBlocks: u32,
|
|
1161
|
+
inputChannels: u32,
|
|
1162
|
+
}
|
|
1163
|
+
${i}
|
|
1164
|
+
@group(0) @binding(${o}) var<storage, read_write> outputData: array<${e}>;
|
|
1165
|
+
@group(0) @binding(${a}) var<uniform> params: Params;
|
|
1166
|
+
|
|
1167
|
+
fn readChannel(pixel: u32, channel: u32) -> f32 {
|
|
1168
|
+
${n}
|
|
1169
|
+
return 0.0;
|
|
1170
|
+
}
|
|
1171
|
+
|
|
1172
|
+
@compute @workgroup_size(${tt}, ${tt}, 1)
|
|
1173
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
1174
|
+
if (gid.x >= params.width || gid.y >= params.height || gid.z >= params.outputBlocks) {
|
|
1175
|
+
return;
|
|
1176
|
+
}
|
|
1177
|
+
let pixel = gid.y * params.width + gid.x;
|
|
1178
|
+
let firstChannel = gid.z * 4u;
|
|
1179
|
+
let value = vec4<f32>(
|
|
1180
|
+
readChannel(pixel, firstChannel),
|
|
1181
|
+
readChannel(pixel, firstChannel + 1u),
|
|
1182
|
+
readChannel(pixel, firstChannel + 2u),
|
|
1183
|
+
readChannel(pixel, firstChannel + 3u)
|
|
1184
|
+
);
|
|
1185
|
+
outputData[pixel * params.outputBlocks + gid.z] = ${ut("value",r)};
|
|
1186
|
+
}
|
|
1187
|
+
`}const Hn="chromium-experimental-subgroup-matrix",Vn="wgpu-cooperative-matrix",Ft=32,Gt=8,Ot=4,jn=1024;function ar(r,t="f32"){var s;const e=r.features.has(Hn)?"tint":r.features.has(Vn)?"naga":void 0;if(!e||!r.features.has("subgroups"))return;const i=r.adapterInfo,n=(i==null?void 0:i.subgroupMaxSize)??0;if(!i||!n||i.subgroupMinSize!==n)return;const o=t==="f16"?[["f16","f16"],["f16","f32"],["f32","f32"]]:[["f16","f32"],["f32","f32"],["f16","f16"]],a=r.features.has("shader-f16");for(const[u,l]of o)if(!(!a&&(u==="f16"||l==="f16"))){for(const p of[8,16])if((s=i.subgroupMatrixConfigs)==null?void 0:s.find(d=>d.componentType===u&&d.resultComponentType===l&&d.M===p&&d.N===p&&d.K===p))return{componentType:u,resultType:l,size:p,subgroupSize:n,dialect:e}}}function Kn(r,t){return t.outputFragments??2}function qn(r,t,e,i){const n=r.componentType==="f16"?2:4,o=r.resultType==="f16"?2:4,a=r.size,s=(Gt+2)*(Ft+2)*a,u=e?16:9;return n*(i?2:1)*(s+u*a*a*t)+o*Ot*Ft*a}function ur(r){const t=r.size,e=r.dialect==="naga",i=(o,a,s)=>e?`coop_mat${t}x${t}<${o}, ${s}>`:`subgroup_matrix_${a}<${o}, ${t}, ${t}>`,n=i(r.resultType,"result","C");return{left:i(r.componentType,"left","A"),right:i(r.componentType,"right","B"),result:n,load:(o,a,s,u)=>e?`coopLoadT<${o}>(&${a}[${s}], ${u})`:`subgroupMatrixLoad<${o}, row_major>(&${a}, ${s}, ${u})`,store:(o,a,s,u)=>e?`coopStoreT(${s}, &${o}[${a}], ${u});`:`subgroupMatrixStore<row_major>(&${o}, ${a}, ${s}, ${u});`,multiplyAdd:e?"coopMultiplyAdd":"subgroupMatrixMultiplyAccumulate",declare:o=>e?`var ${o}: ${n};`:`var ${o} = ${n}();`,extensions:e?"enable wgpu_cooperative_matrix;":`enable subgroups;
|
|
1188
|
+
enable chromium_experimental_subgroup_matrix;
|
|
1189
|
+
diagnostic(off, chromium_experimental.subgroup_matrix_uniformity);`}}function ei(r,t,e,i,n,o,a){const s=(g,f,_)=>`{
|
|
1190
|
+
let staged = vec4<${o}>(${f});
|
|
1191
|
+
let base = ${_} + i * 4u;
|
|
1192
|
+
${g.array}[base] = staged.x;
|
|
1193
|
+
${g.array}[base + 1u] = staged.y;
|
|
1194
|
+
${g.array}[base + 2u] = staged.z;
|
|
1195
|
+
${g.array}[base + 3u] = staged.w;
|
|
1196
|
+
}`,u=g=>t.map(f=>`
|
|
1197
|
+
for (var batch = 0u; batch < ${Math.ceil(f.vectors/i)}u; batch++) {
|
|
1198
|
+
let i = localIndex + batch * ${i}u;
|
|
1199
|
+
if (i < ${f.vectors}u) {
|
|
1200
|
+
var value = ${n}(0.0);
|
|
1201
|
+
${f.read}
|
|
1202
|
+
${s(f,"value",g(f))}
|
|
1203
|
+
}
|
|
1204
|
+
}`).join(""),l=g=>t.map(f=>`let ${f.array}Base = ${g} * ${(f.stride??f.vectors)*4}u;`).join(`
|
|
1205
|
+
`);if(!a||r<2)return`
|
|
1206
|
+
for (var chunk = 0u; chunk < ${r}u; chunk++) {
|
|
1207
|
+
let stageChunk = chunk;
|
|
1208
|
+
${l("0u")}
|
|
1209
|
+
${u(g=>`${g.array}Base`)}
|
|
1210
|
+
workgroupBarrier();
|
|
1211
|
+
${e}
|
|
1212
|
+
workgroupBarrier();
|
|
1213
|
+
}`;const p=t.map(g=>Math.ceil(g.vectors/i)),c=t.map((g,f)=>Array.from({length:p[f]},(_,m)=>`
|
|
1214
|
+
{
|
|
1215
|
+
let i = localIndex + ${m*i}u;
|
|
1216
|
+
var value = ${n}(0.0);
|
|
1217
|
+
if (i < ${g.vectors}u) {
|
|
1218
|
+
${g.read}
|
|
1219
|
+
}
|
|
1220
|
+
prefetch${f}_${m} = value;
|
|
1221
|
+
}`).join("")).join(""),d=t.map((g,f)=>Array.from({length:p[f]},(_,m)=>`var prefetch${f}_${m} = ${n}(0.0);`).join(`
|
|
1222
|
+
`)).join(`
|
|
1223
|
+
`),h=t.map((g,f)=>Array.from({length:p[f]},(_,m)=>`
|
|
1224
|
+
{
|
|
1225
|
+
let i = localIndex + ${m*i}u;
|
|
1226
|
+
if (i < ${g.vectors}u) ${s(g,`prefetch${f}_${m}`,`next * ${(g.stride??g.vectors)*4}u`)}
|
|
1227
|
+
}`).join("")).join("");return`
|
|
1228
|
+
{
|
|
1229
|
+
let stageChunk = 0u;
|
|
1230
|
+
${u(()=>"0u")}
|
|
1231
|
+
}
|
|
1232
|
+
workgroupBarrier();
|
|
1233
|
+
for (var chunk = 0u; chunk < ${r}u; chunk++) {
|
|
1234
|
+
let buffer = chunk & 1u;
|
|
1235
|
+
let next = 1u - buffer;
|
|
1236
|
+
${l("buffer")}
|
|
1237
|
+
${d}
|
|
1238
|
+
if (chunk + 1u < ${r}u) {
|
|
1239
|
+
let stageChunk = chunk + 1u;
|
|
1240
|
+
${c}
|
|
1241
|
+
}
|
|
1242
|
+
${e}
|
|
1243
|
+
if (chunk + 1u < ${r}u) {
|
|
1244
|
+
${h}
|
|
1245
|
+
}
|
|
1246
|
+
workgroupBarrier();
|
|
1247
|
+
}`}function sr(r,t){return r.dialect==="naga"?Array.from({length:9},(e,i)=>`{ let tap = ${i}u; ${t} }`).join(`
|
|
1248
|
+
`):`for (var tap = 0u; tap < 9u; tap++) { ${t} }`}function Zn(r,t,e,i,n,o,a,s,u,l=!1){const p=o.size,c=o.componentType,d=o.resultType,h=Ft,g=h+2,f=Gt+2,_=s?2:1,m=p/4,v=Math.ceil(i/m),w=a*p,T=h/p,$=Ot*o.subgroupSize,b=D(r),x=D(t),{left:B,right:k,load:I,store:P,multiplyAdd:O,declare:H,extensions:rt}=ur(o),at=[];for(let N=0;N<2;N++)for(let V=0;V<T;V++)for(let z=0;z<a;z++)at.push(`acc${N}_${V}_${z}`);const $t=at.map(H).join(`
|
|
1249
|
+
`),Z=Array.from({length:a},(N,V)=>`let b${V} = ${I(k,"weightTile",`weightTileBase + tap * ${p*w}u + ${V*p}u`,`${w}u`)};`).join(`
|
|
1250
|
+
`),nt=[];for(let N=0;N<2;N++)for(let V=0;V<T;V++)nt.push(`{
|
|
1251
|
+
let a = ${I(B,"inputTile",`inputTileBase + ((rowBase + ${N}u + ky) * ${g}u + ${V*p}u + kx) * ${p}u`,`${p}u`)};
|
|
1252
|
+
${Array.from({length:a},(z,K)=>`acc${N}_${V}_${K} = ${O}(a, b${K}, acc${N}_${V}_${K});`).join(`
|
|
1253
|
+
`)}
|
|
1254
|
+
}`);const St=`let ky = tap / 3u;
|
|
1255
|
+
let kx = tap % 3u;
|
|
1256
|
+
${Z}
|
|
1257
|
+
${nt.join(`
|
|
1258
|
+
`)}`,Qt=sr(o,St),Y=[];if(l)for(let N=0;N<a;N++)Y.push(`{
|
|
1259
|
+
var pooled = vec4<f32>(-3.402823e38);
|
|
1260
|
+
let pixel = lane / ${m}u;
|
|
1261
|
+
let outputBlock = blockBase + ${N*m}u + lane % ${m}u;
|
|
1262
|
+
${Array.from({length:2},(V,z)=>`
|
|
1263
|
+
${Array.from({length:T},(K,Lt)=>P("resultTile",`resultBase + ${Lt*p*p}u`,`acc${z}_${Lt}_${N}`,`${p}u`)).join(`
|
|
1264
|
+
`)}
|
|
1265
|
+
workgroupBarrier();
|
|
1266
|
+
{
|
|
1267
|
+
let base = resultBase + pixel * ${p*2}u + (lane % ${m}u) * 4u;
|
|
1268
|
+
let first = vec4<f32>(resultTile[base], resultTile[base + 1u], resultTile[base + 2u], resultTile[base + 3u]);
|
|
1269
|
+
let second = vec4<f32>(resultTile[base + ${p}u], resultTile[base + ${p+1}u], resultTile[base + ${p+2}u], resultTile[base + ${p+3}u]);
|
|
1270
|
+
if (tileY + rowBase + ${z}u < params.inputHeight) {
|
|
1271
|
+
pooled = max(pooled, first);
|
|
1272
|
+
if (tileX + pixel * 2u + 1u < params.inputWidth) { pooled = max(pooled, second); }
|
|
1273
|
+
}
|
|
1274
|
+
}
|
|
1275
|
+
workgroupBarrier();`).join(`
|
|
1276
|
+
`)}
|
|
1277
|
+
let outputX = tileX / 2u + pixel;
|
|
1278
|
+
let outputY = (tileY + rowBase) / 2u;
|
|
1279
|
+
if (outputX < params.outputWidth && outputY < params.outputHeight && outputBlock < ${n}u) {
|
|
1280
|
+
let acc = pooled + bias[outputBlock];
|
|
1281
|
+
outputData[(outputY * params.outputWidth + outputX) * ${n}u + outputBlock] = ${ut(ht("acc",e),t)};
|
|
1282
|
+
}
|
|
1283
|
+
}`);else for(let N=0;N<2;N++)for(let V=0;V<a;V++)Y.push(`{
|
|
1284
|
+
${Array.from({length:T},(z,K)=>P("resultTile",`resultBase + ${K*p*p}u`,`acc${N}_${K}_${V}`,`${p}u`)).join(`
|
|
1285
|
+
`)}
|
|
1286
|
+
workgroupBarrier();
|
|
1287
|
+
let outputY = tileY + rowBase + ${N}u;
|
|
1288
|
+
for (var e = lane; e < ${h*m}u; e += ${o.subgroupSize}u) {
|
|
1289
|
+
let pixel = e / ${m}u;
|
|
1290
|
+
let outputX = tileX + pixel;
|
|
1291
|
+
let outputBlock = blockBase + ${V*m}u + e % ${m}u;
|
|
1292
|
+
if (outputX < params.outputWidth && outputY < params.outputHeight && outputBlock < ${n}u) {
|
|
1293
|
+
let base = resultBase + pixel * ${p}u + (e % ${m}u) * 4u;
|
|
1294
|
+
let acc = vec4<f32>(f32(resultTile[base]), f32(resultTile[base + 1u]),
|
|
1295
|
+
f32(resultTile[base + 2u]), f32(resultTile[base + 3u])) + bias[outputBlock];
|
|
1296
|
+
outputData[(outputY * params.outputWidth + outputX) * ${n}u + outputBlock] = ${ut(ht("acc",e),t)};
|
|
1297
|
+
}
|
|
1298
|
+
}
|
|
1299
|
+
workgroupBarrier();
|
|
1300
|
+
}`);const q=u?(()=>{const[N,V]=u.sourceBlocks,z=K=>K===u.upsampledSource?`min(u32(inputY) / 2u, params.source${K}Height - 1u) * params.source${K}Width + min(u32(inputX) / 2u, params.source${K}Width - 1u)`:`u32(inputY) * params.source${K}Width + u32(inputX)`;return`if (block < ${N}u) {
|
|
1301
|
+
value = input0[(${z(0)}) * ${N}u + block];
|
|
1302
|
+
} else {
|
|
1303
|
+
value = input1[(${z(1)}) * ${V}u + block - ${N}u];
|
|
1304
|
+
}`})():`value = inputData[(u32(inputY) * params.inputWidth + u32(inputX)) * ${i}u + block];`,Wt=u?`struct Params {
|
|
1305
|
+
outputWidth: u32,
|
|
1306
|
+
outputHeight: u32,
|
|
1307
|
+
outputBlocks: u32,
|
|
1308
|
+
inputBlocks: u32,
|
|
1309
|
+
source0Width: u32,
|
|
1310
|
+
source0Height: u32,
|
|
1311
|
+
source1Width: u32,
|
|
1312
|
+
source1Height: u32,
|
|
1313
|
+
}
|
|
1314
|
+
@group(0) @binding(0) var<storage, read> input0: array<${b}>;
|
|
1315
|
+
@group(0) @binding(1) var<storage, read> input1: array<${b}>;
|
|
1316
|
+
@group(0) @binding(2) var<storage, read> weights: array<${b}>;
|
|
1317
|
+
@group(0) @binding(3) var<storage, read> bias: array<vec4<f32>>;
|
|
1318
|
+
@group(0) @binding(4) var<storage, read_write> outputData: array<${x}>;
|
|
1319
|
+
@group(0) @binding(5) var<uniform> params: Params;`:`struct Params {
|
|
1320
|
+
inputWidth: u32,
|
|
1321
|
+
inputHeight: u32,
|
|
1322
|
+
outputWidth: u32,
|
|
1323
|
+
outputHeight: u32,
|
|
1324
|
+
inputBlocks: u32,
|
|
1325
|
+
outputBlocks: u32,
|
|
1326
|
+
}
|
|
1327
|
+
@group(0) @binding(0) var<storage, read> inputData: array<${b}>;
|
|
1328
|
+
@group(0) @binding(1) var<storage, read> weights: array<${b}>;
|
|
1329
|
+
@group(0) @binding(2) var<storage, read> bias: array<vec4<f32>>;
|
|
1330
|
+
@group(0) @binding(3) var<storage, read_write> outputData: array<${x}>;
|
|
1331
|
+
@group(0) @binding(4) var<uniform> params: Params;`,Jt=u?"params.outputWidth":"params.inputWidth",te=u?"params.outputHeight":"params.inputHeight";return`${r==="fp16"||t==="fp16"||c==="f16"||d==="f16"?`enable f16;
|
|
1332
|
+
`:""}${rt}
|
|
1333
|
+
${Wt}
|
|
1334
|
+
|
|
1335
|
+
var<workgroup> inputTile: array<${c}, ${_*f*g*p}>;
|
|
1336
|
+
var<workgroup> weightTile: array<${c}, ${_*9*p*w}>;
|
|
1337
|
+
var<workgroup> resultTile: array<${d}, ${Ot*h*p}>;
|
|
1338
|
+
|
|
1339
|
+
@compute @workgroup_size(${$})
|
|
1340
|
+
fn main(
|
|
1341
|
+
@builtin(local_invocation_index) localIndex: u32,
|
|
1342
|
+
@builtin(workgroup_id) workgroupId: vec3<u32>,
|
|
1343
|
+
@builtin(subgroup_id) subgroup: u32,
|
|
1344
|
+
@builtin(subgroup_invocation_id) lane: u32
|
|
1345
|
+
) {
|
|
1346
|
+
let tileX = workgroupId.x * ${h}u;
|
|
1347
|
+
let tileY = workgroupId.y * ${Gt}u;
|
|
1348
|
+
let blockBase = workgroupId.z * ${w/4}u;
|
|
1349
|
+
let rowBase = subgroup * 2u;
|
|
1350
|
+
let resultBase = subgroup * ${h*p}u;
|
|
1351
|
+
${$t}
|
|
1352
|
+
|
|
1353
|
+
${ei(v,[{array:"inputTile",vectors:f*g*m,read:`let pixel = i / ${m}u;
|
|
1354
|
+
let block = stageChunk * ${m}u + i % ${m}u;
|
|
1355
|
+
let inputX = i32(tileX + pixel % ${g}u) - 1;
|
|
1356
|
+
let inputY = i32(tileY + pixel / ${g}u) - 1;
|
|
1357
|
+
if (block < ${i}u && inputX >= 0 && inputY >= 0 &&
|
|
1358
|
+
inputX < i32(${Jt}) && inputY < i32(${te})) {
|
|
1359
|
+
${q}
|
|
1360
|
+
}`},{array:"weightTile",vectors:9*p*w/4,read:`let quad = i % ${w/4}u;
|
|
1361
|
+
let row = i / ${w/4}u;
|
|
1362
|
+
let tap = row / ${p}u;
|
|
1363
|
+
let inputChannel = stageChunk * ${p}u + row % ${p}u;
|
|
1364
|
+
let outputBlock = blockBase + quad;
|
|
1365
|
+
if (inputChannel < ${i*4}u && outputBlock < ${n}u) {
|
|
1366
|
+
value = weights[((tap * ${i}u + inputChannel / 4u) * ${n}u + outputBlock) * 4u + inputChannel % 4u];
|
|
1367
|
+
}`}],Qt,$,b,c,s)}
|
|
1368
|
+
|
|
1369
|
+
${Y.join(`
|
|
1370
|
+
`)}
|
|
1371
|
+
}
|
|
1372
|
+
`}const ke=[[-1,0],[0,1]],cr=[[[0],[1,2]],[[0,1],[2]]],ue=16;function Qn(r,t,e,i,n,o){const a=M(e.outputChannels),s=n[o],u=n[1-o],l=o===0?0:n[0]*4,p=o===0?n[0]*4:0,c=ue*s*a*16,d=new Float32Array(c+9*u*a*16),h=xe(e.weight),g=(v,w,T,$)=>v<e.outputChannels&&w<e.inputChannels?h[((v*e.inputChannels+w)*3+T)*3+$]:0;for(let v=0;v<a*4;v++){const w=v>>2,T=v&3;for(let $=0;$<s*4;$++)for(let b=0;b<4;b++)for(let x=0;x<4;x++){let B=0;for(const I of cr[b>>1][x>>1])for(const P of cr[b&1][x&1])B+=g(v,l+$,I,P);const k=((b*4+x)*s+($>>2))*a+w;d[k*16+($&3)*4+T]=B}for(let $=0;$<u*4;$++)for(let b=0;b<9;b++){const x=(b*u+($>>2))*a+w;d[c+x*16+($&3)*4+T]=g(v,p+$,b/3|0,b%3)}}let f;if(i==="fp16"){const v=new Uint16Array(d.length);for(let w=0;w<d.length;w++)v[w]=er(d[w]);f=new Uint8Array(v.buffer)}else f=new Uint8Array(d.buffer);const _=new Float32Array(a*4);_.set(xe(e.bias));const m=oe(r,`oidn/${t}/weights/${i}/subpixel`,f,ot.STORAGE);try{return{weights:m,weightLayout:"k-major",bias:oe(r,`oidn/${t}/bias`,new Uint8Array(_.buffer),ot.STORAGE)}}catch(v){throw m.destroy(),v}}function Jn(r,t,e,i,n,o,a,s,u){const l=n.size,p=u?2:1,c=n.componentType,d=n.resultType,h=Ft,g=Gt,f=h+2,_=g+2,m=h/2+2,v=g/2+2,w=l/4,T=a[s],$=a[1-s],b=Math.ceil(T/w),x=Math.ceil($/w),B=o*l,k=g/2,I=h/2/l,P=Ot*n.subgroupSize,O=D(r),H=D(t),{left:rt,right:at,load:$t,store:Z,multiplyAdd:nt,declare:St,extensions:Qt}=ur(n),Y=(U,L,Q)=>`acc${U}_${L}_${Q}`,q=[];for(let U=0;U<k;U++)for(let L=0;L<I;L++)for(let Q=0;Q<o;Q++)q.push(St(Y(U,L,Q)));const Wt=(U,L)=>{const Q=[];for(let xt=0;xt<k;xt++)for(let ge=0;ge<I;ge++)Q.push(`{
|
|
1373
|
+
let a = ${$t(rt,"inputTile",`inputTileBase + ${U(xt,ge)}`,L)};
|
|
1374
|
+
${Array.from({length:o},(da,xi)=>`${Y(xt,ge,xi)} = ${nt}(a, b${xi}, ${Y(xt,ge,xi)});`).join(`
|
|
1375
|
+
`)}
|
|
1376
|
+
}`);return Q.join(`
|
|
1377
|
+
`)},Jt=U=>Array.from({length:o},(L,Q)=>`let b${Q} = ${$t(at,"weightTile",`weightTileBase + ${U} * ${l*B}u + ${Q*l}u`,`${B}u`)};`).join(`
|
|
1378
|
+
`),te=(U,L,Q)=>`(${Q} * params.source${U}Width + ${L}) * ${a[U]}u`,ee=Math.max(_*f,v*m)*w,N=ue*l*B/4,V=(U,L,Q)=>({array:"weightTile",vectors:U*l*B/4,stride:N,read:`let quad = i % ${B/4}u;
|
|
1379
|
+
let row = i / ${B/4}u;
|
|
1380
|
+
let matrix = row / ${l}u;
|
|
1381
|
+
let inputChannel = stageChunk * ${l}u + row % ${l}u;
|
|
1382
|
+
let outputBlock = blockBase + quad;
|
|
1383
|
+
if (inputChannel < ${L*4}u && outputBlock < ${i}u) {
|
|
1384
|
+
value = weights[${Q}u + ((matrix * ${L}u + inputChannel / 4u) * ${i}u + outputBlock) * 4u + inputChannel % 4u];
|
|
1385
|
+
}`}),z=(U,L,Q)=>({array:"inputTile",vectors:U*L*w,stride:ee,read:`let pixel = i / ${w}u;
|
|
1386
|
+
let block = stageChunk * ${w}u + i % ${w}u;
|
|
1387
|
+
${Q}`}),K=`let inputX = i32(tileX / 2u + pixel % ${m}u) - 1;
|
|
1388
|
+
let inputY = i32(tileY / 2u + pixel / ${m}u) - 1;
|
|
1389
|
+
if (block < ${T}u && inputX >= 0 && inputY >= 0 &&
|
|
1390
|
+
inputX < i32(params.source${s}Width) && inputY < i32(params.source${s}Height)) {
|
|
1391
|
+
value = input${s}[${te(s,"u32(inputX)","u32(inputY)")} + block];
|
|
1392
|
+
}`,Lt=1-s,Mt=`let inputX = i32(tileX + pixel % ${f}u) - 1;
|
|
1393
|
+
let inputY = i32(tileY + pixel / ${f}u) - 1;
|
|
1394
|
+
if (block < ${$}u && inputX >= 0 && inputY >= 0 &&
|
|
1395
|
+
inputX < i32(params.outputWidth) && inputY < i32(params.outputHeight)) {
|
|
1396
|
+
value = input${Lt}[${te(Lt,"u32(inputX)","u32(inputY)")} + block];
|
|
1397
|
+
}`,vt=[];for(let U=0;U<2;U++)for(let L=0;L<2;L++)vt.push(`{
|
|
1398
|
+
let matrix = phase * 4u + ${U*2+L}u;
|
|
1399
|
+
let offsetY = select(${ke[0][U]+1}u, ${ke[1][U]+1}u, phaseY == 1u);
|
|
1400
|
+
let offsetX = select(${ke[0][L]+1}u, ${ke[1][L]+1}u, phaseX == 1u);
|
|
1401
|
+
${Jt("matrix")}
|
|
1402
|
+
${Wt((Q,xt)=>`((${Q}u + offsetY) * ${m}u + ${xt*l}u + offsetX) * ${l}u`,`${l}u`)}
|
|
1403
|
+
}`);const fe=`let ky = tap / 3u;
|
|
1404
|
+
let kx = tap % 3u;
|
|
1405
|
+
${Jt("tap")}
|
|
1406
|
+
${Wt((U,L)=>`((phaseY + ${U*2}u + ky) * ${f}u + phaseX + kx + ${L*l*2}u) * ${l}u`,`${l*2}u`)}`,ua=sr(n,fe),nn=[];for(let U=0;U<k;U++)for(let L=0;L<o;L++)nn.push(`{
|
|
1407
|
+
${Array.from({length:I},(Q,xt)=>Z("resultTile",`resultBase + ${xt*l*l}u`,Y(U,xt,L),`${l}u`)).join(`
|
|
1408
|
+
`)}
|
|
1409
|
+
workgroupBarrier();
|
|
1410
|
+
let outputY = tileY + phaseY + ${U*2}u;
|
|
1411
|
+
for (var e = lane; e < ${h/2*w}u; e += ${n.subgroupSize}u) {
|
|
1412
|
+
let pixel = e / ${w}u;
|
|
1413
|
+
let outputX = tileX + phaseX + pixel * 2u;
|
|
1414
|
+
let outputBlock = blockBase + ${L*w}u + e % ${w}u;
|
|
1415
|
+
if (outputX < params.outputWidth && outputY < params.outputHeight && outputBlock < ${i}u) {
|
|
1416
|
+
let base = resultBase + pixel * ${l}u + (e % ${w}u) * 4u;
|
|
1417
|
+
let acc = vec4<f32>(f32(resultTile[base]), f32(resultTile[base + 1u]),
|
|
1418
|
+
f32(resultTile[base + 2u]), f32(resultTile[base + 3u])) + bias[outputBlock];
|
|
1419
|
+
outputData[(outputY * params.outputWidth + outputX) * ${i}u + outputBlock] = ${ut(ht("acc",e),t)};
|
|
1420
|
+
}
|
|
1421
|
+
}
|
|
1422
|
+
workgroupBarrier();
|
|
1423
|
+
}`);return`${r==="fp16"||t==="fp16"||c==="f16"||d==="f16"?`enable f16;
|
|
1424
|
+
`:""}${Qt}
|
|
1425
|
+
struct Params {
|
|
1426
|
+
outputWidth: u32,
|
|
1427
|
+
outputHeight: u32,
|
|
1428
|
+
outputBlocks: u32,
|
|
1429
|
+
inputBlocks: u32,
|
|
1430
|
+
source0Width: u32,
|
|
1431
|
+
source0Height: u32,
|
|
1432
|
+
source1Width: u32,
|
|
1433
|
+
source1Height: u32,
|
|
1434
|
+
}
|
|
1435
|
+
@group(0) @binding(0) var<storage, read> input0: array<${O}>;
|
|
1436
|
+
@group(0) @binding(1) var<storage, read> input1: array<${O}>;
|
|
1437
|
+
@group(0) @binding(2) var<storage, read> weights: array<${O}>;
|
|
1438
|
+
@group(0) @binding(3) var<storage, read> bias: array<vec4<f32>>;
|
|
1439
|
+
@group(0) @binding(4) var<storage, read_write> outputData: array<${H}>;
|
|
1440
|
+
@group(0) @binding(5) var<uniform> params: Params;
|
|
1441
|
+
|
|
1442
|
+
var<workgroup> inputTile: array<${c}, ${p*ee*4}>;
|
|
1443
|
+
var<workgroup> weightTile: array<${c}, ${p*N*4}>;
|
|
1444
|
+
var<workgroup> resultTile: array<${d}, ${Ot*(h/2)*l}>;
|
|
1445
|
+
|
|
1446
|
+
@compute @workgroup_size(${P})
|
|
1447
|
+
fn main(
|
|
1448
|
+
@builtin(local_invocation_index) localIndex: u32,
|
|
1449
|
+
@builtin(workgroup_id) workgroupId: vec3<u32>,
|
|
1450
|
+
@builtin(subgroup_id) subgroup: u32,
|
|
1451
|
+
@builtin(subgroup_invocation_id) lane: u32
|
|
1452
|
+
) {
|
|
1453
|
+
let tileX = workgroupId.x * ${h}u;
|
|
1454
|
+
let tileY = workgroupId.y * ${g}u;
|
|
1455
|
+
let blockBase = workgroupId.z * ${B/4}u;
|
|
1456
|
+
let phase = subgroup;
|
|
1457
|
+
let phaseY = subgroup >> 1u;
|
|
1458
|
+
let phaseX = subgroup & 1u;
|
|
1459
|
+
let resultBase = subgroup * ${h/2*l}u;
|
|
1460
|
+
${q.join(`
|
|
1461
|
+
`)}
|
|
1462
|
+
|
|
1463
|
+
${ei(b,[z(v,m,K),V(ue,T,0)],vt.join(`
|
|
1464
|
+
`),P,O,c,u)}
|
|
1465
|
+
|
|
1466
|
+
${ei(x,[z(_,f,Mt),V(9,$,ue*T*i*4)],ua,P,O,c,u)}
|
|
1467
|
+
|
|
1468
|
+
${nn.join(`
|
|
1469
|
+
`)}
|
|
1470
|
+
}
|
|
1471
|
+
`}function to(r){return r.op==="concat"?r.inputs:r.op==="fusedUpsampleConcatConv2d"?r.inputs.map(t=>t.value):[r.input]}function lr(r,t="auto"){const e=r.features.has("shader-f16");if(t==="fp16"&&!e)throw new Error("OIDN FP16 was requested but the OIDNGPUDevice does not have shader-f16 enabled");return t==="auto"?e?"fp16":"fp32":t}class eo{constructor(t,e,i={}){y(this,"usedKernels",new Set);y(this,"precision");y(this,"kernelSetting");y(this,"gemm");y(this,"maxSpatialInputBlocks");y(this,"subgroupsAvailable");y(this,"subgroupMatrix");y(this,"subgroupMatrixOutputFragments");y(this,"subgroupMatrixSubPixel");y(this,"subgroupMatrixDoubleBuffer");y(this,"_subgroupMatrixFixedFragments");y(this,"_model");y(this,"_gemmByOutputBlocks",new Map);y(this,"_packedConvs",new Map);y(this,"_pipelineCache");y(this,"_executionCache",new Map);y(this,"_retiredExecutions",new Set);y(this,"_clock",0);y(this,"_shapeCacheSize");y(this,"_profileNextExecution",!1);y(this,"_profileOperations",0);y(this,"_resources",new qi);y(this,"_disposed",!1);var c,d,h,g,f,_,m,v,w,T,$,b,x,B;this._device=t,this._pipelineCache=Sn(t),this.precision=lr(t,i.precision??"auto"),this.kernelSetting=i.kernel??"auto";const n=((c=i.gemm)==null?void 0:c.workgroupSize)??[8,8];if(this.gemm=Object.freeze({tilePolicy:((d=i.gemm)==null?void 0:d.tilePolicy)??((h=i.gemm)!=null&&h.workgroupSize?"fixed":"output-aligned"),decoderLoad:((g=i.gemm)==null?void 0:g.decoderLoad)??"per-load",decoderSubPixel:((f=i.gemm)==null?void 0:f.decoderSubPixel)??this.precision==="fp16",unrollRows:((_=i.gemm)==null?void 0:_.unrollRows)??!1,poolLayout:((m=i.gemm)==null?void 0:m.poolLayout)??"channels",sharedLayout:((v=i.gemm)==null?void 0:v.sharedLayout)??(this.precision==="fp16"?"padded-input":"padded"),accumulationOrder:((w=i.gemm)==null?void 0:w.accumulationOrder)??"k-major",finalLayer:((T=i.gemm)==null?void 0:T.finalLayer)??"shared-auto",loadMode:(($=i.gemm)==null?void 0:$.loadMode)??"native",addressMode:((b=i.gemm)==null?void 0:b.addressMode)??"incremental",weightLayout:((x=i.gemm)==null?void 0:x.weightLayout)??"k-major",rowsPerThread:((B=i.gemm)==null?void 0:B.rowsPerThread)??8,workgroupSize:Object.freeze([n[0],n[1]])}),!["fixed","output-aligned"].includes(this.gemm.tilePolicy))throw new Error(`Unsupported GEMM tile policy: ${this.gemm.tilePolicy}`);if(!["per-load","source-first"].includes(this.gemm.decoderLoad))throw new Error(`Unsupported GEMM decoder load: ${this.gemm.decoderLoad}`);if(!["spatial","channels"].includes(this.gemm.poolLayout))throw new Error(`Unsupported GEMM pool layout: ${this.gemm.poolLayout}`);if(!["linear","padded","padded-input","padded-weights"].includes(this.gemm.sharedLayout))throw new Error(`Unsupported GEMM shared layout: ${this.gemm.sharedLayout}`);if(!["k-major","row-major"].includes(this.gemm.accumulationOrder))throw new Error(`Unsupported GEMM accumulation order: ${this.gemm.accumulationOrder}`);if(!["direct","shared-input","shared-input-weights","shared-auto"].includes(this.gemm.finalLayer))throw new Error(`Unsupported GEMM final layer: ${this.gemm.finalLayer}`);if(!["native","packed-weights","packed-all"].includes(this.gemm.loadMode))throw new Error(`Unsupported GEMM load mode: ${this.gemm.loadMode}`);if(!["analytic","incremental","base-offset"].includes(this.gemm.addressMode))throw new Error(`Unsupported GEMM address mode: ${this.gemm.addressMode}`);if(!["output-major","k-major"].includes(this.gemm.weightLayout))throw new Error(`Unsupported GEMM weight layout: ${this.gemm.weightLayout}`);const o=dt(this.gemm),a=(bt(this.gemm).input+bt(this.gemm).weights)*4*(this.precision==="fp16"?2:4);if(![2,4,8].includes(o.rowsPerThread)||![4,8,16].includes(o.workgroupX)||![4,8].includes(o.workgroupY)||n.length!==2||o.workgroupX>t.limits.maxComputeWorkgroupSizeX||o.workgroupY>t.limits.maxComputeWorkgroupSizeY||o.workgroupX*o.workgroupY>t.limits.maxComputeInvocationsPerWorkgroup||a>t.limits.maxComputeWorkgroupStorageSize)throw new Error("Unsupported GEMM tile configuration for this OIDNGPUDevice");this.subgroupsAvailable=t.features.has("subgroups");const s=i.subgroupMatrix??{},u=this.kernelSetting==="subgroup-matrix"?ar(t,s.accumulation):void 0;this.subgroupMatrixOutputFragments=u?Kn(u,s):0,this.subgroupMatrixSubPixel=s.subPixel??!0,this._subgroupMatrixFixedFragments=s.outputFragments!==void 0;const l=k=>u!==void 0&&qn(u,this.subgroupMatrixOutputFragments,this.subgroupMatrixSubPixel,k)<=t.limits.maxComputeWorkgroupStorageSize;if(u&&s.outputFragments===void 0&&!l(!1)&&(this.subgroupMatrixOutputFragments=1),this.subgroupMatrixDoubleBuffer=(s.doubleBuffer??!1)&&l(!0),u&&Ot*u.subgroupSize<=t.limits.maxComputeInvocationsPerWorkgroup&&Ot*u.subgroupSize<=t.limits.maxComputeWorkgroupSizeX&&l(this.subgroupMatrixDoubleBuffer)&&(this.subgroupMatrix=u),this.maxSpatialInputBlocks=this.precision==="fp16"&&t.limits.maxComputeInvocationsPerWorkgroup>=j*j&&t.limits.maxComputeWorkgroupSizeX>=j&&t.limits.maxComputeWorkgroupSizeY>=j?Math.floor(t.limits.maxComputeWorkgroupStorageSize/(pt*pt*4*2)):0,this._shapeCacheSize=Math.max(1,i.shapeCacheSize??2),e.inputChannels%3!==0||e.inputChannels<3||e.inputChannels>9)throw new Error(`Native OIDN expects 3, 6, or 9 input channels, got ${e.inputChannels}`);this._model={spec:e.spec,inputChannels:e.inputChannels,outputChannels:e.outputChannels,channelsByValue:new Map(e.channelsByValue),convChannels:new Map(e.convChannels)};const p=new Map;for(const k of ve(this._model,{fuseConvPool:this._matrixPool()}).nodes){if(k.op==="fusedUpsampleConcatConv2d"){const I=this._subPixelDecoder(k);I&&p.set(k.conv.id,I)}if(k.op!=="conv2d"&&k.op!=="maxPool2d"&&k.op!=="fusedConvReluMaxPool2d"&&k.op!=="fusedUpsampleConcatConv2d")throw new Error(`Native OIDN descriptor ${e.spec.id} leaves unsupported ${k.op} node ${k.id} after graph optimization`)}try{for(const[k,I]of e.convTensors){const P=this._selectConvKernel(M(I.inputChannels),k===e.spec.output),O=P==="implicit-gemm"?this.gemm.weightLayout:P==="subgroup-matrix"?"k-major":"output-major",H=p.get(k),rt=H?Qn(t,k,I,this.precision,H.sourceBlocks,H.upsampledSource):An(t,k,I,this.precision,O);this._resources.track("gpu-buffer",rt.weights),this._resources.track("gpu-buffer",rt.bias),this._packedConvs.set(k,rt)}}catch(k){for(const I of this._packedConvs.values())this._releaseBuffer(I.weights),this._releaseBuffer(I.bias);throw this._packedConvs.clear(),k}}get weightBuffers(){return[...this._packedConvs.values()].flatMap(t=>[t.weights,t.bias])}_pipeline(t,e){let i=this._pipelineCache.get(t);return i||(i=this._device.createComputePipeline({label:`oidn/${t}`,layout:"auto",compute:{module:this._device.createShaderModule({label:`oidn/${t}`,code:e}),entryPoint:"main"}}),this._pipelineCache.set(t,i)),i}_nodePipelineSpec(t,e,i=this.subgroupMatrixOutputFragments){if(t.op==="conv2d"){const n=e?"fp32":this.precision,o=M(this._model.convChannels.get(t.id).inputChannels),a=M(this._model.convChannels.get(t.id).outputChannels),s=this._gemmForOutput(a),u=this._selectConvKernel(o,e),l=s.finalLayer==="shared-input-weights"||s.finalLayer==="shared-auto"&&Ki(this.precision,o,!0)<=this._device.limits.maxComputeWorkgroupStorageSize;return u!=="subgroup-matrix"&&e&&this._model.convChannels.get(t.id).outputChannels===3&&(this.kernelSetting==="auto"||this.kernelSetting==="implicit-gemm"||this.kernelSetting==="subgroup-matrix")&&s.finalLayer!=="direct"&&this._device.limits.maxComputeWorkgroupSizeX>=8&&this._device.limits.maxComputeWorkgroupSizeY>=8&&this._device.limits.maxComputeInvocationsPerWorkgroup>=64&&Ki(this.precision,o,l)<=this._device.limits.maxComputeWorkgroupStorageSize?{key:`conv-final-rgb/${this.precision}/${t.activation}/in${o}/weights-${l?"shared":"storage"}`,kernel:"direct",code:Cn(this.precision,t.activation,o,l)}:u==="subgroup-matrix"?this._subgroupMatrixSpec(`conv/${t.activation}/in${o}/out${a}`,n,t.activation,o,a,i):{key:`conv-${u}/${this.precision}/${n}/${t.activation}/in${o}/out${a}`+(u==="implicit-gemm"?`/address-${s.addressMode}/weights-${s.weightLayout}-v1/tile-${s.workgroupSize.join("x")}-r${s.rowsPerThread}/loads-${s.loadMode}/shared-${s.sharedLayout}/acc-${s.accumulationOrder}/unroll-${Number(s.unrollRows)}`:""),kernel:u,code:u==="implicit-gemm"?Ln(this.precision,n,t.activation,o,a,s):u==="spatial"?Wn(this.precision,n,t.activation,o,a):u==="subgroup"?Un(this.precision,n,t.activation,o,a):zn(this.precision,n,t.activation,o,a)}}if(t.op==="maxPool2d"){const n=M(this._model.channelsByValue.get(t.id)),o=this._coalescedPool();return{key:`max-pool/${this.precision}/out${n}/${o?"channels":"spatial"}`,kernel:"direct",code:Rn(this.precision,n,o)}}if(t.op==="fusedConvReluMaxPool2d"){const n=M(this._model.convChannels.get(t.conv.id).inputChannels),o=M(this._model.convChannels.get(t.conv.id).outputChannels);return this._matrixPool()?this._subgroupMatrixSpec(`conv-pool/${t.conv.activation}/in${n}/out${o}`,this.precision,t.conv.activation,n,o,i,void 0,!0):{key:`conv-pool/${this.precision}/${t.conv.activation}/in${n}/out${o}`,kernel:"direct",code:Dn(this.precision,t.conv.activation,n,o)}}if(t.op==="fusedUpsampleConcatConv2d"){if(t.inputs.length!==2)throw new Error(`Native fused decoder ${t.id} requires two inputs`);const n=t.inputs.map(h=>M(this._model.channelsByValue.get(h.value))),o=t.inputs.findIndex(h=>h.upsample);if(o!==0&&o!==1)throw new Error(`Native fused decoder ${t.id} has no upsample input`);const a=n[0]+n[1],s=e?"fp32":this.precision,u=this._selectConvKernel(a,e),l=u==="subgroup"?"direct":u,p=M(this._model.convChannels.get(t.conv.id).outputChannels);if(l==="subgroup-matrix"&&this._subPixelDecoder(t)){const h=this.subgroupMatrix;return{key:`subgroup-matrix-subpixel/${this.precision}/${s}/${h.dialect}/${h.componentType}-${h.resultType}-${h.size}-sg${h.subgroupSize}/n${i}${this.subgroupMatrixDoubleBuffer?"-db":""}/${t.conv.activation}/${n.join("+")}/out${p}/up${o}`,kernel:"subgroup-matrix",code:Jn(this.precision,s,t.conv.activation,p,h,i,n,o,this.subgroupMatrixDoubleBuffer)}}if(l==="subgroup-matrix")return this._subgroupMatrixSpec(`decoder/${t.conv.activation}/${n.join("+")}/out${p}/up${o}`,s,t.conv.activation,a,p,i,{sourceBlocks:n,upsampledSource:o});const c=this._gemmForOutput(p);return{key:`decoder-${l}${l==="implicit-gemm"&&this._subPixelDecoder(t)?"-subpixel":""}/${this.precision}/${s}/${t.conv.activation}/${n.join("+")}/out${p}/up${o}`+(l==="implicit-gemm"?`/address-${c.addressMode}/weights-${c.weightLayout}-v1/tile-${c.workgroupSize.join("x")}-r${c.rowsPerThread}/loads-${c.loadMode}/shared-${c.sharedLayout}/acc-${c.accumulationOrder}/unroll-${Number(c.unrollRows)}/decoder-${c.decoderLoad}`:""),kernel:l,code:l==="implicit-gemm"?(this._subPixelDecoder(t)?Yn:Xn)(this.precision,s,t.conv.activation,n,o,p,c):l==="spatial"?Gn(this.precision,s,t.conv.activation,n,o,p):Fn(this.precision,s,t.conv.activation,n,o,p)}}throw new Error(`Native OIDN does not implement unfused ${t.op} node ${t.id}`)}_subPixelDecoder(t){if(t.id===this._model.spec.output||t.inputs.length!==2)return;const e=this._selectConvKernel(M(this._model.convChannels.get(t.conv.id).inputChannels),!1);if(!(e==="subgroup-matrix"&&this.subgroupMatrixSubPixel)&&!(e==="implicit-gemm"&&this.gemm.decoderSubPixel&&this.gemm.weightLayout==="k-major"))return;const i=t.inputs[0].upsample?0:t.inputs[1].upsample?1:-1;if(i!==-1)return{sourceBlocks:t.inputs.map(n=>M(this._model.channelsByValue.get(n.value))),upsampledSource:i}}_subgroupMatrixSpec(t,e,i,n,o,a,s,u=!1){const l=this.subgroupMatrix;return{key:`subgroup-matrix/${this.precision}/${e}/${l.dialect}/${l.componentType}-${l.resultType}-${l.size}-sg${l.subgroupSize}/n${a}${this.subgroupMatrixDoubleBuffer?"-db":""}/${t}`,kernel:"subgroup-matrix",code:Zn(this.precision,e,i,n,o,l,a,this.subgroupMatrixDoubleBuffer,s,u)}}_subgroupMatrixFragmentsFor(t){const e=Math.min(this.subgroupMatrixOutputFragments,Math.ceil(M(t.channels)*4/this.subgroupMatrix.size));return this._subgroupMatrixFixedFragments||e===1?e:Math.ceil(t.width/Ft)*Math.ceil(t.height/Gt)*Math.ceil(M(t.channels)*4/(e*this.subgroupMatrix.size))<jn?1:e}_gemmForOutput(t){if(this.gemm.tilePolicy!=="output-aligned"||this.gemm.workgroupSize[0]!==8||t<=0||t%16!==0)return this.gemm;const e=this._gemmByOutputBlocks.get(t);if(e)return e;const i=Object.freeze({...this.gemm,workgroupSize:Object.freeze([16,this.gemm.workgroupSize[1]])}),n=dt(i),o=bt(i),a=this._device.limits,u=n.workgroupX<=a.maxComputeWorkgroupSizeX&&n.workgroupY<=a.maxComputeWorkgroupSizeY&&n.workgroupX*n.workgroupY<=a.maxComputeInvocationsPerWorkgroup&&(o.input+o.weights)*4*(this.precision==="fp16"?2:4)<=a.maxComputeWorkgroupStorageSize?i:this.gemm;return this._gemmByOutputBlocks.set(t,u),u}_matrixPool(){const t=this.subgroupMatrix;return this._model.spec.family==="oidn-unet-large"&&this.precision==="fp16"&&(t==null?void 0:t.dialect)==="naga"&&t.size===8&&t.subgroupSize===32&&t.componentType==="f16"&&t.resultType==="f32"}_coalescedPool(){return this.gemm.poolLayout==="channels"&&(this.kernelSetting==="auto"||this.kernelSetting==="implicit-gemm"||this.kernelSetting==="subgroup-matrix")}_selectConvKernel(t,e){var n;const i=this.precision==="fp16"&&t<=this.maxSpatialInputBlocks;return this.kernelSetting==="direct"?"direct":this.kernelSetting==="spatial"?i?"spatial":"direct":this.kernelSetting==="implicit-gemm"?e?"direct":"implicit-gemm":this.kernelSetting==="subgroup"?this.subgroupsAvailable?"subgroup":"direct":this.kernelSetting==="subgroup-matrix"?e?((n=this.subgroupMatrix)==null?void 0:n.dialect)==="naga"?"subgroup-matrix":"direct":this.subgroupMatrix?"subgroup-matrix":"implicit-gemm":e?"direct":"implicit-gemm"}_nodePipeline(t,e){const{key:i,code:n}=this._nodePipelineSpec(t,e);return this._pipeline(i,n)}_createExecution(t,e){const i=this._matrixPool(),n=Hi(this._model,t,e,{fuseConvPool:i}),o=new Map,a=[],s=[],u=new Map;n.nodes.forEach((c,d)=>{for(const h of to(c))u.set(h,d)}),u.set(n.spec.output,n.nodes.length);const l=[],p=c=>(l.push(c),this._resources.track("gpu-buffer",c));try{const c=($,b,x,B)=>{const k=Mn(b,x),I=s.filter(O=>O.lastUse<B);let P=I.filter(O=>O.capacity>=k).sort((O,H)=>O.capacity-H.capacity)[0];!P&&i&&(P=I.sort((O,H)=>H.capacity-O.capacity)[0]),P||(P={capacity:0,lastUse:-2,values:[]},s.push(P)),P.capacity=Math.max(P.capacity,k),P.lastUse=u.get($)??-1,P.values.push($)};c(n.spec.input,n.inputShape,this.precision==="fp16"?2:4,-1),n.plannedNodes.forEach(({node:$,outputShape:b},x)=>{const B=$.id===n.spec.output;c($.id,b,B||this.precision==="fp32"?4:2,x)}),s.forEach(($,b)=>{const x=p(this._device.createBuffer({label:`oidn/activation/${t}x${e}/${b}`,size:Je($.capacity,4),usage:ot.STORAGE|ot.COPY_SRC|ot.COPY_DST}));a.push({buffer:x,capacity:$.capacity});for(const B of $.values)o.set(B,x)});const d=this._model.inputChannels/3,h=`pack/${this.precision}/${d}`,g=this._pipeline(h,or(this.precision,d)),f=[],_=[],m=[],v=[],w=[];n.plannedNodes.forEach(({node:$,outputShape:b},x)=>{const B=$.id===n.spec.output,k=this.subgroupMatrix?this._subgroupMatrixFragmentsFor(b):0,I=this._nodePipelineSpec($,B,k);I.kernel&&this.usedKernels.add(I.kernel);const P=this._pipeline(I.key,I.code);f.push(P),_.push(I.kernel??"direct"),m.push(I.kernel==="subgroup-matrix"?k:0);const O=o.get($.id);let H,rt,at;if($.op==="conv2d"){const Z=n.valueShapes.get($.input);at=$.id,rt=[Z.width,Z.height,b.width,b.height,M(Z.channels),M(b.channels)];const nt=this._packedConvs.get(at);H=[{binding:0,resource:{buffer:o.get($.input)}},{binding:1,resource:{buffer:nt.weights}},{binding:2,resource:{buffer:nt.bias}},{binding:3,resource:{buffer:O}}]}else if($.op==="maxPool2d"){const Z=n.valueShapes.get($.input);rt=[Z.width,Z.height,b.width,b.height,M(b.channels)],H=[{binding:0,resource:{buffer:o.get($.input)}},{binding:1,resource:{buffer:O}}]}else if($.op==="fusedConvReluMaxPool2d"){const Z=n.valueShapes.get($.input);at=$.conv.id,rt=[Z.width,Z.height,b.width,b.height,M(Z.channels),M(b.channels)];const nt=this._packedConvs.get(at);H=[{binding:0,resource:{buffer:o.get($.input)}},{binding:1,resource:{buffer:nt.weights}},{binding:2,resource:{buffer:nt.bias}},{binding:3,resource:{buffer:O}}]}else if($.op==="fusedUpsampleConcatConv2d"){at=$.conv.id;const Z=n.valueShapes.get($.inputs[0].value),nt=n.valueShapes.get($.inputs[1].value);rt=[b.width,b.height,M(b.channels),M(this._model.convChannels.get(at).inputChannels),Z.width,Z.height,nt.width,nt.height];const St=this._packedConvs.get(at);H=[{binding:0,resource:{buffer:o.get($.inputs[0].value)}},{binding:1,resource:{buffer:o.get($.inputs[1].value)}},{binding:2,resource:{buffer:St.weights}},{binding:3,resource:{buffer:St.bias}},{binding:4,resource:{buffer:O}}]}else throw new Error(`Unexpected native node ${$.op}`);const $t=p(ir(this._device,`oidn/${$.id}/params/${t}x${e}`,rt));w.push($t),H.push({binding:H.length,resource:{buffer:$t}}),v.push(this._device.createBindGroup({label:`oidn/${$.id}/bindings`,layout:P.getBindGroupLayout(0),entries:H}))});const T=p(ir(this._device,`oidn/input/params/${t}x${e}`,[t,e,M(this._model.inputChannels),this._model.inputChannels]));return w.push(T),{plan:n,valueBuffers:o,slots:a,nodeBindings:v,nodePipelines:f,nodeKernels:_,nodeFragments:m,inputPipeline:g,inputUniform:T,ownedBuffers:w,lastUsed:++this._clock}}catch(c){for(const d of l)this._releaseBuffer(d);throw c}}_execution(t,e){if(this._disposed)throw new Error("Native OIDN executor is disposed");const i=`${t}x${e}`;let n=this._executionCache.get(i);if(!n&&(n=this._createExecution(t,e),this._executionCache.set(i,n),this._executionCache.size>this._shapeCacheSize)){const o=[...this._executionCache.entries()].filter(([a])=>a!==i).sort((a,s)=>a[1].lastUsed-s[1].lastUsed)[0];o&&(this._executionCache.delete(o[0]),this._retireExecution(o[1]))}return n.lastUsed=++this._clock,n}_retireExecution(t){this._destroyExecution(t)}prewarm(t){for(const e of t)this._execution(e.width,e.height)}profileNextExecution(){return this._device.features.has("timestamp-query")?(this._profileNextExecution=!0,!0):!1}packedInputBuffer(t,e){const i=this._execution(t,e);return i.valueBuffers.get(i.plan.spec.input)}executePacked(t,e){return this._execute([],t,e,!0)}execute(t,e,i){return this._execute(t,e,i,!1)}_execute(t,e,i,n){const o=this._model.inputChannels/3;if(!n&&t.length!==o)throw new Error(`Native OIDN expected ${o} input buffers, got ${t.length}`);const a=this._execution(e,i),s=[...n?[]:["input-pack"],...a.plan.nodes.map(f=>f.id)],u=this._profileNextExecution&&this._device.features.has("timestamp-query");this._profileNextExecution=!1;const l=s.length*2,p=u?this._resources.track("gpu-query-set",this._device.createQuerySet({type:"timestamp",count:l})):void 0,c=l*8;let d,h;try{d=u?this._resources.track("gpu-buffer",this._device.createBuffer({label:`oidn/profile/resolve/${e}x${i}`,size:c,usage:ot.QUERY_RESOLVE|ot.COPY_SRC})):void 0,h=u?this._resources.track("gpu-buffer",this._device.createBuffer({label:`oidn/profile/readback/${e}x${i}`,size:c,usage:ot.COPY_DST|ot.MAP_READ})):void 0}catch(f){throw this._releaseBuffer(h),this._releaseBuffer(d),this._releaseQuerySet(p),f}const g=(f,_)=>({label:f,...p?{timestampWrites:{querySet:p,beginningOfPassWriteIndex:_*2,endOfPassWriteIndex:_*2+1}}:{}});try{const f=this._device.createCommandEncoder({label:`oidn/native/${e}x${i}`});if(!n){const _=t.map((w,T)=>({binding:T,resource:{buffer:w}}));_.push({binding:o,resource:{buffer:a.valueBuffers.get(a.plan.spec.input)}}),_.push({binding:o+1,resource:{buffer:a.inputUniform}});const m=this._device.createBindGroup({label:"oidn/input/bindings",layout:a.inputPipeline.getBindGroupLayout(0),entries:_}),v=f.beginComputePass(g("oidn/input-pack",0));v.setPipeline(a.inputPipeline),v.setBindGroup(0,m),v.dispatchWorkgroups(Math.ceil(e/tt),Math.ceil(i/tt),M(this._model.inputChannels)),v.end()}a.plan.plannedNodes.forEach(({node:_,outputShape:m},v)=>{const w=f.beginComputePass(g(`oidn/${a.plan.nodes[v].id}`,v+(n?0:1)));if(w.setPipeline(a.nodePipelines[v]),w.setBindGroup(0,a.nodeBindings[v]),_.op==="maxPool2d"&&this._coalescedPool()){const T=Math.ceil(m.width*M(m.channels)/tt),$=this._device.limits.maxComputeWorkgroupsPerDimension;w.dispatchWorkgroups(Math.min(T,$),Math.ceil(m.height/tt),Math.ceil(T/$))}else if(a.nodeKernels[v]==="subgroup-matrix")w.dispatchWorkgroups(Math.ceil(m.width*(_.op==="fusedConvReluMaxPool2d"?2:1)/Ft),Math.ceil(m.height*(_.op==="fusedConvReluMaxPool2d"?2:1)/Gt),Math.ceil(M(m.channels)/(a.nodeFragments[v]*this.subgroupMatrix.size/4)));else if(a.nodeKernels[v]==="implicit-gemm"){const T=dt(this._gemmForOutput(M(m.channels))),$=_.op==="fusedUpsampleConcatConv2d"&&this._subPixelDecoder(_)!==void 0;w.dispatchWorkgroups(Math.ceil(($?Math.ceil(m.width/2)*Math.ceil(m.height/2):m.width*m.height)/T.tileM),Math.ceil(M(m.channels)/T.tileNBlocks),$?4:1)}else w.dispatchWorkgroups(Math.ceil(m.width/tt),Math.ceil(m.height/tt),M(m.channels));w.end()}),p&&(f.resolveQuerySet(p,0,l,d,0),f.copyBufferToBuffer(d,0,h,0,c)),this._device.queue.submit([f.finish()])}catch(f){throw this._releaseBuffer(h),this._releaseBuffer(d),this._releaseQuerySet(p),f}return p&&(this._profileOperations++,this._readProfile(p,d,h,s)),a.valueBuffers.get(a.plan.spec.output)}_readProfile(t,e,i,n){this._releaseQuerySet(t),this._releaseBuffer(e),this._releaseBuffer(i),this._profileOperations--}_releaseBuffer(t){this._resources.release("gpu-buffer",t,()=>t.destroy())}_releaseQuerySet(t){this._resources.release("gpu-query-set",t,()=>t.destroy())}_destroyExecution(t){t.slots.forEach(e=>this._releaseBuffer(e.buffer)),t.ownedBuffers.forEach(e=>this._releaseBuffer(e))}getResourceInfo(){return this._resources.snapshot(this._retiredExecutions.size+this._profileOperations)}dispose(){if(!this._disposed){this._disposed=!0;for(const t of this._packedConvs.values())this._releaseBuffer(t.weights),this._releaseBuffer(t.bias);this._packedConvs.clear();for(const t of this._executionCache.values())this._destroyExecution(t);this._executionCache.clear();for(const t of this._retiredExecutions)this._destroyExecution(t);this._retiredExecutions.clear()}}}function la(r){return r}function io(r,t){return lr(r,t)}function ro(r,t){return ar(r,t)}const pr=new WeakMap;class dr extends eo{constructor(e,i,n){super(e,i,n);y(this,"_lastExecutionProfile")}execute(e,i,n){return super.execute(e,i,n)}get _asyncDevice(){return this._device}get _pipelinePromises(){let e=pr.get(this._device);return e||(e=new Map,pr.set(this._device,e)),e}getLastExecutionProfile(){return this._lastExecutionProfile}_pipelineAsync(e,i){const n=this._pipelineCache.get(e);if(n)return Promise.resolve(n);const o=this._pipelinePromises.get(e);if(o)return o;const a=this._asyncDevice.createComputePipelineAsync({label:`oidn/${e}`,layout:"auto",compute:{module:this._device.createShaderModule({label:`oidn/${e}`,code:i}),entryPoint:"main"}}).then(s=>(this._pipelineCache.set(e,s),this._pipelinePromises.delete(e),s),s=>{throw this._pipelinePromises.delete(e),s});return this._pipelinePromises.set(e,a),a}_retireExecution(e){this._retiredExecutions.add(e),this._asyncDevice.queue.onSubmittedWorkDone().catch(()=>{}).then(()=>{this._retiredExecutions.delete(e),this._destroyExecution(e)})}async prepare(){if(this._disposed)throw new Error("Native OIDN executor is disposed");const e=ve(this._model,{fuseConvPool:!1}),i=this._model.inputChannels/3,n=[{key:`pack/${this.precision}/${i}`,code:or(this.precision,i)},...e.nodes.map(o=>this._nodePipelineSpec(o,o.id===e.spec.output))];await Promise.all(n.map(({key:o,code:a})=>this._pipelineAsync(o,a)))}_readProfile(e,i,n,o){this._lastExecutionProfile=(async()=>{try{await n.mapAsync(Vi.READ);const a=new Uint32Array(n.getMappedRange()),s=l=>a[l*2+1]*4294967296+a[l*2],u=o.map((l,p)=>({id:l,durationMs:(s(p*2+1)-s(p*2))/1e6}));return{totalMs:u.reduce((l,p)=>l+p.durationMs,0),layers:u}}finally{n.mapState==="mapped"&&n.unmap(),this._releaseQuerySet(e),this._releaseBuffer(i),this._releaseBuffer(n),this._profileOperations--}})()}async executeCPU(e,i,n){const o=i*n*this._model.inputChannels;if(e.length!==o)throw new Error(`Native OIDN CPU input has ${e.length} values, expected ${o}`);const a=this._execution(i,n),s=this._model.inputChannels/3,u=i*n;a.cpuInputBuffers||(a.cpuInputBuffers=Array.from({length:s},(h,g)=>{const f=this._device.createBuffer({label:`oidn/cpu-input/${i}x${n}/${g}`,size:u*16,usage:ot.STORAGE|ot.COPY_DST});return this._resources.track("gpu-buffer",f),a.ownedBuffers.push(f),f}),a.cpuReadbackBuffer=this._device.createBuffer({label:`oidn/cpu-readback/${i}x${n}`,size:u*16,usage:ot.COPY_DST|ot.MAP_READ}),this._resources.track("gpu-buffer",a.cpuReadbackBuffer),a.ownedBuffers.push(a.cpuReadbackBuffer));for(let h=0;h<s;h++){const g=new Float32Array(u*4);for(let f=0;f<u;f++){const _=f*this._model.inputChannels+h*3,m=f*4;g[m]=e[_],g[m+1]=e[_+1],g[m+2]=e[_+2]}this._device.queue.writeBuffer(a.cpuInputBuffers[h],0,new Uint8Array(g.buffer,g.byteOffset,g.byteLength))}const l=super.execute(a.cpuInputBuffers,i,n),p=this._device.createCommandEncoder({label:`oidn/cpu-readback/${i}x${n}`});p.copyBufferToBuffer(l,0,a.cpuReadbackBuffer,0,u*16),this._device.queue.submit([p.finish()]),await a.cpuReadbackBuffer.mapAsync(Vi.READ);const c=new Float32Array(a.cpuReadbackBuffer.getMappedRange()),d=new Float32Array(u*3);for(let h=0;h<u;h++)d[h*3]=c[h*4],d[h*3+1]=c[h*4+1],d[h*3+2]=c[h*4+2];return a.cpuReadbackBuffer.unmap(),d}}const no="This is not an object",oo="This is not a Float16Array object",hr="This constructor is not a subclass of Float16Array",fr="The constructor property value is not an object",ao="Species constructor didn't return TypedArray object",uo="Derived constructor created TypedArray object which was too small length",se="Attempting to access detached ArrayBuffer",ii="Cannot convert undefined or null to object",ri="Cannot mix BigInt and other types, use explicit conversions",gr="@@iterator property is not callable",mr="Reduce of empty array with no initial value",so="The comparison function must be either a function or undefined",ni="Offset is out of bounds";function F(r){return(t,...e)=>st(r,t,e)}function Ht(r,t){return F(Vt(r,t).get)}const{apply:st,construct:ce,defineProperty:$r,get:oi,getOwnPropertyDescriptor:Vt,getPrototypeOf:le,has:ai,ownKeys:yr,set:_r,setPrototypeOf:wr}=Reflect,co=Proxy,{EPSILON:lo,MAX_SAFE_INTEGER:vr,isFinite:xr,isNaN:jt}=Number,{iterator:_t,species:po,toStringTag:ui,for:ho}=Symbol,Kt=Object,{create:Te,defineProperty:pe,freeze:fo,is:br}=Kt,si=Kt.prototype,go=si.__lookupGetter__?F(si.__lookupGetter__):(r,t)=>{if(r==null)throw G(ii);let e=Kt(r);do{const i=Vt(e,t);if(i!==void 0)return kt(i,"get")?i.get:void 0}while((e=le(e))!==null)},kt=Kt.hasOwn||F(si.hasOwnProperty),kr=Array,Tr=kr.isArray,Be=kr.prototype,mo=F(Be.join),$o=F(Be.push),yo=F(Be.toLocaleString),ci=Be[_t],_o=F(ci),{abs:wo,trunc:Br}=Math,Ie=ArrayBuffer,vo=Ie.isView,Ir=Ie.prototype,xo=F(Ir.slice),bo=Ht(Ir,"byteLength"),li=typeof SharedArrayBuffer<"u"?SharedArrayBuffer:null,ko=li&&Ht(li.prototype,"byteLength"),pi=le(Uint8Array),To=pi.from,et=pi.prototype,Bo=et[_t],Io=F(et.keys),Po=F(et.values),Co=F(et.entries),So=F(et.set),Pr=F(et.reverse),Mo=F(et.fill),Eo=F(et.copyWithin),Cr=F(et.sort),de=F(et.slice),Ao=F(et.subarray),it=Ht(et,"buffer"),Nt=Ht(et,"byteOffset"),E=Ht(et,"length"),Sr=Ht(et,ui),Oo=Uint8Array,lt=Uint16Array,Mr=(...r)=>st(To,lt,r),di=Uint32Array,No=Float32Array,zt=le([][_t]()),Pe=F(zt.next),zo=F(function*(){}().next),Uo=le(zt),G=TypeError,hi=RangeError,Er=WeakSet,Ar=Er.prototype,Wo=F(Ar.add),Lo=F(Ar.has),Ce=WeakMap,fi=Ce.prototype,Se=F(fi.get),Do=F(fi.has),gi=F(fi.set),Or=new Ce,Ro=Te(null,{next:{value:function(){const t=Se(Or,this);return Pe(t)}},[_t]:{value:function(){return this}}});function Me(r){if(r[_t]===ci&&zt.next===Pe)return r;const t=Te(Ro);return gi(Or,t,_o(r)),t}const Nr=new Ce,zr=Te(Uo,{next:{value:function(){const t=Se(Nr,this);return zo(t)},writable:!0,configurable:!0}});for(const r of yr(zt))r!=="next"&&pe(zr,r,Vt(zt,r));function Ur(r){const t=Te(zr);return gi(Nr,t,r),t}function Ee(r){return r!==null&&typeof r=="object"||typeof r=="function"}function Wr(r){return r!==null&&typeof r=="object"}function Ae(r){return Sr(r)!==void 0}function mi(r){const t=Sr(r);return t==="BigInt64Array"||t==="BigUint64Array"}function Xo(r){try{return Tr(r)?!1:(bo(r),!0)}catch{return!1}}function Lr(r){if(li===null)return!1;try{return ko(r),!0}catch{return!1}}function Yo(r){return Xo(r)||Lr(r)}function Dr(r){return Tr(r)?r[_t]===ci&&zt.next===Pe:!1}function Fo(r){return Ae(r)?r[_t]===Bo&&zt.next===Pe:!1}function Oe(r){if(typeof r!="string")return!1;const t=+r;return r!==t+""||!xr(t)?!1:t===Br(t)}const Ne=ho("__Float16Array__");function Go(r){if(!Wr(r))return!1;const t=le(r);if(!Wr(t))return!1;const e=t.constructor;if(e===void 0)return!1;if(!Ee(e))throw G(fr);return ai(e,Ne)}const $i=1/lo;function Ho(r){return r+$i-$i}const Rr=6103515625e-14,Vo=65504,Xr=.0009765625,Yr=Xr*Rr,jo=Xr*$i;function Ko(r){const t=+r;if(!xr(t)||t===0)return t;const e=t>0?1:-1,i=wo(t);if(i<Rr)return e*Ho(i/Yr)*Yr;const n=(1+jo)*i,o=n-(n-i);return o>Vo||jt(o)?e*(1/0):e*o}const Fr=new Ie(4),Gr=new No(Fr),Hr=new di(Fr),gt=new lt(512),mt=new Oo(512);for(let r=0;r<256;++r){const t=r-127;t<-24?(gt[r]=0,gt[r|256]=32768,mt[r]=24,mt[r|256]=24):t<-14?(gt[r]=1024>>-t-14,gt[r|256]=1024>>-t-14|32768,mt[r]=-t-1,mt[r|256]=-t-1):t<=15?(gt[r]=t+15<<10,gt[r|256]=t+15<<10|32768,mt[r]=13,mt[r|256]=13):t<128?(gt[r]=31744,gt[r|256]=64512,mt[r]=24,mt[r|256]=24):(gt[r]=31744,gt[r|256]=64512,mt[r]=13,mt[r|256]=13)}function wt(r){Gr[0]=Ko(r);const t=Hr[0],e=t>>23&511;return gt[e]+((t&8388607)>>mt[e])}const yi=new di(2048);for(let r=1;r<1024;++r){let t=r<<13,e=0;for(;!(t&8388608);)t<<=1,e-=8388608;t&=-8388609,e+=947912704,yi[r]=t|e}for(let r=1024;r<2048;++r)yi[r]=939524096+(r-1024<<13);const qt=new di(64);for(let r=1;r<31;++r)qt[r]=r<<23;qt[31]=1199570944,qt[32]=2147483648;for(let r=33;r<63;++r)qt[r]=2147483648+(r-32<<23);qt[63]=3347054592;const Vr=new lt(64);for(let r=1;r<64;++r)r!==32&&(Vr[r]=1024);function W(r){const t=r>>10;return Hr[0]=yi[Vr[t]+(r&1023)]+qt[t],Gr[0]}function Tt(r){const t=+r;return jt(t)||t===0?0:Br(t)}function _i(r){const t=Tt(r);return t<0?0:t<vr?t:vr}function ze(r,t){if(!Ee(r))throw G(no);const e=r.constructor;if(e===void 0)return t;if(!Ee(e))throw G(fr);const i=e[po];return i??t}function he(r){if(Lr(r))return!1;try{return xo(r,0,0),!1}catch{}return!0}function jr(r,t){const e=jt(r),i=jt(t);if(e&&i)return 0;if(e)return 1;if(i||r<t)return-1;if(r>t)return 1;if(r===0&&t===0){const n=br(r,0),o=br(t,0);if(!n&&o)return-1;if(n&&!o)return 1}return 0}const wi=2,Ue=new Ce;function Zt(r){return Do(Ue,r)||!vo(r)&&Go(r)}function A(r){if(!Zt(r))throw G(oo)}function We(r,t){const e=Zt(r),i=Ae(r);if(!e&&!i)throw G(ao);if(typeof t=="number"){let n;if(e){const o=C(r);n=E(o)}else n=E(r);if(n<t)throw G(uo)}if(mi(r))throw G(ri)}function C(r){const t=Se(Ue,r);if(t!==void 0){const n=it(t);if(he(n))throw G(se);return t}const e=r.buffer;if(he(e))throw G(se);const i=ce(R,[e,r.byteOffset,r.length],r.constructor);return Se(Ue,i)}function Kr(r){const t=E(r),e=[];for(let i=0;i<t;++i)e[i]=W(r[i]);return e}const qr=new Er;for(const r of yr(et)){if(r===ui)continue;const t=Vt(et,r);kt(t,"get")&&typeof t.get=="function"&&Wo(qr,t.get)}const qo=fo({get(r,t,e){return Oe(t)&&kt(r,t)?W(oi(r,t)):Lo(qr,go(r,t))?oi(r,t):oi(r,t,e)},set(r,t,e,i){return Oe(t)&&kt(r,t)?_r(r,t,wt(e)):_r(r,t,e,i)},getOwnPropertyDescriptor(r,t){if(Oe(t)&&kt(r,t)){const e=Vt(r,t);return e.value=W(e.value),e}return Vt(r,t)},defineProperty(r,t,e){return Oe(t)&&kt(r,t)&&kt(e,"value")&&(e.value=wt(e.value)),$r(r,t,e)}});class R{constructor(t,e,i){let n;if(Zt(t))n=ce(lt,[C(t)],new.target);else if(Ee(t)&&!Yo(t)){let a,s;if(Ae(t)){a=t,s=E(t);const u=it(t);if(he(u))throw G(se);if(mi(t))throw G(ri);const l=new Ie(s*wi);n=ce(lt,[l],new.target)}else{const u=t[_t];if(u!=null&&typeof u!="function")throw G(gr);u!=null?Dr(t)?(a=t,s=t.length):(a=[...t],s=a.length):(a=t,s=_i(a.length)),n=ce(lt,[s],new.target)}for(let u=0;u<s;++u)n[u]=wt(a[u])}else n=ce(lt,arguments,new.target);const o=new co(n,qo);return gi(Ue,o,n),o}static from(t,...e){const i=this;if(!ai(i,Ne))throw G(hr);if(i===R){if(Zt(t)&&e.length===0){const p=C(t),c=new lt(it(p),Nt(p),E(p));return new R(it(de(c)))}if(e.length===0)return new R(it(Mr(t,wt)));const u=e[0],l=e[1];return new R(it(Mr(t,function(p,...c){return wt(st(u,this,[p,...Me(c)]))},l)))}let n,o;const a=t[_t];if(a!=null&&typeof a!="function")throw G(gr);if(a!=null)Dr(t)?(n=t,o=t.length):Fo(t)?(n=t,o=E(t)):(n=[...t],o=n.length);else{if(t==null)throw G(ii);n=Kt(t),o=_i(n.length)}const s=new i(o);if(e.length===0)for(let u=0;u<o;++u)s[u]=n[u];else{const u=e[0],l=e[1];for(let p=0;p<o;++p)s[p]=st(u,l,[n[p],p])}return s}static of(...t){const e=this;if(!ai(e,Ne))throw G(hr);const i=t.length;if(e===R){const o=new R(i),a=C(o);for(let s=0;s<i;++s)a[s]=wt(t[s]);return o}const n=new e(i);for(let o=0;o<i;++o)n[o]=t[o];return n}keys(){A(this);const t=C(this);return Io(t)}values(){A(this);const t=C(this);return Ur(function*(){for(const e of Po(t))yield W(e)}())}entries(){A(this);const t=C(this);return Ur(function*(){for(const[e,i]of Co(t))yield[e,W(i)]}())}at(t){A(this);const e=C(this),i=E(e),n=Tt(t),o=n>=0?n:i+n;if(!(o<0||o>=i))return W(e[o])}with(t,e){A(this);const i=C(this),n=E(i),o=Tt(t),a=o>=0?o:n+o,s=+e;if(a<0||a>=n)throw hi(ni);const u=new lt(it(i),Nt(i),E(i)),l=new R(it(de(u))),p=C(l);return p[a]=wt(s),l}map(t,...e){A(this);const i=C(this),n=E(i),o=e[0],a=ze(i,R);if(a===R){const u=new R(n),l=C(u);for(let p=0;p<n;++p){const c=W(i[p]);l[p]=wt(st(t,o,[c,p,this]))}return u}const s=new a(n);We(s,n);for(let u=0;u<n;++u){const l=W(i[u]);s[u]=st(t,o,[l,u,this])}return s}filter(t,...e){A(this);const i=C(this),n=E(i),o=e[0],a=[];for(let l=0;l<n;++l){const p=W(i[l]);st(t,o,[p,l,this])&&$o(a,p)}const s=ze(i,R),u=new s(a);return We(u),u}reduce(t,...e){A(this);const i=C(this),n=E(i);if(n===0&&e.length===0)throw G(mr);let o,a;e.length===0?(o=W(i[0]),a=1):(o=e[0],a=0);for(let s=a;s<n;++s)o=t(o,W(i[s]),s,this);return o}reduceRight(t,...e){A(this);const i=C(this),n=E(i);if(n===0&&e.length===0)throw G(mr);let o,a;e.length===0?(o=W(i[n-1]),a=n-2):(o=e[0],a=n-1);for(let s=a;s>=0;--s)o=t(o,W(i[s]),s,this);return o}forEach(t,...e){A(this);const i=C(this),n=E(i),o=e[0];for(let a=0;a<n;++a)st(t,o,[W(i[a]),a,this])}find(t,...e){A(this);const i=C(this),n=E(i),o=e[0];for(let a=0;a<n;++a){const s=W(i[a]);if(st(t,o,[s,a,this]))return s}}findIndex(t,...e){A(this);const i=C(this),n=E(i),o=e[0];for(let a=0;a<n;++a){const s=W(i[a]);if(st(t,o,[s,a,this]))return a}return-1}findLast(t,...e){A(this);const i=C(this),n=E(i),o=e[0];for(let a=n-1;a>=0;--a){const s=W(i[a]);if(st(t,o,[s,a,this]))return s}}findLastIndex(t,...e){A(this);const i=C(this),n=E(i),o=e[0];for(let a=n-1;a>=0;--a){const s=W(i[a]);if(st(t,o,[s,a,this]))return a}return-1}every(t,...e){A(this);const i=C(this),n=E(i),o=e[0];for(let a=0;a<n;++a)if(!st(t,o,[W(i[a]),a,this]))return!1;return!0}some(t,...e){A(this);const i=C(this),n=E(i),o=e[0];for(let a=0;a<n;++a)if(st(t,o,[W(i[a]),a,this]))return!0;return!1}set(t,...e){A(this);const i=C(this),n=Tt(e[0]);if(n<0)throw hi(ni);if(t==null)throw G(ii);if(mi(t))throw G(ri);if(Zt(t))return So(C(this),C(t),n);if(Ae(t)){const u=it(t);if(he(u))throw G(se)}const o=E(i),a=Kt(t),s=_i(a.length);if(n===1/0||s+n>o)throw hi(ni);for(let u=0;u<s;++u)i[u+n]=wt(a[u])}reverse(){A(this);const t=C(this);return Pr(t),this}toReversed(){A(this);const t=C(this),e=new lt(it(t),Nt(t),E(t)),i=new R(it(de(e))),n=C(i);return Pr(n),i}fill(t,...e){A(this);const i=C(this);return Mo(i,wt(t),...Me(e)),this}copyWithin(t,e,...i){A(this);const n=C(this);return Eo(n,t,e,...Me(i)),this}sort(t){A(this);const e=C(this),i=t!==void 0?t:jr;return Cr(e,(n,o)=>i(W(n),W(o))),this}toSorted(t){A(this);const e=C(this);if(t!==void 0&&typeof t!="function")throw new G(so);const i=t!==void 0?t:jr,n=new lt(it(e),Nt(e),E(e)),o=new R(it(de(n))),a=C(o);return Cr(a,(s,u)=>i(W(s),W(u))),o}slice(t,e){A(this);const i=C(this),n=ze(i,R);if(n===R){const g=new lt(it(i),Nt(i),E(i));return new R(it(de(g,t,e)))}const o=E(i),a=Tt(t),s=e===void 0?o:Tt(e);let u;a===-1/0?u=0:a<0?u=o+a>0?o+a:0:u=o<a?o:a;let l;s===-1/0?l=0:s<0?l=o+s>0?o+s:0:l=o<s?o:s;const p=l-u>0?l-u:0,c=new n(p);if(We(c,p),p===0)return c;const d=it(i);if(he(d))throw G(se);let h=0;for(;u<l;)c[h]=W(i[u]),++u,++h;return c}subarray(t,e){A(this);const i=C(this),n=ze(i,R),o=new lt(it(i),Nt(i),E(i)),a=Ao(o,t,e),s=new n(it(a),Nt(a),E(a));return We(s),s}indexOf(t,...e){A(this);const i=C(this),n=E(i);let o=Tt(e[0]);if(o===1/0)return-1;o<0&&(o+=n,o<0&&(o=0));for(let a=o;a<n;++a)if(kt(i,a)&&W(i[a])===t)return a;return-1}lastIndexOf(t,...e){A(this);const i=C(this),n=E(i);let o=e.length>=1?Tt(e[0]):n-1;if(o===-1/0)return-1;o>=0?o=o<n-1?o:n-1:o+=n;for(let a=o;a>=0;--a)if(kt(i,a)&&W(i[a])===t)return a;return-1}includes(t,...e){A(this);const i=C(this),n=E(i);let o=Tt(e[0]);if(o===1/0)return!1;o<0&&(o+=n,o<0&&(o=0));const a=jt(t);for(let s=o;s<n;++s){const u=W(i[s]);if(a&&jt(u)||u===t)return!0}return!1}join(t){A(this);const e=C(this),i=Kr(e);return mo(i,t)}toLocaleString(...t){A(this);const e=C(this),i=Kr(e);return yo(i,...Me(t))}get[ui](){if(Zt(this))return"Float16Array"}}pe(R,"BYTES_PER_ELEMENT",{value:wi}),pe(R,Ne,{}),wr(R,pi);const Le=R.prototype;pe(Le,"BYTES_PER_ELEMENT",{value:wi}),pe(Le,_t,{value:Le.values,writable:!0,configurable:!0}),wr(Le,et);const Zr=new WeakMap;function Zo(r,t){let e=Zr.get(r);e||(e=new Map,Zr.set(r,e));let i=e.get(t);if(!i){const n=r.createShaderModule({label:`oidn/webnn/input-pack/${t}`,code:ta(t)}),o=r.createShaderModule({label:"oidn/webnn/output-unpack",code:ea()});i={input:r.createComputePipeline({label:`oidn/webnn/input-pack/${t}`,layout:"auto",compute:{module:n,entryPoint:"main"}}),output:r.createComputePipeline({label:"oidn/webnn/output-unpack",layout:"auto",compute:{module:o,entryPoint:"main"}})},e.set(t,i)}return i}const Ct=8;function Qr(r,t){return Math.ceil(r/t)*t}function Qo(r,t,e,i){const n=r.createBuffer({label:t,size:Qr(e.byteLength,4),usage:i,mappedAtCreation:!0});return new Uint8Array(n.getMappedRange()).set(new Uint8Array(e.buffer,e.byteOffset,e.byteLength)),n.unmap(),n}function Jr(r,t,e){const i=new Uint32Array(Qr(e.length,4));return i.set(e),Qo(r,t,i,GPUBufferUsage.UNIFORM)}function vi(r,t,e,i){var n,o,a,s;return!!((s=(a=(o=(n=r==null?void 0:r[t])==null?void 0:n[e])==null?void 0:o.dataTypes)==null?void 0:a.includes)!=null&&s.call(a,i))}function Jo(r,t){if(t==="fp16"&&r.desc.dataType==="Float16")return new Uint8Array(r.data.buffer,r.data.byteOffset,r.data.byteLength);if(t==="fp32"&&r.desc.dataType==="Float32")return new Uint8Array(r.data.buffer,r.data.byteOffset,r.data.byteLength);const e=r.desc.dataType==="Float32"?new Float32Array(r.data.buffer,r.data.byteOffset,r.data.byteLength/4):new R(r.data.buffer,r.data.byteOffset,r.data.byteLength/2),i=t==="fp16"?new R(e):new Float32Array(e);return new Uint8Array(i.buffer,i.byteOffset,i.byteLength)}function ta(r){const t=Array.from({length:r},(i,n)=>`@group(0) @binding(${n}) var<storage, read> input${n}: array<vec4<f32>>;`).join(`
|
|
1472
|
+
`),e=Array.from({length:r},(i,n)=>{const o=n*3;return`if (channel < ${o+3}u) {
|
|
1473
|
+
return input${n}[pixel][channel - ${o}u];
|
|
1474
|
+
}`}).join(`
|
|
1475
|
+
`);return`enable f16;
|
|
1476
|
+
struct Params { width: u32, height: u32, channels: u32, padding: u32 }
|
|
1477
|
+
${t}
|
|
1478
|
+
@group(0) @binding(${r}) var<storage, read_write> outputData: array<f16>;
|
|
1479
|
+
@group(0) @binding(${r+1}) var<uniform> params: Params;
|
|
1480
|
+
|
|
1481
|
+
fn readChannel(pixel: u32, channel: u32) -> f32 {
|
|
1482
|
+
${e}
|
|
1483
|
+
return 0.0;
|
|
1484
|
+
}
|
|
1485
|
+
|
|
1486
|
+
@compute @workgroup_size(${Ct}, ${Ct}, 1)
|
|
1487
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
1488
|
+
if (gid.x >= params.width || gid.y >= params.height || gid.z >= params.channels) {
|
|
1489
|
+
return;
|
|
1490
|
+
}
|
|
1491
|
+
let pixel = gid.y * params.width + gid.x;
|
|
1492
|
+
let outputIndex = (gid.z * params.height + gid.y) * params.width + gid.x;
|
|
1493
|
+
outputData[outputIndex] = f16(readChannel(pixel, gid.z));
|
|
1494
|
+
}
|
|
1495
|
+
`}function ea(){return`enable f16;
|
|
1496
|
+
struct Params { width: u32, height: u32, padding0: u32, padding1: u32 }
|
|
1497
|
+
@group(0) @binding(0) var<storage, read> inputData: array<f16>;
|
|
1498
|
+
@group(0) @binding(1) var<storage, read_write> outputData: array<vec4<f32>>;
|
|
1499
|
+
@group(0) @binding(2) var<uniform> params: Params;
|
|
1500
|
+
|
|
1501
|
+
@compute @workgroup_size(${Ct}, ${Ct}, 1)
|
|
1502
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
1503
|
+
if (gid.x >= params.width || gid.y >= params.height) { return; }
|
|
1504
|
+
let pixel = gid.y * params.width + gid.x;
|
|
1505
|
+
let plane = params.width * params.height;
|
|
1506
|
+
outputData[pixel] = vec4<f32>(
|
|
1507
|
+
f32(inputData[pixel]),
|
|
1508
|
+
f32(inputData[plane + pixel]),
|
|
1509
|
+
f32(inputData[plane * 2u + pixel]),
|
|
1510
|
+
0.0
|
|
1511
|
+
);
|
|
1512
|
+
}
|
|
1513
|
+
`}function ia(r,t){if(t==="fp32")throw new Error("OIDN WebNN GPU interop currently requires FP16 exportable tensors");if(!r.features.has("shader-f16"))throw new Error("OIDN WebNN requires shader-f16 on the shared GPUDevice");return"fp16"}function ra(r,t,e){return e==="relu"?r.relu(t):t}class tn{constructor(t,e,i={}){y(this,"precision");y(this,"support");y(this,"_context");y(this,"_builderConstructor");y(this,"_shapeCache",new Map);y(this,"_shapePromises",new Map);y(this,"_retiredExecutions",new Set);y(this,"_pendingCreationCount",0);y(this,"_shapeCacheSize");y(this,"_clock",0);y(this,"_inputPipeline");y(this,"_outputPipeline");y(this,"_resources",new qi);y(this,"_disposed",!1);this._device=t,this._model=e,this.precision=ia(t,i.precision??"auto"),this._shapeCacheSize=Math.max(1,i.shapeCacheSize??2),this.support={available:!1,fp16Conv:!1,gpuInterop:!1};const n=Zo(t,e.inputChannels/3);this._inputPipeline=n.input,this._outputPipeline=n.output}async prepare(){var i,n,o;if(this._disposed)throw new Error("OIDN WebNN executor is disposed");const t=(i=globalThis.navigator)==null?void 0:i.ml,e=globalThis.MLGraphBuilder;if(!(t!=null&&t.createContext)||typeof e!="function")throw this.support.reason="WebNN is not exposed by this browser",new Error(this.support.reason);this._builderConstructor=e;try{try{this._context=await t.createContext({deviceType:"gpu",powerPreference:"high-performance"})}catch{this._context=await t.createContext({deviceType:"gpu"})}if(this._resources.track("ml-context",this._context),this._disposed)throw new Error("OIDN WebNN executor is disposed");if(typeof this._context.createExportableTensor!="function"||typeof this._context.exportToGPU!="function")throw this.support.reason="WebNN WebGPU tensor interop is unavailable",new Error(this.support.reason);const a=((o=(n=this._context).opSupportLimits)==null?void 0:o.call(n))??{};if(this.support.fp16Conv=vi(a,"conv2d","input","float16")&&vi(a,"conv2d","filter","float16")&&vi(a,"conv2d","output","float16"),!this.support.fp16Conv)throw this.support.reason="WebNN does not support FP16 conv2d",new Error(this.support.reason);let s,u;try{s=this._resources.track("ml-tensor",await this._context.createExportableTensor({dataType:"float16",shape:[4]},this._device)),u=this._resources.track("gpu-buffer",await this._context.exportToGPU(s)),this.support.gpuInterop=!0}catch(l){throw this.support.reason=`WebNN FP16 WebGPU interop failed: ${String(l)}`,new Error(this.support.reason)}finally{this._releaseBuffer(u),this._releaseTensor(s)}this.support.available=!0}catch(a){throw this._releaseContext(),a}}_constant(t,e){return t.constant({dataType:"float16",shape:[...e.desc.dims]},Jo(e,this.precision))}async _createExecution(t,e){if(this._disposed)throw new Error("OIDN WebNN executor is disposed");const i=new this._builderConstructor(this._context),n=new Map,o=new Map;n.set(this._model.spec.input,i.input("input",{dataType:"float16",shape:[1,this._model.inputChannels,e,t]})),o.set(this._model.spec.input,[this._model.inputChannels,e,t]);for(const d of this._model.spec.nodes){let h,g;if(d.op==="conv2d"){const f=o.get(d.input),_=this._model.convTensors.get(d.id),m=i.conv2d(n.get(d.input),this._constant(i,_.weight),{bias:this._constant(i,_.bias),padding:[1,1,1,1],inputLayout:"nchw",filterLayout:"oihw"});h=ra(i,m,d.activation),g=[_.outputChannels,f[1],f[2]]}else if(d.op==="maxPool2d"){const f=o.get(d.input);h=i.maxPool2d(n.get(d.input),{windowDimensions:[2,2],strides:[2,2],padding:[0,f[1]%2,0,f[2]%2],layout:"nchw"}),g=[f[0],Math.ceil(f[1]/2),Math.ceil(f[2]/2)]}else if(d.op==="upsample2d"){const f=o.get(d.input);h=i.resample2d(n.get(d.input),{mode:"nearest-neighbor",axes:[2,3],scales:[2,2]}),g=[f[0],f[1]*2,f[2]*2]}else{const f=d.inputs.map(_=>o.get(_));if(f.some(_=>_[1]!==f[0][1]||_[2]!==f[0][2]))throw new Error(`WebNN concat ${d.id} has mismatched spatial shapes`);h=i.concat(d.inputs.map(_=>n.get(_)),1),g=[f.reduce((_,m)=>_+m[0],0),f[0][1],f[0][2]]}n.set(d.id,h),o.set(d.id,g)}let a,s,u,l,p,c;try{return a=this._resources.track("ml-graph",await i.build({output:n.get(this._model.spec.output)})),s=this._resources.track("ml-tensor",await this._context.createExportableTensor({dataType:"float16",shape:[1,this._model.inputChannels,e,t],writable:!0},this._device)),u=this._resources.track("ml-tensor",await this._context.createExportableTensor({dataType:"float16",shape:[1,this._model.outputChannels,e,t],readable:!0},this._device)),l=this._resources.track("gpu-buffer",this._device.createBuffer({label:`oidn/webnn/output/${t}x${e}`,size:t*e*4*4,usage:GPUBufferUsage.STORAGE|GPUBufferUsage.COPY_SRC|GPUBufferUsage.COPY_DST})),p=this._resources.track("gpu-buffer",Jr(this._device,`oidn/webnn/input/${t}x${e}`,[t,e,this._model.inputChannels])),c=this._resources.track("gpu-buffer",Jr(this._device,`oidn/webnn/output/${t}x${e}`,[t,e])),{graph:a,inputTensor:s,outputTensor:u,outputBuffer:l,inputUniform:p,outputUniform:c,width:t,height:e,lastUsed:++this._clock}}catch(d){throw this._releaseBuffer(c),this._releaseBuffer(p),this._releaseBuffer(l),this._releaseTensor(u),this._releaseTensor(s),this._releaseGraph(a),d}}async _execution(t,e){const i=`${t}x${e}`;let n=this._shapeCache.get(i);if(!n){let o=this._shapePromises.get(i);o||(o=(async()=>{this._pendingCreationCount++;try{return await this._createExecution(t,e)}finally{this._pendingCreationCount--}})(),this._shapePromises.set(i,o));try{if(n=await o,this._disposed)throw this._destroyExecution(n),new Error("OIDN WebNN executor is disposed");this._shapeCache.set(i,n)}finally{this._shapePromises.get(i)===o&&this._shapePromises.delete(i)}if(this._shapeCache.size>this._shapeCacheSize){const a=[...this._shapeCache.entries()].filter(([s])=>s!==i).sort((s,u)=>s[1].lastUsed-u[1].lastUsed)[0];a&&(this._shapeCache.delete(a[0]),this._retireExecution(a[1]))}}return n.lastUsed=++this._clock,n}async prewarm(t){for(const e of t)await this._execution(e.width,e.height)}async execute(t,e,i){const n=this._model.inputChannels/3;if(t.length!==n)throw new Error(`OIDN WebNN expected ${n} input buffers, got ${t.length}`);const o=await this._execution(e,i),a=this._resources.track("gpu-buffer",await this._context.exportToGPU(o.inputTensor));try{const u=t.map((d,h)=>({binding:h,resource:{buffer:d}}));u.push({binding:n,resource:{buffer:a}}),u.push({binding:n+1,resource:{buffer:o.inputUniform}});const l=this._device.createBindGroup({label:"oidn/webnn/input-bindings",layout:this._inputPipeline.getBindGroupLayout(0),entries:u}),p=this._device.createCommandEncoder({label:"oidn/webnn/input-pack"}),c=p.beginComputePass();c.setPipeline(this._inputPipeline),c.setBindGroup(0,l),c.dispatchWorkgroups(Math.ceil(e/Ct),Math.ceil(i/Ct),this._model.inputChannels),c.end(),this._device.queue.submit([p.finish()])}finally{this._releaseBuffer(a)}this._context.dispatch(o.graph,{input:o.inputTensor},{output:o.outputTensor});const s=this._resources.track("gpu-buffer",await this._context.exportToGPU(o.outputTensor));try{const u=this._device.createBindGroup({label:"oidn/webnn/output-bindings",layout:this._outputPipeline.getBindGroupLayout(0),entries:[{binding:0,resource:{buffer:s}},{binding:1,resource:{buffer:o.outputBuffer}},{binding:2,resource:{buffer:o.outputUniform}}]}),l=this._device.createCommandEncoder({label:"oidn/webnn/output-unpack"}),p=l.beginComputePass();p.setPipeline(this._outputPipeline),p.setBindGroup(0,u),p.dispatchWorkgroups(Math.ceil(e/Ct),Math.ceil(i/Ct)),p.end(),this._device.queue.submit([l.finish()])}finally{this._releaseBuffer(s)}return o.outputBuffer}async executeCPU(t,e,i){const n=await this._execution(e,i),o=e*i,a=new R(o*this._model.inputChannels);for(let l=0;l<o;l++)for(let p=0;p<this._model.inputChannels;p++)a[p*o+l]=t[l*this._model.inputChannels+p];this._context.writeTensor(n.inputTensor,a),this._context.dispatch(n.graph,{input:n.inputTensor},{output:n.outputTensor});const s=new R(await this._context.readTensor(n.outputTensor)),u=new Float32Array(o*this._model.outputChannels);for(let l=0;l<o;l++)for(let p=0;p<this._model.outputChannels;p++)u[l*this._model.outputChannels+p]=s[p*o+l];return u}_destroyExecution(t){this._releaseGraph(t.graph),this._releaseTensor(t.inputTensor),this._releaseTensor(t.outputTensor),this._releaseBuffer(t.outputBuffer),this._releaseBuffer(t.inputUniform),this._releaseBuffer(t.outputUniform)}_retireExecution(t){this._retiredExecutions.add(t),this._device.queue.onSubmittedWorkDone().catch(()=>{}).then(()=>{this._retiredExecutions.delete(t),this._destroyExecution(t)})}_releaseBuffer(t){this._resources.release("gpu-buffer",t,()=>t.destroy())}_releaseTensor(t){this._resources.release("ml-tensor",t,()=>t.destroy())}_releaseGraph(t){this._resources.release("ml-graph",t,()=>{var e;return(e=t.destroy)==null?void 0:e.call(t)})}_releaseContext(){this._resources.release("ml-context",this._context,()=>{var t,e;return(e=(t=this._context).destroy)==null?void 0:e.call(t)})}getResourceInfo(){return this._resources.snapshot(this._pendingCreationCount+this._retiredExecutions.size)}dispose(){if(!this._disposed){this._disposed=!0;for(const t of this._shapeCache.values())this._destroyExecution(t);this._shapeCache.clear();for(const t of this._retiredExecutions)this._destroyExecution(t);this._retiredExecutions.clear(),this._shapePromises.clear(),this._releaseContext()}}}const na=100;function Ut(r,t){return Math.ceil(r/t)*t}function De(r){return r.data instanceof GPUBuffer||r.data instanceof GPUTexture}class en{constructor(t,e,i={}){y(this,"_device");y(this,"_aux");y(this,"_hdr");y(this,"_hdrTransfer");y(this,"_dataProcessGPU");y(this,"_nativeExecutor");y(this,"_webNNExecutor");y(this,"_modelSpec");y(this,"_inputChannels");y(this,"_engine");y(this,"_dynamicTileController");y(this,"_lastExecution");y(this,"_profileNextImage",!1);y(this,"_lastImageProfile");y(this,"_activeExecutionFailures",new Set);y(this,"_deviceLostObserved",!1);y(this,"_deviceLostSettled",!1);y(this,"_deviceLostReason");this._aux=i.aux||!1,this._hdr=i.hdr||!1,this._hdrTransfer=i.hdrTransfer??"pu",this._engine=i.engine??"auto";const n=i.modelSpec??qe(t),o=Fi(t,n);this._modelSpec=o.spec,this._inputChannels=o.inputChannels;const a=this._aux?9:3;if(o.inputChannels!==a)throw new Error(`OIDN model expects ${o.inputChannels} input channels, but aux=${this._aux} provides ${a}`);this._dynamicTileController=new vn(i.maxTileSize??512,i.dynamicTile),this._device=e,this._observeDeviceLoss(),this._engine==="webnn"?this._webNNExecutor=new tn(this._device,o,{precision:i.precision}):this._nativeExecutor=new dr(this._device,o,{precision:i.precision,kernel:i.kernel,gemm:i.gemm,subgroupMatrix:i.subgroupMatrix})}getDevice(){return this._device}async prepare(){if(this._webNNExecutor){await this._webNNExecutor.prepare();const t=Ut(this._modelSpec.receptiveField/2,J),e=[this._dynamicTileController.tileSize,this._dynamicTileController.minTileSize];await this._webNNExecutor.prewarm([...new Set(e)].map(i=>({width:i+2*t,height:i+2*t})));return}await this._nativeExecutor.prepare()}async prepareForImage(t,e,i={}){var u,l;const n=Ut(this._modelSpec.receptiveField/2,J),o=i.tileOverlap===void 0?n:Ut(Math.max(0,i.tileOverlap),J),a=Ke(t,e,i.wholeImage?Ut(Math.max(t,e),J):this._dynamicTileController.tileSize,o),s=[...new Map(a.tiles.map(({input:p})=>[`${p.width}x${p.height}`,{width:p.width,height:p.height}])).values()];await((u=this._webNNExecutor)==null?void 0:u.prewarm(s)),(l=this._nativeExecutor)==null||l.prewarm(s)}getRuntimeInfo(){var t;return{configuredEngine:this._engine,gpuEngine:this._webNNExecutor?"webnn":"wgsl",precision:(this._webNNExecutor??this._nativeExecutor).precision,kernel:this._nativeExecutor?{configured:this._nativeExecutor.kernelSetting,gemm:this._nativeExecutor.gemm,maxSpatialInputBlocks:this._nativeExecutor.maxSpatialInputBlocks,subgroupsAvailable:this._nativeExecutor.subgroupsAvailable,subgroupMatrix:this._nativeExecutor.subgroupMatrix}:void 0,webnn:(t=this._webNNExecutor)==null?void 0:t.support,resources:(this._webNNExecutor??this._nativeExecutor).getResourceInfo(),model:this._modelSpec.id,modelFamily:this._modelSpec.family,inputChannels:this._inputChannels,hdrTransfer:this._hdrTransfer,dynamicTile:{enabled:this._dynamicTileController.enabled,currentTileSize:this._dynamicTileController.tileSize,minTileSize:this._dynamicTileController.minTileSize,maxTileSize:this._dynamicTileController.maxTileSize,targetTileTimeMs:this._dynamicTileController.targetTileTimeMs},lastExecution:this._lastExecution,activeExecutionCount:this._activeExecutionFailures.size}}_observeDeviceLoss(){if(this._activeExecutionFailures??(this._activeExecutionFailures=new Set),this._deviceLostObserved)return;this._deviceLostObserved=!0;const t=this._device.lost;if(!t)return;const e=i=>{if(this._deviceLostSettled)return;this._deviceLostSettled=!0,this._deviceLostReason=i;const n=[...this._activeExecutionFailures];this._activeExecutionFailures.clear();for(const o of n)o(i)};t.then(i=>e(new Error(`WebGPU device lost: ${i.message}`)),e)}_registerExecutionFailure(t){return this._observeDeviceLoss(),this._activeExecutionFailures.add(t),this._deviceLostSettled&&(this._activeExecutionFailures.delete(t),queueMicrotask(()=>t(this._deviceLostReason))),()=>this._activeExecutionFailures.delete(t)}profileNextExecution(){return!this._nativeExecutor||!this._device.features.has("timestamp-query")?!1:(this._profileNextImage=!0,!0)}getLastExecutionProfile(){return this._lastImageProfile?Promise.resolve(this._lastImageProfile):void 0}_processImageData(t,e,i,n){const o=t.data,a=o.length/4,s=this._aux?9:3,u=new Float32Array(a*s);if(e&&!i||i&&!e)throw new Error("Normal map and albedo map are both required");if(e&&i&&(e.width!==i.width||e.height!==i.height||t.width!==e.width||t.height!==e.height))throw new Error("Image size mismatch");const l=e==null?void 0:e.data,p=i==null?void 0:i.data,c=l instanceof Float32Array?1:255,d=p instanceof Float32Array?1:255;for(let h=0;h<o.length;h+=4){const g=h/4*s;for(let f=0;f<3;f++)n?u[g+f]=o[h+f]:u[g+f]=o[h+f]/255,l&&(u[g+f+3]=l[h+f]/c),p&&(u[g+f+6]=p[h+f]/d)}return u}_readTile(t,e,i,n){const o=new Float32Array(i.width*i.height*e),a=t.length/(n*e);for(let s=0;s<i.height;s++)for(let u=0;u<i.width;u++){const l=Math.min(n-1,u+i.x),c=(Math.min(a-1,s+i.y)*n+l)*e,d=(s*i.width+u)*e;for(let h=0;h<e;h++)o[d+h]=t[c+h]}return o}_writeTile(t,e,i,n,o,a){const{data:s,width:u}=t,l=i.x-e.x,p=i.y-e.y;for(let c=0;c<i.height;c++)for(let d=0;d<i.width;d++){const h=((c+p)*o+d+l)*3,g=((c+i.y)*u+(d+i.x))*4;for(let f=0;f<3;f++)a?s[g+f]=n[h+f]:s[g+f]=Math.min(Math.max(n[h+f]*255,0),255);t.data[g+3]=a?1:255}}async _executeTile(t,e,i,n,o,a,s,u,l){const p=this._aux?9:3,c=new je(n.input.x,n.input.y,n.input.width,n.input.height),d=new je(n.output.x,n.output.y,n.output.width,n.output.height),h=c.width,g=c.height;let f,_,m=1;const v=this._device;let w=this._dataProcessGPU;if(t instanceof Float32Array){let $=this._readTile(t,p,c,a);u&&(m=dn({data:$,channels:p}),$=ln({data:$,channels:p,inputScale:m,transfer:this._hdrTransfer})),_=await(this._webNNExecutor??this._nativeExecutor).executeCPU($,h,g)}else{w||(w=this._dataProcessGPU=new gn(v,u,this._hdrTransfer)),w.setImageSize(a,s),w.setInputTile(c),o&&w.copyInputDataToOutput(t.color);const{color:$,albedo:b,normal:x}=w.forward(t.color,this._aux?t.albedo:void 0,this._aux?t.normal:void 0,l);f=await(this._webNNExecutor??this._nativeExecutor).execute(this._aux?[$,b,x]:[$],h,g)}let T;if(t instanceof Float32Array){u&&(_=pn({data:_,channels:3,inputScale:m,transfer:this._hdrTransfer})),this._writeTile(i,c,d,_,c.width,u);for(let $=0;$<d.height;$++)for(let b=0;b<d.width;b++){const x=($*d.width+b)*4,B=(($+d.y)*a+(b+d.x))*4;for(let k=0;k<4;k++)e.data[x+k]=i.data[B+k]}}else w.setOutputTile(d,c),T=w.inverse(f,t.color);return T}tileExecute({color:t,albedo:e,normal:i,done:n,progress:o,denoiseAlpha:a,tileOverlap:s,wholeImage:u,scheduling:l="event-loop",error:p}){if(this._aux&&(!e||!i))throw new Error("Normal map and albedo map are both required");if(!this._aux&&(e||i))throw new Error("Normal map and albedo map are not required");const c=t.width,d=t.height,h=this._dynamicTileController.tileSize,g=u?Ut(Math.max(c,d),J):h,f=Ut(this._modelSpec.receptiveField/2,J),_=s===void 0?f:Ut(Math.max(0,s),J),m=Ke(c,d,g,_),v=m.tiles.length>1,w=this._hdr||!1;let T;De(t)||(T=this._processImageData(t,e,i,w));function $(Y,q){return w?{data:new Float32Array(Y*q*4),width:Y,height:q}:new ImageData(Y,q)}const b=De(t)?void 0:$(c,d);let x="active",B,k,I=()=>!1;const P=()=>typeof performance>"u"?Date.now():performance.now(),O=P(),H=[],rt=this._profileNextImage===!0;this._profileNextImage=!1;const at=new Map,$t=()=>{B!==void 0&&(clearTimeout(B),B=void 0),k!==void 0&&typeof cancelAnimationFrame<"u"&&(cancelAnimationFrame(k),k=void 0)},Z=Y=>{console.error("OIDN error callback failed",Y)},nt=Y=>{if(x==="active")if(x="settled",$t(),I(),p)try{Promise.resolve(p(Y)).catch(Z)}catch(q){Z(q)}else console.error("OIDN execution failed",Y)},St=Y=>{if(l==="event-loop"||typeof requestAnimationFrame>"u")B=setTimeout(()=>{B=void 0,Y()},0);else{const q=()=>{$t(),Y()};k=requestAnimationFrame(q),B=setTimeout(q,na)}},Qt=async Y=>{if(x!=="active")return;const q=m.tiles[Y],Wt=De(t)?void 0:$(q.output.width,q.output.height),Jt=P();rt&&this._nativeExecutor.profileNextExecution();const te=await this._executeTile(De(t)?{color:t.data,albedo:e==null?void 0:e.data,normal:i==null?void 0:i.data}:T,Wt,b,q,Y===0,c,d,w,a);if(rt){const z=await this._nativeExecutor.getLastExecutionProfile();for(const K of(z==null?void 0:z.layers)??[])at.set(K.id,(at.get(K.id)??0)+K.durationMs)}if(x!=="active")return;const ee=b||{data:te,width:c,height:d};if(o&&await o(ee,Wt,new je(q.output.x,q.output.y,q.output.width,q.output.height),Y,m.tiles.length),x!=="active")return;const N=Y+1<m.tiles.length;if(await this._device.queue.onSubmittedWorkDone(),x!=="active")return;await(async()=>{if(H.push(P()-Jt),x==="active")if(N)St(()=>{x==="active"&&Qt(Y+1).catch(nt)});else{const z=[...H].sort((Mt,vt)=>Mt-vt),K=Math.floor(z.length/2),Lt=z.length%2?z[K]:(z[K-1]+z[K])/2;if(this._lastExecution={width:c,height:d,tileCount:m.tiles.length,tileColumns:m.columns,tileRows:m.rows,tileOverlap:m.overlap,inputPixelCount:m.inputPixelCount,inputShapeCount:m.inputShapeCount,durationMs:P()-O,tileTimeMs:{min:z[0],median:Lt,mean:z.reduce((Mt,vt)=>Mt+vt,0)/z.length,max:z[z.length-1]}},v&&this._dynamicTileController.observe(H),I(),rt){const Mt=[...at].map(([vt,fe])=>({id:vt,durationMs:fe}));this._lastImageProfile={totalMs:Mt.reduce((vt,fe)=>vt+fe.durationMs,0),layers:Mt}}await n(ee),x==="active"&&(x="settled")}})()};return I=this._registerExecutionFailure(nt),Qt(0).catch(nt),()=>{x==="active"&&(x="aborted",$t(),I())}}dispose(){var t,e,i;(t=this._dataProcessGPU)==null||t.dispose(),(e=this._nativeExecutor)==null||e.dispose(),(i=this._webNNExecutor)==null||i.dispose()}}async function oa(){if(!navigator.gpu)throw new Error("WebGPU is not available");const r={powerPreference:"high-performance"},t=await navigator.gpu.requestAdapter(r);if(!t)throw new Error("No WebGPU adapter is available");const e={},i=[];t.features.has("timestamp-query")&&i.push("timestamp-query"),t.features.has("bgra8unorm-storage")&&i.push("bgra8unorm-storage"),t.features.has("shader-f16")&&i.push("shader-f16");for(const o of["subgroups","chromium-experimental-subgroup-matrix"])t.features.has(o)&&i.push(o);e.requiredFeatures=i;const n=t.limits;return e.requiredLimits={maxComputeWorkgroupStorageSize:n.maxComputeWorkgroupStorageSize,maxComputeWorkgroupsPerDimension:n.maxComputeWorkgroupsPerDimension,maxStorageBufferBindingSize:n.maxStorageBufferBindingSize,maxBufferSize:n.maxBufferSize,maxComputeWorkgroupSizeX:n.maxComputeWorkgroupSizeX,maxComputeInvocationsPerWorkgroup:n.maxComputeInvocationsPerWorkgroup},t.requestDevice(e)}async function rn(r,t,e){const i=(t==null?void 0:t.device)??await oa(),n=bi(r),o=new en(n,i,e);return await o.prepare(),o}async function aa(r,t,e){return fetch(r).then(i=>i.arrayBuffer()).then(i=>rn(i,t,e))}X.NativeUNetExecutor=dr,X.OIDN_UNET_LARGE_SPEC=Li,X.OIDN_UNET_SMALL_SPEC=Wi,X.UNet=en,X.WebNNUNetExecutor=tn,X.detectUNetModelSpec=qe,X.initUNetFromBuffer=rn,X.initUNetFromURL=aa,X.optimizeModelGraph=ve,X.parseTZA=bi,X.planModelExecution=Hi,X.planTileGrid=Ke,X.resolveNativeUNetPrecision=io,X.selectSubgroupMatrixConfig=ro,X.validateUNetModel=Fi,Object.defineProperty(X,Symbol.toStringTag,{value:"Module"})});
|