longpipe 0.0.1
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/LICENSE +21 -0
- package/README.md +193 -0
- package/dist/index.cjs +1731 -0
- package/dist/index.d.cts +226 -0
- package/dist/index.d.ts +226 -0
- package/dist/index.js +1731 -0
- package/package.json +61 -0
package/dist/index.js
ADDED
|
@@ -0,0 +1,1731 @@
|
|
|
1
|
+
var v=[{model:"xxs",dtype:"f16",resolution:{w:128,h:72},skipFrames:3},{model:"xs",dtype:"f16",resolution:{w:192,h:108},skipFrames:1},{model:"small",dtype:"f16",resolution:{w:256,h:144},skipFrames:1},{model:"compact",dtype:"f16",resolution:{w:256,h:144},skipFrames:1},{model:"medium",dtype:"f16",resolution:{w:256,h:144},skipFrames:1},{model:"large",dtype:"f32",resolution:{w:256,h:144},skipFrames:0},{model:"xl",dtype:"f32",resolution:{w:512,h:288},skipFrames:0}];var w=(...i)=>console.log("[longpipe/adaptive]",...i),Qe=2e3,Ze=1e4,et=30,tt=20,rt=5e3,at=.3,it=15e3,C=class{constructor(t){this.opts=t;this.currentPresetIdx=v.findIndex(e=>e.model===t.initialModel),this.lastSwapAt=performance.now();}opts;timer=null;currentPresetIdx;lastSwapAt=0;overshootStart=0;headroomStart=0;start(){w(`started; backend=${this.opts.backendKind} initial=${this.opts.initialModel} idx=${this.currentPresetIdx}`),this.timer=setInterval(()=>{this.tick().catch(t=>w("tick failed:",t));},Qe);}stop(){this.timer&&clearInterval(this.timer),this.timer=null;}async tick(){if(typeof document<"u"&&document.hidden)return;let t=performance.now();if(t-this.lastSwapAt<Ze)return;let e=await this.opts.getStats();if(e.fps<tt?(this.overshootStart===0&&(this.overshootStart=t),t-this.overshootStart>=rt&&this.tryDowngrade(e.fps)):this.overshootStart=0,this.opts.backendKind==="webgpu"&&e.modelMs>0){let r=1e3/et*at;e.modelMs<r?(this.headroomStart===0&&(this.headroomStart=t),t-this.headroomStart>=it&&this.tryUpgrade(e.modelMs,r)):this.headroomStart=0;}}tryDowngrade(t){if(this.currentPresetIdx<=0)return;let e=this.currentPresetIdx-1;w(`downgrade ${v[this.currentPresetIdx].model} \u2192 ${v[e].model} (fps=${t})`),this.swap(e);}tryUpgrade(t,e){if(this.currentPresetIdx>=v.length-1)return;let r=this.currentPresetIdx+1;w(`upgrade ${v[this.currentPresetIdx].model} \u2192 ${v[r].model} (modelMs=${t.toFixed(1)} < ${e.toFixed(1)})`),this.swap(r);}async swap(t){let e=v[t];this.lastSwapAt=performance.now(),this.overshootStart=0,this.headroomStart=0;try{let r=this.opts.buildWeightsUrl(this.opts.weightsBaseUrl,e.model),a=await fetch(r);if(!a.ok)throw new Error(`weights fetch failed: ${a.status} ${r}`);let o=await a.arrayBuffer();await this.opts.swapPreset(e.model,o),this.currentPresetIdx=t,w(`swap to ${e.model} done`);}catch(r){w("swap failed:",r),this.opts.onError?.({message:`adaptive preset swap failed: ${r.message??String(r)}`,source:"adaptive",recoverable:true,cause:r});}}};async function re(i){let t,e=false;if(i instanceof HTMLVideoElement){if(t=i.currentSrc||i.src,!t)throw new Error("background video: HTMLVideoElement has no src")}else if(i instanceof Blob)t=URL.createObjectURL(i),e=true;else if(typeof i=="string")t=i;else throw new Error("background video: unsupported input shape");let r=document.createElement("video");r.crossOrigin="anonymous",r.muted=true,r.loop=true,r.playsInline=true,r.src=t,await new Promise((p,l)=>{r.addEventListener("loadeddata",()=>p(),{once:true}),r.addEventListener("error",()=>l(new Error(`background video: failed to load ${t}`)),{once:true});});try{await r.play();}catch(p){throw new Error(`background video: play() rejected \u2014 ${p.message}`)}let a=new MessageChannel,o=a.port1,s=a.port2,n=false,u=r.requestVideoFrameCallback?.bind(r);if(u){let p=(l,f)=>{if(!n){try{let b=new VideoFrame(r,{timestamp:f.mediaTime*1e6});o.postMessage({frame:b},[b]);}catch{}u(p);}};u(p);}else {let p=()=>{if(!n){try{let l=new VideoFrame(r,{timestamp:performance.now()*1e3});o.postMessage({frame:l},[l]);}catch{}requestAnimationFrame(p);}};requestAnimationFrame(p);}return {port:s,transferList:[s],cleanup:()=>{n=true,o.close(),r.pause(),r.removeAttribute("src"),r.load(),e&&URL.revokeObjectURL(t);}}}var J=4,X=8,ae=16;async function Q(i){if(typeof i=="string")return i==="none"?{background:{kind:"none"}}:i==="blur"?{background:{kind:"blur",sigma:X}}:{background:{kind:"image",bitmap:await oe(i)}};if(i instanceof ImageBitmap)return {background:{kind:"image",bitmap:i}};if(typeof HTMLImageElement<"u"&&i instanceof HTMLImageElement)return {background:{kind:"image",bitmap:await createImageBitmap(i)}};if(typeof HTMLVideoElement<"u"&&i instanceof HTMLVideoElement)return await ie(i);let t=i;if(t.color)return {background:{kind:"color",rgb:t.color}};if(t.blur!==void 0)return {background:{kind:"blur",sigma:ot(t.blur)}};if(t.image!==void 0)return {background:{kind:"image",bitmap:await st(t.image)}};if(t.video!==void 0)return await ie(t.video);throw new Error(`background: unrecognized input shape \u2014 ${P(i)}`)}async function ie(i){let t;if(typeof i=="string")t=i;else if(typeof HTMLVideoElement<"u"&&i instanceof HTMLVideoElement)t=i;else if(i&&typeof i=="object"&&"data"in i&&"type"in i)t=i.data instanceof Blob?i.data:new Blob([i.data],{type:i.type});else throw new Error(`background.video: unrecognized shape \u2014 ${P(i)}`);let e=await re(t);return {background:{kind:"video",port:e.port},transferList:e.transferList,cleanup:e.cleanup}}function ot(i){if(i===true)return X;if("sigma"in i){if(typeof i.sigma!="number"||!Number.isFinite(i.sigma)||i.sigma<0)throw new Error(`background.blur.sigma must be a non-negative number, got ${i.sigma}`);return i.sigma}if("strength"in i){let t=i.strength;if(typeof t=="number"){let e=Math.max(0,Math.min(1,t));return J+(ae-J)*e}if(t==="low")return J;if(t==="medium")return X;if(t==="high")return ae;throw new Error(`background.blur.strength must be 'low' | 'medium' | 'high' | number, got ${P(t)}`)}throw new Error(`background.blur: must be true, { strength }, or { sigma } \u2014 got ${P(i)}`)}async function st(i){if(i instanceof ImageBitmap)return i;if(typeof HTMLImageElement<"u"&&i instanceof HTMLImageElement)return createImageBitmap(i);if(typeof i=="string")return oe(i);if(i&&typeof i=="object"&&"data"in i&&"type"in i){let t=i.data instanceof Blob?i.data:new Blob([i.data],{type:i.type});return createImageBitmap(t)}throw new Error(`background.image: unrecognized shape \u2014 ${P(i)}`)}async function oe(i){let t=await fetch(i);if(!t.ok)throw new Error(`background: failed to load image from ${i} (HTTP ${t.status})`);let e=await t.blob();return createImageBitmap(e)}function P(i){return i===null?"null":typeof i!="object"?typeof i+"("+String(i)+")":`${i.constructor?.name??"Object"} ${JSON.stringify(i).slice(0,80)}`}function se(){return {input:"MediaStreamTrackProcessor"in self?"mstp":"rvfc-postmessage",output:"MediaStreamTrackGenerator"in self?"mstg":nt()?"transfer-capture":"bitmap-shuttle"}}var T=null;function nt(){return T!==null||(T=!/Firefox/.test(navigator.userAgent)),T}function ne(i,t,e){let r=[i];return e==="passthrough"&&r.push(...t.getAudioTracks()),new MediaStream(r)}var W=(...i)=>console.log("[longpipe/worker_controller]",...i),M=class{worker;listeners=new Map;persistentEvents=new Set;constructor(t){this.worker=t,this.worker.addEventListener("message",this.handleMessage.bind(this)),W("constructed; message listener attached");}handleMessage(t){let{request_id:e,res:r}=t.data,a=this.listeners.get(e);W("msg received: request_id=",e,"handler?",!!a,"persistent?",this.persistentEvents.has(e)),a&&(a(r),this.persistentEvents.has(e)||this.listeners.delete(e));}addPersistentListener(t,e){this.persistentEvents.add(t),this.listeners.set(t,e),W("addPersistentListener:",t);}removePersistentListener(t){this.persistentEvents.delete(t),this.listeners.delete(t);}sendMessage(t,e,r=[]){let a=crypto.randomUUID();return W("sendMessage:",t,"request_id=",a,"transferables=",r.length),new Promise(o=>{this.listeners.set(a,n=>o(n));let s={cmd:t,data:e,request_id:a};this.worker.postMessage(s,r);})}terminate(){this.worker.terminate(),this.listeners.clear(),this.persistentEvents.clear();}};function ue(i){let t=document.createElement("video");t.srcObject=i,t.muted=true,t.autoplay=true,t.playsInline=true,t.play().catch(c=>{console.warn("[input-postmessage/main] video.play() rejected:",c);});let e=t.requestVideoFrameCallback?.bind(t);if(!e)throw new Error("setupPostMessageInput: requestVideoFrameCallback not supported on this browser");let r=new MessageChannel,a=r.port1,o=r.port2,s=false,n=(c,p)=>{if(!s){try{let l=new VideoFrame(t,{timestamp:p.mediaTime*1e6});a.postMessage({frame:l},[l]);}catch(l){console.warn("[input-postmessage/main] VideoFrame() failed; skipping frame:",l);}e(n);}};return e(n),{port:o,transferList:[o],cleanup:()=>{s=true,a.close(),t.srcObject=null,t.pause();}}}function pe(i){let t=i.getVideoTracks();if(t.length===0)throw new Error("setupMstpInput: input MediaStream has no video tracks");let e=t[0],r=self.MediaStreamTrackProcessor;if(!r)throw new Error("setupMstpInput: MediaStreamTrackProcessor not supported on this browser");let o=new r({track:e}).readable;return {readable:o,transferList:[o],cleanup:()=>{o.cancel().catch(()=>{});}}}function le(i,t){switch(i){case "mstp":{let e=pe(t);return {initFields:{inputReadable:e.readable},transferList:e.transferList,cleanup:e.cleanup}}case "rvfc-postmessage":{let e=ue(t);return {initFields:{inputPort:e.port},transferList:e.transferList,cleanup:e.cleanup}}}}function ce(i,t=30){let e=document.createElement("canvas");e.width=i.w,e.height=i.h;let r=e.getContext("bitmaprenderer");if(!r)throw new Error("setupBitmapShuttleOutput: bitmaprenderer context not available");let o=e.captureStream(t).getVideoTracks();if(o.length===0)throw new Error("setupBitmapShuttleOutput: captureStream produced no video tracks");let s=o[0],n=new MessageChannel,u=n.port1,c=n.port2,p=false,l=null,f=()=>{},b=3,U=0;return u.onmessage=I=>{if(p){if(U++,U<b){I.data.bmp.close();return}f();}try{r.transferFromImageBitmap(I.data.bmp);}catch(g){console.warn("[bitmap-shuttle/main] transferFromImageBitmap failed:",g);}},u.start(),{videoTrack:s,port:c,transferList:[c],startPassthrough:I=>{if(p)return;p=true;let g=document.createElement("video");g.srcObject=I,g.muted=true,g.playsInline=true,g.play().catch(()=>{}),l=g;let G="requestVideoFrameCallback"in g?y=>{g.requestVideoFrameCallback(y);}:y=>{requestAnimationFrame(y);},B=()=>{if(p){if(g.videoWidth===0){G(B);return}createImageBitmap(g,{resizeWidth:e.width,resizeHeight:e.height,resizeQuality:"medium"}).then(y=>{if(!p){y.close();return}try{r.transferFromImageBitmap(y);}catch{}G(B);}).catch(()=>{G(B);});}};G(B),f=()=>{p=false,l&&(l.srcObject=null,l.pause(),l=null);};},cleanup:()=>{f(),u.onmessage=null,u.close(),s.stop();}}}function de(){let i=self.MediaStreamTrackGenerator;if(!i)throw new Error("setupMstgOutput: MediaStreamTrackGenerator not supported on this browser");let t=new i({kind:"video"}),e=t.writable;return {videoTrack:t,writable:e,transferList:[e],cleanup:()=>{e.abort().catch(()=>{}),t.stop();}}}function fe(i,t=30){let e=document.createElement("canvas");e.width=i.w,e.height=i.h;let a=e.captureStream(t).getVideoTracks();if(a.length===0)throw new Error("setupTransferCaptureOutput: captureStream produced no video tracks");let o=a[0],s=e.transferControlToOffscreen();return {videoTrack:o,canvas:s,transferList:[s],cleanup:()=>{o.stop();}}}function me(i,t){switch(i){case "mstg":{let e=de();return {videoTrack:e.videoTrack,initFields:{outputWritable:e.writable},transferList:e.transferList,cleanup:e.cleanup}}case "transfer-capture":{let e=fe(t);return {videoTrack:e.videoTrack,initFields:{outputCanvas:e.canvas},transferList:e.transferList,cleanup:e.cleanup}}case "bitmap-shuttle":{let e=ce(t);return {videoTrack:e.videoTrack,initFields:{outputPort:e.port},transferList:e.transferList,startPassthrough:e.startPassthrough,cleanup:e.cleanup}}}}var ut="https://cdn.longpipe.dev/models/v/0.0.1/",pt={background:"blur",preset:"auto",weightsBaseUrl:ut,audio:"passthrough",enabled:true,adaptive:true},ge={w:1280,h:720};function lt(i,t){if(t)return t;let r=i.getVideoTracks()[0]?.getSettings()??{};return {w:r.width||ge.w,h:r.height||ge.h}}function be(i,t){return `${i.replace(/\/$/,"")}/model_${t}.bin`}var Z=class{stream;ready;controller;worker;inputCleanup;outputCleanup;adaptive=null;bgCleanup=null;constructor(t,e){let r={...pt,...e},a=se(),o=le(a.input,t),s=lt(t,r.outputResolution),n=me(a.output,s);this.inputCleanup=o.cleanup,this.outputCleanup=n.cleanup,this.stream=ne(n.videoTrack,t,r.audio),n.startPassthrough?.(t),this.worker=new Worker(new URL("./worker/index.ts",import.meta.url),{type:"module"}),this.controller=new M(this.worker),this.ready=new Promise((p,l)=>{this.controller.addPersistentListener("ready",()=>{console.log("[longpipe/pipeline] ready handler invoked; resolving .ready"),r.onReady?.(),p();}),this.controller.addPersistentListener("error",f=>{console.error(`[longpipe/pipeline] error (${f.source}):`,f.message),f.recoverable||l(new Error(f.message)),r.onError?.(f);});});let u={topology:a,preset:r.preset,enabled:r.enabled,backend:"auto",dtype:"f16",canvasSize:s,...o.initFields,...n.initFields},c=[...o.transferList,...n.transferList];this.bootstrap(u,r.background,c,r.weightsBaseUrl,r.adaptive,r.onError);}async bootstrap(t,e,r,a,o,s){try{console.log("[longpipe/pipeline] normalizing background\u2026");let n=await Q(e);this.bgCleanup=n.cleanup??null;let u={...t,background:n.background},c=[...r,...n.transferList??[]];console.log("[longpipe/pipeline] sending init\u2026");let p=await this.controller.sendMessage("init",u,c);console.log("[longpipe/pipeline] init resolved:",p);let l=be(a,p.resolvedPreset.model);console.log("[longpipe/pipeline] fetching weights:",l);let f=await fetch(l);if(!f.ok)throw new Error(`weights fetch failed: ${f.status} ${l}`);let b=await f.arrayBuffer();console.log("[longpipe/pipeline] weights bytes:",b.byteLength),console.log("[longpipe/pipeline] sending startRender\u2026"),await this.controller.sendMessage("startRender",{weights:b},[b]),console.log("[longpipe/pipeline] startRender resolved; awaiting first frame"),u.preset==="auto"&&o&&(this.adaptive=new C({backendKind:p.resolvedBackend,initialModel:p.resolvedPreset.model,weightsBaseUrl:a,buildWeightsUrl:be,getStats:()=>this.getStats(),swapPreset:async(U,K)=>{await this.controller.sendMessage("setPreset",{preset:U,weights:K},[K]);},onError:s}),this.adaptive.start());}catch(n){console.error("[longpipe/pipeline] bootstrap failed:",n),this.controller.handleMessage.call(this.controller,{data:{request_id:"error",res:{message:n.message??String(n),source:"worker",recoverable:false,cause:n}}});}}then(t,e){return this.ready.then(()=>t?t(this):this,e)}async setBackground(t){let e=await Q(t),r=this.bgCleanup;this.bgCleanup=e.cleanup??null,await this.controller.sendMessage("setBackground",e.background,e.transferList),r?.();}setPreset(t,e){this.controller.sendMessage("setPreset",{preset:t,weights:e});}setEnabled(t){this.controller.sendMessage("setEnabled",{enabled:t});}async getStats(){return this.controller.sendMessage("getStats",{})}destroy(){this.adaptive?.stop(),this.adaptive=null,this.bgCleanup?.(),this.bgCleanup=null,this.controller.sendMessage("destroy",{}),this.controller.terminate(),this.inputCleanup(),this.outputCleanup();}};var x=new Float32Array(1),k=new Uint32Array(x.buffer);function ct(i){x[0]=i;let t=k[0],e=t>>>16&32768,r=t&2147483647,a=(r>>>23)-127+15;if((t&2139095040)===2139095040)return (t&8388607)!==0?e|32256:e|31744;if(a>=31)return e|31744;if(a<=0){if(a<-10)return e;let c=r&8388607|8388608,p=14-a,l=c>>>p,f=c>>>p-1&1,b=(c&(1<<p-1)-1)!==0?1:0;return f&&(b||l&1)&&(l+=1),e|l}let o=r>>>13&1023,s=r>>>12&1,n=(r&4095)!==0?1:0,u=a<<10|o;return s&&(n||o&1)&&(u+=1,(u>>>10&31)===31)?e|31744:e|u}function dt(i){let t=(i&32768)<<16,e=(i&31744)>>>10,r=i&1023;if(e===0){if(r===0)return k[0]=t,x[0];let a=r,o=1;for(;(a&1024)===0;)a<<=1,o-=1;return a&=1023,k[0]=t|o+127-15<<23|a<<13,x[0]}return e===31?(k[0]=t|2139095040|r<<13,x[0]):(k[0]=t|e+127-15<<23|r<<13,x[0])}function he(i){let t=new Uint16Array(i.length);for(let e=0;e<i.length;e++)t[e]=ct(i[e]);return t}function ee(i){let t=new Float32Array(i.length);for(let e=0;e<i.length;e++)t[e]=dt(i[e]);return t}var d=class{constructor(t){this.backend=t;}backend;shader="";pipeline;bindGroup;uniformDefs=[];uniformBuffers={};createUniform(t,e){this.uniformDefs.push({name:t,type:e});}setUniform(t,e){let r=this.backend.device.createBuffer({size:e.byteLength,usage:GPUBufferUsage.UNIFORM|GPUBufferUsage.COPY_DST,mappedAtCreation:true});e instanceof Uint32Array?new Uint32Array(r.getMappedRange()).set(e):new Float32Array(r.getMappedRange()).set(e),r.unmap(),this.uniformBuffers[t]=r;}defaultBindGroup(){let t=[],e=0;for(let r of this.inputs)t.push({binding:e++,resource:{buffer:r.buffer}});for(let r of this.weights)t.push({binding:e++,resource:{buffer:r.buffer}});for(let r of this.uniformDefs)t.push({binding:e++,resource:{buffer:this.uniformBuffers[r.name]}});return t.push({binding:e,resource:{buffer:this.output.buffer}}),this.backend.device.createBindGroup({layout:this.pipeline.getBindGroupLayout(0),entries:t})}defaultSetup(){let t=this.backend.device.createShaderModule({code:this.shader});this.pipeline=this.backend.device.createComputePipeline({layout:"auto",compute:{module:t,entryPoint:"main"}}),this.bindGroup=this.defaultBindGroup();}run(){let t=this.backend.device.createCommandEncoder(),e=t.beginComputePass();e.setPipeline(this.pipeline),e.setBindGroup(0,this.bindGroup),e.dispatchWorkgroups(...this.dispatch),e.end(),this.backend.device.queue.submit([t.finish()]);}};var _e=`// Conv2d compute shader \u2014 handles all variants:
|
|
2
|
+
// 1\xD71 (pointwise), 3\xD73, 5\xD75, strided, BN-fused
|
|
3
|
+
//
|
|
4
|
+
// Tensor layout: NHWC with channels in vec4 groups
|
|
5
|
+
// buffer index = y * out_w * out_groups + x * out_groups + o
|
|
6
|
+
//
|
|
7
|
+
// Weight layout: [K*K][out_groups][in_groups] array of mat4x4
|
|
8
|
+
// weight index = (ky*kernel_w + kx) * out_groups * in_groups + o * in_groups + i
|
|
9
|
+
// mat4x4[col][row] = weight(in_channel=col, out_channel=row)
|
|
10
|
+
// so (mat4x4 * vec4)[r] = sum_c(w(in=c, out=r) * input[c]) \u2014 correct matmul direction
|
|
11
|
+
//
|
|
12
|
+
// Padding model: only \`pad_top\` and \`pad_left\` are applied to the input offset.
|
|
13
|
+
// Asymmetric SAME padding (top<bottom or left<right) is handled implicitly:
|
|
14
|
+
// out-of-bounds reads on the right/bottom side are skipped via the bounds
|
|
15
|
+
// check, so the runner only needs to specify the top/left pad.
|
|
16
|
+
|
|
17
|
+
struct Params {
|
|
18
|
+
in_h : u32,
|
|
19
|
+
in_w : u32,
|
|
20
|
+
out_h : u32,
|
|
21
|
+
out_w : u32,
|
|
22
|
+
in_groups : u32, // in_channels / 4
|
|
23
|
+
out_groups : u32, // out_channels / 4
|
|
24
|
+
kernel_h : u32,
|
|
25
|
+
kernel_w : u32,
|
|
26
|
+
stride : u32,
|
|
27
|
+
pad_top : u32,
|
|
28
|
+
pad_left : u32,
|
|
29
|
+
activation : u32, // 0 = none, 1 = relu6, 2 = relu
|
|
30
|
+
}
|
|
31
|
+
|
|
32
|
+
@group(0) @binding(0) var<storage, read> input_buf : array<vec4<f32>>;
|
|
33
|
+
@group(0) @binding(1) var<storage, read> weight_buf : array<mat4x4<f32>>;
|
|
34
|
+
@group(0) @binding(2) var<storage, read> bias_buf : array<vec4<f32>>;
|
|
35
|
+
@group(0) @binding(3) var<uniform> params : Params;
|
|
36
|
+
@group(0) @binding(4) var<storage, read_write> output_buf : array<vec4<f32>>;
|
|
37
|
+
|
|
38
|
+
@compute @workgroup_size(8, 8, 1)
|
|
39
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
40
|
+
let x = gid.x; // output column
|
|
41
|
+
let y = gid.y; // output row
|
|
42
|
+
let o = gid.z; // output channel group
|
|
43
|
+
|
|
44
|
+
if (x >= params.out_w || y >= params.out_h || o >= params.out_groups) {
|
|
45
|
+
return;
|
|
46
|
+
}
|
|
47
|
+
|
|
48
|
+
let I = params.in_groups;
|
|
49
|
+
let O = params.out_groups;
|
|
50
|
+
|
|
51
|
+
var result = bias_buf[o];
|
|
52
|
+
|
|
53
|
+
for (var ky = 0u; ky < params.kernel_h; ky++) {
|
|
54
|
+
for (var kx = 0u; kx < params.kernel_w; kx++) {
|
|
55
|
+
let in_y_s = i32(y * params.stride + ky) - i32(params.pad_top);
|
|
56
|
+
let in_x_s = i32(x * params.stride + kx) - i32(params.pad_left);
|
|
57
|
+
|
|
58
|
+
if (in_y_s < 0 || in_x_s < 0 ||
|
|
59
|
+
u32(in_y_s) >= params.in_h || u32(in_x_s) >= params.in_w) {
|
|
60
|
+
continue;
|
|
61
|
+
}
|
|
62
|
+
|
|
63
|
+
let in_y = u32(in_y_s);
|
|
64
|
+
let in_x = u32(in_x_s);
|
|
65
|
+
let z = ky * params.kernel_w + kx;
|
|
66
|
+
|
|
67
|
+
for (var i = 0u; i < I; i++) {
|
|
68
|
+
let in_idx = in_y * params.in_w * I + in_x * I + i;
|
|
69
|
+
let w_idx = z * O * I + o * I + i;
|
|
70
|
+
result += weight_buf[w_idx] * input_buf[in_idx];
|
|
71
|
+
}
|
|
72
|
+
}
|
|
73
|
+
}
|
|
74
|
+
|
|
75
|
+
if (params.activation == 1u) {
|
|
76
|
+
result = clamp(result, vec4<f32>(0.0), vec4<f32>(6.0));
|
|
77
|
+
} else if (params.activation == 2u) {
|
|
78
|
+
result = max(result, vec4<f32>(0.0));
|
|
79
|
+
}
|
|
80
|
+
|
|
81
|
+
let out_idx = y * params.out_w * O + x * O + o;
|
|
82
|
+
output_buf[out_idx] = result;
|
|
83
|
+
}
|
|
84
|
+
`;var ve=`enable f16;
|
|
85
|
+
|
|
86
|
+
// Conv2d \u2014 full f16 variant.
|
|
87
|
+
// All buffers (input, output, weights, bias) are f16. Weights are converted
|
|
88
|
+
// from f32\u2192f16 in JavaScript before upload. mat4x4<f16> * vec4<f16> runs
|
|
89
|
+
// entirely in f16 on hardware that has native f16 ALUs.
|
|
90
|
+
|
|
91
|
+
struct Params {
|
|
92
|
+
in_h : u32,
|
|
93
|
+
in_w : u32,
|
|
94
|
+
out_h : u32,
|
|
95
|
+
out_w : u32,
|
|
96
|
+
in_groups : u32,
|
|
97
|
+
out_groups : u32,
|
|
98
|
+
kernel_h : u32,
|
|
99
|
+
kernel_w : u32,
|
|
100
|
+
stride : u32,
|
|
101
|
+
pad_top : u32,
|
|
102
|
+
pad_left : u32,
|
|
103
|
+
activation : u32,
|
|
104
|
+
}
|
|
105
|
+
|
|
106
|
+
@group(0) @binding(0) var<storage, read> input_buf : array<vec4<f16>>;
|
|
107
|
+
@group(0) @binding(1) var<storage, read> weight_buf : array<mat4x4<f16>>;
|
|
108
|
+
@group(0) @binding(2) var<storage, read> bias_buf : array<vec4<f16>>;
|
|
109
|
+
@group(0) @binding(3) var<uniform> params : Params;
|
|
110
|
+
@group(0) @binding(4) var<storage, read_write> output_buf : array<vec4<f16>>;
|
|
111
|
+
|
|
112
|
+
@compute @workgroup_size(8, 8, 1)
|
|
113
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
114
|
+
let x = gid.x;
|
|
115
|
+
let y = gid.y;
|
|
116
|
+
let o = gid.z;
|
|
117
|
+
|
|
118
|
+
if (x >= params.out_w || y >= params.out_h || o >= params.out_groups) {
|
|
119
|
+
return;
|
|
120
|
+
}
|
|
121
|
+
|
|
122
|
+
let I = params.in_groups;
|
|
123
|
+
let O = params.out_groups;
|
|
124
|
+
|
|
125
|
+
var result = bias_buf[o];
|
|
126
|
+
|
|
127
|
+
for (var ky = 0u; ky < params.kernel_h; ky++) {
|
|
128
|
+
for (var kx = 0u; kx < params.kernel_w; kx++) {
|
|
129
|
+
let in_y_s = i32(y * params.stride + ky) - i32(params.pad_top);
|
|
130
|
+
let in_x_s = i32(x * params.stride + kx) - i32(params.pad_left);
|
|
131
|
+
|
|
132
|
+
if (in_y_s < 0 || in_x_s < 0 ||
|
|
133
|
+
u32(in_y_s) >= params.in_h || u32(in_x_s) >= params.in_w) {
|
|
134
|
+
continue;
|
|
135
|
+
}
|
|
136
|
+
|
|
137
|
+
let in_y = u32(in_y_s);
|
|
138
|
+
let in_x = u32(in_x_s);
|
|
139
|
+
let z = ky * params.kernel_w + kx;
|
|
140
|
+
|
|
141
|
+
for (var i = 0u; i < I; i++) {
|
|
142
|
+
let in_idx = in_y * params.in_w * I + in_x * I + i;
|
|
143
|
+
let w_idx = z * O * I + o * I + i;
|
|
144
|
+
result += weight_buf[w_idx] * input_buf[in_idx];
|
|
145
|
+
}
|
|
146
|
+
}
|
|
147
|
+
}
|
|
148
|
+
|
|
149
|
+
if (params.activation == 1u) {
|
|
150
|
+
result = clamp(result, vec4<f16>(0.0h), vec4<f16>(6.0h));
|
|
151
|
+
} else if (params.activation == 2u) {
|
|
152
|
+
result = max(result, vec4<f16>(0.0h));
|
|
153
|
+
}
|
|
154
|
+
|
|
155
|
+
output_buf[y * params.out_w * O + x * O + o] = result;
|
|
156
|
+
}
|
|
157
|
+
`;function h(i,t,e,r){return typeof r=="number"?Math.floor((i+2*r-t)/e)+1:r==="same"?Math.ceil(i/e):Math.floor((i-t)/e)+1}function gt(i,t,e,r){return Math.floor(Math.max((t-1)*r+e-i,0)/2)}function _(i,t,e,r,a){return typeof i=="number"?i:i==="same"?gt(t,e,r,a):0}function m(i){return i instanceof Float32Array||i instanceof Uint16Array?i:new Float32Array(i)}var S=class extends d{inputs;weights;output;dispatch;shader;constructor(t,e,r,a){super(t),this.shader=t.dtype==="f16"?ve:_e;let o=h(e.h,a.kernel,a.stride,a.padding),s=h(e.w,a.kernel,a.stride,a.padding),n=e.c/4,u=a.outChannels/4,c=_(a.padding,e.h,o,a.kernel,a.stride),p=_(a.padding,e.w,s,a.kernel,a.stride);this.output=t.tensor(o,s,a.outChannels),this.inputs=[e],this.weights=[t.upload(m(r.weights)),t.upload(m(r.bias))],this.createUniform("params","Params"),this.setUniform("params",new Uint32Array([e.h,e.w,o,s,n,u,a.kernel,a.kernel,a.stride,c,p,a.activation==="relu6"?1:0])),this.defaultSetup(),this.dispatch=[Math.ceil(s/8),Math.ceil(o/8),u];}};var ye=`// Depthwise Conv2d \u2014 groups = in_channels (each channel convolved independently).
|
|
158
|
+
//
|
|
159
|
+
// Weight layout: [K*K][channel_groups] array of vec4
|
|
160
|
+
// weight index = (ky*kernel_w + kx) * channel_groups + c
|
|
161
|
+
// Each vec4 holds the kernel weight for 4 consecutive channels at one spatial position.
|
|
162
|
+
// Operation: element-wise multiply (each input channel multiplied by its own weight).
|
|
163
|
+
//
|
|
164
|
+
// Contrast with conv2d.wgsl which uses mat4x4 (dense cross-channel mixing).
|
|
165
|
+
// 4\xD7 smaller weight buffer than a diagonal mat4x4 representation (4 floats vs 16 per group).
|
|
166
|
+
//
|
|
167
|
+
// Padding model: only \`pad_top\` and \`pad_left\` are applied to the input offset.
|
|
168
|
+
// Asymmetric SAME padding is handled implicitly via the in_h/in_w bounds check.
|
|
169
|
+
|
|
170
|
+
struct Params {
|
|
171
|
+
in_h : u32,
|
|
172
|
+
in_w : u32,
|
|
173
|
+
out_h : u32,
|
|
174
|
+
out_w : u32,
|
|
175
|
+
channel_groups : u32, // channels / 4
|
|
176
|
+
kernel_h : u32,
|
|
177
|
+
kernel_w : u32,
|
|
178
|
+
stride : u32,
|
|
179
|
+
pad_top : u32,
|
|
180
|
+
pad_left : u32,
|
|
181
|
+
apply_relu6 : u32,
|
|
182
|
+
_pad0 : u32,
|
|
183
|
+
}
|
|
184
|
+
|
|
185
|
+
@group(0) @binding(0) var<storage, read> input_buf : array<vec4<f32>>;
|
|
186
|
+
@group(0) @binding(1) var<storage, read> weight_buf : array<vec4<f32>>;
|
|
187
|
+
@group(0) @binding(2) var<storage, read> bias_buf : array<vec4<f32>>;
|
|
188
|
+
@group(0) @binding(3) var<uniform> params : Params;
|
|
189
|
+
@group(0) @binding(4) var<storage, read_write> output_buf : array<vec4<f32>>;
|
|
190
|
+
|
|
191
|
+
@compute @workgroup_size(8, 8, 1)
|
|
192
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
193
|
+
let x = gid.x; // output column
|
|
194
|
+
let y = gid.y; // output row
|
|
195
|
+
let c = gid.z; // channel group
|
|
196
|
+
|
|
197
|
+
if (x >= params.out_w || y >= params.out_h || c >= params.channel_groups) {
|
|
198
|
+
return;
|
|
199
|
+
}
|
|
200
|
+
|
|
201
|
+
let C = params.channel_groups;
|
|
202
|
+
|
|
203
|
+
var result = bias_buf[c];
|
|
204
|
+
|
|
205
|
+
for (var ky = 0u; ky < params.kernel_h; ky++) {
|
|
206
|
+
for (var kx = 0u; kx < params.kernel_w; kx++) {
|
|
207
|
+
let in_y_s = i32(y * params.stride + ky) - i32(params.pad_top);
|
|
208
|
+
let in_x_s = i32(x * params.stride + kx) - i32(params.pad_left);
|
|
209
|
+
|
|
210
|
+
if (in_y_s < 0 || in_x_s < 0 ||
|
|
211
|
+
u32(in_y_s) >= params.in_h || u32(in_x_s) >= params.in_w) {
|
|
212
|
+
continue;
|
|
213
|
+
}
|
|
214
|
+
|
|
215
|
+
let in_y = u32(in_y_s);
|
|
216
|
+
let in_x = u32(in_x_s);
|
|
217
|
+
let z = ky * params.kernel_w + kx;
|
|
218
|
+
|
|
219
|
+
let in_idx = in_y * params.in_w * C + in_x * C + c;
|
|
220
|
+
let w_idx = z * C + c;
|
|
221
|
+
result += weight_buf[w_idx] * input_buf[in_idx];
|
|
222
|
+
}
|
|
223
|
+
}
|
|
224
|
+
|
|
225
|
+
if (params.apply_relu6 == 1u) {
|
|
226
|
+
result = clamp(result, vec4<f32>(0.0), vec4<f32>(6.0));
|
|
227
|
+
}
|
|
228
|
+
|
|
229
|
+
let out_idx = y * params.out_w * C + x * C + c;
|
|
230
|
+
output_buf[out_idx] = result;
|
|
231
|
+
}
|
|
232
|
+
`;var we=`enable f16;
|
|
233
|
+
|
|
234
|
+
// Depthwise conv2d \u2014 full f16 variant. All buffers f16.
|
|
235
|
+
|
|
236
|
+
struct Params {
|
|
237
|
+
in_h : u32,
|
|
238
|
+
in_w : u32,
|
|
239
|
+
out_h : u32,
|
|
240
|
+
out_w : u32,
|
|
241
|
+
channel_groups : u32,
|
|
242
|
+
kernel_h : u32,
|
|
243
|
+
kernel_w : u32,
|
|
244
|
+
stride : u32,
|
|
245
|
+
pad_top : u32,
|
|
246
|
+
pad_left : u32,
|
|
247
|
+
apply_relu6 : u32,
|
|
248
|
+
_pad0 : u32,
|
|
249
|
+
}
|
|
250
|
+
|
|
251
|
+
@group(0) @binding(0) var<storage, read> input_buf : array<vec4<f16>>;
|
|
252
|
+
@group(0) @binding(1) var<storage, read> weight_buf : array<vec4<f16>>;
|
|
253
|
+
@group(0) @binding(2) var<storage, read> bias_buf : array<vec4<f16>>;
|
|
254
|
+
@group(0) @binding(3) var<uniform> params : Params;
|
|
255
|
+
@group(0) @binding(4) var<storage, read_write> output_buf : array<vec4<f16>>;
|
|
256
|
+
|
|
257
|
+
@compute @workgroup_size(8, 8, 1)
|
|
258
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
259
|
+
let x = gid.x;
|
|
260
|
+
let y = gid.y;
|
|
261
|
+
let c = gid.z;
|
|
262
|
+
|
|
263
|
+
if (x >= params.out_w || y >= params.out_h || c >= params.channel_groups) {
|
|
264
|
+
return;
|
|
265
|
+
}
|
|
266
|
+
|
|
267
|
+
let C = params.channel_groups;
|
|
268
|
+
|
|
269
|
+
var result = bias_buf[c];
|
|
270
|
+
|
|
271
|
+
for (var ky = 0u; ky < params.kernel_h; ky++) {
|
|
272
|
+
for (var kx = 0u; kx < params.kernel_w; kx++) {
|
|
273
|
+
let in_y_s = i32(y * params.stride + ky) - i32(params.pad_top);
|
|
274
|
+
let in_x_s = i32(x * params.stride + kx) - i32(params.pad_left);
|
|
275
|
+
|
|
276
|
+
if (in_y_s < 0 || in_x_s < 0 ||
|
|
277
|
+
u32(in_y_s) >= params.in_h || u32(in_x_s) >= params.in_w) {
|
|
278
|
+
continue;
|
|
279
|
+
}
|
|
280
|
+
|
|
281
|
+
let in_y = u32(in_y_s);
|
|
282
|
+
let in_x = u32(in_x_s);
|
|
283
|
+
let z = ky * params.kernel_w + kx;
|
|
284
|
+
|
|
285
|
+
let in_idx = in_y * params.in_w * C + in_x * C + c;
|
|
286
|
+
let w_idx = z * C + c;
|
|
287
|
+
result += weight_buf[w_idx] * input_buf[in_idx];
|
|
288
|
+
}
|
|
289
|
+
}
|
|
290
|
+
|
|
291
|
+
if (params.apply_relu6 == 1u) {
|
|
292
|
+
result = clamp(result, vec4<f16>(0.0h), vec4<f16>(6.0h));
|
|
293
|
+
}
|
|
294
|
+
|
|
295
|
+
output_buf[y * params.out_w * C + x * C + c] = result;
|
|
296
|
+
}
|
|
297
|
+
`;var A=class extends d{inputs;weights;output;dispatch;shader;constructor(t,e,r,a){super(t),this.shader=t.dtype==="f16"?we:ye;let o=h(e.h,a.kernel,a.stride,a.padding),s=h(e.w,a.kernel,a.stride,a.padding),n=e.c/4,u=_(a.padding,e.h,o,a.kernel,a.stride),c=_(a.padding,e.w,s,a.kernel,a.stride);this.output=t.tensor(o,s,e.c),this.inputs=[e],this.weights=[t.upload(m(r.weights)),t.upload(m(r.bias))],this.createUniform("params","Params"),this.setUniform("params",new Uint32Array([e.h,e.w,o,s,n,a.kernel,a.kernel,a.stride,u,c,a.activation==="relu6"?1:0,0])),this.defaultSetup(),this.dispatch=[Math.ceil(s/8),Math.ceil(o/8),n];}};var xe=`// Element-wise add of two tensors \u2014 used for residual connections in MBConv blocks.
|
|
298
|
+
// Operates on the flat float buffer directly; layout is irrelevant for a pure element-wise op.
|
|
299
|
+
|
|
300
|
+
struct Params {
|
|
301
|
+
size : u32, // total number of f32 elements
|
|
302
|
+
_pad0 : u32,
|
|
303
|
+
_pad1 : u32,
|
|
304
|
+
_pad2 : u32,
|
|
305
|
+
}
|
|
306
|
+
|
|
307
|
+
@group(0) @binding(0) var<storage, read> input_a : array<f32>;
|
|
308
|
+
@group(0) @binding(1) var<storage, read> input_b : array<f32>;
|
|
309
|
+
@group(0) @binding(2) var<uniform> params : Params;
|
|
310
|
+
@group(0) @binding(3) var<storage, read_write> output : array<f32>;
|
|
311
|
+
|
|
312
|
+
@compute @workgroup_size(256)
|
|
313
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
314
|
+
let idx = gid.x;
|
|
315
|
+
if (idx >= params.size) { return; }
|
|
316
|
+
output[idx] = input_a[idx] + input_b[idx];
|
|
317
|
+
}
|
|
318
|
+
`;var Pe=`enable f16;
|
|
319
|
+
|
|
320
|
+
// Element-wise add \u2014 full f16 variant.
|
|
321
|
+
// array<f16> is binary-compatible with array<vec4<f16>> written by the f16 conv shaders.
|
|
322
|
+
|
|
323
|
+
struct Params {
|
|
324
|
+
size : u32,
|
|
325
|
+
_pad0 : u32,
|
|
326
|
+
_pad1 : u32,
|
|
327
|
+
_pad2 : u32,
|
|
328
|
+
}
|
|
329
|
+
|
|
330
|
+
@group(0) @binding(0) var<storage, read> input_a : array<f16>;
|
|
331
|
+
@group(0) @binding(1) var<storage, read> input_b : array<f16>;
|
|
332
|
+
@group(0) @binding(2) var<uniform> params : Params;
|
|
333
|
+
@group(0) @binding(3) var<storage, read_write> output : array<f16>;
|
|
334
|
+
|
|
335
|
+
@compute @workgroup_size(256)
|
|
336
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
337
|
+
let idx = gid.x;
|
|
338
|
+
if (idx >= params.size) { return; }
|
|
339
|
+
output[idx] = input_a[idx] + input_b[idx];
|
|
340
|
+
}
|
|
341
|
+
`;var O=class extends d{inputs;weights=[];output;dispatch;shader;constructor(t,e,r){super(t),this.shader=t.dtype==="f16"?Pe:xe;let a=e.h*e.w*e.c;this.output=t.tensor(e.h,e.w,e.c),this.inputs=[e,r],this.createUniform("params","Params"),this.setUniform("params",new Uint32Array([a,0,0,0])),this.defaultSetup(),this.dispatch=[Math.ceil(a/256),1,1];}};var ke=`// Element-wise sigmoid: output = 1 / (1 + exp(-x)).
|
|
342
|
+
// Operates on packed vec4 buffers (NHWC layout). exp() and arithmetic are element-wise on vec4.
|
|
343
|
+
|
|
344
|
+
struct Params {
|
|
345
|
+
n_groups : u32, // total vec4 elements (H * W * channel_groups)
|
|
346
|
+
_pad0 : u32,
|
|
347
|
+
_pad1 : u32,
|
|
348
|
+
_pad2 : u32,
|
|
349
|
+
}
|
|
350
|
+
|
|
351
|
+
@group(0) @binding(0) var<storage, read> input_buf : array<vec4<f32>>;
|
|
352
|
+
@group(0) @binding(1) var<uniform> params : Params;
|
|
353
|
+
@group(0) @binding(2) var<storage, read_write> output_buf : array<vec4<f32>>;
|
|
354
|
+
|
|
355
|
+
@compute @workgroup_size(256)
|
|
356
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
357
|
+
let idx = gid.x;
|
|
358
|
+
if (idx >= params.n_groups) { return; }
|
|
359
|
+
output_buf[idx] = 1.0 / (1.0 + exp(-input_buf[idx]));
|
|
360
|
+
}
|
|
361
|
+
`;var Ue=`enable f16;
|
|
362
|
+
|
|
363
|
+
// Sigmoid \u2014 full f16 variant.
|
|
364
|
+
|
|
365
|
+
struct Params {
|
|
366
|
+
n_groups : u32,
|
|
367
|
+
_pad0 : u32,
|
|
368
|
+
_pad1 : u32,
|
|
369
|
+
_pad2 : u32,
|
|
370
|
+
}
|
|
371
|
+
|
|
372
|
+
@group(0) @binding(0) var<storage, read> input_buf : array<vec4<f16>>;
|
|
373
|
+
@group(0) @binding(1) var<uniform> params : Params;
|
|
374
|
+
@group(0) @binding(2) var<storage, read_write> output_buf : array<vec4<f16>>;
|
|
375
|
+
|
|
376
|
+
@compute @workgroup_size(256)
|
|
377
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
378
|
+
let idx = gid.x;
|
|
379
|
+
if (idx >= params.n_groups) { return; }
|
|
380
|
+
let x = vec4<f32>(input_buf[idx]);
|
|
381
|
+
output_buf[idx] = vec4<f16>(1.0 / (1.0 + exp(-x)));
|
|
382
|
+
}
|
|
383
|
+
`;var E=class extends d{inputs;weights=[];output;dispatch;shader;constructor(t,e){super(t),this.shader=t.dtype==="f16"?Ue:ke;let r=e.h*e.w*(e.c/4);this.output=t.tensor(e.h,e.w,e.c),this.inputs=[e],this.createUniform("params","Params"),this.setUniform("params",new Uint32Array([r,0,0,0])),this.defaultSetup(),this.dispatch=[Math.ceil(r/256),1,1];}};var Ie=`// Bilinear upsample (arbitrary ratio), align_corners=False (matches PyTorch default).
|
|
384
|
+
// Input/output in NHWC vec4 format: index = y*W*(C/4) + x*(C/4) + c_group.
|
|
385
|
+
// Each thread computes one output pixel for one channel group (vec4 = 4 channels).
|
|
386
|
+
|
|
387
|
+
struct Params {
|
|
388
|
+
in_h : u32,
|
|
389
|
+
in_w : u32,
|
|
390
|
+
out_h : u32,
|
|
391
|
+
out_w : u32,
|
|
392
|
+
channel_groups : u32,
|
|
393
|
+
_pad0 : u32,
|
|
394
|
+
_pad1 : u32,
|
|
395
|
+
_pad2 : u32,
|
|
396
|
+
}
|
|
397
|
+
|
|
398
|
+
@group(0) @binding(0) var<storage, read> input_buf : array<vec4<f32>>;
|
|
399
|
+
@group(0) @binding(1) var<uniform> params : Params;
|
|
400
|
+
@group(0) @binding(2) var<storage, read_write> output_buf : array<vec4<f32>>;
|
|
401
|
+
|
|
402
|
+
@compute @workgroup_size(8, 8, 1)
|
|
403
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
404
|
+
let ox = gid.x;
|
|
405
|
+
let oy = gid.y;
|
|
406
|
+
let c = gid.z;
|
|
407
|
+
|
|
408
|
+
if (ox >= params.out_w || oy >= params.out_h || c >= params.channel_groups) {
|
|
409
|
+
return;
|
|
410
|
+
}
|
|
411
|
+
|
|
412
|
+
let IH = params.in_h;
|
|
413
|
+
let IW = params.in_w;
|
|
414
|
+
let C = params.channel_groups;
|
|
415
|
+
|
|
416
|
+
// align_corners=False: src = (out + 0.5) * (in / out) - 0.5
|
|
417
|
+
let src_x = (f32(ox) + 0.5) * (f32(IW) / f32(params.out_w)) - 0.5;
|
|
418
|
+
let src_y = (f32(oy) + 0.5) * (f32(IH) / f32(params.out_h)) - 0.5;
|
|
419
|
+
|
|
420
|
+
// Clamp to [0, in-1] for border replication. Use i32 intermediates so that
|
|
421
|
+
// floor() returning -1.0 doesn't produce an invalid u32 conversion.
|
|
422
|
+
let x0 = u32(clamp(i32(floor(src_x)), 0, i32(IW) - 1));
|
|
423
|
+
let x1 = u32(clamp(i32(floor(src_x)) + 1, 0, i32(IW) - 1));
|
|
424
|
+
let y0 = u32(clamp(i32(floor(src_y)), 0, i32(IH) - 1));
|
|
425
|
+
let y1 = u32(clamp(i32(floor(src_y)) + 1, 0, i32(IH) - 1));
|
|
426
|
+
|
|
427
|
+
let wx = src_x - floor(src_x);
|
|
428
|
+
let wy = src_y - floor(src_y);
|
|
429
|
+
|
|
430
|
+
let tl = input_buf[y0 * IW * C + x0 * C + c];
|
|
431
|
+
let tr = input_buf[y0 * IW * C + x1 * C + c];
|
|
432
|
+
let bl = input_buf[y1 * IW * C + x0 * C + c];
|
|
433
|
+
let br = input_buf[y1 * IW * C + x1 * C + c];
|
|
434
|
+
|
|
435
|
+
// Bilinear blend \u2014 vec4 ops are element-wise, so all 4 channels blend identically.
|
|
436
|
+
let result = (1.0 - wy) * ((1.0 - wx) * tl + wx * tr)
|
|
437
|
+
+ wy * ((1.0 - wx) * bl + wx * br);
|
|
438
|
+
|
|
439
|
+
output_buf[oy * params.out_w * C + ox * C + c] = result;
|
|
440
|
+
}
|
|
441
|
+
`;var Ge=`enable f16;
|
|
442
|
+
|
|
443
|
+
// Bilinear upsample 2\xD7 \u2014 full f16 variant.
|
|
444
|
+
// Interpolation weights (wx, wy) and intermediate blends computed in f32 for
|
|
445
|
+
// accuracy; result cast to f16 on write.
|
|
446
|
+
|
|
447
|
+
struct Params {
|
|
448
|
+
in_h : u32,
|
|
449
|
+
in_w : u32,
|
|
450
|
+
out_h : u32,
|
|
451
|
+
out_w : u32,
|
|
452
|
+
channel_groups : u32,
|
|
453
|
+
_pad0 : u32,
|
|
454
|
+
_pad1 : u32,
|
|
455
|
+
_pad2 : u32,
|
|
456
|
+
}
|
|
457
|
+
|
|
458
|
+
@group(0) @binding(0) var<storage, read> input_buf : array<vec4<f16>>;
|
|
459
|
+
@group(0) @binding(1) var<uniform> params : Params;
|
|
460
|
+
@group(0) @binding(2) var<storage, read_write> output_buf : array<vec4<f16>>;
|
|
461
|
+
|
|
462
|
+
@compute @workgroup_size(8, 8, 1)
|
|
463
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
464
|
+
let ox = gid.x;
|
|
465
|
+
let oy = gid.y;
|
|
466
|
+
let c = gid.z;
|
|
467
|
+
|
|
468
|
+
if (ox >= params.out_w || oy >= params.out_h || c >= params.channel_groups) {
|
|
469
|
+
return;
|
|
470
|
+
}
|
|
471
|
+
|
|
472
|
+
let IH = params.in_h;
|
|
473
|
+
let IW = params.in_w;
|
|
474
|
+
let C = params.channel_groups;
|
|
475
|
+
|
|
476
|
+
let src_x = (f32(ox) + 0.5) * (f32(IW) / f32(params.out_w)) - 0.5;
|
|
477
|
+
let src_y = (f32(oy) + 0.5) * (f32(IH) / f32(params.out_h)) - 0.5;
|
|
478
|
+
|
|
479
|
+
let x0 = u32(clamp(i32(floor(src_x)), 0, i32(IW) - 1));
|
|
480
|
+
let x1 = u32(clamp(i32(floor(src_x)) + 1, 0, i32(IW) - 1));
|
|
481
|
+
let y0 = u32(clamp(i32(floor(src_y)), 0, i32(IH) - 1));
|
|
482
|
+
let y1 = u32(clamp(i32(floor(src_y)) + 1, 0, i32(IH) - 1));
|
|
483
|
+
|
|
484
|
+
let wx = src_x - floor(src_x);
|
|
485
|
+
let wy = src_y - floor(src_y);
|
|
486
|
+
|
|
487
|
+
let tl = vec4<f32>(input_buf[y0 * IW * C + x0 * C + c]);
|
|
488
|
+
let tr = vec4<f32>(input_buf[y0 * IW * C + x1 * C + c]);
|
|
489
|
+
let bl = vec4<f32>(input_buf[y1 * IW * C + x0 * C + c]);
|
|
490
|
+
let br = vec4<f32>(input_buf[y1 * IW * C + x1 * C + c]);
|
|
491
|
+
|
|
492
|
+
let result = (1.0 - wy) * ((1.0 - wx) * tl + wx * tr)
|
|
493
|
+
+ wy * ((1.0 - wx) * bl + wx * br);
|
|
494
|
+
|
|
495
|
+
output_buf[oy * params.out_w * C + ox * C + c] = vec4<f16>(result);
|
|
496
|
+
}
|
|
497
|
+
`;var F=class extends d{inputs;weights=[];output;dispatch;shader;constructor(t,e,r){super(t),this.shader=t.dtype==="f16"?Ge:Ie;let a=e.c/4;this.output=t.tensor(r.outH,r.outW,e.c),this.inputs=[e],this.createUniform("params","Params"),this.setUniform("params",new Uint32Array([e.h,e.w,r.outH,r.outW,a,0,0,0])),this.defaultSetup(),this.dispatch=[Math.ceil(r.outW/8),Math.ceil(r.outH/8),a];}};var Be=`// Bicubic upsample (arbitrary scale) \u2014 Keys cubic, a=-0.75 (PyTorch default
|
|
498
|
+
// for mode='bicubic', align_corners=False). Direct 2D, 4\xD74 = 16 taps per
|
|
499
|
+
// output pixel. NHWC vec4 layout: index = y*W*(C/4) + x*(C/4) + c_group.
|
|
500
|
+
|
|
501
|
+
struct Params {
|
|
502
|
+
in_h : u32,
|
|
503
|
+
in_w : u32,
|
|
504
|
+
out_h : u32,
|
|
505
|
+
out_w : u32,
|
|
506
|
+
channel_groups : u32,
|
|
507
|
+
_pad0 : u32,
|
|
508
|
+
_pad1 : u32,
|
|
509
|
+
_pad2 : u32,
|
|
510
|
+
}
|
|
511
|
+
|
|
512
|
+
@group(0) @binding(0) var<storage, read> input_buf : array<vec4<f32>>;
|
|
513
|
+
@group(0) @binding(1) var<uniform> params : Params;
|
|
514
|
+
@group(0) @binding(2) var<storage, read_write> output_buf : array<vec4<f32>>;
|
|
515
|
+
|
|
516
|
+
const A: f32 = -0.75;
|
|
517
|
+
|
|
518
|
+
fn wcubic(d: f32) -> f32 {
|
|
519
|
+
let ad = abs(d);
|
|
520
|
+
if (ad <= 1.0) { return ((A + 2.0) * ad - (A + 3.0)) * ad * ad + 1.0; }
|
|
521
|
+
if (ad < 2.0) { return ((A * ad - 5.0 * A) * ad + 8.0 * A) * ad - 4.0 * A; }
|
|
522
|
+
return 0.0;
|
|
523
|
+
}
|
|
524
|
+
|
|
525
|
+
@compute @workgroup_size(8, 8, 1)
|
|
526
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
527
|
+
let ox = gid.x;
|
|
528
|
+
let oy = gid.y;
|
|
529
|
+
let c = gid.z;
|
|
530
|
+
|
|
531
|
+
if (ox >= params.out_w || oy >= params.out_h || c >= params.channel_groups) {
|
|
532
|
+
return;
|
|
533
|
+
}
|
|
534
|
+
|
|
535
|
+
let IH = params.in_h;
|
|
536
|
+
let IW = params.in_w;
|
|
537
|
+
let C = params.channel_groups;
|
|
538
|
+
|
|
539
|
+
let src_x = (f32(ox) + 0.5) * (f32(IW) / f32(params.out_w)) - 0.5;
|
|
540
|
+
let src_y = (f32(oy) + 0.5) * (f32(IH) / f32(params.out_h)) - 0.5;
|
|
541
|
+
|
|
542
|
+
let x0 = i32(floor(src_x));
|
|
543
|
+
let y0 = i32(floor(src_y));
|
|
544
|
+
let fx = src_x - f32(x0);
|
|
545
|
+
let fy = src_y - f32(y0);
|
|
546
|
+
|
|
547
|
+
var wx: array<f32, 4>;
|
|
548
|
+
var wy: array<f32, 4>;
|
|
549
|
+
wx[0] = wcubic(1.0 + fx); wx[1] = wcubic(fx); wx[2] = wcubic(1.0 - fx); wx[3] = wcubic(2.0 - fx);
|
|
550
|
+
wy[0] = wcubic(1.0 + fy); wy[1] = wcubic(fy); wy[2] = wcubic(1.0 - fy); wy[3] = wcubic(2.0 - fy);
|
|
551
|
+
|
|
552
|
+
var acc: vec4<f32> = vec4<f32>(0.0);
|
|
553
|
+
for (var j: i32 = 0; j < 4; j = j + 1) {
|
|
554
|
+
let sy = u32(clamp(y0 + j - 1, 0, i32(IH) - 1));
|
|
555
|
+
for (var i: i32 = 0; i < 4; i = i + 1) {
|
|
556
|
+
let sx = u32(clamp(x0 + i - 1, 0, i32(IW) - 1));
|
|
557
|
+
let v = input_buf[sy * IW * C + sx * C + c];
|
|
558
|
+
acc = acc + (wx[i] * wy[j]) * v;
|
|
559
|
+
}
|
|
560
|
+
}
|
|
561
|
+
|
|
562
|
+
output_buf[oy * params.out_w * C + ox * C + c] = acc;
|
|
563
|
+
}
|
|
564
|
+
`;var Ce=`enable f16;
|
|
565
|
+
|
|
566
|
+
// Bicubic upsample \u2014 full f16 storage variant. Bicubic weights and the per-
|
|
567
|
+
// pixel accumulator are computed in f32 (cheap to keep precision around the
|
|
568
|
+
// kernel arithmetic), then demoted to f16 when written to the output buffer.
|
|
569
|
+
|
|
570
|
+
struct Params {
|
|
571
|
+
in_h : u32,
|
|
572
|
+
in_w : u32,
|
|
573
|
+
out_h : u32,
|
|
574
|
+
out_w : u32,
|
|
575
|
+
channel_groups : u32,
|
|
576
|
+
_pad0 : u32,
|
|
577
|
+
_pad1 : u32,
|
|
578
|
+
_pad2 : u32,
|
|
579
|
+
}
|
|
580
|
+
|
|
581
|
+
@group(0) @binding(0) var<storage, read> input_buf : array<vec4<f16>>;
|
|
582
|
+
@group(0) @binding(1) var<uniform> params : Params;
|
|
583
|
+
@group(0) @binding(2) var<storage, read_write> output_buf : array<vec4<f16>>;
|
|
584
|
+
|
|
585
|
+
const A: f32 = -0.75;
|
|
586
|
+
|
|
587
|
+
fn wcubic(d: f32) -> f32 {
|
|
588
|
+
let ad = abs(d);
|
|
589
|
+
if (ad <= 1.0) { return ((A + 2.0) * ad - (A + 3.0)) * ad * ad + 1.0; }
|
|
590
|
+
if (ad < 2.0) { return ((A * ad - 5.0 * A) * ad + 8.0 * A) * ad - 4.0 * A; }
|
|
591
|
+
return 0.0;
|
|
592
|
+
}
|
|
593
|
+
|
|
594
|
+
@compute @workgroup_size(8, 8, 1)
|
|
595
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
596
|
+
let ox = gid.x;
|
|
597
|
+
let oy = gid.y;
|
|
598
|
+
let c = gid.z;
|
|
599
|
+
|
|
600
|
+
if (ox >= params.out_w || oy >= params.out_h || c >= params.channel_groups) {
|
|
601
|
+
return;
|
|
602
|
+
}
|
|
603
|
+
|
|
604
|
+
let IH = params.in_h;
|
|
605
|
+
let IW = params.in_w;
|
|
606
|
+
let C = params.channel_groups;
|
|
607
|
+
|
|
608
|
+
let src_x = (f32(ox) + 0.5) * (f32(IW) / f32(params.out_w)) - 0.5;
|
|
609
|
+
let src_y = (f32(oy) + 0.5) * (f32(IH) / f32(params.out_h)) - 0.5;
|
|
610
|
+
|
|
611
|
+
let x0 = i32(floor(src_x));
|
|
612
|
+
let y0 = i32(floor(src_y));
|
|
613
|
+
let fx = src_x - f32(x0);
|
|
614
|
+
let fy = src_y - f32(y0);
|
|
615
|
+
|
|
616
|
+
var wx: array<f32, 4>;
|
|
617
|
+
var wy: array<f32, 4>;
|
|
618
|
+
wx[0] = wcubic(1.0 + fx); wx[1] = wcubic(fx); wx[2] = wcubic(1.0 - fx); wx[3] = wcubic(2.0 - fx);
|
|
619
|
+
wy[0] = wcubic(1.0 + fy); wy[1] = wcubic(fy); wy[2] = wcubic(1.0 - fy); wy[3] = wcubic(2.0 - fy);
|
|
620
|
+
|
|
621
|
+
var acc: vec4<f32> = vec4<f32>(0.0);
|
|
622
|
+
for (var j: i32 = 0; j < 4; j = j + 1) {
|
|
623
|
+
let sy = u32(clamp(y0 + j - 1, 0, i32(IH) - 1));
|
|
624
|
+
for (var i: i32 = 0; i < 4; i = i + 1) {
|
|
625
|
+
let sx = u32(clamp(x0 + i - 1, 0, i32(IW) - 1));
|
|
626
|
+
let v = vec4<f32>(input_buf[sy * IW * C + sx * C + c]);
|
|
627
|
+
acc = acc + (wx[i] * wy[j]) * v;
|
|
628
|
+
}
|
|
629
|
+
}
|
|
630
|
+
|
|
631
|
+
output_buf[oy * params.out_w * C + ox * C + c] = vec4<f16>(acc);
|
|
632
|
+
}
|
|
633
|
+
`;var V=class extends d{inputs;weights=[];output;dispatch;shader;constructor(t,e,r){super(t),this.shader=t.dtype==="f16"?Ce:Be;let a=e.c/4;this.output=t.tensor(r.outH,r.outW,e.c),this.inputs=[e],this.createUniform("params","Params"),this.setUniform("params",new Uint32Array([e.h,e.w,r.outH,r.outW,a,0,0,0])),this.defaultSetup(),this.dispatch=[Math.ceil(r.outW/8),Math.ceil(r.outH/8),a];}};var Te=`// Channel concatenation: output = cat(A, B, dim=channel).
|
|
634
|
+
// Both A and B must share the same H\xD7W and be in NHWC vec4 format.
|
|
635
|
+
// Output layout: for each (y, x), the first a_groups vec4s come from A,
|
|
636
|
+
// followed by b_groups vec4s from B.
|
|
637
|
+
|
|
638
|
+
struct Params {
|
|
639
|
+
height : u32,
|
|
640
|
+
width : u32,
|
|
641
|
+
a_groups : u32, // Ca / 4
|
|
642
|
+
b_groups : u32, // Cb / 4
|
|
643
|
+
out_groups: u32, // a_groups + b_groups
|
|
644
|
+
_pad0 : u32,
|
|
645
|
+
_pad1 : u32,
|
|
646
|
+
_pad2 : u32,
|
|
647
|
+
}
|
|
648
|
+
|
|
649
|
+
@group(0) @binding(0) var<storage, read> input_a : array<vec4<f32>>;
|
|
650
|
+
@group(0) @binding(1) var<storage, read> input_b : array<vec4<f32>>;
|
|
651
|
+
@group(0) @binding(2) var<uniform> params : Params;
|
|
652
|
+
@group(0) @binding(3) var<storage, read_write> output_buf : array<vec4<f32>>;
|
|
653
|
+
|
|
654
|
+
@compute @workgroup_size(8, 8, 1)
|
|
655
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
656
|
+
let x = gid.x;
|
|
657
|
+
let y = gid.y;
|
|
658
|
+
let c = gid.z; // output channel group
|
|
659
|
+
|
|
660
|
+
let W = params.width;
|
|
661
|
+
let Ag = params.a_groups;
|
|
662
|
+
let Bg = params.b_groups;
|
|
663
|
+
let Cg = params.out_groups;
|
|
664
|
+
|
|
665
|
+
if (x >= W || y >= params.height || c >= Cg) {
|
|
666
|
+
return;
|
|
667
|
+
}
|
|
668
|
+
|
|
669
|
+
let out_idx = y * W * Cg + x * Cg + c;
|
|
670
|
+
|
|
671
|
+
if (c < Ag) {
|
|
672
|
+
output_buf[out_idx] = input_a[y * W * Ag + x * Ag + c];
|
|
673
|
+
} else {
|
|
674
|
+
let c_b = c - Ag;
|
|
675
|
+
output_buf[out_idx] = input_b[y * W * Bg + x * Bg + c_b];
|
|
676
|
+
}
|
|
677
|
+
}
|
|
678
|
+
`;var We=`enable f16;
|
|
679
|
+
|
|
680
|
+
// Channel concatenation \u2014 full f16 variant.
|
|
681
|
+
|
|
682
|
+
struct Params {
|
|
683
|
+
height : u32,
|
|
684
|
+
width : u32,
|
|
685
|
+
a_groups : u32,
|
|
686
|
+
b_groups : u32,
|
|
687
|
+
out_groups: u32,
|
|
688
|
+
_pad0 : u32,
|
|
689
|
+
_pad1 : u32,
|
|
690
|
+
_pad2 : u32,
|
|
691
|
+
}
|
|
692
|
+
|
|
693
|
+
@group(0) @binding(0) var<storage, read> input_a : array<vec4<f16>>;
|
|
694
|
+
@group(0) @binding(1) var<storage, read> input_b : array<vec4<f16>>;
|
|
695
|
+
@group(0) @binding(2) var<uniform> params : Params;
|
|
696
|
+
@group(0) @binding(3) var<storage, read_write> output_buf : array<vec4<f16>>;
|
|
697
|
+
|
|
698
|
+
@compute @workgroup_size(8, 8, 1)
|
|
699
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
700
|
+
let x = gid.x;
|
|
701
|
+
let y = gid.y;
|
|
702
|
+
let c = gid.z;
|
|
703
|
+
|
|
704
|
+
let W = params.width;
|
|
705
|
+
let Ag = params.a_groups;
|
|
706
|
+
let Bg = params.b_groups;
|
|
707
|
+
let Cg = params.out_groups;
|
|
708
|
+
|
|
709
|
+
if (x >= W || y >= params.height || c >= Cg) { return; }
|
|
710
|
+
|
|
711
|
+
let out_idx = y * W * Cg + x * Cg + c;
|
|
712
|
+
if (c < Ag) {
|
|
713
|
+
output_buf[out_idx] = input_a[y * W * Ag + x * Ag + c];
|
|
714
|
+
} else {
|
|
715
|
+
let c_b = c - Ag;
|
|
716
|
+
output_buf[out_idx] = input_b[y * W * Bg + x * Bg + c_b];
|
|
717
|
+
}
|
|
718
|
+
}
|
|
719
|
+
`;var L=class extends d{inputs;weights=[];output;dispatch;shader;constructor(t,e,r){super(t),this.shader=t.dtype==="f16"?We:Te;let a=e.c/4,o=r.c/4,s=a+o;this.output=t.tensor(e.h,e.w,e.c+r.c),this.inputs=[e,r],this.createUniform("params","Params"),this.setUniform("params",new Uint32Array([e.h,e.w,a,o,s,0,0,0])),this.defaultSetup(),this.dispatch=[Math.ceil(e.w/8),Math.ceil(e.h/8),s];}};var Me=`// Conv2d + skip add fused.
|
|
720
|
+
// Identical to conv2d.wgsl except skip is an activation input (binding 1),
|
|
721
|
+
// added element-wise to the conv result at write time.
|
|
722
|
+
// Eliminates the separate add dispatch and its intermediate buffer round-trip.
|
|
723
|
+
// Binding order: input(0), skip(1), weights(2), bias(3), params(4), output(5)
|
|
724
|
+
|
|
725
|
+
struct Params {
|
|
726
|
+
in_h : u32,
|
|
727
|
+
in_w : u32,
|
|
728
|
+
out_h : u32,
|
|
729
|
+
out_w : u32,
|
|
730
|
+
in_groups : u32,
|
|
731
|
+
out_groups : u32,
|
|
732
|
+
kernel_h : u32,
|
|
733
|
+
kernel_w : u32,
|
|
734
|
+
stride : u32,
|
|
735
|
+
pad_top : u32,
|
|
736
|
+
pad_left : u32,
|
|
737
|
+
activation : u32,
|
|
738
|
+
}
|
|
739
|
+
|
|
740
|
+
@group(0) @binding(0) var<storage, read> input_buf : array<vec4<f32>>;
|
|
741
|
+
@group(0) @binding(1) var<storage, read> skip_buf : array<vec4<f32>>;
|
|
742
|
+
@group(0) @binding(2) var<storage, read> weight_buf : array<mat4x4<f32>>;
|
|
743
|
+
@group(0) @binding(3) var<storage, read> bias_buf : array<vec4<f32>>;
|
|
744
|
+
@group(0) @binding(4) var<uniform> params : Params;
|
|
745
|
+
@group(0) @binding(5) var<storage, read_write> output_buf : array<vec4<f32>>;
|
|
746
|
+
|
|
747
|
+
@compute @workgroup_size(8, 8, 1)
|
|
748
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
749
|
+
let x = gid.x;
|
|
750
|
+
let y = gid.y;
|
|
751
|
+
let o = gid.z;
|
|
752
|
+
|
|
753
|
+
if (x >= params.out_w || y >= params.out_h || o >= params.out_groups) {
|
|
754
|
+
return;
|
|
755
|
+
}
|
|
756
|
+
|
|
757
|
+
let I = params.in_groups;
|
|
758
|
+
let O = params.out_groups;
|
|
759
|
+
|
|
760
|
+
var result = bias_buf[o];
|
|
761
|
+
|
|
762
|
+
for (var ky = 0u; ky < params.kernel_h; ky++) {
|
|
763
|
+
for (var kx = 0u; kx < params.kernel_w; kx++) {
|
|
764
|
+
let in_y_s = i32(y * params.stride + ky) - i32(params.pad_top);
|
|
765
|
+
let in_x_s = i32(x * params.stride + kx) - i32(params.pad_left);
|
|
766
|
+
|
|
767
|
+
if (in_y_s < 0 || in_x_s < 0 ||
|
|
768
|
+
u32(in_y_s) >= params.in_h || u32(in_x_s) >= params.in_w) {
|
|
769
|
+
continue;
|
|
770
|
+
}
|
|
771
|
+
|
|
772
|
+
let in_y = u32(in_y_s);
|
|
773
|
+
let in_x = u32(in_x_s);
|
|
774
|
+
let z = ky * params.kernel_w + kx;
|
|
775
|
+
|
|
776
|
+
for (var i = 0u; i < I; i++) {
|
|
777
|
+
let in_idx = in_y * params.in_w * I + in_x * I + i;
|
|
778
|
+
let w_idx = z * O * I + o * I + i;
|
|
779
|
+
result += weight_buf[w_idx] * input_buf[in_idx];
|
|
780
|
+
}
|
|
781
|
+
}
|
|
782
|
+
}
|
|
783
|
+
|
|
784
|
+
if (params.activation == 1u) {
|
|
785
|
+
result = clamp(result, vec4<f32>(0.0), vec4<f32>(6.0));
|
|
786
|
+
} else if (params.activation == 2u) {
|
|
787
|
+
result = max(result, vec4<f32>(0.0));
|
|
788
|
+
}
|
|
789
|
+
|
|
790
|
+
let out_idx = y * params.out_w * O + x * O + o;
|
|
791
|
+
output_buf[out_idx] = result + skip_buf[out_idx];
|
|
792
|
+
}
|
|
793
|
+
`;var Se=`enable f16;
|
|
794
|
+
|
|
795
|
+
// Conv2d + skip add fused \u2014 full f16 variant.
|
|
796
|
+
|
|
797
|
+
struct Params {
|
|
798
|
+
in_h : u32,
|
|
799
|
+
in_w : u32,
|
|
800
|
+
out_h : u32,
|
|
801
|
+
out_w : u32,
|
|
802
|
+
in_groups : u32,
|
|
803
|
+
out_groups : u32,
|
|
804
|
+
kernel_h : u32,
|
|
805
|
+
kernel_w : u32,
|
|
806
|
+
stride : u32,
|
|
807
|
+
pad_top : u32,
|
|
808
|
+
pad_left : u32,
|
|
809
|
+
activation : u32,
|
|
810
|
+
}
|
|
811
|
+
|
|
812
|
+
// Binding order matches conv2d_add.wgsl: input(0), skip(1), weight(2), bias(3),
|
|
813
|
+
// params(4), output(5).
|
|
814
|
+
@group(0) @binding(0) var<storage, read> input_buf : array<vec4<f16>>;
|
|
815
|
+
@group(0) @binding(1) var<storage, read> skip_buf : array<vec4<f16>>;
|
|
816
|
+
@group(0) @binding(2) var<storage, read> weight_buf : array<mat4x4<f16>>;
|
|
817
|
+
@group(0) @binding(3) var<storage, read> bias_buf : array<vec4<f16>>;
|
|
818
|
+
@group(0) @binding(4) var<uniform> params : Params;
|
|
819
|
+
@group(0) @binding(5) var<storage, read_write> output_buf : array<vec4<f16>>;
|
|
820
|
+
|
|
821
|
+
@compute @workgroup_size(8, 8, 1)
|
|
822
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
823
|
+
let x = gid.x;
|
|
824
|
+
let y = gid.y;
|
|
825
|
+
let o = gid.z;
|
|
826
|
+
|
|
827
|
+
if (x >= params.out_w || y >= params.out_h || o >= params.out_groups) {
|
|
828
|
+
return;
|
|
829
|
+
}
|
|
830
|
+
|
|
831
|
+
let I = params.in_groups;
|
|
832
|
+
let O = params.out_groups;
|
|
833
|
+
|
|
834
|
+
var result = bias_buf[o];
|
|
835
|
+
|
|
836
|
+
for (var ky = 0u; ky < params.kernel_h; ky++) {
|
|
837
|
+
for (var kx = 0u; kx < params.kernel_w; kx++) {
|
|
838
|
+
let in_y_s = i32(y * params.stride + ky) - i32(params.pad_top);
|
|
839
|
+
let in_x_s = i32(x * params.stride + kx) - i32(params.pad_left);
|
|
840
|
+
|
|
841
|
+
if (in_y_s < 0 || in_x_s < 0 ||
|
|
842
|
+
u32(in_y_s) >= params.in_h || u32(in_x_s) >= params.in_w) {
|
|
843
|
+
continue;
|
|
844
|
+
}
|
|
845
|
+
|
|
846
|
+
let in_y = u32(in_y_s);
|
|
847
|
+
let in_x = u32(in_x_s);
|
|
848
|
+
let z = ky * params.kernel_w + kx;
|
|
849
|
+
|
|
850
|
+
for (var i = 0u; i < I; i++) {
|
|
851
|
+
let in_idx = in_y * params.in_w * I + in_x * I + i;
|
|
852
|
+
let w_idx = z * O * I + o * I + i;
|
|
853
|
+
result += weight_buf[w_idx] * input_buf[in_idx];
|
|
854
|
+
}
|
|
855
|
+
}
|
|
856
|
+
}
|
|
857
|
+
|
|
858
|
+
if (params.activation == 1u) {
|
|
859
|
+
result = clamp(result, vec4<f16>(0.0h), vec4<f16>(6.0h));
|
|
860
|
+
} else if (params.activation == 2u) {
|
|
861
|
+
result = max(result, vec4<f16>(0.0h));
|
|
862
|
+
}
|
|
863
|
+
|
|
864
|
+
let out_idx = y * params.out_w * O + x * O + o;
|
|
865
|
+
output_buf[out_idx] = result + skip_buf[out_idx];
|
|
866
|
+
}
|
|
867
|
+
`;var H=class extends d{inputs;weights;output;dispatch;shader;constructor(t,e,r,a,o){super(t),this.shader=t.dtype==="f16"?Se:Me;let s=h(e.h,o.kernel,o.stride,o.padding),n=h(e.w,o.kernel,o.stride,o.padding),u=e.c/4,c=o.outChannels/4,p=_(o.padding,e.h,s,o.kernel,o.stride),l=_(o.padding,e.w,n,o.kernel,o.stride);this.output=t.tensor(s,n,o.outChannels),this.inputs=[e,r],this.weights=[t.upload(m(a.weights)),t.upload(m(a.bias))],this.createUniform("params","Params"),this.setUniform("params",new Uint32Array([e.h,e.w,s,n,u,c,o.kernel,o.kernel,o.stride,p,l,o.activation==="relu6"?1:0])),this.defaultSetup(),this.dispatch=[Math.ceil(n/8),Math.ceil(s/8),c];}};var Ae=`// Bilinear upsample + channel concat fused.
|
|
868
|
+
// input_a is the decoder tensor at small spatial resolution (in_h \xD7 in_w).
|
|
869
|
+
// input_b is the encoder skip feature already at output resolution (out_h \xD7 out_w).
|
|
870
|
+
// For output channels 0..a_groups-1: bilinearly interpolate from input_a.
|
|
871
|
+
// For output channels a_groups..out_groups-1: copy directly from input_b.
|
|
872
|
+
// Eliminates the intermediate upsample buffer and the separate concat dispatch.
|
|
873
|
+
|
|
874
|
+
struct Params {
|
|
875
|
+
in_h : u32, // input_a spatial height
|
|
876
|
+
in_w : u32, // input_a spatial width
|
|
877
|
+
out_h : u32,
|
|
878
|
+
out_w : u32,
|
|
879
|
+
a_groups : u32, // input_a channel groups (upsampled)
|
|
880
|
+
b_groups : u32, // input_b channel groups (encoder feature)
|
|
881
|
+
out_groups : u32, // a_groups + b_groups
|
|
882
|
+
_pad0 : u32,
|
|
883
|
+
}
|
|
884
|
+
|
|
885
|
+
@group(0) @binding(0) var<storage, read> input_a : array<vec4<f32>>;
|
|
886
|
+
@group(0) @binding(1) var<storage, read> input_b : array<vec4<f32>>;
|
|
887
|
+
@group(0) @binding(2) var<uniform> params : Params;
|
|
888
|
+
@group(0) @binding(3) var<storage, read_write> output_buf : array<vec4<f32>>;
|
|
889
|
+
|
|
890
|
+
@compute @workgroup_size(8, 8, 1)
|
|
891
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
892
|
+
let ox = gid.x;
|
|
893
|
+
let oy = gid.y;
|
|
894
|
+
let c = gid.z;
|
|
895
|
+
|
|
896
|
+
if (ox >= params.out_w || oy >= params.out_h || c >= params.out_groups) { return; }
|
|
897
|
+
|
|
898
|
+
let out_idx = oy * params.out_w * params.out_groups + ox * params.out_groups + c;
|
|
899
|
+
|
|
900
|
+
if (c < params.a_groups) {
|
|
901
|
+
let IH = params.in_h;
|
|
902
|
+
let IW = params.in_w;
|
|
903
|
+
let AG = params.a_groups;
|
|
904
|
+
|
|
905
|
+
let src_x = (f32(ox) + 0.5) * (f32(IW) / f32(params.out_w)) - 0.5;
|
|
906
|
+
let src_y = (f32(oy) + 0.5) * (f32(IH) / f32(params.out_h)) - 0.5;
|
|
907
|
+
|
|
908
|
+
let x0 = u32(clamp(i32(floor(src_x)), 0, i32(IW) - 1));
|
|
909
|
+
let x1 = u32(clamp(i32(floor(src_x)) + 1, 0, i32(IW) - 1));
|
|
910
|
+
let y0 = u32(clamp(i32(floor(src_y)), 0, i32(IH) - 1));
|
|
911
|
+
let y1 = u32(clamp(i32(floor(src_y)) + 1, 0, i32(IH) - 1));
|
|
912
|
+
|
|
913
|
+
let wx = src_x - floor(src_x);
|
|
914
|
+
let wy = src_y - floor(src_y);
|
|
915
|
+
|
|
916
|
+
let tl = input_a[y0 * IW * AG + x0 * AG + c];
|
|
917
|
+
let tr = input_a[y0 * IW * AG + x1 * AG + c];
|
|
918
|
+
let bl = input_a[y1 * IW * AG + x0 * AG + c];
|
|
919
|
+
let br = input_a[y1 * IW * AG + x1 * AG + c];
|
|
920
|
+
|
|
921
|
+
output_buf[out_idx] = (1.0 - wy) * ((1.0 - wx) * tl + wx * tr)
|
|
922
|
+
+ wy * ((1.0 - wx) * bl + wx * br);
|
|
923
|
+
} else {
|
|
924
|
+
let c_b = c - params.a_groups;
|
|
925
|
+
output_buf[out_idx] = input_b[oy * params.out_w * params.b_groups + ox * params.b_groups + c_b];
|
|
926
|
+
}
|
|
927
|
+
}
|
|
928
|
+
`;var Oe=`enable f16;
|
|
929
|
+
|
|
930
|
+
// Bilinear upsample + channel concat fused \u2014 full f16 variant.
|
|
931
|
+
// Bilinear weights (wx, wy) and intermediate blends computed in f32; result cast to f16.
|
|
932
|
+
|
|
933
|
+
struct Params {
|
|
934
|
+
in_h : u32,
|
|
935
|
+
in_w : u32,
|
|
936
|
+
out_h : u32,
|
|
937
|
+
out_w : u32,
|
|
938
|
+
a_groups : u32,
|
|
939
|
+
b_groups : u32,
|
|
940
|
+
out_groups : u32,
|
|
941
|
+
_pad0 : u32,
|
|
942
|
+
}
|
|
943
|
+
|
|
944
|
+
@group(0) @binding(0) var<storage, read> input_a : array<vec4<f16>>;
|
|
945
|
+
@group(0) @binding(1) var<storage, read> input_b : array<vec4<f16>>;
|
|
946
|
+
@group(0) @binding(2) var<uniform> params : Params;
|
|
947
|
+
@group(0) @binding(3) var<storage, read_write> output_buf : array<vec4<f16>>;
|
|
948
|
+
|
|
949
|
+
@compute @workgroup_size(8, 8, 1)
|
|
950
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
951
|
+
let ox = gid.x;
|
|
952
|
+
let oy = gid.y;
|
|
953
|
+
let c = gid.z;
|
|
954
|
+
|
|
955
|
+
if (ox >= params.out_w || oy >= params.out_h || c >= params.out_groups) { return; }
|
|
956
|
+
|
|
957
|
+
let out_idx = oy * params.out_w * params.out_groups + ox * params.out_groups + c;
|
|
958
|
+
|
|
959
|
+
if (c < params.a_groups) {
|
|
960
|
+
let IH = params.in_h;
|
|
961
|
+
let IW = params.in_w;
|
|
962
|
+
let AG = params.a_groups;
|
|
963
|
+
|
|
964
|
+
let src_x = (f32(ox) + 0.5) * (f32(IW) / f32(params.out_w)) - 0.5;
|
|
965
|
+
let src_y = (f32(oy) + 0.5) * (f32(IH) / f32(params.out_h)) - 0.5;
|
|
966
|
+
|
|
967
|
+
let x0 = u32(clamp(i32(floor(src_x)), 0, i32(IW) - 1));
|
|
968
|
+
let x1 = u32(clamp(i32(floor(src_x)) + 1, 0, i32(IW) - 1));
|
|
969
|
+
let y0 = u32(clamp(i32(floor(src_y)), 0, i32(IH) - 1));
|
|
970
|
+
let y1 = u32(clamp(i32(floor(src_y)) + 1, 0, i32(IH) - 1));
|
|
971
|
+
|
|
972
|
+
let wx = src_x - floor(src_x);
|
|
973
|
+
let wy = src_y - floor(src_y);
|
|
974
|
+
|
|
975
|
+
let tl = vec4<f32>(input_a[y0 * IW * AG + x0 * AG + c]);
|
|
976
|
+
let tr = vec4<f32>(input_a[y0 * IW * AG + x1 * AG + c]);
|
|
977
|
+
let bl = vec4<f32>(input_a[y1 * IW * AG + x0 * AG + c]);
|
|
978
|
+
let br = vec4<f32>(input_a[y1 * IW * AG + x1 * AG + c]);
|
|
979
|
+
|
|
980
|
+
let result = (1.0 - wy) * ((1.0 - wx) * tl + wx * tr)
|
|
981
|
+
+ wy * ((1.0 - wx) * bl + wx * br);
|
|
982
|
+
output_buf[out_idx] = vec4<f16>(result);
|
|
983
|
+
} else {
|
|
984
|
+
let c_b = c - params.a_groups;
|
|
985
|
+
output_buf[out_idx] = input_b[oy * params.out_w * params.b_groups + ox * params.b_groups + c_b];
|
|
986
|
+
}
|
|
987
|
+
}
|
|
988
|
+
`;var R=class extends d{inputs;weights=[];output;dispatch;shader;constructor(t,e,r,a){super(t),this.shader=t.dtype==="f16"?Oe:Ae;let o=e.c/4,s=r.c/4,n=o+s;this.output=t.tensor(a.outH,a.outW,e.c+r.c),this.inputs=[e,r],this.createUniform("params","Params"),this.setUniform("params",new Uint32Array([e.h,e.w,a.outH,a.outW,o,s,n,0])),this.defaultSetup(),this.dispatch=[Math.ceil(a.outW/8),Math.ceil(a.outH/8),n];}};var Ee=`// Bilinear upsample + 1\xD71 pointwise conv fused.
|
|
989
|
+
// For each output pixel, bilinearly samples the small input for each in_group,
|
|
990
|
+
// immediately applies the 1\xD71 conv weights, and writes the activated result.
|
|
991
|
+
// Eliminates the intermediate full-resolution upsample buffer.
|
|
992
|
+
|
|
993
|
+
struct Params {
|
|
994
|
+
in_h : u32,
|
|
995
|
+
in_w : u32,
|
|
996
|
+
out_h : u32,
|
|
997
|
+
out_w : u32,
|
|
998
|
+
in_groups : u32,
|
|
999
|
+
out_groups : u32,
|
|
1000
|
+
activation : u32,
|
|
1001
|
+
_pad0 : u32,
|
|
1002
|
+
}
|
|
1003
|
+
|
|
1004
|
+
@group(0) @binding(0) var<storage, read> input_buf : array<vec4<f32>>;
|
|
1005
|
+
@group(0) @binding(1) var<storage, read> weight_buf : array<mat4x4<f32>>;
|
|
1006
|
+
@group(0) @binding(2) var<storage, read> bias_buf : array<vec4<f32>>;
|
|
1007
|
+
@group(0) @binding(3) var<uniform> params : Params;
|
|
1008
|
+
@group(0) @binding(4) var<storage, read_write> output_buf : array<vec4<f32>>;
|
|
1009
|
+
|
|
1010
|
+
@compute @workgroup_size(8, 8, 1)
|
|
1011
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
1012
|
+
let ox = gid.x;
|
|
1013
|
+
let oy = gid.y;
|
|
1014
|
+
let og = gid.z;
|
|
1015
|
+
|
|
1016
|
+
if (ox >= params.out_w || oy >= params.out_h || og >= params.out_groups) { return; }
|
|
1017
|
+
|
|
1018
|
+
let IH = params.in_h;
|
|
1019
|
+
let IW = params.in_w;
|
|
1020
|
+
let IG = params.in_groups;
|
|
1021
|
+
|
|
1022
|
+
let src_x = (f32(ox) + 0.5) * (f32(IW) / f32(params.out_w)) - 0.5;
|
|
1023
|
+
let src_y = (f32(oy) + 0.5) * (f32(IH) / f32(params.out_h)) - 0.5;
|
|
1024
|
+
|
|
1025
|
+
let x0 = u32(clamp(i32(floor(src_x)), 0, i32(IW) - 1));
|
|
1026
|
+
let x1 = u32(clamp(i32(floor(src_x)) + 1, 0, i32(IW) - 1));
|
|
1027
|
+
let y0 = u32(clamp(i32(floor(src_y)), 0, i32(IH) - 1));
|
|
1028
|
+
let y1 = u32(clamp(i32(floor(src_y)) + 1, 0, i32(IH) - 1));
|
|
1029
|
+
|
|
1030
|
+
let wx = src_x - floor(src_x);
|
|
1031
|
+
let wy = src_y - floor(src_y);
|
|
1032
|
+
|
|
1033
|
+
var result = bias_buf[og];
|
|
1034
|
+
|
|
1035
|
+
for (var ig = 0u; ig < IG; ig++) {
|
|
1036
|
+
let tl = input_buf[y0 * IW * IG + x0 * IG + ig];
|
|
1037
|
+
let tr = input_buf[y0 * IW * IG + x1 * IG + ig];
|
|
1038
|
+
let bl = input_buf[y1 * IW * IG + x0 * IG + ig];
|
|
1039
|
+
let br = input_buf[y1 * IW * IG + x1 * IG + ig];
|
|
1040
|
+
|
|
1041
|
+
let sampled = (1.0 - wy) * ((1.0 - wx) * tl + wx * tr)
|
|
1042
|
+
+ wy * ((1.0 - wx) * bl + wx * br);
|
|
1043
|
+
|
|
1044
|
+
result += weight_buf[og * IG + ig] * sampled;
|
|
1045
|
+
}
|
|
1046
|
+
|
|
1047
|
+
if (params.activation == 1u) {
|
|
1048
|
+
result = clamp(result, vec4<f32>(0.0), vec4<f32>(6.0));
|
|
1049
|
+
} else if (params.activation == 2u) {
|
|
1050
|
+
result = result * clamp(result + 3.0, vec4<f32>(0.0), vec4<f32>(6.0)) / 6.0;
|
|
1051
|
+
}
|
|
1052
|
+
|
|
1053
|
+
output_buf[oy * params.out_w * params.out_groups + ox * params.out_groups + og] = result;
|
|
1054
|
+
}
|
|
1055
|
+
`;var Fe=`enable f16;
|
|
1056
|
+
|
|
1057
|
+
// Bilinear upsample + 1\xD71 pointwise conv fused \u2014 full f16 variant.
|
|
1058
|
+
// Bilinear weights and intermediate blends computed in f32; conv accumulation in f16.
|
|
1059
|
+
|
|
1060
|
+
struct Params {
|
|
1061
|
+
in_h : u32,
|
|
1062
|
+
in_w : u32,
|
|
1063
|
+
out_h : u32,
|
|
1064
|
+
out_w : u32,
|
|
1065
|
+
in_groups : u32,
|
|
1066
|
+
out_groups : u32,
|
|
1067
|
+
activation : u32,
|
|
1068
|
+
_pad0 : u32,
|
|
1069
|
+
}
|
|
1070
|
+
|
|
1071
|
+
@group(0) @binding(0) var<storage, read> input_buf : array<vec4<f16>>;
|
|
1072
|
+
@group(0) @binding(1) var<storage, read> weight_buf : array<mat4x4<f16>>;
|
|
1073
|
+
@group(0) @binding(2) var<storage, read> bias_buf : array<vec4<f16>>;
|
|
1074
|
+
@group(0) @binding(3) var<uniform> params : Params;
|
|
1075
|
+
@group(0) @binding(4) var<storage, read_write> output_buf : array<vec4<f16>>;
|
|
1076
|
+
|
|
1077
|
+
@compute @workgroup_size(8, 8, 1)
|
|
1078
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
1079
|
+
let ox = gid.x;
|
|
1080
|
+
let oy = gid.y;
|
|
1081
|
+
let og = gid.z;
|
|
1082
|
+
|
|
1083
|
+
if (ox >= params.out_w || oy >= params.out_h || og >= params.out_groups) { return; }
|
|
1084
|
+
|
|
1085
|
+
let IH = params.in_h;
|
|
1086
|
+
let IW = params.in_w;
|
|
1087
|
+
let IG = params.in_groups;
|
|
1088
|
+
|
|
1089
|
+
let src_x = (f32(ox) + 0.5) * (f32(IW) / f32(params.out_w)) - 0.5;
|
|
1090
|
+
let src_y = (f32(oy) + 0.5) * (f32(IH) / f32(params.out_h)) - 0.5;
|
|
1091
|
+
|
|
1092
|
+
let x0 = u32(clamp(i32(floor(src_x)), 0, i32(IW) - 1));
|
|
1093
|
+
let x1 = u32(clamp(i32(floor(src_x)) + 1, 0, i32(IW) - 1));
|
|
1094
|
+
let y0 = u32(clamp(i32(floor(src_y)), 0, i32(IH) - 1));
|
|
1095
|
+
let y1 = u32(clamp(i32(floor(src_y)) + 1, 0, i32(IH) - 1));
|
|
1096
|
+
|
|
1097
|
+
let wx = src_x - floor(src_x);
|
|
1098
|
+
let wy = src_y - floor(src_y);
|
|
1099
|
+
|
|
1100
|
+
var result = bias_buf[og];
|
|
1101
|
+
|
|
1102
|
+
for (var ig = 0u; ig < IG; ig++) {
|
|
1103
|
+
let tl = vec4<f32>(input_buf[y0 * IW * IG + x0 * IG + ig]);
|
|
1104
|
+
let tr = vec4<f32>(input_buf[y0 * IW * IG + x1 * IG + ig]);
|
|
1105
|
+
let bl = vec4<f32>(input_buf[y1 * IW * IG + x0 * IG + ig]);
|
|
1106
|
+
let br = vec4<f32>(input_buf[y1 * IW * IG + x1 * IG + ig]);
|
|
1107
|
+
|
|
1108
|
+
let sampled = vec4<f16>((1.0 - wy) * ((1.0 - wx) * tl + wx * tr)
|
|
1109
|
+
+ wy * ((1.0 - wx) * bl + wx * br));
|
|
1110
|
+
|
|
1111
|
+
result += weight_buf[og * IG + ig] * sampled;
|
|
1112
|
+
}
|
|
1113
|
+
|
|
1114
|
+
if (params.activation == 1u) {
|
|
1115
|
+
result = clamp(result, vec4<f16>(0.0h), vec4<f16>(6.0h));
|
|
1116
|
+
} else if (params.activation == 2u) {
|
|
1117
|
+
result = result * clamp(result + 3.0h, vec4<f16>(0.0h), vec4<f16>(6.0h)) / 6.0h;
|
|
1118
|
+
}
|
|
1119
|
+
|
|
1120
|
+
output_buf[oy * params.out_w * params.out_groups + ox * params.out_groups + og] = result;
|
|
1121
|
+
}
|
|
1122
|
+
`;var z=class extends d{inputs;weights;output;dispatch;shader;constructor(t,e,r,a){super(t),this.shader=t.dtype==="f16"?Fe:Ee;let o=e.c/4,s=a.outChannels/4;this.output=t.tensor(a.outH,a.outW,a.outChannels),this.inputs=[e],this.weights=[t.upload(m(r.weights)),t.upload(m(r.bias))],this.createUniform("params","Params"),this.setUniform("params",new Uint32Array([e.h,e.w,a.outH,a.outW,o,s,a.activation==="relu6"?1:0,0])),this.defaultSetup(),this.dispatch=[Math.ceil(a.outW/8),Math.ceil(a.outH/8),s];}};var Ve=`// Bilinear upsample + sigmoid fused.
|
|
1123
|
+
// Identical to bilinear_upsample.wgsl except sigmoid is applied at write time.
|
|
1124
|
+
// Eliminates the intermediate full-resolution buffer and the separate sigmoid dispatch.
|
|
1125
|
+
|
|
1126
|
+
struct Params {
|
|
1127
|
+
in_h : u32,
|
|
1128
|
+
in_w : u32,
|
|
1129
|
+
out_h : u32,
|
|
1130
|
+
out_w : u32,
|
|
1131
|
+
channel_groups : u32,
|
|
1132
|
+
_pad0 : u32,
|
|
1133
|
+
_pad1 : u32,
|
|
1134
|
+
_pad2 : u32,
|
|
1135
|
+
}
|
|
1136
|
+
|
|
1137
|
+
@group(0) @binding(0) var<storage, read> input_buf : array<vec4<f32>>;
|
|
1138
|
+
@group(0) @binding(1) var<uniform> params : Params;
|
|
1139
|
+
@group(0) @binding(2) var<storage, read_write> output_buf : array<vec4<f32>>;
|
|
1140
|
+
|
|
1141
|
+
@compute @workgroup_size(8, 8, 1)
|
|
1142
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
1143
|
+
let ox = gid.x;
|
|
1144
|
+
let oy = gid.y;
|
|
1145
|
+
let c = gid.z;
|
|
1146
|
+
|
|
1147
|
+
if (ox >= params.out_w || oy >= params.out_h || c >= params.channel_groups) {
|
|
1148
|
+
return;
|
|
1149
|
+
}
|
|
1150
|
+
|
|
1151
|
+
let IH = params.in_h;
|
|
1152
|
+
let IW = params.in_w;
|
|
1153
|
+
let C = params.channel_groups;
|
|
1154
|
+
|
|
1155
|
+
let src_x = (f32(ox) + 0.5) * (f32(IW) / f32(params.out_w)) - 0.5;
|
|
1156
|
+
let src_y = (f32(oy) + 0.5) * (f32(IH) / f32(params.out_h)) - 0.5;
|
|
1157
|
+
|
|
1158
|
+
let x0 = u32(clamp(i32(floor(src_x)), 0, i32(IW) - 1));
|
|
1159
|
+
let x1 = u32(clamp(i32(floor(src_x)) + 1, 0, i32(IW) - 1));
|
|
1160
|
+
let y0 = u32(clamp(i32(floor(src_y)), 0, i32(IH) - 1));
|
|
1161
|
+
let y1 = u32(clamp(i32(floor(src_y)) + 1, 0, i32(IH) - 1));
|
|
1162
|
+
|
|
1163
|
+
let wx = src_x - floor(src_x);
|
|
1164
|
+
let wy = src_y - floor(src_y);
|
|
1165
|
+
|
|
1166
|
+
let tl = input_buf[y0 * IW * C + x0 * C + c];
|
|
1167
|
+
let tr = input_buf[y0 * IW * C + x1 * C + c];
|
|
1168
|
+
let bl = input_buf[y1 * IW * C + x0 * C + c];
|
|
1169
|
+
let br = input_buf[y1 * IW * C + x1 * C + c];
|
|
1170
|
+
|
|
1171
|
+
let result = (1.0 - wy) * ((1.0 - wx) * tl + wx * tr)
|
|
1172
|
+
+ wy * ((1.0 - wx) * bl + wx * br);
|
|
1173
|
+
|
|
1174
|
+
output_buf[oy * params.out_w * C + ox * C + c] = 1.0 / (1.0 + exp(-result));
|
|
1175
|
+
}
|
|
1176
|
+
`;var Le=`enable f16;
|
|
1177
|
+
|
|
1178
|
+
// Bilinear upsample + sigmoid fused \u2014 full f16 variant.
|
|
1179
|
+
// Interpolation computed in f32 for accuracy; sigmoid and result cast to f16.
|
|
1180
|
+
|
|
1181
|
+
struct Params {
|
|
1182
|
+
in_h : u32,
|
|
1183
|
+
in_w : u32,
|
|
1184
|
+
out_h : u32,
|
|
1185
|
+
out_w : u32,
|
|
1186
|
+
channel_groups : u32,
|
|
1187
|
+
_pad0 : u32,
|
|
1188
|
+
_pad1 : u32,
|
|
1189
|
+
_pad2 : u32,
|
|
1190
|
+
}
|
|
1191
|
+
|
|
1192
|
+
@group(0) @binding(0) var<storage, read> input_buf : array<vec4<f16>>;
|
|
1193
|
+
@group(0) @binding(1) var<uniform> params : Params;
|
|
1194
|
+
@group(0) @binding(2) var<storage, read_write> output_buf : array<vec4<f16>>;
|
|
1195
|
+
|
|
1196
|
+
@compute @workgroup_size(8, 8, 1)
|
|
1197
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
1198
|
+
let ox = gid.x;
|
|
1199
|
+
let oy = gid.y;
|
|
1200
|
+
let c = gid.z;
|
|
1201
|
+
|
|
1202
|
+
if (ox >= params.out_w || oy >= params.out_h || c >= params.channel_groups) {
|
|
1203
|
+
return;
|
|
1204
|
+
}
|
|
1205
|
+
|
|
1206
|
+
let IH = params.in_h;
|
|
1207
|
+
let IW = params.in_w;
|
|
1208
|
+
let C = params.channel_groups;
|
|
1209
|
+
|
|
1210
|
+
let src_x = (f32(ox) + 0.5) * (f32(IW) / f32(params.out_w)) - 0.5;
|
|
1211
|
+
let src_y = (f32(oy) + 0.5) * (f32(IH) / f32(params.out_h)) - 0.5;
|
|
1212
|
+
|
|
1213
|
+
let x0 = u32(clamp(i32(floor(src_x)), 0, i32(IW) - 1));
|
|
1214
|
+
let x1 = u32(clamp(i32(floor(src_x)) + 1, 0, i32(IW) - 1));
|
|
1215
|
+
let y0 = u32(clamp(i32(floor(src_y)), 0, i32(IH) - 1));
|
|
1216
|
+
let y1 = u32(clamp(i32(floor(src_y)) + 1, 0, i32(IH) - 1));
|
|
1217
|
+
|
|
1218
|
+
let wx = src_x - floor(src_x);
|
|
1219
|
+
let wy = src_y - floor(src_y);
|
|
1220
|
+
|
|
1221
|
+
let tl = vec4<f32>(input_buf[y0 * IW * C + x0 * C + c]);
|
|
1222
|
+
let tr = vec4<f32>(input_buf[y0 * IW * C + x1 * C + c]);
|
|
1223
|
+
let bl = vec4<f32>(input_buf[y1 * IW * C + x0 * C + c]);
|
|
1224
|
+
let br = vec4<f32>(input_buf[y1 * IW * C + x1 * C + c]);
|
|
1225
|
+
|
|
1226
|
+
let result = (1.0 - wy) * ((1.0 - wx) * tl + wx * tr)
|
|
1227
|
+
+ wy * ((1.0 - wx) * bl + wx * br);
|
|
1228
|
+
|
|
1229
|
+
output_buf[oy * params.out_w * C + ox * C + c] = vec4<f16>(1.0 / (1.0 + exp(-result)));
|
|
1230
|
+
}
|
|
1231
|
+
`;var D=class extends d{inputs;weights=[];output;dispatch;shader;constructor(t,e,r){super(t),this.shader=t.dtype==="f16"?Le:Ve;let a=e.c/4;this.output=t.tensor(r.outH,r.outW,e.c),this.inputs=[e],this.createUniform("params","Params"),this.setUniform("params",new Uint32Array([e.h,e.w,r.outH,r.outW,a,0,0,0])),this.defaultSetup(),this.dispatch=[Math.ceil(r.outW/8),Math.ceil(r.outH/8),a];}};var He=`// Composite an RGBA image over a solid background, gated by a 1-ch alpha.
|
|
1232
|
+
// Fragment writes the canvas's swapchain texture (premultiplied output).
|
|
1233
|
+
//
|
|
1234
|
+
// Caller invariants (matched in CompositeSolidWebGPU):
|
|
1235
|
+
// - image and alpha are NHWC vec4 storage buffers, same h \xD7 w
|
|
1236
|
+
// - canvas.width === image.w, canvas.height === image.h (no resampling)
|
|
1237
|
+
|
|
1238
|
+
struct VertexOut {
|
|
1239
|
+
@builtin(position) pos: vec4<f32>,
|
|
1240
|
+
};
|
|
1241
|
+
|
|
1242
|
+
@vertex
|
|
1243
|
+
fn vs(@builtin(vertex_index) vi: u32) -> VertexOut {
|
|
1244
|
+
let verts = array<vec2<f32>, 6>(
|
|
1245
|
+
vec2(-1.0, -1.0), vec2( 1.0, -1.0), vec2(-1.0, 1.0),
|
|
1246
|
+
vec2(-1.0, 1.0), vec2( 1.0, -1.0), vec2( 1.0, 1.0),
|
|
1247
|
+
);
|
|
1248
|
+
var out: VertexOut;
|
|
1249
|
+
out.pos = vec4<f32>(verts[vi], 0.0, 1.0);
|
|
1250
|
+
return out;
|
|
1251
|
+
}
|
|
1252
|
+
|
|
1253
|
+
struct Params {
|
|
1254
|
+
width: u32, // image width in pixels (= canvas width)
|
|
1255
|
+
_pad0: u32,
|
|
1256
|
+
_pad1: u32,
|
|
1257
|
+
_pad2: u32,
|
|
1258
|
+
bgColor: vec4<f32>, // .rgb used; .a ignored
|
|
1259
|
+
};
|
|
1260
|
+
|
|
1261
|
+
@group(0) @binding(0) var<storage, read> image: array<vec4<f32>>;
|
|
1262
|
+
@group(0) @binding(1) var<storage, read> alpha: array<vec4<f32>>;
|
|
1263
|
+
@group(0) @binding(2) var<uniform> params: Params;
|
|
1264
|
+
|
|
1265
|
+
@fragment
|
|
1266
|
+
fn fs(in: VertexOut) -> @location(0) vec4<f32> {
|
|
1267
|
+
let x = u32(in.pos.x);
|
|
1268
|
+
let y = u32(in.pos.y);
|
|
1269
|
+
let i = y * params.width + x;
|
|
1270
|
+
|
|
1271
|
+
let fg = image[i].rgb;
|
|
1272
|
+
let a = alpha[i].r;
|
|
1273
|
+
let rgb = fg * a + params.bgColor.rgb * (1.0 - a);
|
|
1274
|
+
return vec4<f32>(rgb, 1.0);
|
|
1275
|
+
}
|
|
1276
|
+
`;var Re=`enable f16;
|
|
1277
|
+
|
|
1278
|
+
// Composite an RGBA image over a solid background, gated by a 1-ch alpha.
|
|
1279
|
+
// f16 variant: image and alpha are stored as f16; values promote to f32 on
|
|
1280
|
+
// read and the fragment writes f32 to the canvas swapchain (color attachment
|
|
1281
|
+
// format is the swapchain's preferred format, always f32-equivalent).
|
|
1282
|
+
//
|
|
1283
|
+
// Caller invariants (matched in CompositeSolidWebGPU):
|
|
1284
|
+
// - image and alpha are NHWC vec4 storage buffers, same h \xD7 w
|
|
1285
|
+
// - canvas.width === image.w, canvas.height === image.h (no resampling)
|
|
1286
|
+
|
|
1287
|
+
struct VertexOut {
|
|
1288
|
+
@builtin(position) pos: vec4<f32>,
|
|
1289
|
+
};
|
|
1290
|
+
|
|
1291
|
+
@vertex
|
|
1292
|
+
fn vs(@builtin(vertex_index) vi: u32) -> VertexOut {
|
|
1293
|
+
let verts = array<vec2<f32>, 6>(
|
|
1294
|
+
vec2(-1.0, -1.0), vec2( 1.0, -1.0), vec2(-1.0, 1.0),
|
|
1295
|
+
vec2(-1.0, 1.0), vec2( 1.0, -1.0), vec2( 1.0, 1.0),
|
|
1296
|
+
);
|
|
1297
|
+
var out: VertexOut;
|
|
1298
|
+
out.pos = vec4<f32>(verts[vi], 0.0, 1.0);
|
|
1299
|
+
return out;
|
|
1300
|
+
}
|
|
1301
|
+
|
|
1302
|
+
struct Params {
|
|
1303
|
+
width: u32,
|
|
1304
|
+
_pad0: u32,
|
|
1305
|
+
_pad1: u32,
|
|
1306
|
+
_pad2: u32,
|
|
1307
|
+
bgColor: vec4<f32>,
|
|
1308
|
+
};
|
|
1309
|
+
|
|
1310
|
+
@group(0) @binding(0) var<storage, read> image: array<vec4<f16>>;
|
|
1311
|
+
@group(0) @binding(1) var<storage, read> alpha: array<vec4<f16>>;
|
|
1312
|
+
@group(0) @binding(2) var<uniform> params: Params;
|
|
1313
|
+
|
|
1314
|
+
@fragment
|
|
1315
|
+
fn fs(in: VertexOut) -> @location(0) vec4<f32> {
|
|
1316
|
+
let x = u32(in.pos.x);
|
|
1317
|
+
let y = u32(in.pos.y);
|
|
1318
|
+
let i = y * params.width + x;
|
|
1319
|
+
|
|
1320
|
+
let fg = vec3<f32>(image[i].rgb);
|
|
1321
|
+
let a = f32(alpha[i].r);
|
|
1322
|
+
let rgb = fg * a + params.bgColor.rgb * (1.0 - a);
|
|
1323
|
+
return vec4<f32>(rgb, 1.0);
|
|
1324
|
+
}
|
|
1325
|
+
`;var N=class{constructor(t,e,r,a){this.backend=t;if(e.h!==r.h||e.w!==r.w)throw new Error(`CompositeSolid: image (${e.h}\xD7${e.w}) and alpha (${r.h}\xD7${r.w}) must match. Run the upscaler first.`);let o=t.device;this.uniformBuffer=o.createBuffer({size:32,usage:GPUBufferUsage.UNIFORM|GPUBufferUsage.COPY_DST});let s=new ArrayBuffer(32);new Uint32Array(s,0,1)[0]=e.w,new Float32Array(s,16,4).set([a[0],a[1],a[2],0]),o.queue.writeBuffer(this.uniformBuffer,0,s);let n=t.dtype==="f16"?Re:He,u=o.createShaderModule({code:n});this.pipeline=o.createRenderPipeline({layout:"auto",vertex:{module:u,entryPoint:"vs"},fragment:{module:u,entryPoint:"fs",targets:[{format:t.canvasFormat}]},primitive:{topology:"triangle-list"}}),this.bindGroup=o.createBindGroup({layout:this.pipeline.getBindGroupLayout(0),entries:[{binding:0,resource:{buffer:e.buffer}},{binding:1,resource:{buffer:r.buffer}},{binding:2,resource:{buffer:this.uniformBuffer}}]});}backend;pipeline;bindGroup;uniformBuffer;outputView=null;setOutput(t){this.outputView=t.createView();}run(){if(!this.outputView)throw new Error("CompositeSolidWebGPU.run() called before setOutput()");let t=this.backend.device.createCommandEncoder(),e=t.beginRenderPass({colorAttachments:[{view:this.outputView,clearValue:[0,0,0,1],loadOp:"clear",storeOp:"store"}]});e.setPipeline(this.pipeline),e.setBindGroup(0,this.bindGroup),e.draw(6),e.end(),this.backend.device.queue.submit([t.finish()]),this.outputView=null;}};var ze=`// Like composite_solid but bg is an NHWC vec4 storage buffer (e.g. virtual
|
|
1326
|
+
// background image, or a blurred copy of the input).
|
|
1327
|
+
|
|
1328
|
+
struct VertexOut {
|
|
1329
|
+
@builtin(position) pos: vec4<f32>,
|
|
1330
|
+
};
|
|
1331
|
+
|
|
1332
|
+
@vertex
|
|
1333
|
+
fn vs(@builtin(vertex_index) vi: u32) -> VertexOut {
|
|
1334
|
+
let verts = array<vec2<f32>, 6>(
|
|
1335
|
+
vec2(-1.0, -1.0), vec2( 1.0, -1.0), vec2(-1.0, 1.0),
|
|
1336
|
+
vec2(-1.0, 1.0), vec2( 1.0, -1.0), vec2( 1.0, 1.0),
|
|
1337
|
+
);
|
|
1338
|
+
var out: VertexOut;
|
|
1339
|
+
out.pos = vec4<f32>(verts[vi], 0.0, 1.0);
|
|
1340
|
+
return out;
|
|
1341
|
+
}
|
|
1342
|
+
|
|
1343
|
+
struct Params {
|
|
1344
|
+
width: u32,
|
|
1345
|
+
_pad0: u32,
|
|
1346
|
+
_pad1: u32,
|
|
1347
|
+
_pad2: u32,
|
|
1348
|
+
};
|
|
1349
|
+
|
|
1350
|
+
@group(0) @binding(0) var<storage, read> image: array<vec4<f32>>;
|
|
1351
|
+
@group(0) @binding(1) var<storage, read> alpha: array<vec4<f32>>;
|
|
1352
|
+
@group(0) @binding(2) var<uniform> params: Params;
|
|
1353
|
+
@group(0) @binding(3) var<storage, read> bg: array<vec4<f32>>;
|
|
1354
|
+
|
|
1355
|
+
@fragment
|
|
1356
|
+
fn fs(in: VertexOut) -> @location(0) vec4<f32> {
|
|
1357
|
+
let x = u32(in.pos.x);
|
|
1358
|
+
let y = u32(in.pos.y);
|
|
1359
|
+
let i = y * params.width + x;
|
|
1360
|
+
|
|
1361
|
+
let fg = image[i].rgb;
|
|
1362
|
+
let a = alpha[i].r;
|
|
1363
|
+
let bgc = bg[i].rgb;
|
|
1364
|
+
let rgb = fg * a + bgc * (1.0 - a);
|
|
1365
|
+
return vec4<f32>(rgb, 1.0);
|
|
1366
|
+
}
|
|
1367
|
+
`;var De=`enable f16;
|
|
1368
|
+
|
|
1369
|
+
// Like composite_solid_f16 but bg is an NHWC vec4<f16> storage buffer (e.g. a
|
|
1370
|
+
// virtual background image, or a blurred copy of the input).
|
|
1371
|
+
|
|
1372
|
+
struct VertexOut {
|
|
1373
|
+
@builtin(position) pos: vec4<f32>,
|
|
1374
|
+
};
|
|
1375
|
+
|
|
1376
|
+
@vertex
|
|
1377
|
+
fn vs(@builtin(vertex_index) vi: u32) -> VertexOut {
|
|
1378
|
+
let verts = array<vec2<f32>, 6>(
|
|
1379
|
+
vec2(-1.0, -1.0), vec2( 1.0, -1.0), vec2(-1.0, 1.0),
|
|
1380
|
+
vec2(-1.0, 1.0), vec2( 1.0, -1.0), vec2( 1.0, 1.0),
|
|
1381
|
+
);
|
|
1382
|
+
var out: VertexOut;
|
|
1383
|
+
out.pos = vec4<f32>(verts[vi], 0.0, 1.0);
|
|
1384
|
+
return out;
|
|
1385
|
+
}
|
|
1386
|
+
|
|
1387
|
+
struct Params {
|
|
1388
|
+
width: u32,
|
|
1389
|
+
_pad0: u32,
|
|
1390
|
+
_pad1: u32,
|
|
1391
|
+
_pad2: u32,
|
|
1392
|
+
};
|
|
1393
|
+
|
|
1394
|
+
@group(0) @binding(0) var<storage, read> image: array<vec4<f16>>;
|
|
1395
|
+
@group(0) @binding(1) var<storage, read> alpha: array<vec4<f16>>;
|
|
1396
|
+
@group(0) @binding(2) var<uniform> params: Params;
|
|
1397
|
+
@group(0) @binding(3) var<storage, read> bg: array<vec4<f16>>;
|
|
1398
|
+
|
|
1399
|
+
@fragment
|
|
1400
|
+
fn fs(in: VertexOut) -> @location(0) vec4<f32> {
|
|
1401
|
+
let x = u32(in.pos.x);
|
|
1402
|
+
let y = u32(in.pos.y);
|
|
1403
|
+
let i = y * params.width + x;
|
|
1404
|
+
|
|
1405
|
+
let fg = vec3<f32>(image[i].rgb);
|
|
1406
|
+
let a = f32(alpha[i].r);
|
|
1407
|
+
let bgc = vec3<f32>(bg[i].rgb);
|
|
1408
|
+
let rgb = fg * a + bgc * (1.0 - a);
|
|
1409
|
+
return vec4<f32>(rgb, 1.0);
|
|
1410
|
+
}
|
|
1411
|
+
`;var $=class{constructor(t,e,r,a){this.backend=t;if(e.h!==r.h||e.w!==r.w||e.h!==a.h||e.w!==a.w)throw new Error(`CompositeImage: image (${e.h}\xD7${e.w}), alpha (${r.h}\xD7${r.w}), and bg (${a.h}\xD7${a.w}) must all match. Run upscaler / resizer first.`);let o=t.device;this.uniformBuffer=o.createBuffer({size:16,usage:GPUBufferUsage.UNIFORM|GPUBufferUsage.COPY_DST});let s=new ArrayBuffer(16);new Uint32Array(s,0,1)[0]=e.w,o.queue.writeBuffer(this.uniformBuffer,0,s);let n=t.dtype==="f16"?De:ze,u=o.createShaderModule({code:n});this.pipeline=o.createRenderPipeline({layout:"auto",vertex:{module:u,entryPoint:"vs"},fragment:{module:u,entryPoint:"fs",targets:[{format:t.canvasFormat}]},primitive:{topology:"triangle-list"}}),this.bindGroup=o.createBindGroup({layout:this.pipeline.getBindGroupLayout(0),entries:[{binding:0,resource:{buffer:e.buffer}},{binding:1,resource:{buffer:r.buffer}},{binding:2,resource:{buffer:this.uniformBuffer}},{binding:3,resource:{buffer:a.buffer}}]});}backend;pipeline;bindGroup;uniformBuffer;outputView=null;setOutput(t){this.outputView=t.createView();}run(){if(!this.outputView)throw new Error("CompositeImageWebGPU.run() called before setOutput()");let t=this.backend.device.createCommandEncoder(),e=t.beginRenderPass({colorAttachments:[{view:this.outputView,clearValue:[0,0,0,1],loadOp:"clear",storeOp:"store"}]});e.setPipeline(this.pipeline),e.setBindGroup(0,this.bindGroup),e.draw(6),e.end(),this.backend.device.queue.submit([t.finish()]),this.outputView=null;}};var Ne=`// Like composite_image but bg is sampled bilinearly \u2014 bg may be smaller than
|
|
1412
|
+
// (image, alpha). Used by CompositorBlur to skip the final full-res upsample
|
|
1413
|
+
// in the blur pyramid and let this shader's own per-pixel scan do the
|
|
1414
|
+
// expansion (the work is already happening here for the composite anyway).
|
|
1415
|
+
//
|
|
1416
|
+
// Layout invariant: image, alpha share output dims (canvas h \xD7 w); bg is
|
|
1417
|
+
// at a smaller resolution (bg_h \xD7 bg_w). All NHWC vec4 storage buffers,
|
|
1418
|
+
// channelGroups = 1 (RGB padded to vec4).
|
|
1419
|
+
|
|
1420
|
+
struct VertexOut {
|
|
1421
|
+
@builtin(position) pos: vec4<f32>,
|
|
1422
|
+
};
|
|
1423
|
+
|
|
1424
|
+
@vertex
|
|
1425
|
+
fn vs(@builtin(vertex_index) vi: u32) -> VertexOut {
|
|
1426
|
+
let verts = array<vec2<f32>, 6>(
|
|
1427
|
+
vec2(-1.0, -1.0), vec2( 1.0, -1.0), vec2(-1.0, 1.0),
|
|
1428
|
+
vec2(-1.0, 1.0), vec2( 1.0, -1.0), vec2( 1.0, 1.0),
|
|
1429
|
+
);
|
|
1430
|
+
var out: VertexOut;
|
|
1431
|
+
out.pos = vec4<f32>(verts[vi], 0.0, 1.0);
|
|
1432
|
+
return out;
|
|
1433
|
+
}
|
|
1434
|
+
|
|
1435
|
+
struct Params {
|
|
1436
|
+
out_w: u32,
|
|
1437
|
+
out_h: u32,
|
|
1438
|
+
bg_w: u32,
|
|
1439
|
+
bg_h: u32,
|
|
1440
|
+
};
|
|
1441
|
+
|
|
1442
|
+
@group(0) @binding(0) var<storage, read> image: array<vec4<f32>>;
|
|
1443
|
+
@group(0) @binding(1) var<storage, read> alpha: array<vec4<f32>>;
|
|
1444
|
+
@group(0) @binding(2) var<uniform> params: Params;
|
|
1445
|
+
@group(0) @binding(3) var<storage, read> bg: array<vec4<f32>>;
|
|
1446
|
+
|
|
1447
|
+
@fragment
|
|
1448
|
+
fn fs(in: VertexOut) -> @location(0) vec4<f32> {
|
|
1449
|
+
let x = u32(in.pos.x);
|
|
1450
|
+
let y = u32(in.pos.y);
|
|
1451
|
+
let i = y * params.out_w + x;
|
|
1452
|
+
|
|
1453
|
+
let fg = image[i].rgb;
|
|
1454
|
+
let a = alpha[i].r;
|
|
1455
|
+
|
|
1456
|
+
// Bilinear sample bg at the corresponding location. align_corners=False.
|
|
1457
|
+
let src_x = (f32(x) + 0.5) * (f32(params.bg_w) / f32(params.out_w)) - 0.5;
|
|
1458
|
+
let src_y = (f32(y) + 0.5) * (f32(params.bg_h) / f32(params.out_h)) - 0.5;
|
|
1459
|
+
let x0 = u32(clamp(i32(floor(src_x)), 0, i32(params.bg_w) - 1));
|
|
1460
|
+
let x1 = u32(clamp(i32(floor(src_x)) + 1, 0, i32(params.bg_w) - 1));
|
|
1461
|
+
let y0 = u32(clamp(i32(floor(src_y)), 0, i32(params.bg_h) - 1));
|
|
1462
|
+
let y1 = u32(clamp(i32(floor(src_y)) + 1, 0, i32(params.bg_h) - 1));
|
|
1463
|
+
let wx = src_x - floor(src_x);
|
|
1464
|
+
let wy = src_y - floor(src_y);
|
|
1465
|
+
|
|
1466
|
+
let tl = bg[y0 * params.bg_w + x0].rgb;
|
|
1467
|
+
let tr = bg[y0 * params.bg_w + x1].rgb;
|
|
1468
|
+
let bl = bg[y1 * params.bg_w + x0].rgb;
|
|
1469
|
+
let br = bg[y1 * params.bg_w + x1].rgb;
|
|
1470
|
+
let bgc = (1.0 - wy) * ((1.0 - wx) * tl + wx * tr)
|
|
1471
|
+
+ wy * ((1.0 - wx) * bl + wx * br);
|
|
1472
|
+
|
|
1473
|
+
let rgb = fg * a + bgc * (1.0 - a);
|
|
1474
|
+
return vec4<f32>(rgb, 1.0);
|
|
1475
|
+
}
|
|
1476
|
+
`;var $e=`enable f16;
|
|
1477
|
+
|
|
1478
|
+
// f16 storage variant of composite_image_bilinear. bg is at a smaller
|
|
1479
|
+
// resolution and is bilinearly sampled to match (image, alpha) at full res.
|
|
1480
|
+
// Computation in f32; storage in f16.
|
|
1481
|
+
|
|
1482
|
+
struct VertexOut {
|
|
1483
|
+
@builtin(position) pos: vec4<f32>,
|
|
1484
|
+
};
|
|
1485
|
+
|
|
1486
|
+
@vertex
|
|
1487
|
+
fn vs(@builtin(vertex_index) vi: u32) -> VertexOut {
|
|
1488
|
+
let verts = array<vec2<f32>, 6>(
|
|
1489
|
+
vec2(-1.0, -1.0), vec2( 1.0, -1.0), vec2(-1.0, 1.0),
|
|
1490
|
+
vec2(-1.0, 1.0), vec2( 1.0, -1.0), vec2( 1.0, 1.0),
|
|
1491
|
+
);
|
|
1492
|
+
var out: VertexOut;
|
|
1493
|
+
out.pos = vec4<f32>(verts[vi], 0.0, 1.0);
|
|
1494
|
+
return out;
|
|
1495
|
+
}
|
|
1496
|
+
|
|
1497
|
+
struct Params {
|
|
1498
|
+
out_w: u32,
|
|
1499
|
+
out_h: u32,
|
|
1500
|
+
bg_w: u32,
|
|
1501
|
+
bg_h: u32,
|
|
1502
|
+
};
|
|
1503
|
+
|
|
1504
|
+
@group(0) @binding(0) var<storage, read> image: array<vec4<f16>>;
|
|
1505
|
+
@group(0) @binding(1) var<storage, read> alpha: array<vec4<f16>>;
|
|
1506
|
+
@group(0) @binding(2) var<uniform> params: Params;
|
|
1507
|
+
@group(0) @binding(3) var<storage, read> bg: array<vec4<f16>>;
|
|
1508
|
+
|
|
1509
|
+
@fragment
|
|
1510
|
+
fn fs(in: VertexOut) -> @location(0) vec4<f32> {
|
|
1511
|
+
let x = u32(in.pos.x);
|
|
1512
|
+
let y = u32(in.pos.y);
|
|
1513
|
+
let i = y * params.out_w + x;
|
|
1514
|
+
|
|
1515
|
+
let fg = vec3<f32>(image[i].rgb);
|
|
1516
|
+
let a = f32(alpha[i].r);
|
|
1517
|
+
|
|
1518
|
+
let src_x = (f32(x) + 0.5) * (f32(params.bg_w) / f32(params.out_w)) - 0.5;
|
|
1519
|
+
let src_y = (f32(y) + 0.5) * (f32(params.bg_h) / f32(params.out_h)) - 0.5;
|
|
1520
|
+
let x0 = u32(clamp(i32(floor(src_x)), 0, i32(params.bg_w) - 1));
|
|
1521
|
+
let x1 = u32(clamp(i32(floor(src_x)) + 1, 0, i32(params.bg_w) - 1));
|
|
1522
|
+
let y0 = u32(clamp(i32(floor(src_y)), 0, i32(params.bg_h) - 1));
|
|
1523
|
+
let y1 = u32(clamp(i32(floor(src_y)) + 1, 0, i32(params.bg_h) - 1));
|
|
1524
|
+
let wx = src_x - floor(src_x);
|
|
1525
|
+
let wy = src_y - floor(src_y);
|
|
1526
|
+
|
|
1527
|
+
let tl = vec3<f32>(bg[y0 * params.bg_w + x0].rgb);
|
|
1528
|
+
let tr = vec3<f32>(bg[y0 * params.bg_w + x1].rgb);
|
|
1529
|
+
let bl = vec3<f32>(bg[y1 * params.bg_w + x0].rgb);
|
|
1530
|
+
let br = vec3<f32>(bg[y1 * params.bg_w + x1].rgb);
|
|
1531
|
+
let bgc = (1.0 - wy) * ((1.0 - wx) * tl + wx * tr)
|
|
1532
|
+
+ wy * ((1.0 - wx) * bl + wx * br);
|
|
1533
|
+
|
|
1534
|
+
let rgb = fg * a + bgc * (1.0 - a);
|
|
1535
|
+
return vec4<f32>(rgb, 1.0);
|
|
1536
|
+
}
|
|
1537
|
+
`;var q=class{constructor(t,e,r,a){this.backend=t;if(e.h!==r.h||e.w!==r.w)throw new Error(`CompositeImageBilinear: image (${e.h}\xD7${e.w}) and alpha (${r.h}\xD7${r.w}) must match.`);let o=t.device;this.uniformBuffer=o.createBuffer({size:16,usage:GPUBufferUsage.UNIFORM|GPUBufferUsage.COPY_DST});let s=new ArrayBuffer(16),n=new Uint32Array(s);n[0]=e.w,n[1]=e.h,n[2]=a.w,n[3]=a.h,o.queue.writeBuffer(this.uniformBuffer,0,s);let u=t.dtype==="f16"?$e:Ne,c=o.createShaderModule({code:u});this.pipeline=o.createRenderPipeline({layout:"auto",vertex:{module:c,entryPoint:"vs"},fragment:{module:c,entryPoint:"fs",targets:[{format:t.canvasFormat}]},primitive:{topology:"triangle-list"}}),this.bindGroup=o.createBindGroup({layout:this.pipeline.getBindGroupLayout(0),entries:[{binding:0,resource:{buffer:e.buffer}},{binding:1,resource:{buffer:r.buffer}},{binding:2,resource:{buffer:this.uniformBuffer}},{binding:3,resource:{buffer:a.buffer}}]});}backend;pipeline;bindGroup;uniformBuffer;outputView=null;setOutput(t){this.outputView=t.createView();}run(){if(!this.outputView)throw new Error("CompositeImageBilinearWebGPU.run() called before setOutput()");let t=this.backend.device.createCommandEncoder(),e=t.beginRenderPass({colorAttachments:[{view:this.outputView,clearValue:[0,0,0,1],loadOp:"clear",storeOp:"store"}]});e.setPipeline(this.pipeline),e.setBindGroup(0,this.bindGroup),e.draw(6),e.end(),this.backend.device.queue.submit([t.finish()]),this.outputView=null;}};var qe=`// Passthrough "compositor" \u2014 writes the image directly to the canvas
|
|
1538
|
+
// swapchain texture with no alpha math and no background. Used by RenderOp
|
|
1539
|
+
// when the renderer is in disabled state, so the output canvas reflects
|
|
1540
|
+
// the unmodified input frame instead of a stale matted result.
|
|
1541
|
+
//
|
|
1542
|
+
// Caller invariants (matched in CompositePassthroughWebGPU):
|
|
1543
|
+
// - image is an NHWC vec4 storage buffer
|
|
1544
|
+
// - canvas.width === image.w, canvas.height === image.h (no resampling)
|
|
1545
|
+
|
|
1546
|
+
struct VertexOut {
|
|
1547
|
+
@builtin(position) pos: vec4<f32>,
|
|
1548
|
+
};
|
|
1549
|
+
|
|
1550
|
+
@vertex
|
|
1551
|
+
fn vs(@builtin(vertex_index) vi: u32) -> VertexOut {
|
|
1552
|
+
let verts = array<vec2<f32>, 6>(
|
|
1553
|
+
vec2(-1.0, -1.0), vec2( 1.0, -1.0), vec2(-1.0, 1.0),
|
|
1554
|
+
vec2(-1.0, 1.0), vec2( 1.0, -1.0), vec2( 1.0, 1.0),
|
|
1555
|
+
);
|
|
1556
|
+
var out: VertexOut;
|
|
1557
|
+
out.pos = vec4<f32>(verts[vi], 0.0, 1.0);
|
|
1558
|
+
return out;
|
|
1559
|
+
}
|
|
1560
|
+
|
|
1561
|
+
struct Params {
|
|
1562
|
+
width: u32,
|
|
1563
|
+
};
|
|
1564
|
+
|
|
1565
|
+
@group(0) @binding(0) var<storage, read> image: array<vec4<f32>>;
|
|
1566
|
+
@group(0) @binding(1) var<uniform> params: Params;
|
|
1567
|
+
|
|
1568
|
+
@fragment
|
|
1569
|
+
fn fs(in: VertexOut) -> @location(0) vec4<f32> {
|
|
1570
|
+
let x = u32(in.pos.x);
|
|
1571
|
+
let y = u32(in.pos.y);
|
|
1572
|
+
let i = y * params.width + x;
|
|
1573
|
+
return vec4<f32>(image[i].rgb, 1.0);
|
|
1574
|
+
}
|
|
1575
|
+
`;var je=`enable f16;
|
|
1576
|
+
|
|
1577
|
+
// Passthrough "compositor" \u2014 f16 variant. Image is stored as f16; values
|
|
1578
|
+
// promote to f32 on read and the fragment writes f32 to the canvas
|
|
1579
|
+
// swapchain (color attachment is the swapchain's preferred f32-equivalent
|
|
1580
|
+
// format). See composite_passthrough.wgsl for the f32 version.
|
|
1581
|
+
|
|
1582
|
+
struct VertexOut {
|
|
1583
|
+
@builtin(position) pos: vec4<f32>,
|
|
1584
|
+
};
|
|
1585
|
+
|
|
1586
|
+
@vertex
|
|
1587
|
+
fn vs(@builtin(vertex_index) vi: u32) -> VertexOut {
|
|
1588
|
+
let verts = array<vec2<f32>, 6>(
|
|
1589
|
+
vec2(-1.0, -1.0), vec2( 1.0, -1.0), vec2(-1.0, 1.0),
|
|
1590
|
+
vec2(-1.0, 1.0), vec2( 1.0, -1.0), vec2( 1.0, 1.0),
|
|
1591
|
+
);
|
|
1592
|
+
var out: VertexOut;
|
|
1593
|
+
out.pos = vec4<f32>(verts[vi], 0.0, 1.0);
|
|
1594
|
+
return out;
|
|
1595
|
+
}
|
|
1596
|
+
|
|
1597
|
+
struct Params {
|
|
1598
|
+
width: u32,
|
|
1599
|
+
};
|
|
1600
|
+
|
|
1601
|
+
@group(0) @binding(0) var<storage, read> image: array<vec4<f16>>;
|
|
1602
|
+
@group(0) @binding(1) var<uniform> params: Params;
|
|
1603
|
+
|
|
1604
|
+
@fragment
|
|
1605
|
+
fn fs(in: VertexOut) -> @location(0) vec4<f32> {
|
|
1606
|
+
let x = u32(in.pos.x);
|
|
1607
|
+
let y = u32(in.pos.y);
|
|
1608
|
+
let i = y * params.width + x;
|
|
1609
|
+
return vec4<f32>(vec3<f32>(image[i].rgb), 1.0);
|
|
1610
|
+
}
|
|
1611
|
+
`;var j=class{constructor(t,e){this.backend=t;let r=t.device;this.uniformBuffer=r.createBuffer({size:16,usage:GPUBufferUsage.UNIFORM|GPUBufferUsage.COPY_DST});let a=new ArrayBuffer(16);new Uint32Array(a,0,1)[0]=e.w,r.queue.writeBuffer(this.uniformBuffer,0,a);let o=t.dtype==="f16"?je:qe,s=r.createShaderModule({code:o});this.pipeline=r.createRenderPipeline({layout:"auto",vertex:{module:s,entryPoint:"vs"},fragment:{module:s,entryPoint:"fs",targets:[{format:t.canvasFormat}]},primitive:{topology:"triangle-list"}}),this.bindGroup=r.createBindGroup({layout:this.pipeline.getBindGroupLayout(0),entries:[{binding:0,resource:{buffer:e.buffer}},{binding:1,resource:{buffer:this.uniformBuffer}}]});}backend;pipeline;bindGroup;uniformBuffer;outputView=null;setOutput(t){this.outputView=t.createView();}run(){if(!this.outputView)throw new Error("CompositePassthroughWebGPU.run() called before setOutput()");let t=this.backend.device.createCommandEncoder(),e=t.beginRenderPass({colorAttachments:[{view:this.outputView,clearValue:[0,0,0,1],loadOp:"clear",storeOp:"store"}]});e.setPipeline(this.pipeline),e.setBindGroup(0,this.bindGroup),e.draw(6),e.end(),this.backend.device.queue.submit([t.finish()]),this.outputView=null;}};var Ye=`// Input op \u2014 sample a regular 2D source texture (RGBA8 unorm) into the NHWC
|
|
1612
|
+
// vec4<f32> output buffer at the target resolution. Used for ImageBitmap
|
|
1613
|
+
// sources, which are uploaded into a persistent staging texture via
|
|
1614
|
+
// copyExternalImageToTexture before each dispatch.
|
|
1615
|
+
|
|
1616
|
+
struct Params {
|
|
1617
|
+
out_w : u32,
|
|
1618
|
+
out_h : u32,
|
|
1619
|
+
_pad0 : u32,
|
|
1620
|
+
_pad1 : u32,
|
|
1621
|
+
}
|
|
1622
|
+
|
|
1623
|
+
@group(0) @binding(0) var src_tex : texture_2d<f32>;
|
|
1624
|
+
@group(0) @binding(1) var src_sampler : sampler;
|
|
1625
|
+
@group(0) @binding(2) var<uniform> params : Params;
|
|
1626
|
+
@group(0) @binding(3) var<storage, read_write> output_buf : array<vec4<f32>>;
|
|
1627
|
+
|
|
1628
|
+
@compute @workgroup_size(8, 8, 1)
|
|
1629
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
1630
|
+
let x = gid.x;
|
|
1631
|
+
let y = gid.y;
|
|
1632
|
+
if (x >= params.out_w || y >= params.out_h) { return; }
|
|
1633
|
+
|
|
1634
|
+
let uv = vec2<f32>(
|
|
1635
|
+
(f32(x) + 0.5) / f32(params.out_w),
|
|
1636
|
+
(f32(y) + 0.5) / f32(params.out_h),
|
|
1637
|
+
);
|
|
1638
|
+
let rgba = textureSampleLevel(src_tex, src_sampler, uv, 0.0);
|
|
1639
|
+
output_buf[y * params.out_w + x] = rgba;
|
|
1640
|
+
}
|
|
1641
|
+
`;var Ke=`enable f16;
|
|
1642
|
+
|
|
1643
|
+
// Input op \u2014 f16 storage variant. Source is sampled as vec4<f32>; only the
|
|
1644
|
+
// store is demoted to vec4<f16>.
|
|
1645
|
+
|
|
1646
|
+
struct Params {
|
|
1647
|
+
out_w : u32,
|
|
1648
|
+
out_h : u32,
|
|
1649
|
+
_pad0 : u32,
|
|
1650
|
+
_pad1 : u32,
|
|
1651
|
+
}
|
|
1652
|
+
|
|
1653
|
+
@group(0) @binding(0) var src_tex : texture_2d<f32>;
|
|
1654
|
+
@group(0) @binding(1) var src_sampler : sampler;
|
|
1655
|
+
@group(0) @binding(2) var<uniform> params : Params;
|
|
1656
|
+
@group(0) @binding(3) var<storage, read_write> output_buf : array<vec4<f16>>;
|
|
1657
|
+
|
|
1658
|
+
@compute @workgroup_size(8, 8, 1)
|
|
1659
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
1660
|
+
let x = gid.x;
|
|
1661
|
+
let y = gid.y;
|
|
1662
|
+
if (x >= params.out_w || y >= params.out_h) { return; }
|
|
1663
|
+
|
|
1664
|
+
let uv = vec2<f32>(
|
|
1665
|
+
(f32(x) + 0.5) / f32(params.out_w),
|
|
1666
|
+
(f32(y) + 0.5) / f32(params.out_h),
|
|
1667
|
+
);
|
|
1668
|
+
let rgba = textureSampleLevel(src_tex, src_sampler, uv, 0.0);
|
|
1669
|
+
output_buf[y * params.out_w + x] = vec4<f16>(rgba);
|
|
1670
|
+
}
|
|
1671
|
+
`;var Je=`// Input op \u2014 sample a GPUExternalTexture (zero-copy VideoFrame import) into
|
|
1672
|
+
// the NHWC vec4<f32> output buffer at the target resolution. Used when the
|
|
1673
|
+
// caller passes a VideoFrame to setSource(); requires importExternalTexture
|
|
1674
|
+
// to be called in the same task as the dispatch (the texture is invalidated
|
|
1675
|
+
// after the current task completes).
|
|
1676
|
+
|
|
1677
|
+
struct Params {
|
|
1678
|
+
out_w : u32,
|
|
1679
|
+
out_h : u32,
|
|
1680
|
+
_pad0 : u32,
|
|
1681
|
+
_pad1 : u32,
|
|
1682
|
+
}
|
|
1683
|
+
|
|
1684
|
+
@group(0) @binding(0) var src_tex : texture_external;
|
|
1685
|
+
@group(0) @binding(1) var src_sampler : sampler;
|
|
1686
|
+
@group(0) @binding(2) var<uniform> params : Params;
|
|
1687
|
+
@group(0) @binding(3) var<storage, read_write> output_buf : array<vec4<f32>>;
|
|
1688
|
+
|
|
1689
|
+
@compute @workgroup_size(8, 8, 1)
|
|
1690
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
1691
|
+
let x = gid.x;
|
|
1692
|
+
let y = gid.y;
|
|
1693
|
+
if (x >= params.out_w || y >= params.out_h) { return; }
|
|
1694
|
+
|
|
1695
|
+
let uv = vec2<f32>(
|
|
1696
|
+
(f32(x) + 0.5) / f32(params.out_w),
|
|
1697
|
+
(f32(y) + 0.5) / f32(params.out_h),
|
|
1698
|
+
);
|
|
1699
|
+
let rgba = textureSampleBaseClampToEdge(src_tex, src_sampler, uv);
|
|
1700
|
+
output_buf[y * params.out_w + x] = rgba;
|
|
1701
|
+
}
|
|
1702
|
+
`;var Xe=`enable f16;
|
|
1703
|
+
|
|
1704
|
+
// Input op \u2014 texture_external + f16 storage. Sample in f32, store as f16.
|
|
1705
|
+
|
|
1706
|
+
struct Params {
|
|
1707
|
+
out_w : u32,
|
|
1708
|
+
out_h : u32,
|
|
1709
|
+
_pad0 : u32,
|
|
1710
|
+
_pad1 : u32,
|
|
1711
|
+
}
|
|
1712
|
+
|
|
1713
|
+
@group(0) @binding(0) var src_tex : texture_external;
|
|
1714
|
+
@group(0) @binding(1) var src_sampler : sampler;
|
|
1715
|
+
@group(0) @binding(2) var<uniform> params : Params;
|
|
1716
|
+
@group(0) @binding(3) var<storage, read_write> output_buf : array<vec4<f16>>;
|
|
1717
|
+
|
|
1718
|
+
@compute @workgroup_size(8, 8, 1)
|
|
1719
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
1720
|
+
let x = gid.x;
|
|
1721
|
+
let y = gid.y;
|
|
1722
|
+
if (x >= params.out_w || y >= params.out_h) { return; }
|
|
1723
|
+
|
|
1724
|
+
let uv = vec2<f32>(
|
|
1725
|
+
(f32(x) + 0.5) / f32(params.out_w),
|
|
1726
|
+
(f32(y) + 0.5) / f32(params.out_h),
|
|
1727
|
+
);
|
|
1728
|
+
let rgba = textureSampleBaseClampToEdge(src_tex, src_sampler, uv);
|
|
1729
|
+
output_buf[y * params.out_w + x] = vec4<f16>(rgba);
|
|
1730
|
+
}
|
|
1731
|
+
`;var Yt=i=>typeof VideoFrame<"u"&&i instanceof VideoFrame,Y=class{output;device;dtype;sampler;uniformBuffer;dispatch;pipeline2d=null;pipelineExternal=null;stagingTex=null;stagingW=0;stagingH=0;source=null;constructor(t,e,r){this.device=t.device,this.dtype=t.dtype,this.output=t.tensor(e,r,4),this.sampler=this.device.createSampler({magFilter:"linear",minFilter:"linear",addressModeU:"clamp-to-edge",addressModeV:"clamp-to-edge"}),this.uniformBuffer=this.device.createBuffer({size:16,usage:GPUBufferUsage.UNIFORM|GPUBufferUsage.COPY_DST});let a=new ArrayBuffer(16);new Uint32Array(a,0,2).set([r,e]),this.device.queue.writeBuffer(this.uniformBuffer,0,a),this.dispatch=[Math.ceil(r/8),Math.ceil(e/8),1],this.pipeline2d=this.buildPipeline(this.dtype==="f16"?Ke:Ye);}setSource(t){this.source=t;}run(){if(!this.source)throw new Error("InputWebGPU.run() called before setSource()");Yt(this.source)?this.runExternal(this.source):this.run2d(this.source);}run2d(t){this.ensureStagingTexture(t.width,t.height),this.device.queue.copyExternalImageToTexture({source:t,flipY:false},{texture:this.stagingTex},[t.width,t.height]);let e=this.device.createBindGroup({layout:this.pipeline2d.getBindGroupLayout(0),entries:[{binding:0,resource:this.stagingTex.createView()},{binding:1,resource:this.sampler},{binding:2,resource:{buffer:this.uniformBuffer}},{binding:3,resource:{buffer:this.output.buffer}}]});this.dispatchOnce(this.pipeline2d,e);}runExternal(t){this.pipelineExternal||(this.pipelineExternal=this.buildPipeline(this.dtype==="f16"?Xe:Je));let e=this.device.importExternalTexture({source:t}),r=this.device.createBindGroup({layout:this.pipelineExternal.getBindGroupLayout(0),entries:[{binding:0,resource:e},{binding:1,resource:this.sampler},{binding:2,resource:{buffer:this.uniformBuffer}},{binding:3,resource:{buffer:this.output.buffer}}]});this.dispatchOnce(this.pipelineExternal,r);}ensureStagingTexture(t,e){this.stagingTex&&this.stagingW===t&&this.stagingH===e||(this.stagingTex?.destroy(),this.stagingTex=this.device.createTexture({size:[t,e,1],format:"rgba8unorm",usage:GPUTextureUsage.TEXTURE_BINDING|GPUTextureUsage.COPY_DST|GPUTextureUsage.RENDER_ATTACHMENT}),this.stagingW=t,this.stagingH=e);}buildPipeline(t){let e=this.device.createShaderModule({code:t});return this.device.createComputePipeline({layout:"auto",compute:{module:e,entryPoint:"main"}})}dispatchOnce(t,e){let r=this.device.createCommandEncoder(),a=r.beginComputePass();a.setPipeline(t),a.setBindGroup(0,e),a.dispatchWorkgroups(...this.dispatch),a.end(),this.device.queue.submit([r.finish()]);}};var Kt=navigator.gpu?GPUBufferUsage.STORAGE|GPUBufferUsage.COPY_SRC|GPUBufferUsage.COPY_DST:0,te=class i{constructor(t,e,r){this.device=t;this.canvas=e;this.dtype=r;this.bytesPerElement=r==="f16"?2:4;let a=e.getContext("webgpu");if(!a)throw new Error("Failed to get WebGPU context from canvas");this.canvasContext=a,this.canvasFormat=navigator.gpu.getPreferredCanvasFormat(),this.canvasContext.configure({device:t,format:this.canvasFormat,alphaMode:"premultiplied"}),this.ops={Conv2d:(o,s,n)=>new S(this,o,s,n),DepthwiseConv2d:(o,s,n)=>new A(this,o,s,n),Add:(o,s)=>new O(this,o,s),Sigmoid:o=>new E(this,o),BilinearUpsample:(o,s)=>new F(this,o,s),BicubicUpsample:(o,s)=>new V(this,o,s),ChannelConcat:(o,s)=>new L(this,o,s),Conv2dAdd:(o,s,n,u)=>new H(this,o,s,n,u),UpsampleConcat:(o,s,n)=>new R(this,o,s,n),UpsampleConv1x1:(o,s,n)=>new z(this,o,s,n),UpsampleSigmoid:(o,s)=>new D(this,o,s),Input:(o,s)=>new Y(this,o,s)},this.presenters={CompositeSolid:(o,s,n)=>{let u=new N(this,o,s,n);return {run:()=>{u.setOutput(this.getCurrentDisplayTexture()),u.run();}}},CompositeImage:(o,s,n)=>{let u=new $(this,o,s,n);return {run:()=>{u.setOutput(this.getCurrentDisplayTexture()),u.run();}}},CompositeImageBilinear:(o,s,n)=>{let u=new q(this,o,s,n);return {run:()=>{u.setOutput(this.getCurrentDisplayTexture()),u.run();}}},CompositePassthrough:o=>{let s=new j(this,o);return {run:()=>{s.setOutput(this.getCurrentDisplayTexture()),s.run();}}}};}device;canvas;dtype;ops;presenters;canvasContext;canvasFormat;bytesPerElement;static async isAvailable(){return navigator.gpu?await navigator.gpu.requestAdapter()!==null:false}static async hasF16Support(){if(!navigator.gpu)return false;let t=await navigator.gpu.requestAdapter();return t?t.features.has("shader-f16"):false}static async create(t){let e=t.dtype??"f32",r=t.device;if(r){if(e==="f16"&&!r.features.has("shader-f16"))throw new Error("WebGPU dtype='f16' requested but supplied device lacks `shader-f16`")}else {let a=await navigator.gpu.requestAdapter();if(!a)throw new Error("WebGPU adapter not available");if(e==="f16"&&!a.features.has("shader-f16"))throw new Error("WebGPU dtype='f16' requested but adapter lacks `shader-f16` feature");r=await a.requestDevice({requiredFeatures:e==="f16"?["shader-f16"]:[]});}return new i(r,t.canvas,e)}getCurrentDisplayTexture(){return this.canvasContext.getCurrentTexture()}tensor(t,e,r,a){let s=t*e*r*this.bytesPerElement,n=this.device.createBuffer({size:s,usage:Kt,mappedAtCreation:a!==void 0});if(a!==void 0){let u=n.getMappedRange();this.writeView(u,a),n.unmap();}return {h:t,w:e,c:r,buffer:n}}upload(t){let r=t.length*this.bytesPerElement,a=this.device.createBuffer({size:r,usage:GPUBufferUsage.STORAGE|GPUBufferUsage.COPY_DST,mappedAtCreation:true});return this.writeView(a.getMappedRange(),t),a.unmap(),{buffer:a}}writeView(t,e){let r=this.dtype==="f16",a=e instanceof Uint16Array;if(r===a){r?new Uint16Array(t).set(e):new Float32Array(t).set(e);return}r?new Uint16Array(t).set(he(e)):new Float32Array(t).set(ee(e));}async readback(t){let e=this.device.createBuffer({size:t.buffer.size,usage:GPUBufferUsage.COPY_DST|GPUBufferUsage.MAP_READ}),r=this.device.createCommandEncoder();r.copyBufferToBuffer(t.buffer,0,e,0,t.buffer.size),this.device.queue.submit([r.finish()]),await e.mapAsync(GPUMapMode.READ);let a=e.getMappedRange(),o=this.dtype==="f16"?ee(new Uint16Array(a.slice(0))):new Float32Array(a.slice(0));return e.unmap(),e.destroy(),o}async sync(){await this.device.queue.onSubmittedWorkDone();}destroy(){this.device.destroy();}};export{Z as EffectsPipeline,te as WebGPUBackend};
|