diff --git a/.gitignore b/.gitignore index edb60d0..f6c65be 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ node_modules dist dist-example +dist-benchmark .vite coverage *.log diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..70c1228 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,46 @@ +# AGENTS.md + +## Project Overview +- **wgpu-executorch**: A high-performance WebGPU runtime and WGSL compute engine for ExecuTorch ML inference in the browser. +- **Tech Stack**: TypeScript, WebGPU / WGSL, Vitest, Vite, ESLint, Prettier, Lefthook. + +## Commands & Workflows +- **Build / Bundle**: `npm run build` +- **Typecheck**: `npm run typecheck` (`tsc --noEmit`) +- **Lint & Fix**: `npm run lint` / `npx eslint --fix .` +- **Full Test Suite**: `npm test` (`vitest run` — 27 test suites, 112+ tests) +- **Run Single Test**: `npx vitest run ` +- **Benchmarks**: `npm run benchmark:web` (starts local benchmark server on `http://localhost:5178/`) +- **Process Safety Guardrail**: NEVER run `taskkill`, `pkill`, or kill external host browser processes (`chrome.exe`, `msedge.exe`). To test or benchmark in browser environments, spin up a dedicated local HTTP / Vite server (`npm run benchmark:web`) with an auto-reporting endpoint (`POST /api/results`), allowing the user to navigate to the URL without process disruption. +- **Git Branching**: All performance and kernel work must be developed on the active feature branch (e.g., `perf/kernel-optimizations`) and pushed to remote origin regularly. + +## Architectural Boundaries +- **Core vs. Shaders**: + - Code changes outside `src/shaders/` must be minimal, focused, and well justified. + - Preserve core API style and existing contracts (`src/tensor.ts`, `src/command.ts`, `src/program.ts`, `src/context.ts`). +- **Shader Modularization**: + - Every operator resides in a dedicated operator subdirectory (`src/shaders//index.ts`). + - Submodules must follow the naming pattern `.wgsl.ts` (e.g. `conv2d-tile64.wgsl.ts`, `mm-64x64.wgsl.ts`). + - Submodule files must export strictly the WGSL code template string (`export const SHADER_... = /* wgsl */ \`...\`;`) and include an MIT license and citation note (`https://github.com/barhanc/wgpu-executorch`). + - Driver files (`src/shaders//index.ts`) must have an empty line between standard imports and shader imports, and between shader imports and the `const name = ...` definition. + - In WGSL compute shaders, loop index variables must be declared with mutable `var` (e.g., `for (var i = 0u; ...)`), never `let`. + +## WebGPU Alignment & Buffer Rules +- **4-Byte Alignment Invariant**: + - WebGPU `copyBufferToBuffer` requires `size`, `sourceOffset`, and `destinationOffset` to be multiples of 4 bytes. + - When copying float16 or odd-sized tensor buffers, size must be padded using `Math.ceil(size / 4) * 4`. +- **Node.js FP16 Support**: + - Node.js LTS environments lack global `Float16Array`; polyfill via `@petamoriken/float16` conditionally in test setups. + +## WebGPU Kernel Performance Guidelines +- **2D Register Tiling over Naive Spatial Barrier Loops**: + - Avoid per-channel workgroup loops with frequent `workgroupBarrier()` calls. + - Prefer 2D register-tiled GEMM architectures (e.g., 64x64 output blocks with 8x8 register accumulators per thread) to maximize arithmetic intensity and minimize barrier synchronization. +- **128-bit Vectorized Memory Transactions**: + - Use `vec4` (or `vec4`) storage buffer reads for aligned tensor buffers to saturate the GPU memory bus with 128-bit memory transactions. +- **Bitwise Indexing**: + - Replace hardware integer division (`/`) and modulo (`%`) with bitwise shifts (`>>`, `<<`) and bitwise masks (`&`) whenever tile or tensor dimensions are powers of two. + +## Benchmarking & Verification Discipline +- **Zero Regressions**: All 25 test suites must pass 100% green against PyTorch reference outputs before and after optimizations. +- **Repeatable Benchmarking**: In `WgpuExecutionContext`, record the program/shaders ONCE outside the benchmark loop. Inside the measurement loop, only execute `ctx.submit()`, followed by `await ctx.sync()`. Never call `recordShader` or `addDispatch` inside the timed loop to avoid command list accumulation. diff --git a/benchmark-results.json b/benchmark-results.json new file mode 100644 index 0000000..f85288c --- /dev/null +++ b/benchmark-results.json @@ -0,0 +1,220 @@ +{ + "gpu": { + "vendor": "nvidia", + "architecture": "turing", + "device": "", + "hasF16": true + }, + "results": [ + { + "category": "End-to-End Model", + "name": "DepthART Monocular Depth", + "shape": "[1,3,448,448] → [1,1,448,448]", + "baselineMs": 28.299999952316284, + "fp32Ms": 20.5, + "fp16Ms": 18.200000047683716, + "speedupFp32": "1.38x", + "speedupFp16": "1.55x" + }, + { + "category": "End-to-End Model", + "name": "YOLO26-N Object Detection", + "shape": "[1,3,640,640] → [1,84,8400]", + "baselineMs": 24.699999928474426, + "fp32Ms": 15.5, + "fp16Ms": 14, + "speedupFp32": "1.59x", + "speedupFp16": "1.76x" + }, + { + "category": "End-to-End Model", + "name": "Selfie Segmentation", + "shape": "[1,3,256,256] → [1,1,256,256]", + "baselineMs": 3.100000023841858, + "fp32Ms": 3.200000047683716, + "fp16Ms": 3.100000023841858, + "speedupFp32": "0.97x", + "speedupFp16": "1.00x" + }, + { + "category": "Linear Algebra", + "name": "MatMul (MM Vec4)", + "shape": "[256, 256] * [256, 256]", + "baselineMs": 0.2, + "fp32Ms": 0.18, + "fp16Ms": 0.16399999618530273, + "speedupFp32": "1.11x", + "speedupFp16": "1.22x" + }, + { + "category": "Linear Algebra", + "name": "MatMul (MM Vec4)", + "shape": "[512, 512] * [512, 512]", + "baselineMs": 0.275, + "fp32Ms": 0.2300000011920929, + "fp16Ms": 0.1800000011920929, + "speedupFp32": "1.20x", + "speedupFp16": "1.53x" + }, + { + "category": "Linear Algebra", + "name": "MatMul (MM 64x64)", + "shape": "[1024, 1024] * [1024, 1024]", + "baselineMs": 1.6, + "fp32Ms": 1.5866666634877522, + "fp16Ms": 1.0866666714350381, + "speedupFp32": "1.01x", + "speedupFp16": "1.47x" + }, + { + "category": "Linear Algebra", + "name": "MatMul (MM 64x64)", + "shape": "[2048, 2048] * [2048, 2048]", + "baselineMs": 12.849999994039536, + "fp32Ms": 11.162499994039536, + "fp16Ms": 6.587499991059303, + "speedupFp32": "1.15x", + "speedupFp16": "1.95x" + }, + { + "category": "Linear Algebra", + "name": "Batch MatMul (BMM)", + "shape": "[4, 64, 64] * [4, 64, 64]", + "baselineMs": 0.13999999761581422, + "fp32Ms": 0.16333333253860474, + "fp16Ms": 0.13333333333333333, + "speedupFp32": "0.86x", + "speedupFp16": "1.05x" + }, + { + "category": "Linear Algebra", + "name": "Batch MatMul (BMM)", + "shape": "[8, 128, 128] * [8, 128, 128]", + "baselineMs": 0.21000000238418579, + "fp32Ms": 0.1550000011920929, + "fp16Ms": 0.15, + "speedupFp32": "1.35x", + "speedupFp16": "1.40x" + }, + { + "category": "Convolution", + "name": "Conv2D 3x3 (Tile64)", + "shape": "[1,64,56,56] * [64,64,3,3]", + "baselineMs": 0.575, + "fp32Ms": 0.3449999988079071, + "fp16Ms": 0.2199999988079071, + "speedupFp32": "1.67x", + "speedupFp16": "2.61x" + }, + { + "category": "Convolution", + "name": "Conv2D 1x1 (Tile64)", + "shape": "[1,128,56,56] * [128,128,1,1]", + "baselineMs": 0.2949999988079071, + "fp32Ms": 0.2, + "fp16Ms": 0.2, + "speedupFp32": "1.47x", + "speedupFp16": "1.47x" + }, + { + "category": "Convolution", + "name": "Conv2D 3x3 Depthwise", + "shape": "[1,64,56,56] * [64,1,3,3]", + "baselineMs": 0.16799999713897706, + "fp32Ms": 0.2, + "fp16Ms": 0.19599999904632567, + "speedupFp32": "0.84x", + "speedupFp16": "0.86x" + }, + { + "category": "Convolution", + "name": "ConvWithClamp 3x3 (Tile64)", + "shape": "[1,64,56,56] * [64,64,3,3]", + "baselineMs": 0.525, + "fp32Ms": 0.425, + "fp16Ms": 0.2, + "speedupFp32": "1.24x", + "speedupFp16": "2.63x" + }, + { + "category": "Convolution", + "name": "ConvWithClamp 1x1 (Tile64)", + "shape": "[1,128,56,56] * [128,128,1,1]", + "baselineMs": 0.275, + "fp32Ms": 0.1799999952316284, + "fp16Ms": 0.1550000011920929, + "speedupFp32": "1.53x", + "speedupFp16": "1.77x" + }, + { + "category": "Activation", + "name": "ReLU (Vec4)", + "shape": "[1,64,56,56]", + "baselineMs": 0.1, + "fp32Ms": 0.10750000178813934, + "fp16Ms": 0.12249999940395355, + "speedupFp32": "0.93x", + "speedupFp16": "0.82x" + }, + { + "category": "Activation", + "name": "Sigmoid (Vec4)", + "shape": "[1,64,56,56]", + "baselineMs": 0.1324999988079071, + "fp32Ms": 0.15, + "fp16Ms": 0.12250000238418579, + "speedupFp32": "0.88x", + "speedupFp16": "1.08x" + }, + { + "category": "Activation", + "name": "Hardswish (Vec4)", + "shape": "[1,64,56,56]", + "baselineMs": 0.12250000238418579, + "fp32Ms": 0.11499999761581421, + "fp16Ms": 0.125, + "speedupFp32": "1.07x", + "speedupFp16": "0.98x" + }, + { + "category": "Reduction", + "name": "Softmax (Last-Dim Vec4)", + "shape": "[128,1024]", + "baselineMs": 0.5266666650772095, + "fp32Ms": 0.36666666666666664, + "fp16Ms": 0.6033333341280619, + "speedupFp32": "1.44x", + "speedupFp16": "0.87x" + }, + { + "category": "Pooling & Resampling", + "name": "MaxPool2D (2D Spatial Grid)", + "shape": "[1,64,112,112]", + "baselineMs": 0.22399999618530272, + "fp32Ms": 0.16799999713897706, + "fp16Ms": 0.18, + "speedupFp32": "1.33x", + "speedupFp16": "1.24x" + }, + { + "category": "Elementwise Math", + "name": "Elementwise Mul (Vec4)", + "shape": "[1,64,56,56]", + "baselineMs": 0.12750000059604644, + "fp32Ms": 0.1125, + "fp16Ms": 0.1375, + "speedupFp32": "1.13x", + "speedupFp16": "0.93x" + }, + { + "category": "Data Movement", + "name": "Matrix Transpose (Shared-Memory 16x16)", + "shape": "[1024,1024]", + "baselineMs": 0.20399999618530273, + "fp32Ms": 0.14799999713897705, + "fp16Ms": 0.14, + "speedupFp32": "1.38x", + "speedupFp16": "1.46x" + } + ] +} diff --git a/example/index.html b/example/index.html index 94933a7..109eae9 100644 --- a/example/index.html +++ b/example/index.html @@ -30,6 +30,22 @@
← all demos

+
+ Precision + +
diff --git a/example/public/assets/depthart/program-fp16.json b/example/public/assets/depthart/program-fp16.json new file mode 100644 index 0000000..550d031 --- /dev/null +++ b/example/public/assets/depthart/program-fp16.json @@ -0,0 +1,3234 @@ +{ + "version": "1", + "chain": [ + { + "name": "et_vk.conv_with_clamp.default", + "args": [136, 0, 1, 137, 138, 139, 140, 141, 142, 143, 144, 145] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [145, 2, 3, 146, 147, 148, 140, 149, 142, 143, 144, 150] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [150, 4, 5, 151, 152, 153, 140, 154, 142, 143, 144, 155] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [155, 6, 7, 156, 157, 158, 140, 159, 142, 143, 144, 160] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [160, 8, 9, 161, 162, 163, 140, 164, 142, 143, 144, 165] + }, + { + "name": "aten.convolution.default", + "args": [165, 10, 11, 166, 167, 168, 140, 169, 170, 171] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [171, 12, 13, 172, 173, 174, 140, 175, 142, 143, 144, 176] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [176, 14, 15, 177, 178, 179, 140, 180, 181, 143, 144, 182] + }, + { + "name": "aten.convolution.default", + "args": [182, 16, 17, 183, 184, 185, 140, 186, 142, 187] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [187, 18, 19, 188, 189, 190, 140, 191, 142, 143, 144, 192] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [192, 20, 21, 193, 194, 195, 140, 196, 181, 143, 144, 197] + }, + { + "name": "aten.convolution.default", + "args": [197, 22, 23, 198, 199, 200, 140, 201, 142, 202] + }, + { + "name": "aten.add.Tensor", + "args": [202, 187, 142, 203] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [203, 24, 25, 204, 205, 206, 140, 207, 142, 143, 144, 208] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [208, 26, 27, 209, 210, 211, 140, 212, 181, 143, 144, 213] + }, + { + "name": "aten.convolution.default", + "args": [213, 28, 29, 214, 215, 216, 140, 217, 142, 218] + }, + { + "name": "aten.add.Tensor", + "args": [218, 203, 142, 219] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [219, 30, 31, 220, 221, 222, 140, 223, 142, 143, 144, 224] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [224, 32, 33, 225, 226, 227, 140, 228, 181, 143, 144, 229] + }, + { + "name": "aten.convolution.default", + "args": [229, 34, 35, 230, 231, 232, 140, 233, 142, 234] + }, + { + "name": "aten.add.Tensor", + "args": [234, 219, 142, 235] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [235, 36, 37, 236, 237, 238, 140, 239, 142, 143, 144, 240] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [240, 38, 39, 241, 242, 243, 140, 244, 181, 143, 144, 245] + }, + { + "name": "aten.convolution.default", + "args": [245, 40, 41, 246, 247, 248, 140, 249, 142, 250] + }, + { + "name": "aten.add.Tensor", + "args": [250, 235, 142, 251] + }, + { + "name": "aten.convolution.default", + "args": [251, 42, 43, 252, 253, 254, 140, 255, 256, 257] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [257, 44, 45, 258, 259, 260, 140, 261, 142, 143, 144, 262] + }, + { + "name": "aten.convolution.default", + "args": [262, 46, 47, 263, 264, 265, 140, 266, 142, 267] + }, + { + "name": "aten.add.Tensor", + "args": [267, 251, 142, 268] + }, + { + "name": "aten.convolution.default", + "args": [268, 48, 49, 269, 270, 271, 140, 272, 256, 273] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [273, 50, 51, 274, 275, 276, 140, 277, 142, 143, 144, 278] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [278, 52, 53, 279, 280, 281, 140, 282, 283, 143, 144, 284] + }, + { + "name": "aten.convolution.default", + "args": [284, 54, 55, 285, 286, 287, 140, 288, 142, 289] + }, + { + "name": "aten.convolution.default", + "args": [289, 56, 57, 290, 291, 292, 140, 293, 294, 295] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [295, 58, 59, 296, 297, 298, 140, 299, 142, 143, 144, 300] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [300, 60, 61, 301, 302, 303, 140, 304, 305, 143, 144, 306] + }, + { + "name": "aten.convolution.default", + "args": [306, 62, 63, 307, 308, 309, 140, 310, 142, 311] + }, + { + "name": "aten.add.Tensor", + "args": [311, 289, 142, 312] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [312, 64, 65, 313, 314, 315, 140, 316, 142, 143, 144, 317] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [317, 66, 67, 318, 319, 320, 140, 321, 305, 143, 144, 322] + }, + { + "name": "aten.convolution.default", + "args": [322, 68, 69, 323, 324, 325, 140, 326, 142, 327] + }, + { + "name": "aten.add.Tensor", + "args": [327, 312, 142, 328] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [328, 70, 71, 329, 330, 331, 140, 332, 142, 143, 144, 333] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [333, 72, 73, 334, 335, 336, 140, 337, 338, 143, 144, 339] + }, + { + "name": "aten.convolution.default", + "args": [339, 74, 75, 340, 341, 342, 140, 343, 142, 344] + }, + { + "name": "aten.add.Tensor", + "args": [344, 328, 142, 345] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [345, 76, 77, 346, 347, 348, 140, 349, 142, 143, 144, 350] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [350, 78, 79, 351, 352, 353, 140, 354, 305, 143, 144, 355] + }, + { + "name": "aten.convolution.default", + "args": [355, 80, 81, 356, 357, 358, 140, 359, 142, 360] + }, + { + "name": "aten.add.Tensor", + "args": [360, 345, 142, 361] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [361, 82, 83, 362, 363, 364, 140, 365, 142, 143, 144, 366] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [366, 84, 85, 367, 368, 369, 140, 370, 305, 143, 144, 371] + }, + { + "name": "aten.convolution.default", + "args": [371, 86, 87, 372, 373, 374, 140, 375, 142, 376] + }, + { + "name": "aten.add.Tensor", + "args": [376, 361, 142, 377] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [377, 88, 89, 378, 379, 380, 140, 381, 142, 143, 144, 382] + }, + { + "name": "aten.convolution.default", + "args": [155, 90, 383, 384, 385, 386, 140, 387, 142, 388] + }, + { + "name": "aten.convolution.default", + "args": [165, 91, 389, 390, 391, 392, 140, 393, 142, 394] + }, + { + "name": "aten.convolution.default", + "args": [268, 92, 395, 396, 397, 398, 140, 399, 142, 400] + }, + { + "name": "aten.convolution.default", + "args": [382, 93, 401, 402, 403, 404, 140, 405, 142, 406] + }, + { + "name": "aten.relu.default", + "args": [406, 407] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [407, 126, 127, 408, 409, 410, 140, 411, 142, 143, 144, 412] + }, + { + "name": "aten.convolution.default", + "args": [412, 128, 129, 413, 414, 415, 140, 416, 142, 417] + }, + { + "name": "aten.add.Tensor", + "args": [417, 406, 142, 418] + }, + { + "name": "aten.upsample_bilinear2d.vec", + "args": [418, 419, 420, 421, 422] + }, + { + "name": "aten.convolution.default", + "args": [422, 124, 125, 423, 424, 425, 140, 426, 142, 427] + }, + { + "name": "aten.relu.default", + "args": [400, 428] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [428, 116, 117, 429, 430, 431, 140, 432, 142, 143, 144, 433] + }, + { + "name": "aten.convolution.default", + "args": [433, 118, 119, 434, 435, 436, 140, 437, 142, 438] + }, + { + "name": "aten.add.Tensor", + "args": [438, 400, 142, 439] + }, + { + "name": "aten.add.Tensor", + "args": [427, 439, 142, 440] + }, + { + "name": "aten.relu.default", + "args": [440, 441] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [441, 120, 121, 442, 443, 444, 140, 445, 142, 143, 144, 446] + }, + { + "name": "aten.convolution.default", + "args": [446, 122, 123, 447, 448, 449, 140, 450, 142, 451] + }, + { + "name": "aten.add.Tensor", + "args": [451, 440, 142, 452] + }, + { + "name": "aten.upsample_bilinear2d.vec", + "args": [452, 453, 420, 454, 455] + }, + { + "name": "aten.convolution.default", + "args": [455, 114, 115, 456, 457, 458, 140, 459, 142, 460] + }, + { + "name": "aten.relu.default", + "args": [394, 461] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [461, 106, 107, 462, 463, 464, 140, 465, 142, 143, 144, 466] + }, + { + "name": "aten.convolution.default", + "args": [466, 108, 109, 467, 468, 469, 140, 470, 142, 471] + }, + { + "name": "aten.add.Tensor", + "args": [471, 394, 142, 472] + }, + { + "name": "aten.add.Tensor", + "args": [460, 472, 142, 473] + }, + { + "name": "aten.relu.default", + "args": [473, 474] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [474, 110, 111, 475, 476, 477, 140, 478, 142, 143, 144, 479] + }, + { + "name": "aten.convolution.default", + "args": [479, 112, 113, 480, 481, 482, 140, 483, 142, 484] + }, + { + "name": "aten.add.Tensor", + "args": [484, 473, 142, 485] + }, + { + "name": "aten.upsample_bilinear2d.vec", + "args": [485, 486, 420, 487, 488] + }, + { + "name": "aten.convolution.default", + "args": [488, 104, 105, 489, 490, 491, 140, 492, 142, 493] + }, + { + "name": "aten.relu.default", + "args": [388, 494] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [494, 96, 97, 495, 496, 497, 140, 498, 142, 143, 144, 499] + }, + { + "name": "aten.convolution.default", + "args": [499, 98, 99, 500, 501, 502, 140, 503, 142, 504] + }, + { + "name": "aten.add.Tensor", + "args": [504, 388, 142, 505] + }, + { + "name": "aten.add.Tensor", + "args": [493, 505, 142, 506] + }, + { + "name": "aten.relu.default", + "args": [506, 507] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [507, 100, 101, 508, 509, 510, 140, 511, 142, 143, 144, 512] + }, + { + "name": "aten.convolution.default", + "args": [512, 102, 103, 513, 514, 515, 140, 516, 142, 517] + }, + { + "name": "aten.add.Tensor", + "args": [517, 506, 142, 518] + }, + { + "name": "aten.upsample_bilinear2d.vec", + "args": [518, 519, 420, 520, 521] + }, + { + "name": "aten.convolution.default", + "args": [521, 94, 95, 522, 523, 524, 140, 525, 142, 526] + }, + { + "name": "aten.convolution.default", + "args": [526, 130, 131, 527, 528, 529, 140, 530, 142, 531] + }, + { + "name": "aten.upsample_bilinear2d.vec", + "args": [531, 532, 420, 533, 534] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [534, 132, 133, 535, 536, 537, 140, 538, 142, 143, 144, 539] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [539, 134, 135, 540, 541, 542, 140, 543, 142, 143, 144, 544] + } + ], + "values": [ + { + "type": "tensor", + "shape": [32, 3, 3, 3], + "dtype": "float16", + "weightsOffset": 0, + "weightsLength": 1728 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 1728, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [32, 32, 3, 3], + "dtype": "float16", + "weightsOffset": 1792, + "weightsLength": 18432 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 20224, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [32, 32, 1, 1], + "dtype": "float16", + "weightsOffset": 20288, + "weightsLength": 2048 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 22336, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [96, 32, 3, 3], + "dtype": "float16", + "weightsOffset": 22400, + "weightsLength": 55296 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 77696, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [64, 96, 1, 1], + "dtype": "float16", + "weightsOffset": 77888, + "weightsLength": 12288 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 90176, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [64, 1, 5, 5], + "dtype": "float16", + "weightsOffset": 90304, + "weightsLength": 3200 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 93504, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [192, 64, 1, 1], + "dtype": "float16", + "weightsOffset": 93632, + "weightsLength": 24576 + }, + { + "type": "tensor", + "shape": [192], + "dtype": "float16", + "weightsOffset": 118208, + "weightsLength": 384 + }, + { + "type": "tensor", + "shape": [192, 1, 5, 5], + "dtype": "float16", + "weightsOffset": 118592, + "weightsLength": 9600 + }, + { + "type": "tensor", + "shape": [192], + "dtype": "float16", + "weightsOffset": 128192, + "weightsLength": 384 + }, + { + "type": "tensor", + "shape": [96, 192, 1, 1], + "dtype": "float16", + "weightsOffset": 128576, + "weightsLength": 36864 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 165440, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [192, 96, 1, 1], + "dtype": "float16", + "weightsOffset": 165632, + "weightsLength": 36864 + }, + { + "type": "tensor", + "shape": [192], + "dtype": "float16", + "weightsOffset": 202496, + "weightsLength": 384 + }, + { + "type": "tensor", + "shape": [192, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 202880, + "weightsLength": 3456 + }, + { + "type": "tensor", + "shape": [192], + "dtype": "float16", + "weightsOffset": 206336, + "weightsLength": 384 + }, + { + "type": "tensor", + "shape": [96, 192, 1, 1], + "dtype": "float16", + "weightsOffset": 206720, + "weightsLength": 36864 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 243584, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [192, 96, 1, 1], + "dtype": "float16", + "weightsOffset": 243776, + "weightsLength": 36864 + }, + { + "type": "tensor", + "shape": [192], + "dtype": "float16", + "weightsOffset": 280640, + "weightsLength": 384 + }, + { + "type": "tensor", + "shape": [192, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 281024, + "weightsLength": 3456 + }, + { + "type": "tensor", + "shape": [192], + "dtype": "float16", + "weightsOffset": 284480, + "weightsLength": 384 + }, + { + "type": "tensor", + "shape": [96, 192, 1, 1], + "dtype": "float16", + "weightsOffset": 284864, + "weightsLength": 36864 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 321728, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [192, 96, 1, 1], + "dtype": "float16", + "weightsOffset": 321920, + "weightsLength": 36864 + }, + { + "type": "tensor", + "shape": [192], + "dtype": "float16", + "weightsOffset": 358784, + "weightsLength": 384 + }, + { + "type": "tensor", + "shape": [192, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 359168, + "weightsLength": 3456 + }, + { + "type": "tensor", + "shape": [192], + "dtype": "float16", + "weightsOffset": 362624, + "weightsLength": 384 + }, + { + "type": "tensor", + "shape": [96, 192, 1, 1], + "dtype": "float16", + "weightsOffset": 363008, + "weightsLength": 36864 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 399872, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [192, 96, 1, 1], + "dtype": "float16", + "weightsOffset": 400064, + "weightsLength": 36864 + }, + { + "type": "tensor", + "shape": [192], + "dtype": "float16", + "weightsOffset": 436928, + "weightsLength": 384 + }, + { + "type": "tensor", + "shape": [192, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 437312, + "weightsLength": 3456 + }, + { + "type": "tensor", + "shape": [192], + "dtype": "float16", + "weightsOffset": 440768, + "weightsLength": 384 + }, + { + "type": "tensor", + "shape": [96, 192, 1, 1], + "dtype": "float16", + "weightsOffset": 441152, + "weightsLength": 36864 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 478016, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [96, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 478208, + "weightsLength": 1728 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 479936, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [384, 96, 1, 1], + "dtype": "float16", + "weightsOffset": 480128, + "weightsLength": 73728 + }, + { + "type": "tensor", + "shape": [384], + "dtype": "float16", + "weightsOffset": 553856, + "weightsLength": 768 + }, + { + "type": "tensor", + "shape": [96, 384, 1, 1], + "dtype": "float16", + "weightsOffset": 554624, + "weightsLength": 73728 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 628352, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [96, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 628544, + "weightsLength": 1728 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 630272, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [576, 96, 1, 1], + "dtype": "float16", + "weightsOffset": 630464, + "weightsLength": 110592 + }, + { + "type": "tensor", + "shape": [576], + "dtype": "float16", + "weightsOffset": 741056, + "weightsLength": 1152 + }, + { + "type": "tensor", + "shape": [576, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 742208, + "weightsLength": 10368 + }, + { + "type": "tensor", + "shape": [576], + "dtype": "float16", + "weightsOffset": 752576, + "weightsLength": 1152 + }, + { + "type": "tensor", + "shape": [128, 576, 1, 1], + "dtype": "float16", + "weightsOffset": 753728, + "weightsLength": 147456 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 901184, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [128, 1, 5, 5], + "dtype": "float16", + "weightsOffset": 901440, + "weightsLength": 6400 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 907840, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [512, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 908096, + "weightsLength": 131072 + }, + { + "type": "tensor", + "shape": [512], + "dtype": "float16", + "weightsOffset": 1039168, + "weightsLength": 1024 + }, + { + "type": "tensor", + "shape": [512, 1, 5, 5], + "dtype": "float16", + "weightsOffset": 1040192, + "weightsLength": 25600 + }, + { + "type": "tensor", + "shape": [512], + "dtype": "float16", + "weightsOffset": 1065792, + "weightsLength": 1024 + }, + { + "type": "tensor", + "shape": [128, 512, 1, 1], + "dtype": "float16", + "weightsOffset": 1066816, + "weightsLength": 131072 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 1197888, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [512, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 1198144, + "weightsLength": 131072 + }, + { + "type": "tensor", + "shape": [512], + "dtype": "float16", + "weightsOffset": 1329216, + "weightsLength": 1024 + }, + { + "type": "tensor", + "shape": [512, 1, 5, 5], + "dtype": "float16", + "weightsOffset": 1330240, + "weightsLength": 25600 + }, + { + "type": "tensor", + "shape": [512], + "dtype": "float16", + "weightsOffset": 1355840, + "weightsLength": 1024 + }, + { + "type": "tensor", + "shape": [128, 512, 1, 1], + "dtype": "float16", + "weightsOffset": 1356864, + "weightsLength": 131072 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 1487936, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [384, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 1488192, + "weightsLength": 98304 + }, + { + "type": "tensor", + "shape": [384], + "dtype": "float16", + "weightsOffset": 1586496, + "weightsLength": 768 + }, + { + "type": "tensor", + "shape": [384, 1, 5, 5], + "dtype": "float16", + "weightsOffset": 1587264, + "weightsLength": 19200 + }, + { + "type": "tensor", + "shape": [384], + "dtype": "float16", + "weightsOffset": 1606464, + "weightsLength": 768 + }, + { + "type": "tensor", + "shape": [128, 384, 1, 1], + "dtype": "float16", + "weightsOffset": 1607232, + "weightsLength": 98304 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 1705536, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [512, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 1705792, + "weightsLength": 131072 + }, + { + "type": "tensor", + "shape": [512], + "dtype": "float16", + "weightsOffset": 1836864, + "weightsLength": 1024 + }, + { + "type": "tensor", + "shape": [512, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 1837888, + "weightsLength": 9216 + }, + { + "type": "tensor", + "shape": [512], + "dtype": "float16", + "weightsOffset": 1847104, + "weightsLength": 1024 + }, + { + "type": "tensor", + "shape": [128, 512, 1, 1], + "dtype": "float16", + "weightsOffset": 1848128, + "weightsLength": 131072 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 1979200, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [512, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 1979456, + "weightsLength": 131072 + }, + { + "type": "tensor", + "shape": [512], + "dtype": "float16", + "weightsOffset": 2110528, + "weightsLength": 1024 + }, + { + "type": "tensor", + "shape": [512, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 2111552, + "weightsLength": 9216 + }, + { + "type": "tensor", + "shape": [512], + "dtype": "float16", + "weightsOffset": 2120768, + "weightsLength": 1024 + }, + { + "type": "tensor", + "shape": [128, 512, 1, 1], + "dtype": "float16", + "weightsOffset": 2121792, + "weightsLength": 131072 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 2252864, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [960, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 2253120, + "weightsLength": 245760 + }, + { + "type": "tensor", + "shape": [960], + "dtype": "float16", + "weightsOffset": 2498880, + "weightsLength": 1920 + }, + { + "type": "tensor", + "shape": [48, 32, 3, 3], + "dtype": "float16", + "weightsOffset": 2500800, + "weightsLength": 27648 + }, + { + "type": "tensor", + "shape": [48, 64, 3, 3], + "dtype": "float16", + "weightsOffset": 2528448, + "weightsLength": 55296 + }, + { + "type": "tensor", + "shape": [48, 96, 3, 3], + "dtype": "float16", + "weightsOffset": 2583744, + "weightsLength": 82944 + }, + { + "type": "tensor", + "shape": [48, 960, 3, 3], + "dtype": "float16", + "weightsOffset": 2666688, + "weightsLength": 829440 + }, + { + "type": "tensor", + "shape": [48, 48, 1, 1], + "dtype": "float16", + "weightsOffset": 3496128, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 3500736, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 3, 3], + "dtype": "float16", + "weightsOffset": 3500832, + "weightsLength": 41472 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 3542304, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 3, 3], + "dtype": "float16", + "weightsOffset": 3542400, + "weightsLength": 41472 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 3583872, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 3, 3], + "dtype": "float16", + "weightsOffset": 3583968, + "weightsLength": 41472 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 3625440, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 3, 3], + "dtype": "float16", + "weightsOffset": 3625536, + "weightsLength": 41472 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 3667008, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 1, 1], + "dtype": "float16", + "weightsOffset": 3667104, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 3671712, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 3, 3], + "dtype": "float16", + "weightsOffset": 3671808, + "weightsLength": 41472 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 3713280, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 3, 3], + "dtype": "float16", + "weightsOffset": 3713376, + "weightsLength": 41472 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 3754848, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 3, 3], + "dtype": "float16", + "weightsOffset": 3754944, + "weightsLength": 41472 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 3796416, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 3, 3], + "dtype": "float16", + "weightsOffset": 3796512, + "weightsLength": 41472 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 3837984, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 1, 1], + "dtype": "float16", + "weightsOffset": 3838080, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 3842688, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 3, 3], + "dtype": "float16", + "weightsOffset": 3842784, + "weightsLength": 41472 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 3884256, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 3, 3], + "dtype": "float16", + "weightsOffset": 3884352, + "weightsLength": 41472 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 3925824, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 3, 3], + "dtype": "float16", + "weightsOffset": 3925920, + "weightsLength": 41472 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 3967392, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 3, 3], + "dtype": "float16", + "weightsOffset": 3967488, + "weightsLength": 41472 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 4008960, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 1, 1], + "dtype": "float16", + "weightsOffset": 4009056, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 4013664, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 3, 3], + "dtype": "float16", + "weightsOffset": 4013760, + "weightsLength": 41472 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 4055232, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [48, 48, 3, 3], + "dtype": "float16", + "weightsOffset": 4055328, + "weightsLength": 41472 + }, + { + "type": "tensor", + "shape": [48], + "dtype": "float16", + "weightsOffset": 4096800, + "weightsLength": 96 + }, + { + "type": "tensor", + "shape": [24, 48, 3, 3], + "dtype": "float16", + "weightsOffset": 4096896, + "weightsLength": 20736 + }, + { + "type": "tensor", + "shape": [24], + "dtype": "float16", + "weightsOffset": 4117632, + "weightsLength": 48 + }, + { + "type": "tensor", + "shape": [16, 24, 3, 3], + "dtype": "float16", + "weightsOffset": 4117680, + "weightsLength": 6912 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 4124592, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [1, 16, 1, 1], + "dtype": "float16", + "weightsOffset": 4124624, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [1], + "dtype": "float16", + "weightsOffset": 4124656, + "weightsLength": 2 + }, + { + "type": "tensor", + "shape": [1, 3, 448, 448], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "scalar", + "value": false + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 1 + }, + { + "type": "scalar", + "value": 0 + }, + { + "type": "scalar", + "value": 1.7976931348623157e308 + }, + { + "type": "tensor", + "shape": [1, 32, 224, 224], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 112, 112], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 112, 112], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 96, 56, 56], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 56, 56], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 64 + }, + { + "type": "tensor", + "shape": [1, 64, 56, 56], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 192, 56, 56], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 192 + }, + { + "type": "tensor", + "shape": [1, 192, 28, 28], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 96, 28, 28], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 192, 28, 28], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 192, 28, 28], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 96, 28, 28], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 96, 28, 28], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 192, 28, 28], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 192, 28, 28], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 96, 28, 28], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 96, 28, 28], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 192, 28, 28], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 192, 28, 28], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 96, 28, 28], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 96, 28, 28], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 192, 28, 28], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 192, 28, 28], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 96, 28, 28], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 96, 28, 28], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 96 + }, + { + "type": "tensor", + "shape": [1, 96, 28, 28], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 384, 28, 28], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 96, 28, 28], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 96, 28, 28], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 96, 28, 28], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 576, 28, 28], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 576 + }, + { + "type": "tensor", + "shape": [1, 576, 14, 14], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 14, 14], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 128 + }, + { + "type": "tensor", + "shape": [1, 128, 14, 14], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 512, 14, 14], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 512 + }, + { + "type": "tensor", + "shape": [1, 512, 14, 14], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 14, 14], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 128, 14, 14], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 512, 14, 14], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 512, 14, 14], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 14, 14], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 128, 14, 14], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 384, 14, 14], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 384 + }, + { + "type": "tensor", + "shape": [1, 384, 14, 14], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 14, 14], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 128, 14, 14], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 512, 14, 14], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 512, 14, 14], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 14, 14], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 128, 14, 14], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 512, 14, 14], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 512, 14, 14], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 14, 14], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "tensor", + "shape": [1, 128, 14, 14], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 960, 14, 14], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "null" + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 112, 112], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "null" + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 56, 56], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "null" + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 28, 28], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "null" + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 14, 14], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 48, 14, 14], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 14, 14], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 14, 14], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 48, 14, 14], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [28, 28] + }, + { + "type": "scalar", + "value": true + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 48, 28, 28], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 28, 28], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 48, 28, 28], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 28, 28], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 28, 28], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 48, 28, 28], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 48, 28, 28], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 48, 28, 28], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 28, 28], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 28, 28], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 48, 28, 28], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [56, 56] + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 48, 56, 56], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 56, 56], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 48, 56, 56], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 56, 56], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 56, 56], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 48, 56, 56], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 48, 56, 56], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 48, 56, 56], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 56, 56], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 56, 56], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 48, 56, 56], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [112, 112] + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 48, 112, 112], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 112, 112], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 48, 112, 112], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 112, 112], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 112, 112], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 48, 112, 112], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 48, 112, 112], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 48, 112, 112], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 112, 112], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 112, 112], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 48, 112, 112], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "null" + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "tensor", + "shape": [1, 48, 224, 224], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 48, 224, 224], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 24, 224, 224], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [448, 448] + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 24, 448, 448], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 448, 448], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 1, 448, 448], + "dtype": "float16", + "memObjId": 0 + } + ], + "inputIds": [136], + "outputIds": [544], + "memoryPlan": { + "pools": [ + { + "id": 0, + "size": 9633792 + }, + { + "id": 1, + "size": 6422528 + }, + { + "id": 2, + "size": 1204224 + }, + { + "id": 3, + "size": 1204224 + }, + { + "id": 4, + "size": 301056 + }, + { + "id": 5, + "size": 301056 + }, + { + "id": 6, + "size": 50176 + } + ] + } +} diff --git a/example/public/assets/depthart/weights-fp16.bin b/example/public/assets/depthart/weights-fp16.bin new file mode 100644 index 0000000..8d5b481 Binary files /dev/null and b/example/public/assets/depthart/weights-fp16.bin differ diff --git a/example/public/assets/selfie-segmentation/program-fp16.json b/example/public/assets/selfie-segmentation/program-fp16.json new file mode 100644 index 0000000..aa88b7f --- /dev/null +++ b/example/public/assets/selfie-segmentation/program-fp16.json @@ -0,0 +1,3028 @@ +{ + "version": "1", + "chain": [ + { + "name": "aten.constant_pad_nd.default", + "args": [110, 111, 112, 113] + }, + { + "name": "aten.convolution.default", + "args": [113, 0, 1, 114, 115, 116, 117, 118, 119, 120] + }, + { + "name": "aten.hardswish.default", + "args": [120, 121] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [121, 2, 3, 122, 123, 124, 117, 125, 119, 112, 126, 127] + }, + { + "name": "aten.constant_pad_nd.default", + "args": [127, 128, 112, 129] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [129, 4, 5, 130, 131, 132, 117, 133, 134, 112, 126, 135] + }, + { + "name": "aten.mean.dim", + "args": [135, 136, 137, 138, 139] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [139, 6, 7, 140, 141, 142, 117, 143, 119, 112, 126, 144] + }, + { + "name": "aten.convolution.default", + "args": [144, 8, 9, 145, 146, 147, 117, 148, 119, 149] + }, + { + "name": "aten.sigmoid.default", + "args": [149, 150] + }, + { + "name": "aten.mul.Tensor", + "args": [135, 150, 151] + }, + { + "name": "aten.convolution.default", + "args": [151, 10, 11, 152, 153, 154, 117, 155, 119, 156] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [156, 12, 13, 157, 158, 159, 117, 160, 119, 112, 126, 161] + }, + { + "name": "aten.constant_pad_nd.default", + "args": [161, 162, 112, 163] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [163, 14, 15, 164, 165, 166, 117, 167, 168, 112, 126, 169] + }, + { + "name": "aten.convolution.default", + "args": [169, 16, 17, 170, 171, 172, 117, 173, 119, 174] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [174, 18, 19, 175, 176, 177, 117, 178, 119, 112, 126, 179] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [179, 20, 21, 180, 181, 182, 117, 183, 184, 112, 126, 185] + }, + { + "name": "aten.convolution.default", + "args": [185, 22, 23, 186, 187, 188, 117, 189, 119, 190] + }, + { + "name": "aten.add.Tensor", + "args": [174, 190, 119, 191] + }, + { + "name": "aten.convolution.default", + "args": [191, 24, 25, 192, 193, 194, 117, 195, 119, 196] + }, + { + "name": "aten.hardswish.default", + "args": [196, 197] + }, + { + "name": "aten.constant_pad_nd.default", + "args": [197, 198, 112, 199] + }, + { + "name": "aten.convolution.default", + "args": [199, 26, 27, 200, 201, 202, 117, 203, 204, 205] + }, + { + "name": "aten.hardswish.default", + "args": [205, 206] + }, + { + "name": "aten.mean.dim", + "args": [206, 207, 137, 208, 209] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [209, 28, 29, 210, 211, 212, 117, 213, 119, 112, 126, 214] + }, + { + "name": "aten.convolution.default", + "args": [214, 30, 31, 215, 216, 217, 117, 218, 119, 219] + }, + { + "name": "aten.sigmoid.default", + "args": [219, 220] + }, + { + "name": "aten.mul.Tensor", + "args": [206, 220, 221] + }, + { + "name": "aten.convolution.default", + "args": [221, 32, 33, 222, 223, 224, 117, 225, 119, 226] + }, + { + "name": "aten.convolution.default", + "args": [226, 34, 35, 227, 228, 229, 117, 230, 119, 231] + }, + { + "name": "aten.hardswish.default", + "args": [231, 232] + }, + { + "name": "aten.convolution.default", + "args": [232, 36, 37, 233, 234, 235, 117, 236, 237, 238] + }, + { + "name": "aten.hardswish.default", + "args": [238, 239] + }, + { + "name": "aten.mean.dim", + "args": [239, 240, 137, 241, 242] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [242, 38, 39, 243, 244, 245, 117, 246, 119, 112, 126, 247] + }, + { + "name": "aten.convolution.default", + "args": [247, 40, 41, 248, 249, 250, 117, 251, 119, 252] + }, + { + "name": "aten.sigmoid.default", + "args": [252, 253] + }, + { + "name": "aten.mul.Tensor", + "args": [239, 253, 254] + }, + { + "name": "aten.convolution.default", + "args": [254, 42, 43, 255, 256, 257, 117, 258, 119, 259] + }, + { + "name": "aten.add.Tensor", + "args": [226, 259, 119, 260] + }, + { + "name": "aten.convolution.default", + "args": [260, 44, 45, 261, 262, 263, 117, 264, 119, 265] + }, + { + "name": "aten.hardswish.default", + "args": [265, 266] + }, + { + "name": "aten.convolution.default", + "args": [266, 46, 47, 267, 268, 269, 117, 270, 237, 271] + }, + { + "name": "aten.hardswish.default", + "args": [271, 272] + }, + { + "name": "aten.mean.dim", + "args": [272, 273, 137, 274, 275] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [275, 48, 49, 276, 277, 278, 117, 279, 119, 112, 126, 280] + }, + { + "name": "aten.convolution.default", + "args": [280, 50, 51, 281, 282, 283, 117, 284, 119, 285] + }, + { + "name": "aten.sigmoid.default", + "args": [285, 286] + }, + { + "name": "aten.mul.Tensor", + "args": [272, 286, 287] + }, + { + "name": "aten.convolution.default", + "args": [287, 52, 53, 288, 289, 290, 117, 291, 119, 292] + }, + { + "name": "aten.add.Tensor", + "args": [260, 292, 119, 293] + }, + { + "name": "aten.convolution.default", + "args": [293, 54, 55, 294, 295, 296, 117, 297, 119, 298] + }, + { + "name": "aten.hardswish.default", + "args": [298, 299] + }, + { + "name": "aten.convolution.default", + "args": [299, 56, 57, 300, 301, 302, 117, 303, 204, 304] + }, + { + "name": "aten.hardswish.default", + "args": [304, 305] + }, + { + "name": "aten.mean.dim", + "args": [305, 306, 137, 307, 308] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [308, 58, 59, 309, 310, 311, 117, 312, 119, 112, 126, 313] + }, + { + "name": "aten.convolution.default", + "args": [313, 60, 61, 314, 315, 316, 117, 317, 119, 318] + }, + { + "name": "aten.sigmoid.default", + "args": [318, 319] + }, + { + "name": "aten.mul.Tensor", + "args": [305, 319, 320] + }, + { + "name": "aten.convolution.default", + "args": [320, 62, 63, 321, 322, 323, 117, 324, 119, 325] + }, + { + "name": "aten.add.Tensor", + "args": [293, 325, 119, 326] + }, + { + "name": "aten.convolution.default", + "args": [326, 64, 65, 327, 328, 329, 117, 330, 119, 331] + }, + { + "name": "aten.hardswish.default", + "args": [331, 332] + }, + { + "name": "aten.convolution.default", + "args": [332, 66, 67, 333, 334, 335, 117, 336, 204, 337] + }, + { + "name": "aten.hardswish.default", + "args": [337, 338] + }, + { + "name": "aten.mean.dim", + "args": [338, 339, 137, 340, 341] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [341, 68, 69, 342, 343, 344, 117, 345, 119, 112, 126, 346] + }, + { + "name": "aten.convolution.default", + "args": [346, 70, 71, 347, 348, 349, 117, 350, 119, 351] + }, + { + "name": "aten.sigmoid.default", + "args": [351, 352] + }, + { + "name": "aten.mul.Tensor", + "args": [338, 352, 353] + }, + { + "name": "aten.convolution.default", + "args": [353, 72, 73, 354, 355, 356, 117, 357, 119, 358] + }, + { + "name": "aten.add.Tensor", + "args": [326, 358, 119, 359] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [359, 74, 75, 360, 361, 362, 117, 363, 119, 112, 126, 364] + }, + { + "name": "aten.mean.dim", + "args": [359, 365, 137, 366, 367] + }, + { + "name": "aten.convolution.default", + "args": [367, 76, 77, 368, 369, 370, 117, 371, 119, 372] + }, + { + "name": "aten.sigmoid.default", + "args": [372, 373] + }, + { + "name": "aten.mul.Tensor", + "args": [364, 373, 374] + }, + { + "name": "aten.upsample_bilinear2d.vec", + "args": [374, 375, 117, 376, 377] + }, + { + "name": "aten.convolution.default", + "args": [377, 78, 79, 378, 379, 380, 117, 381, 119, 382] + }, + { + "name": "aten.add.Tensor", + "args": [382, 191, 119, 383] + }, + { + "name": "aten.mean.dim", + "args": [383, 384, 137, 385, 386] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [386, 80, 81, 387, 388, 389, 117, 390, 119, 112, 126, 391] + }, + { + "name": "aten.convolution.default", + "args": [391, 82, 83, 392, 393, 394, 117, 395, 119, 396] + }, + { + "name": "aten.sigmoid.default", + "args": [396, 397] + }, + { + "name": "aten.mul.Tensor", + "args": [191, 397, 398] + }, + { + "name": "aten.add.Tensor", + "args": [382, 398, 119, 399] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [399, 84, 85, 400, 401, 402, 117, 403, 119, 112, 126, 404] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [404, 86, 87, 405, 406, 407, 117, 408, 409, 112, 126, 410] + }, + { + "name": "aten.add.Tensor", + "args": [404, 410, 119, 411] + }, + { + "name": "aten.upsample_bilinear2d.vec", + "args": [411, 412, 117, 413, 414] + }, + { + "name": "aten.convolution.default", + "args": [414, 88, 89, 415, 416, 417, 117, 418, 119, 419] + }, + { + "name": "aten.add.Tensor", + "args": [419, 156, 119, 420] + }, + { + "name": "aten.mean.dim", + "args": [420, 421, 137, 422, 423] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [423, 90, 91, 424, 425, 426, 117, 427, 119, 112, 126, 428] + }, + { + "name": "aten.convolution.default", + "args": [428, 92, 93, 429, 430, 431, 117, 432, 119, 433] + }, + { + "name": "aten.sigmoid.default", + "args": [433, 434] + }, + { + "name": "aten.mul.Tensor", + "args": [156, 434, 435] + }, + { + "name": "aten.add.Tensor", + "args": [419, 435, 119, 436] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [436, 94, 95, 437, 438, 439, 117, 440, 119, 112, 126, 441] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [441, 96, 97, 442, 443, 444, 117, 445, 134, 112, 126, 446] + }, + { + "name": "aten.add.Tensor", + "args": [441, 446, 119, 447] + }, + { + "name": "aten.upsample_bilinear2d.vec", + "args": [447, 448, 117, 449, 450] + }, + { + "name": "aten.convolution.default", + "args": [450, 98, 99, 451, 452, 453, 117, 454, 119, 455] + }, + { + "name": "aten.add.Tensor", + "args": [455, 121, 119, 456] + }, + { + "name": "aten.mean.dim", + "args": [456, 457, 137, 458, 459] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [459, 100, 101, 460, 461, 462, 117, 463, 119, 112, 126, 464] + }, + { + "name": "aten.convolution.default", + "args": [464, 102, 103, 465, 466, 467, 117, 468, 119, 469] + }, + { + "name": "aten.sigmoid.default", + "args": [469, 470] + }, + { + "name": "aten.mul.Tensor", + "args": [121, 470, 471] + }, + { + "name": "aten.add.Tensor", + "args": [455, 471, 119, 472] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [472, 104, 105, 473, 474, 475, 117, 476, 119, 112, 126, 477] + }, + { + "name": "et_vk.conv_with_clamp.default", + "args": [477, 106, 107, 478, 479, 480, 117, 481, 134, 112, 126, 482] + }, + { + "name": "aten.add.Tensor", + "args": [477, 482, 119, 483] + }, + { + "name": "aten.convolution.default", + "args": [483, 108, 109, 484, 485, 486, 137, 487, 119, 488] + }, + { + "name": "aten.sigmoid.default", + "args": [488, 489] + } + ], + "values": [ + { + "type": "tensor", + "shape": [16, 3, 3, 3], + "dtype": "float16", + "weightsOffset": 0, + "weightsLength": 864 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 864, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 16, 1, 1], + "dtype": "float16", + "weightsOffset": 896, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 1408, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 1440, + "weightsLength": 288 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 1728, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [8, 16, 1, 1], + "dtype": "float16", + "weightsOffset": 1760, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [8], + "dtype": "float16", + "weightsOffset": 2016, + "weightsLength": 16 + }, + { + "type": "tensor", + "shape": [16, 8, 1, 1], + "dtype": "float16", + "weightsOffset": 2032, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 2288, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 16, 1, 1], + "dtype": "float16", + "weightsOffset": 2320, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 2832, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [72, 16, 1, 1], + "dtype": "float16", + "weightsOffset": 2864, + "weightsLength": 2304 + }, + { + "type": "tensor", + "shape": [72], + "dtype": "float16", + "weightsOffset": 5168, + "weightsLength": 144 + }, + { + "type": "tensor", + "shape": [72, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 5312, + "weightsLength": 1296 + }, + { + "type": "tensor", + "shape": [72], + "dtype": "float16", + "weightsOffset": 6608, + "weightsLength": 144 + }, + { + "type": "tensor", + "shape": [24, 72, 1, 1], + "dtype": "float16", + "weightsOffset": 6752, + "weightsLength": 3456 + }, + { + "type": "tensor", + "shape": [24], + "dtype": "float16", + "weightsOffset": 10208, + "weightsLength": 48 + }, + { + "type": "tensor", + "shape": [88, 24, 1, 1], + "dtype": "float16", + "weightsOffset": 10256, + "weightsLength": 4224 + }, + { + "type": "tensor", + "shape": [88], + "dtype": "float16", + "weightsOffset": 14480, + "weightsLength": 176 + }, + { + "type": "tensor", + "shape": [88, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 14656, + "weightsLength": 1584 + }, + { + "type": "tensor", + "shape": [88], + "dtype": "float16", + "weightsOffset": 16240, + "weightsLength": 176 + }, + { + "type": "tensor", + "shape": [24, 88, 1, 1], + "dtype": "float16", + "weightsOffset": 16416, + "weightsLength": 4224 + }, + { + "type": "tensor", + "shape": [24], + "dtype": "float16", + "weightsOffset": 20640, + "weightsLength": 48 + }, + { + "type": "tensor", + "shape": [96, 24, 1, 1], + "dtype": "float16", + "weightsOffset": 20688, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 25296, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [96, 1, 5, 5], + "dtype": "float16", + "weightsOffset": 25488, + "weightsLength": 4800 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 30288, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [24, 96, 1, 1], + "dtype": "float16", + "weightsOffset": 30480, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [24], + "dtype": "float16", + "weightsOffset": 35088, + "weightsLength": 48 + }, + { + "type": "tensor", + "shape": [96, 24, 1, 1], + "dtype": "float16", + "weightsOffset": 35136, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 39744, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [32, 96, 1, 1], + "dtype": "float16", + "weightsOffset": 39936, + "weightsLength": 6144 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 46080, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [128, 32, 1, 1], + "dtype": "float16", + "weightsOffset": 46144, + "weightsLength": 8192 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 54336, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [128, 1, 5, 5], + "dtype": "float16", + "weightsOffset": 54592, + "weightsLength": 6400 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 60992, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [32, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 61248, + "weightsLength": 8192 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 69440, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [128, 32, 1, 1], + "dtype": "float16", + "weightsOffset": 69504, + "weightsLength": 8192 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 77696, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [32, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 77952, + "weightsLength": 8192 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 86144, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [128, 32, 1, 1], + "dtype": "float16", + "weightsOffset": 86208, + "weightsLength": 8192 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 94400, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [128, 1, 5, 5], + "dtype": "float16", + "weightsOffset": 94656, + "weightsLength": 6400 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 101056, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [32, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 101312, + "weightsLength": 8192 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 109504, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [128, 32, 1, 1], + "dtype": "float16", + "weightsOffset": 109568, + "weightsLength": 8192 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 117760, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [32, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 118016, + "weightsLength": 8192 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 126208, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [96, 32, 1, 1], + "dtype": "float16", + "weightsOffset": 126272, + "weightsLength": 6144 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 132416, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [96, 1, 5, 5], + "dtype": "float16", + "weightsOffset": 132608, + "weightsLength": 4800 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 137408, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [24, 96, 1, 1], + "dtype": "float16", + "weightsOffset": 137600, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [24], + "dtype": "float16", + "weightsOffset": 142208, + "weightsLength": 48 + }, + { + "type": "tensor", + "shape": [96, 24, 1, 1], + "dtype": "float16", + "weightsOffset": 142256, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 146864, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [32, 96, 1, 1], + "dtype": "float16", + "weightsOffset": 147056, + "weightsLength": 6144 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 153200, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [96, 32, 1, 1], + "dtype": "float16", + "weightsOffset": 153264, + "weightsLength": 6144 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 159408, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [96, 1, 5, 5], + "dtype": "float16", + "weightsOffset": 159600, + "weightsLength": 4800 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 164400, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [24, 96, 1, 1], + "dtype": "float16", + "weightsOffset": 164592, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [24], + "dtype": "float16", + "weightsOffset": 169200, + "weightsLength": 48 + }, + { + "type": "tensor", + "shape": [96, 24, 1, 1], + "dtype": "float16", + "weightsOffset": 169248, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [96], + "dtype": "float16", + "weightsOffset": 173856, + "weightsLength": 192 + }, + { + "type": "tensor", + "shape": [32, 96, 1, 1], + "dtype": "float16", + "weightsOffset": 174048, + "weightsLength": 6144 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 180192, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [128, 32, 1, 1], + "dtype": "float16", + "weightsOffset": 180256, + "weightsLength": 8192 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 188448, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [128, 32, 1, 1], + "dtype": "float16", + "weightsOffset": 188704, + "weightsLength": 8192 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 196896, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [24, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 197152, + "weightsLength": 6144 + }, + { + "type": "tensor", + "shape": [24], + "dtype": "float16", + "weightsOffset": 203296, + "weightsLength": 48 + }, + { + "type": "tensor", + "shape": [24, 24, 1, 1], + "dtype": "float16", + "weightsOffset": 203344, + "weightsLength": 1152 + }, + { + "type": "tensor", + "shape": [24], + "dtype": "float16", + "weightsOffset": 204496, + "weightsLength": 48 + }, + { + "type": "tensor", + "shape": [24, 24, 1, 1], + "dtype": "float16", + "weightsOffset": 204544, + "weightsLength": 1152 + }, + { + "type": "tensor", + "shape": [24], + "dtype": "float16", + "weightsOffset": 205696, + "weightsLength": 48 + }, + { + "type": "tensor", + "shape": [24, 24, 1, 1], + "dtype": "float16", + "weightsOffset": 205744, + "weightsLength": 1152 + }, + { + "type": "tensor", + "shape": [24], + "dtype": "float16", + "weightsOffset": 206896, + "weightsLength": 48 + }, + { + "type": "tensor", + "shape": [24, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 206944, + "weightsLength": 432 + }, + { + "type": "tensor", + "shape": [24], + "dtype": "float16", + "weightsOffset": 207376, + "weightsLength": 48 + }, + { + "type": "tensor", + "shape": [16, 24, 1, 1], + "dtype": "float16", + "weightsOffset": 207424, + "weightsLength": 768 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 208192, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 16, 1, 1], + "dtype": "float16", + "weightsOffset": 208224, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 208736, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 16, 1, 1], + "dtype": "float16", + "weightsOffset": 208768, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 209280, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 16, 1, 1], + "dtype": "float16", + "weightsOffset": 209312, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 209824, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 209856, + "weightsLength": 288 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 210144, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 16, 1, 1], + "dtype": "float16", + "weightsOffset": 210176, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 210688, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 16, 1, 1], + "dtype": "float16", + "weightsOffset": 210720, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 211232, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 16, 1, 1], + "dtype": "float16", + "weightsOffset": 211264, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 211776, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 16, 1, 1], + "dtype": "float16", + "weightsOffset": 211808, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 212320, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 212352, + "weightsLength": 288 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 212640, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 1, 2, 2], + "dtype": "float16", + "weightsOffset": 212672, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [1], + "dtype": "float16", + "weightsOffset": 212800, + "weightsLength": 2 + }, + { + "type": "tensor", + "shape": [1, 3, 256, 256], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [0, 1, 0, 1] + }, + { + "type": "scalar", + "value": 0 + }, + { + "type": "tensor", + "shape": [1, 3, 257, 257], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "scalar", + "value": false + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 1 + }, + { + "type": "tensor", + "shape": [1, 16, 128, 128], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 16, 128, 128], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 1.7976931348623157e308 + }, + { + "type": "tensor", + "shape": [1, 16, 128, 128], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [0, 1, 0, 1] + }, + { + "type": "tensor", + "shape": [1, 16, 129, 129], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 16 + }, + { + "type": "tensor", + "shape": [1, 16, 64, 64], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [2, 3] + }, + { + "type": "scalar", + "value": true + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 16, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 8, 1, 1], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 16, 1, 1], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 64, 64], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 64, 64], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 72, 64, 64], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [0, 1, 0, 1] + }, + { + "type": "tensor", + "shape": [1, 72, 65, 65], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 72 + }, + { + "type": "tensor", + "shape": [1, 72, 32, 32], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 24, 32, 32], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 88, 32, 32], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 88 + }, + { + "type": "tensor", + "shape": [1, 88, 32, 32], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 24, 32, 32], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 24, 32, 32], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 96, 32, 32], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 96, 32, 32], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 2, 1, 2] + }, + { + "type": "tensor", + "shape": [1, 96, 35, 35], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 96 + }, + { + "type": "tensor", + "shape": [1, 96, 16, 16], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 96, 16, 16], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [2, 3] + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 96, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 24, 1, 1], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 96, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 96, 1, 1], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 96, 16, 16], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 16, 16], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 16, 16], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 128, 16, 16], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 128 + }, + { + "type": "tensor", + "shape": [1, 128, 16, 16], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 128, 16, 16], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [2, 3] + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 128, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 1, 1], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 128, 1, 1], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "tensor", + "shape": [1, 128, 16, 16], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 16, 16], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 16, 16], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 16, 16], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 128, 16, 16], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 16, 16], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 128, 16, 16], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [2, 3] + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 128, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 1, 1], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 128, 1, 1], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 128, 16, 16], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 16, 16], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 16, 16], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 96, 16, 16], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 96, 16, 16], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 96, 16, 16], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 96, 16, 16], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [2, 3] + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 96, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 24, 1, 1], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 96, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 96, 1, 1], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "tensor", + "shape": [1, 96, 16, 16], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 16, 16], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 16, 16], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 96, 16, 16], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 96, 16, 16], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 96, 16, 16], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 96, 16, 16], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [2, 3] + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 96, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 24, 1, 1], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 96, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 96, 1, 1], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 96, 16, 16], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 16, 16], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 32, 16, 16], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 16, 16], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [2, 3] + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 32, 1, 1], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 1, 1], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 1, 1], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 128, 16, 16], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [32, 32] + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 128, 32, 32], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 24, 32, 32], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 24, 32, 32], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [2, 3] + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 24, 1, 1], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 24, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 24, 1, 1], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 24, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 24, 32, 32], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 24, 32, 32], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 24, 32, 32], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 24 + }, + { + "type": "tensor", + "shape": [1, 24, 32, 32], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 24, 32, 32], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [64, 64] + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 24, 64, 64], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 64, 64], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 64, 64], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [2, 3] + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 16, 1, 1], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 1, 1], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 16, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 16, 64, 64], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 16, 64, 64], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 64, 64], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 64, 64], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 16, 64, 64], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [128, 128] + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 16, 128, 128], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 128, 128], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 128, 128], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [2, 3] + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 16, 1, 1], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 1, 1], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 16, 1, 1], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 16, 128, 128], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 16, 128, 128], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 128, 128], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 128, 128], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 16, 128, 128], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 1, 256, 256], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 1, 256, 256], + "dtype": "float16", + "memObjId": 0 + } + ], + "inputIds": [110], + "outputIds": [489], + "memoryPlan": { + "pools": [ + { + "id": 0, + "size": 608400 + }, + { + "id": 1, + "size": 589824 + }, + { + "id": 2, + "size": 524288 + }, + { + "id": 3, + "size": 524288 + }, + { + "id": 4, + "size": 131072 + }, + { + "id": 5, + "size": 49152 + }, + { + "id": 6, + "size": 16384 + } + ] + } +} diff --git a/example/public/assets/selfie-segmentation/weights-fp16.bin b/example/public/assets/selfie-segmentation/weights-fp16.bin new file mode 100644 index 0000000..2a7a7a3 Binary files /dev/null and b/example/public/assets/selfie-segmentation/weights-fp16.bin differ diff --git a/example/public/assets/yolo26n/program-fp16.json b/example/public/assets/yolo26n/program-fp16.json new file mode 100644 index 0000000..3d6aa1e --- /dev/null +++ b/example/public/assets/yolo26n/program-fp16.json @@ -0,0 +1,7770 @@ +{ + "version": "1", + "chain": [ + { + "name": "et_vk.prepack.default", + "args": [204, 205] + }, + { + "name": "et_vk.prepack.default", + "args": [206, 207] + }, + { + "name": "et_vk.prepack.default", + "args": [208, 209] + }, + { + "name": "et_vk.prepack.default", + "args": [210, 211] + }, + { + "name": "aten.convolution.default", + "args": [212, 0, 1, 213, 214, 215, 216, 217, 218, 219] + }, + { + "name": "aten.sigmoid.default", + "args": [219, 220] + }, + { + "name": "aten.mul.Tensor", + "args": [219, 220, 221] + }, + { + "name": "aten.convolution.default", + "args": [221, 2, 3, 222, 223, 224, 216, 225, 218, 226] + }, + { + "name": "aten.sigmoid.default", + "args": [226, 227] + }, + { + "name": "aten.mul.Tensor", + "args": [226, 227, 228] + }, + { + "name": "aten.convolution.default", + "args": [228, 4, 5, 229, 230, 231, 216, 232, 218, 233] + }, + { + "name": "aten.sigmoid.default", + "args": [233, 234] + }, + { + "name": "aten.mul.Tensor", + "args": [233, 234, 235] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [235, 236, 218, 239] + }, + { + "name": "aten.convolution.default", + "args": [238, 8, 9, 240, 241, 242, 216, 243, 218, 244] + }, + { + "name": "aten.sigmoid.default", + "args": [244, 245] + }, + { + "name": "aten.mul.Tensor", + "args": [244, 245, 246] + }, + { + "name": "aten.convolution.default", + "args": [246, 10, 11, 247, 248, 249, 216, 250, 218, 251] + }, + { + "name": "aten.sigmoid.default", + "args": [251, 252] + }, + { + "name": "aten.mul.Tensor", + "args": [251, 252, 253] + }, + { + "name": "aten.add.Tensor", + "args": [238, 253, 218, 254] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [235, 255, 218, 258] + }, + { + "name": "aten.cat.default", + "args": [259, 218, 260] + }, + { + "name": "aten.convolution.default", + "args": [260, 6, 7, 261, 262, 263, 216, 264, 218, 265] + }, + { + "name": "aten.sigmoid.default", + "args": [265, 266] + }, + { + "name": "aten.mul.Tensor", + "args": [265, 266, 267] + }, + { + "name": "aten.convolution.default", + "args": [267, 12, 13, 268, 269, 270, 216, 271, 218, 272] + }, + { + "name": "aten.sigmoid.default", + "args": [272, 273] + }, + { + "name": "aten.mul.Tensor", + "args": [272, 273, 274] + }, + { + "name": "aten.convolution.default", + "args": [274, 14, 15, 275, 276, 277, 216, 278, 218, 279] + }, + { + "name": "aten.sigmoid.default", + "args": [279, 280] + }, + { + "name": "aten.mul.Tensor", + "args": [279, 280, 281] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [281, 282, 218, 285] + }, + { + "name": "aten.convolution.default", + "args": [284, 18, 19, 286, 287, 288, 216, 289, 218, 290] + }, + { + "name": "aten.sigmoid.default", + "args": [290, 291] + }, + { + "name": "aten.mul.Tensor", + "args": [290, 291, 292] + }, + { + "name": "aten.convolution.default", + "args": [292, 20, 21, 293, 294, 295, 216, 296, 218, 297] + }, + { + "name": "aten.sigmoid.default", + "args": [297, 298] + }, + { + "name": "aten.mul.Tensor", + "args": [297, 298, 299] + }, + { + "name": "aten.add.Tensor", + "args": [284, 299, 218, 300] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [281, 301, 218, 304] + }, + { + "name": "aten.cat.default", + "args": [305, 218, 306] + }, + { + "name": "aten.convolution.default", + "args": [306, 16, 17, 307, 308, 309, 216, 310, 218, 311] + }, + { + "name": "aten.sigmoid.default", + "args": [311, 312] + }, + { + "name": "aten.mul.Tensor", + "args": [311, 312, 313] + }, + { + "name": "aten.convolution.default", + "args": [313, 22, 23, 314, 315, 316, 216, 317, 218, 318] + }, + { + "name": "aten.sigmoid.default", + "args": [318, 319] + }, + { + "name": "aten.mul.Tensor", + "args": [318, 319, 320] + }, + { + "name": "aten.convolution.default", + "args": [320, 24, 25, 321, 322, 323, 216, 324, 218, 325] + }, + { + "name": "aten.sigmoid.default", + "args": [325, 326] + }, + { + "name": "aten.mul.Tensor", + "args": [325, 326, 327] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [327, 328, 218, 331] + }, + { + "name": "aten.convolution.default", + "args": [330, 28, 29, 332, 333, 334, 216, 335, 218, 336] + }, + { + "name": "aten.sigmoid.default", + "args": [336, 337] + }, + { + "name": "aten.mul.Tensor", + "args": [336, 337, 338] + }, + { + "name": "aten.convolution.default", + "args": [338, 34, 35, 339, 340, 341, 216, 342, 218, 343] + }, + { + "name": "aten.sigmoid.default", + "args": [343, 344] + }, + { + "name": "aten.mul.Tensor", + "args": [343, 344, 345] + }, + { + "name": "aten.convolution.default", + "args": [345, 36, 37, 346, 347, 348, 216, 349, 218, 350] + }, + { + "name": "aten.sigmoid.default", + "args": [350, 351] + }, + { + "name": "aten.mul.Tensor", + "args": [350, 351, 352] + }, + { + "name": "aten.add.Tensor", + "args": [338, 352, 218, 353] + }, + { + "name": "aten.convolution.default", + "args": [353, 38, 39, 354, 355, 356, 216, 357, 218, 358] + }, + { + "name": "aten.sigmoid.default", + "args": [358, 359] + }, + { + "name": "aten.mul.Tensor", + "args": [358, 359, 360] + }, + { + "name": "aten.convolution.default", + "args": [360, 40, 41, 361, 362, 363, 216, 364, 218, 365] + }, + { + "name": "aten.sigmoid.default", + "args": [365, 366] + }, + { + "name": "aten.mul.Tensor", + "args": [365, 366, 367] + }, + { + "name": "aten.add.Tensor", + "args": [353, 367, 218, 368] + }, + { + "name": "aten.convolution.default", + "args": [330, 30, 31, 369, 370, 371, 216, 372, 218, 373] + }, + { + "name": "aten.sigmoid.default", + "args": [373, 374] + }, + { + "name": "aten.mul.Tensor", + "args": [373, 374, 375] + }, + { + "name": "aten.cat.default", + "args": [376, 218, 377] + }, + { + "name": "aten.convolution.default", + "args": [377, 32, 33, 378, 379, 380, 216, 381, 218, 382] + }, + { + "name": "aten.sigmoid.default", + "args": [382, 383] + }, + { + "name": "aten.mul.Tensor", + "args": [382, 383, 384] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [327, 385, 218, 388] + }, + { + "name": "aten.cat.default", + "args": [389, 218, 390] + }, + { + "name": "aten.convolution.default", + "args": [390, 26, 27, 391, 392, 393, 216, 394, 218, 395] + }, + { + "name": "aten.sigmoid.default", + "args": [395, 396] + }, + { + "name": "aten.mul.Tensor", + "args": [395, 396, 397] + }, + { + "name": "aten.convolution.default", + "args": [397, 42, 43, 398, 399, 400, 216, 401, 218, 402] + }, + { + "name": "aten.sigmoid.default", + "args": [402, 403] + }, + { + "name": "aten.mul.Tensor", + "args": [402, 403, 404] + }, + { + "name": "aten.convolution.default", + "args": [404, 44, 45, 405, 406, 407, 216, 408, 218, 409] + }, + { + "name": "aten.sigmoid.default", + "args": [409, 410] + }, + { + "name": "aten.mul.Tensor", + "args": [409, 410, 411] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [411, 412, 218, 415] + }, + { + "name": "aten.convolution.default", + "args": [414, 48, 49, 416, 417, 418, 216, 419, 218, 420] + }, + { + "name": "aten.sigmoid.default", + "args": [420, 421] + }, + { + "name": "aten.mul.Tensor", + "args": [420, 421, 422] + }, + { + "name": "aten.convolution.default", + "args": [422, 54, 55, 423, 424, 425, 216, 426, 218, 427] + }, + { + "name": "aten.sigmoid.default", + "args": [427, 428] + }, + { + "name": "aten.mul.Tensor", + "args": [427, 428, 429] + }, + { + "name": "aten.convolution.default", + "args": [429, 56, 57, 430, 431, 432, 216, 433, 218, 434] + }, + { + "name": "aten.sigmoid.default", + "args": [434, 435] + }, + { + "name": "aten.mul.Tensor", + "args": [434, 435, 436] + }, + { + "name": "aten.add.Tensor", + "args": [422, 436, 218, 437] + }, + { + "name": "aten.convolution.default", + "args": [437, 58, 59, 438, 439, 440, 216, 441, 218, 442] + }, + { + "name": "aten.sigmoid.default", + "args": [442, 443] + }, + { + "name": "aten.mul.Tensor", + "args": [442, 443, 444] + }, + { + "name": "aten.convolution.default", + "args": [444, 60, 61, 445, 446, 447, 216, 448, 218, 449] + }, + { + "name": "aten.sigmoid.default", + "args": [449, 450] + }, + { + "name": "aten.mul.Tensor", + "args": [449, 450, 451] + }, + { + "name": "aten.add.Tensor", + "args": [437, 451, 218, 452] + }, + { + "name": "aten.convolution.default", + "args": [414, 50, 51, 453, 454, 455, 216, 456, 218, 457] + }, + { + "name": "aten.sigmoid.default", + "args": [457, 458] + }, + { + "name": "aten.mul.Tensor", + "args": [457, 458, 459] + }, + { + "name": "aten.cat.default", + "args": [460, 218, 461] + }, + { + "name": "aten.convolution.default", + "args": [461, 52, 53, 462, 463, 464, 216, 465, 218, 466] + }, + { + "name": "aten.sigmoid.default", + "args": [466, 467] + }, + { + "name": "aten.mul.Tensor", + "args": [466, 467, 468] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [411, 469, 218, 472] + }, + { + "name": "aten.cat.default", + "args": [473, 218, 474] + }, + { + "name": "aten.convolution.default", + "args": [474, 46, 47, 475, 476, 477, 216, 478, 218, 479] + }, + { + "name": "aten.sigmoid.default", + "args": [479, 480] + }, + { + "name": "aten.mul.Tensor", + "args": [479, 480, 481] + }, + { + "name": "aten.convolution.default", + "args": [481, 62, 63, 482, 483, 484, 216, 485, 218, 486] + }, + { + "name": "aten.max_pool2d_with_indices.default", + "args": [486, 487, 488, 489, 490, 216, 493] + }, + { + "name": "aten.max_pool2d_with_indices.default", + "args": [491, 494, 495, 496, 497, 216, 500] + }, + { + "name": "aten.max_pool2d_with_indices.default", + "args": [498, 501, 502, 503, 504, 216, 507] + }, + { + "name": "aten.cat.default", + "args": [508, 218, 509] + }, + { + "name": "aten.convolution.default", + "args": [509, 64, 65, 510, 511, 512, 216, 513, 218, 514] + }, + { + "name": "aten.sigmoid.default", + "args": [514, 515] + }, + { + "name": "aten.mul.Tensor", + "args": [514, 515, 516] + }, + { + "name": "aten.add.Tensor", + "args": [516, 481, 218, 517] + }, + { + "name": "aten.convolution.default", + "args": [517, 66, 67, 518, 519, 520, 216, 521, 218, 522] + }, + { + "name": "aten.sigmoid.default", + "args": [522, 523] + }, + { + "name": "aten.mul.Tensor", + "args": [522, 523, 524] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [524, 525, 218, 528] + }, + { + "name": "aten.convolution.default", + "args": [527, 70, 71, 529, 530, 531, 216, 532, 218, 533] + }, + { + "name": "aten.clone.default", + "args": [533, 534, 535] + }, + { + "name": "aten.view_copy.default", + "args": [535, 536, 537] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [537, 538, 539, 543] + }, + { + "name": "aten.mul.Tensor", + "args": [540, 205, 544] + }, + { + "name": "aten.permute_copy.default", + "args": [544, 545, 546] + }, + { + "name": "aten.squeeze_copy.dims", + "args": [546, 547, 548] + }, + { + "name": "aten.squeeze_copy.dims", + "args": [541, 549, 550] + }, + { + "name": "aten.bmm.default", + "args": [548, 550, 551] + }, + { + "name": "aten.view_copy.default", + "args": [551, 552, 553] + }, + { + "name": "aten._softmax.default", + "args": [553, 554, 216, 555] + }, + { + "name": "aten.permute_copy.default", + "args": [555, 556, 557] + }, + { + "name": "aten.squeeze_copy.dims", + "args": [542, 558, 559] + }, + { + "name": "aten.squeeze_copy.dims", + "args": [557, 560, 561] + }, + { + "name": "aten.bmm.default", + "args": [559, 561, 562] + }, + { + "name": "aten.clone.default", + "args": [562, 563, 564] + }, + { + "name": "aten.view_copy.default", + "args": [564, 565, 566] + }, + { + "name": "aten.view_copy.default", + "args": [542, 567, 568] + }, + { + "name": "aten.clone.default", + "args": [568, 569, 570] + }, + { + "name": "aten.convolution.default", + "args": [570, 74, 75, 571, 572, 573, 216, 574, 575, 576] + }, + { + "name": "aten.add.Tensor", + "args": [566, 576, 218, 577] + }, + { + "name": "aten.convolution.default", + "args": [577, 72, 73, 578, 579, 580, 216, 581, 218, 582] + }, + { + "name": "aten.add.Tensor", + "args": [527, 582, 218, 583] + }, + { + "name": "aten.convolution.default", + "args": [583, 76, 77, 584, 585, 586, 216, 587, 218, 588] + }, + { + "name": "aten.sigmoid.default", + "args": [588, 589] + }, + { + "name": "aten.mul.Tensor", + "args": [588, 589, 590] + }, + { + "name": "aten.convolution.default", + "args": [590, 78, 79, 591, 592, 593, 216, 594, 218, 595] + }, + { + "name": "aten.add.Tensor", + "args": [583, 595, 218, 596] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [524, 597, 218, 600] + }, + { + "name": "aten.cat.default", + "args": [601, 218, 602] + }, + { + "name": "aten.convolution.default", + "args": [602, 68, 69, 603, 604, 605, 216, 606, 218, 607] + }, + { + "name": "aten.sigmoid.default", + "args": [607, 608] + }, + { + "name": "aten.mul.Tensor", + "args": [607, 608, 609] + }, + { + "name": "aten.upsample_nearest2d.vec", + "args": [609, 610, 611, 612] + }, + { + "name": "aten.cat.default", + "args": [613, 218, 614] + }, + { + "name": "aten.convolution.default", + "args": [614, 80, 81, 615, 616, 617, 216, 618, 218, 619] + }, + { + "name": "aten.sigmoid.default", + "args": [619, 620] + }, + { + "name": "aten.mul.Tensor", + "args": [619, 620, 621] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [621, 622, 218, 625] + }, + { + "name": "aten.convolution.default", + "args": [624, 84, 85, 626, 627, 628, 216, 629, 218, 630] + }, + { + "name": "aten.sigmoid.default", + "args": [630, 631] + }, + { + "name": "aten.mul.Tensor", + "args": [630, 631, 632] + }, + { + "name": "aten.convolution.default", + "args": [632, 90, 91, 633, 634, 635, 216, 636, 218, 637] + }, + { + "name": "aten.sigmoid.default", + "args": [637, 638] + }, + { + "name": "aten.mul.Tensor", + "args": [637, 638, 639] + }, + { + "name": "aten.convolution.default", + "args": [639, 92, 93, 640, 641, 642, 216, 643, 218, 644] + }, + { + "name": "aten.sigmoid.default", + "args": [644, 645] + }, + { + "name": "aten.mul.Tensor", + "args": [644, 645, 646] + }, + { + "name": "aten.add.Tensor", + "args": [632, 646, 218, 647] + }, + { + "name": "aten.convolution.default", + "args": [647, 94, 95, 648, 649, 650, 216, 651, 218, 652] + }, + { + "name": "aten.sigmoid.default", + "args": [652, 653] + }, + { + "name": "aten.mul.Tensor", + "args": [652, 653, 654] + }, + { + "name": "aten.convolution.default", + "args": [654, 96, 97, 655, 656, 657, 216, 658, 218, 659] + }, + { + "name": "aten.sigmoid.default", + "args": [659, 660] + }, + { + "name": "aten.mul.Tensor", + "args": [659, 660, 661] + }, + { + "name": "aten.add.Tensor", + "args": [647, 661, 218, 662] + }, + { + "name": "aten.convolution.default", + "args": [624, 86, 87, 663, 664, 665, 216, 666, 218, 667] + }, + { + "name": "aten.sigmoid.default", + "args": [667, 668] + }, + { + "name": "aten.mul.Tensor", + "args": [667, 668, 669] + }, + { + "name": "aten.cat.default", + "args": [670, 218, 671] + }, + { + "name": "aten.convolution.default", + "args": [671, 88, 89, 672, 673, 674, 216, 675, 218, 676] + }, + { + "name": "aten.sigmoid.default", + "args": [676, 677] + }, + { + "name": "aten.mul.Tensor", + "args": [676, 677, 678] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [621, 679, 218, 682] + }, + { + "name": "aten.cat.default", + "args": [683, 218, 684] + }, + { + "name": "aten.convolution.default", + "args": [684, 82, 83, 685, 686, 687, 216, 688, 218, 689] + }, + { + "name": "aten.sigmoid.default", + "args": [689, 690] + }, + { + "name": "aten.mul.Tensor", + "args": [689, 690, 691] + }, + { + "name": "aten.upsample_nearest2d.vec", + "args": [691, 692, 693, 694] + }, + { + "name": "aten.cat.default", + "args": [695, 218, 696] + }, + { + "name": "aten.convolution.default", + "args": [696, 98, 99, 697, 698, 699, 216, 700, 218, 701] + }, + { + "name": "aten.sigmoid.default", + "args": [701, 702] + }, + { + "name": "aten.mul.Tensor", + "args": [701, 702, 703] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [703, 704, 218, 707] + }, + { + "name": "aten.convolution.default", + "args": [706, 102, 103, 708, 709, 710, 216, 711, 218, 712] + }, + { + "name": "aten.sigmoid.default", + "args": [712, 713] + }, + { + "name": "aten.mul.Tensor", + "args": [712, 713, 714] + }, + { + "name": "aten.convolution.default", + "args": [714, 108, 109, 715, 716, 717, 216, 718, 218, 719] + }, + { + "name": "aten.sigmoid.default", + "args": [719, 720] + }, + { + "name": "aten.mul.Tensor", + "args": [719, 720, 721] + }, + { + "name": "aten.convolution.default", + "args": [721, 110, 111, 722, 723, 724, 216, 725, 218, 726] + }, + { + "name": "aten.sigmoid.default", + "args": [726, 727] + }, + { + "name": "aten.mul.Tensor", + "args": [726, 727, 728] + }, + { + "name": "aten.add.Tensor", + "args": [714, 728, 218, 729] + }, + { + "name": "aten.convolution.default", + "args": [729, 112, 113, 730, 731, 732, 216, 733, 218, 734] + }, + { + "name": "aten.sigmoid.default", + "args": [734, 735] + }, + { + "name": "aten.mul.Tensor", + "args": [734, 735, 736] + }, + { + "name": "aten.convolution.default", + "args": [736, 114, 115, 737, 738, 739, 216, 740, 218, 741] + }, + { + "name": "aten.sigmoid.default", + "args": [741, 742] + }, + { + "name": "aten.mul.Tensor", + "args": [741, 742, 743] + }, + { + "name": "aten.add.Tensor", + "args": [729, 743, 218, 744] + }, + { + "name": "aten.convolution.default", + "args": [706, 104, 105, 745, 746, 747, 216, 748, 218, 749] + }, + { + "name": "aten.sigmoid.default", + "args": [749, 750] + }, + { + "name": "aten.mul.Tensor", + "args": [749, 750, 751] + }, + { + "name": "aten.cat.default", + "args": [752, 218, 753] + }, + { + "name": "aten.convolution.default", + "args": [753, 106, 107, 754, 755, 756, 216, 757, 218, 758] + }, + { + "name": "aten.sigmoid.default", + "args": [758, 759] + }, + { + "name": "aten.mul.Tensor", + "args": [758, 759, 760] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [703, 761, 218, 764] + }, + { + "name": "aten.cat.default", + "args": [765, 218, 766] + }, + { + "name": "aten.convolution.default", + "args": [766, 100, 101, 767, 768, 769, 216, 770, 218, 771] + }, + { + "name": "aten.sigmoid.default", + "args": [771, 772] + }, + { + "name": "aten.mul.Tensor", + "args": [771, 772, 773] + }, + { + "name": "aten.convolution.default", + "args": [773, 116, 117, 774, 775, 776, 216, 777, 218, 778] + }, + { + "name": "aten.sigmoid.default", + "args": [778, 779] + }, + { + "name": "aten.mul.Tensor", + "args": [778, 779, 780] + }, + { + "name": "aten.cat.default", + "args": [781, 218, 782] + }, + { + "name": "aten.convolution.default", + "args": [782, 118, 119, 783, 784, 785, 216, 786, 218, 787] + }, + { + "name": "aten.sigmoid.default", + "args": [787, 788] + }, + { + "name": "aten.mul.Tensor", + "args": [787, 788, 789] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [789, 790, 218, 793] + }, + { + "name": "aten.convolution.default", + "args": [792, 122, 123, 794, 795, 796, 216, 797, 218, 798] + }, + { + "name": "aten.sigmoid.default", + "args": [798, 799] + }, + { + "name": "aten.mul.Tensor", + "args": [798, 799, 800] + }, + { + "name": "aten.convolution.default", + "args": [800, 128, 129, 801, 802, 803, 216, 804, 218, 805] + }, + { + "name": "aten.sigmoid.default", + "args": [805, 806] + }, + { + "name": "aten.mul.Tensor", + "args": [805, 806, 807] + }, + { + "name": "aten.convolution.default", + "args": [807, 130, 131, 808, 809, 810, 216, 811, 218, 812] + }, + { + "name": "aten.sigmoid.default", + "args": [812, 813] + }, + { + "name": "aten.mul.Tensor", + "args": [812, 813, 814] + }, + { + "name": "aten.add.Tensor", + "args": [800, 814, 218, 815] + }, + { + "name": "aten.convolution.default", + "args": [815, 132, 133, 816, 817, 818, 216, 819, 218, 820] + }, + { + "name": "aten.sigmoid.default", + "args": [820, 821] + }, + { + "name": "aten.mul.Tensor", + "args": [820, 821, 822] + }, + { + "name": "aten.convolution.default", + "args": [822, 134, 135, 823, 824, 825, 216, 826, 218, 827] + }, + { + "name": "aten.sigmoid.default", + "args": [827, 828] + }, + { + "name": "aten.mul.Tensor", + "args": [827, 828, 829] + }, + { + "name": "aten.add.Tensor", + "args": [815, 829, 218, 830] + }, + { + "name": "aten.convolution.default", + "args": [792, 124, 125, 831, 832, 833, 216, 834, 218, 835] + }, + { + "name": "aten.sigmoid.default", + "args": [835, 836] + }, + { + "name": "aten.mul.Tensor", + "args": [835, 836, 837] + }, + { + "name": "aten.cat.default", + "args": [838, 218, 839] + }, + { + "name": "aten.convolution.default", + "args": [839, 126, 127, 840, 841, 842, 216, 843, 218, 844] + }, + { + "name": "aten.sigmoid.default", + "args": [844, 845] + }, + { + "name": "aten.mul.Tensor", + "args": [844, 845, 846] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [789, 847, 218, 850] + }, + { + "name": "aten.cat.default", + "args": [851, 218, 852] + }, + { + "name": "aten.convolution.default", + "args": [852, 120, 121, 853, 854, 855, 216, 856, 218, 857] + }, + { + "name": "aten.sigmoid.default", + "args": [857, 858] + }, + { + "name": "aten.mul.Tensor", + "args": [857, 858, 859] + }, + { + "name": "aten.convolution.default", + "args": [859, 136, 137, 860, 861, 862, 216, 863, 218, 864] + }, + { + "name": "aten.sigmoid.default", + "args": [864, 865] + }, + { + "name": "aten.mul.Tensor", + "args": [864, 865, 866] + }, + { + "name": "aten.cat.default", + "args": [867, 218, 868] + }, + { + "name": "aten.convolution.default", + "args": [868, 138, 139, 869, 870, 871, 216, 872, 218, 873] + }, + { + "name": "aten.sigmoid.default", + "args": [873, 874] + }, + { + "name": "aten.mul.Tensor", + "args": [873, 874, 875] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [875, 876, 218, 879] + }, + { + "name": "aten.convolution.default", + "args": [878, 142, 143, 880, 881, 882, 216, 883, 218, 884] + }, + { + "name": "aten.sigmoid.default", + "args": [884, 885] + }, + { + "name": "aten.mul.Tensor", + "args": [884, 885, 886] + }, + { + "name": "aten.convolution.default", + "args": [886, 144, 145, 887, 888, 889, 216, 890, 218, 891] + }, + { + "name": "aten.sigmoid.default", + "args": [891, 892] + }, + { + "name": "aten.mul.Tensor", + "args": [891, 892, 893] + }, + { + "name": "aten.add.Tensor", + "args": [878, 893, 218, 894] + }, + { + "name": "aten.convolution.default", + "args": [894, 146, 147, 895, 896, 897, 216, 898, 218, 899] + }, + { + "name": "aten.clone.default", + "args": [899, 900, 901] + }, + { + "name": "aten.view_copy.default", + "args": [901, 902, 903] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [903, 904, 539, 908] + }, + { + "name": "aten.mul.Tensor", + "args": [905, 207, 909] + }, + { + "name": "aten.permute_copy.default", + "args": [909, 910, 911] + }, + { + "name": "aten.squeeze_copy.dims", + "args": [911, 912, 913] + }, + { + "name": "aten.squeeze_copy.dims", + "args": [906, 914, 915] + }, + { + "name": "aten.bmm.default", + "args": [913, 915, 916] + }, + { + "name": "aten.view_copy.default", + "args": [916, 917, 918] + }, + { + "name": "aten._softmax.default", + "args": [918, 554, 216, 919] + }, + { + "name": "aten.permute_copy.default", + "args": [919, 920, 921] + }, + { + "name": "aten.squeeze_copy.dims", + "args": [907, 922, 923] + }, + { + "name": "aten.squeeze_copy.dims", + "args": [921, 924, 925] + }, + { + "name": "aten.bmm.default", + "args": [923, 925, 926] + }, + { + "name": "aten.clone.default", + "args": [926, 927, 928] + }, + { + "name": "aten.view_copy.default", + "args": [928, 929, 930] + }, + { + "name": "aten.view_copy.default", + "args": [907, 931, 932] + }, + { + "name": "aten.clone.default", + "args": [932, 933, 934] + }, + { + "name": "aten.convolution.default", + "args": [934, 150, 151, 935, 936, 937, 216, 938, 575, 939] + }, + { + "name": "aten.add.Tensor", + "args": [930, 939, 218, 940] + }, + { + "name": "aten.convolution.default", + "args": [940, 148, 149, 941, 942, 943, 216, 944, 218, 945] + }, + { + "name": "aten.add.Tensor", + "args": [894, 945, 218, 946] + }, + { + "name": "aten.convolution.default", + "args": [946, 152, 153, 947, 948, 949, 216, 950, 218, 951] + }, + { + "name": "aten.sigmoid.default", + "args": [951, 952] + }, + { + "name": "aten.mul.Tensor", + "args": [951, 952, 953] + }, + { + "name": "aten.convolution.default", + "args": [953, 154, 155, 954, 955, 956, 216, 957, 218, 958] + }, + { + "name": "aten.add.Tensor", + "args": [946, 958, 218, 959] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [875, 960, 218, 963] + }, + { + "name": "aten.cat.default", + "args": [964, 218, 965] + }, + { + "name": "aten.convolution.default", + "args": [965, 140, 141, 966, 967, 968, 216, 969, 218, 970] + }, + { + "name": "aten.sigmoid.default", + "args": [970, 971] + }, + { + "name": "aten.mul.Tensor", + "args": [970, 971, 972] + }, + { + "name": "aten.convolution.default", + "args": [773, 156, 157, 973, 974, 975, 216, 976, 218, 977] + }, + { + "name": "aten.sigmoid.default", + "args": [977, 978] + }, + { + "name": "aten.mul.Tensor", + "args": [977, 978, 979] + }, + { + "name": "aten.convolution.default", + "args": [979, 158, 159, 980, 981, 982, 216, 983, 218, 984] + }, + { + "name": "aten.sigmoid.default", + "args": [984, 985] + }, + { + "name": "aten.mul.Tensor", + "args": [984, 985, 986] + }, + { + "name": "aten.convolution.default", + "args": [986, 160, 161, 987, 988, 989, 216, 990, 218, 991] + }, + { + "name": "aten.view_copy.default", + "args": [991, 992, 993] + }, + { + "name": "aten.convolution.default", + "args": [859, 162, 163, 994, 995, 996, 216, 997, 218, 998] + }, + { + "name": "aten.sigmoid.default", + "args": [998, 999] + }, + { + "name": "aten.mul.Tensor", + "args": [998, 999, 1000] + }, + { + "name": "aten.convolution.default", + "args": [1000, 164, 165, 1001, 1002, 1003, 216, 1004, 218, 1005] + }, + { + "name": "aten.sigmoid.default", + "args": [1005, 1006] + }, + { + "name": "aten.mul.Tensor", + "args": [1005, 1006, 1007] + }, + { + "name": "aten.convolution.default", + "args": [1007, 166, 167, 1008, 1009, 1010, 216, 1011, 218, 1012] + }, + { + "name": "aten.view_copy.default", + "args": [1012, 1013, 1014] + }, + { + "name": "aten.convolution.default", + "args": [972, 168, 169, 1015, 1016, 1017, 216, 1018, 218, 1019] + }, + { + "name": "aten.sigmoid.default", + "args": [1019, 1020] + }, + { + "name": "aten.mul.Tensor", + "args": [1019, 1020, 1021] + }, + { + "name": "aten.convolution.default", + "args": [1021, 170, 171, 1022, 1023, 1024, 216, 1025, 218, 1026] + }, + { + "name": "aten.sigmoid.default", + "args": [1026, 1027] + }, + { + "name": "aten.mul.Tensor", + "args": [1026, 1027, 1028] + }, + { + "name": "aten.convolution.default", + "args": [1028, 172, 173, 1029, 1030, 1031, 216, 1032, 218, 1033] + }, + { + "name": "aten.view_copy.default", + "args": [1033, 1034, 1035] + }, + { + "name": "aten.cat.default", + "args": [1036, 554, 1037] + }, + { + "name": "aten.convolution.default", + "args": [773, 174, 175, 1038, 1039, 1040, 216, 1041, 1042, 1043] + }, + { + "name": "aten.sigmoid.default", + "args": [1043, 1044] + }, + { + "name": "aten.mul.Tensor", + "args": [1043, 1044, 1045] + }, + { + "name": "aten.convolution.default", + "args": [1045, 176, 177, 1046, 1047, 1048, 216, 1049, 218, 1050] + }, + { + "name": "aten.sigmoid.default", + "args": [1050, 1051] + }, + { + "name": "aten.mul.Tensor", + "args": [1050, 1051, 1052] + }, + { + "name": "aten.convolution.default", + "args": [1052, 178, 179, 1053, 1054, 1055, 216, 1056, 1057, 1058] + }, + { + "name": "aten.sigmoid.default", + "args": [1058, 1059] + }, + { + "name": "aten.mul.Tensor", + "args": [1058, 1059, 1060] + }, + { + "name": "aten.convolution.default", + "args": [1060, 180, 181, 1061, 1062, 1063, 216, 1064, 218, 1065] + }, + { + "name": "aten.sigmoid.default", + "args": [1065, 1066] + }, + { + "name": "aten.mul.Tensor", + "args": [1065, 1066, 1067] + }, + { + "name": "aten.convolution.default", + "args": [1067, 182, 183, 1068, 1069, 1070, 216, 1071, 218, 1072] + }, + { + "name": "aten.view_copy.default", + "args": [1072, 1073, 1074] + }, + { + "name": "aten.convolution.default", + "args": [859, 184, 185, 1075, 1076, 1077, 216, 1078, 575, 1079] + }, + { + "name": "aten.sigmoid.default", + "args": [1079, 1080] + }, + { + "name": "aten.mul.Tensor", + "args": [1079, 1080, 1081] + }, + { + "name": "aten.convolution.default", + "args": [1081, 186, 187, 1082, 1083, 1084, 216, 1085, 218, 1086] + }, + { + "name": "aten.sigmoid.default", + "args": [1086, 1087] + }, + { + "name": "aten.mul.Tensor", + "args": [1086, 1087, 1088] + }, + { + "name": "aten.convolution.default", + "args": [1088, 188, 189, 1089, 1090, 1091, 216, 1092, 1057, 1093] + }, + { + "name": "aten.sigmoid.default", + "args": [1093, 1094] + }, + { + "name": "aten.mul.Tensor", + "args": [1093, 1094, 1095] + }, + { + "name": "aten.convolution.default", + "args": [1095, 190, 191, 1096, 1097, 1098, 216, 1099, 218, 1100] + }, + { + "name": "aten.sigmoid.default", + "args": [1100, 1101] + }, + { + "name": "aten.mul.Tensor", + "args": [1100, 1101, 1102] + }, + { + "name": "aten.convolution.default", + "args": [1102, 192, 193, 1103, 1104, 1105, 216, 1106, 218, 1107] + }, + { + "name": "aten.view_copy.default", + "args": [1107, 1108, 1109] + }, + { + "name": "aten.convolution.default", + "args": [972, 194, 195, 1110, 1111, 1112, 216, 1113, 1114, 1115] + }, + { + "name": "aten.sigmoid.default", + "args": [1115, 1116] + }, + { + "name": "aten.mul.Tensor", + "args": [1115, 1116, 1117] + }, + { + "name": "aten.convolution.default", + "args": [1117, 196, 197, 1118, 1119, 1120, 216, 1121, 218, 1122] + }, + { + "name": "aten.sigmoid.default", + "args": [1122, 1123] + }, + { + "name": "aten.mul.Tensor", + "args": [1122, 1123, 1124] + }, + { + "name": "aten.convolution.default", + "args": [1124, 198, 199, 1125, 1126, 1127, 216, 1128, 1057, 1129] + }, + { + "name": "aten.sigmoid.default", + "args": [1129, 1130] + }, + { + "name": "aten.mul.Tensor", + "args": [1129, 1130, 1131] + }, + { + "name": "aten.convolution.default", + "args": [1131, 200, 201, 1132, 1133, 1134, 216, 1135, 218, 1136] + }, + { + "name": "aten.sigmoid.default", + "args": [1136, 1137] + }, + { + "name": "aten.mul.Tensor", + "args": [1136, 1137, 1138] + }, + { + "name": "aten.convolution.default", + "args": [1138, 202, 203, 1139, 1140, 1141, 216, 1142, 218, 1143] + }, + { + "name": "aten.view_copy.default", + "args": [1143, 1144, 1145] + }, + { + "name": "aten.cat.default", + "args": [1146, 554, 1147] + }, + { + "name": "aten.unsqueeze_copy.default", + "args": [209, 1148, 1149] + }, + { + "name": "aten.split_with_sizes_copy.default", + "args": [1037, 1150, 218, 1153] + }, + { + "name": "aten.clone.default", + "args": [1151, 1154, 1155] + }, + { + "name": "aten.sub.Tensor", + "args": [1149, 1155, 218, 1156] + }, + { + "name": "aten.clone.default", + "args": [1152, 1157, 1158] + }, + { + "name": "aten.add.Tensor", + "args": [1149, 1158, 218, 1159] + }, + { + "name": "aten.cat.default", + "args": [1160, 218, 1161] + }, + { + "name": "aten.mul.Tensor", + "args": [1161, 211, 1162] + }, + { + "name": "aten.sigmoid.default", + "args": [1147, 1163] + }, + { + "name": "aten.clone.default", + "args": [1163, 1164, 1165] + }, + { + "name": "aten.cat.default", + "args": [1166, 218, 1167] + } + ], + "values": [ + { + "type": "tensor", + "shape": [16, 3, 3, 3], + "dtype": "float16", + "weightsOffset": 0, + "weightsLength": 864 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 864, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [32, 16, 3, 3], + "dtype": "float16", + "weightsOffset": 896, + "weightsLength": 9216 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 10112, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [32, 32, 1, 1], + "dtype": "float16", + "weightsOffset": 10176, + "weightsLength": 2048 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 12224, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [64, 48, 1, 1], + "dtype": "float16", + "weightsOffset": 12288, + "weightsLength": 6144 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 18432, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [8, 16, 3, 3], + "dtype": "float16", + "weightsOffset": 18560, + "weightsLength": 2304 + }, + { + "type": "tensor", + "shape": [8], + "dtype": "float16", + "weightsOffset": 20864, + "weightsLength": 16 + }, + { + "type": "tensor", + "shape": [16, 8, 3, 3], + "dtype": "float16", + "weightsOffset": 20880, + "weightsLength": 2304 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 23184, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [64, 64, 3, 3], + "dtype": "float16", + "weightsOffset": 23216, + "weightsLength": 73728 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 96944, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [64, 64, 1, 1], + "dtype": "float16", + "weightsOffset": 97072, + "weightsLength": 8192 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 105264, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [128, 96, 1, 1], + "dtype": "float16", + "weightsOffset": 105392, + "weightsLength": 24576 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 129968, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [16, 32, 3, 3], + "dtype": "float16", + "weightsOffset": 130224, + "weightsLength": 9216 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 139440, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [32, 16, 3, 3], + "dtype": "float16", + "weightsOffset": 139472, + "weightsLength": 9216 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 148688, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [128, 128, 3, 3], + "dtype": "float16", + "weightsOffset": 148752, + "weightsLength": 294912 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 443664, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [128, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 443920, + "weightsLength": 32768 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 476688, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [128, 192, 1, 1], + "dtype": "float16", + "weightsOffset": 476944, + "weightsLength": 49152 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 526096, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [32, 64, 1, 1], + "dtype": "float16", + "weightsOffset": 526352, + "weightsLength": 4096 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 530448, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [32, 64, 1, 1], + "dtype": "float16", + "weightsOffset": 530512, + "weightsLength": 4096 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 534608, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [64, 64, 1, 1], + "dtype": "float16", + "weightsOffset": 534672, + "weightsLength": 8192 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 542864, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [32, 32, 3, 3], + "dtype": "float16", + "weightsOffset": 542992, + "weightsLength": 18432 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 561424, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [32, 32, 3, 3], + "dtype": "float16", + "weightsOffset": 561488, + "weightsLength": 18432 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 579920, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [32, 32, 3, 3], + "dtype": "float16", + "weightsOffset": 579984, + "weightsLength": 18432 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 598416, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [32, 32, 3, 3], + "dtype": "float16", + "weightsOffset": 598480, + "weightsLength": 18432 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 616912, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [256, 128, 3, 3], + "dtype": "float16", + "weightsOffset": 616976, + "weightsLength": 589824 + }, + { + "type": "tensor", + "shape": [256], + "dtype": "float16", + "weightsOffset": 1206800, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [256, 256, 1, 1], + "dtype": "float16", + "weightsOffset": 1207312, + "weightsLength": 131072 + }, + { + "type": "tensor", + "shape": [256], + "dtype": "float16", + "weightsOffset": 1338384, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [256, 384, 1, 1], + "dtype": "float16", + "weightsOffset": 1338896, + "weightsLength": 196608 + }, + { + "type": "tensor", + "shape": [256], + "dtype": "float16", + "weightsOffset": 1535504, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [64, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 1536016, + "weightsLength": 16384 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 1552400, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [64, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 1552528, + "weightsLength": 16384 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 1568912, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [128, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 1569040, + "weightsLength": 32768 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 1601808, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [64, 64, 3, 3], + "dtype": "float16", + "weightsOffset": 1602064, + "weightsLength": 73728 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 1675792, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [64, 64, 3, 3], + "dtype": "float16", + "weightsOffset": 1675920, + "weightsLength": 73728 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 1749648, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [64, 64, 3, 3], + "dtype": "float16", + "weightsOffset": 1749776, + "weightsLength": 73728 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 1823504, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [64, 64, 3, 3], + "dtype": "float16", + "weightsOffset": 1823632, + "weightsLength": 73728 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 1897360, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [128, 256, 1, 1], + "dtype": "float16", + "weightsOffset": 1897488, + "weightsLength": 65536 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 1963024, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [256, 512, 1, 1], + "dtype": "float16", + "weightsOffset": 1963280, + "weightsLength": 262144 + }, + { + "type": "tensor", + "shape": [256], + "dtype": "float16", + "weightsOffset": 2225424, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [256, 256, 1, 1], + "dtype": "float16", + "weightsOffset": 2225936, + "weightsLength": 131072 + }, + { + "type": "tensor", + "shape": [256], + "dtype": "float16", + "weightsOffset": 2357008, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [256, 256, 1, 1], + "dtype": "float16", + "weightsOffset": 2357520, + "weightsLength": 131072 + }, + { + "type": "tensor", + "shape": [256], + "dtype": "float16", + "weightsOffset": 2488592, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [256, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 2489104, + "weightsLength": 65536 + }, + { + "type": "tensor", + "shape": [256], + "dtype": "float16", + "weightsOffset": 2554640, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [128, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 2555152, + "weightsLength": 32768 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 2587920, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [128, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 2588176, + "weightsLength": 2304 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 2590480, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [256, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 2590736, + "weightsLength": 65536 + }, + { + "type": "tensor", + "shape": [256], + "dtype": "float16", + "weightsOffset": 2656272, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [128, 256, 1, 1], + "dtype": "float16", + "weightsOffset": 2656784, + "weightsLength": 65536 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 2722320, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [128, 384, 1, 1], + "dtype": "float16", + "weightsOffset": 2722576, + "weightsLength": 98304 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 2820880, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [128, 192, 1, 1], + "dtype": "float16", + "weightsOffset": 2821136, + "weightsLength": 49152 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 2870288, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [32, 64, 1, 1], + "dtype": "float16", + "weightsOffset": 2870544, + "weightsLength": 4096 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 2874640, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [32, 64, 1, 1], + "dtype": "float16", + "weightsOffset": 2874704, + "weightsLength": 4096 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 2878800, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [64, 64, 1, 1], + "dtype": "float16", + "weightsOffset": 2878864, + "weightsLength": 8192 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 2887056, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [32, 32, 3, 3], + "dtype": "float16", + "weightsOffset": 2887184, + "weightsLength": 18432 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 2905616, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [32, 32, 3, 3], + "dtype": "float16", + "weightsOffset": 2905680, + "weightsLength": 18432 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 2924112, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [32, 32, 3, 3], + "dtype": "float16", + "weightsOffset": 2924176, + "weightsLength": 18432 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 2942608, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [32, 32, 3, 3], + "dtype": "float16", + "weightsOffset": 2942672, + "weightsLength": 18432 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 2961104, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [64, 256, 1, 1], + "dtype": "float16", + "weightsOffset": 2961168, + "weightsLength": 32768 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 2993936, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [64, 96, 1, 1], + "dtype": "float16", + "weightsOffset": 2994064, + "weightsLength": 12288 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 3006352, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [16, 32, 1, 1], + "dtype": "float16", + "weightsOffset": 3006480, + "weightsLength": 1024 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 3007504, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 32, 1, 1], + "dtype": "float16", + "weightsOffset": 3007536, + "weightsLength": 1024 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 3008560, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [32, 32, 1, 1], + "dtype": "float16", + "weightsOffset": 3008592, + "weightsLength": 2048 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 3010640, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [16, 16, 3, 3], + "dtype": "float16", + "weightsOffset": 3010704, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 3015312, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 16, 3, 3], + "dtype": "float16", + "weightsOffset": 3015344, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 3019952, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 16, 3, 3], + "dtype": "float16", + "weightsOffset": 3019984, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 3024592, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 16, 3, 3], + "dtype": "float16", + "weightsOffset": 3024624, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 3029232, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [64, 64, 3, 3], + "dtype": "float16", + "weightsOffset": 3029264, + "weightsLength": 73728 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 3102992, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [128, 192, 1, 1], + "dtype": "float16", + "weightsOffset": 3103120, + "weightsLength": 49152 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 3152272, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [128, 192, 1, 1], + "dtype": "float16", + "weightsOffset": 3152528, + "weightsLength": 49152 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 3201680, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [32, 64, 1, 1], + "dtype": "float16", + "weightsOffset": 3201936, + "weightsLength": 4096 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 3206032, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [32, 64, 1, 1], + "dtype": "float16", + "weightsOffset": 3206096, + "weightsLength": 4096 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 3210192, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [64, 64, 1, 1], + "dtype": "float16", + "weightsOffset": 3210256, + "weightsLength": 8192 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 3218448, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [32, 32, 3, 3], + "dtype": "float16", + "weightsOffset": 3218576, + "weightsLength": 18432 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 3237008, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [32, 32, 3, 3], + "dtype": "float16", + "weightsOffset": 3237072, + "weightsLength": 18432 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 3255504, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [32, 32, 3, 3], + "dtype": "float16", + "weightsOffset": 3255568, + "weightsLength": 18432 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 3274000, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [32, 32, 3, 3], + "dtype": "float16", + "weightsOffset": 3274064, + "weightsLength": 18432 + }, + { + "type": "tensor", + "shape": [32], + "dtype": "float16", + "weightsOffset": 3292496, + "weightsLength": 64 + }, + { + "type": "tensor", + "shape": [128, 128, 3, 3], + "dtype": "float16", + "weightsOffset": 3292560, + "weightsLength": 294912 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 3587472, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [256, 384, 1, 1], + "dtype": "float16", + "weightsOffset": 3587728, + "weightsLength": 196608 + }, + { + "type": "tensor", + "shape": [256], + "dtype": "float16", + "weightsOffset": 3784336, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [256, 384, 1, 1], + "dtype": "float16", + "weightsOffset": 3784848, + "weightsLength": 196608 + }, + { + "type": "tensor", + "shape": [256], + "dtype": "float16", + "weightsOffset": 3981456, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [64, 128, 3, 3], + "dtype": "float16", + "weightsOffset": 3981968, + "weightsLength": 147456 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 4129424, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [128, 64, 3, 3], + "dtype": "float16", + "weightsOffset": 4129552, + "weightsLength": 147456 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 4277008, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [256, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 4277264, + "weightsLength": 65536 + }, + { + "type": "tensor", + "shape": [256], + "dtype": "float16", + "weightsOffset": 4342800, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [128, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 4343312, + "weightsLength": 32768 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 4376080, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [128, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 4376336, + "weightsLength": 2304 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 4378640, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [256, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 4378896, + "weightsLength": 65536 + }, + { + "type": "tensor", + "shape": [256], + "dtype": "float16", + "weightsOffset": 4444432, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [128, 256, 1, 1], + "dtype": "float16", + "weightsOffset": 4444944, + "weightsLength": 65536 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 4510480, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [16, 64, 3, 3], + "dtype": "float16", + "weightsOffset": 4510736, + "weightsLength": 18432 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 4529168, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 16, 3, 3], + "dtype": "float16", + "weightsOffset": 4529200, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 4533808, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [4, 16, 1, 1], + "dtype": "float16", + "weightsOffset": 4533840, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [4], + "dtype": "float16", + "weightsOffset": 4533968, + "weightsLength": 8 + }, + { + "type": "tensor", + "shape": [16, 128, 3, 3], + "dtype": "float16", + "weightsOffset": 4533976, + "weightsLength": 36864 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 4570840, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 16, 3, 3], + "dtype": "float16", + "weightsOffset": 4570872, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 4575480, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [4, 16, 1, 1], + "dtype": "float16", + "weightsOffset": 4575512, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [4], + "dtype": "float16", + "weightsOffset": 4575640, + "weightsLength": 8 + }, + { + "type": "tensor", + "shape": [16, 256, 3, 3], + "dtype": "float16", + "weightsOffset": 4575648, + "weightsLength": 73728 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 4649376, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [16, 16, 3, 3], + "dtype": "float16", + "weightsOffset": 4649408, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [16], + "dtype": "float16", + "weightsOffset": 4654016, + "weightsLength": 32 + }, + { + "type": "tensor", + "shape": [4, 16, 1, 1], + "dtype": "float16", + "weightsOffset": 4654048, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [4], + "dtype": "float16", + "weightsOffset": 4654176, + "weightsLength": 8 + }, + { + "type": "tensor", + "shape": [64, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 4654184, + "weightsLength": 1152 + }, + { + "type": "tensor", + "shape": [64], + "dtype": "float16", + "weightsOffset": 4655336, + "weightsLength": 128 + }, + { + "type": "tensor", + "shape": [80, 64, 1, 1], + "dtype": "float16", + "weightsOffset": 4655464, + "weightsLength": 10240 + }, + { + "type": "tensor", + "shape": [80], + "dtype": "float16", + "weightsOffset": 4665704, + "weightsLength": 160 + }, + { + "type": "tensor", + "shape": [80, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 4665864, + "weightsLength": 1440 + }, + { + "type": "tensor", + "shape": [80], + "dtype": "float16", + "weightsOffset": 4667304, + "weightsLength": 160 + }, + { + "type": "tensor", + "shape": [80, 80, 1, 1], + "dtype": "float16", + "weightsOffset": 4667464, + "weightsLength": 12800 + }, + { + "type": "tensor", + "shape": [80], + "dtype": "float16", + "weightsOffset": 4680264, + "weightsLength": 160 + }, + { + "type": "tensor", + "shape": [80, 80, 1, 1], + "dtype": "float16", + "weightsOffset": 4680424, + "weightsLength": 12800 + }, + { + "type": "tensor", + "shape": [80], + "dtype": "float16", + "weightsOffset": 4693224, + "weightsLength": 160 + }, + { + "type": "tensor", + "shape": [128, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 4693384, + "weightsLength": 2304 + }, + { + "type": "tensor", + "shape": [128], + "dtype": "float16", + "weightsOffset": 4695688, + "weightsLength": 256 + }, + { + "type": "tensor", + "shape": [80, 128, 1, 1], + "dtype": "float16", + "weightsOffset": 4695944, + "weightsLength": 20480 + }, + { + "type": "tensor", + "shape": [80], + "dtype": "float16", + "weightsOffset": 4716424, + "weightsLength": 160 + }, + { + "type": "tensor", + "shape": [80, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 4716584, + "weightsLength": 1440 + }, + { + "type": "tensor", + "shape": [80], + "dtype": "float16", + "weightsOffset": 4718024, + "weightsLength": 160 + }, + { + "type": "tensor", + "shape": [80, 80, 1, 1], + "dtype": "float16", + "weightsOffset": 4718184, + "weightsLength": 12800 + }, + { + "type": "tensor", + "shape": [80], + "dtype": "float16", + "weightsOffset": 4730984, + "weightsLength": 160 + }, + { + "type": "tensor", + "shape": [80, 80, 1, 1], + "dtype": "float16", + "weightsOffset": 4731144, + "weightsLength": 12800 + }, + { + "type": "tensor", + "shape": [80], + "dtype": "float16", + "weightsOffset": 4743944, + "weightsLength": 160 + }, + { + "type": "tensor", + "shape": [256, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 4744104, + "weightsLength": 4608 + }, + { + "type": "tensor", + "shape": [256], + "dtype": "float16", + "weightsOffset": 4748712, + "weightsLength": 512 + }, + { + "type": "tensor", + "shape": [80, 256, 1, 1], + "dtype": "float16", + "weightsOffset": 4749224, + "weightsLength": 40960 + }, + { + "type": "tensor", + "shape": [80], + "dtype": "float16", + "weightsOffset": 4790184, + "weightsLength": 160 + }, + { + "type": "tensor", + "shape": [80, 1, 3, 3], + "dtype": "float16", + "weightsOffset": 4790344, + "weightsLength": 1440 + }, + { + "type": "tensor", + "shape": [80], + "dtype": "float16", + "weightsOffset": 4791784, + "weightsLength": 160 + }, + { + "type": "tensor", + "shape": [80, 80, 1, 1], + "dtype": "float16", + "weightsOffset": 4791944, + "weightsLength": 12800 + }, + { + "type": "tensor", + "shape": [80], + "dtype": "float16", + "weightsOffset": 4804744, + "weightsLength": 160 + }, + { + "type": "tensor", + "shape": [80, 80, 1, 1], + "dtype": "float16", + "weightsOffset": 4804904, + "weightsLength": 12800 + }, + { + "type": "tensor", + "shape": [80], + "dtype": "float16", + "weightsOffset": 4817704, + "weightsLength": 160 + }, + { + "type": "tensor", + "shape": [], + "dtype": "float16", + "weightsOffset": 4817864, + "weightsLength": 2 + }, + { + "type": "tensor", + "shape": [], + "dtype": "float16" + }, + { + "type": "tensor", + "shape": [], + "dtype": "float16", + "weightsOffset": 4817868, + "weightsLength": 2 + }, + { + "type": "tensor", + "shape": [], + "dtype": "float16" + }, + { + "type": "tensor", + "shape": [2, 8400], + "dtype": "float16", + "weightsOffset": 4817872, + "weightsLength": 33600 + }, + { + "type": "tensor", + "shape": [2, 8400], + "dtype": "float16" + }, + { + "type": "tensor", + "shape": [1, 8400], + "dtype": "float16", + "weightsOffset": 4851472, + "weightsLength": 16800 + }, + { + "type": "tensor", + "shape": [1, 8400], + "dtype": "float16" + }, + { + "type": "tensor", + "shape": [1, 3, 640, 640], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "scalar", + "value": false + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 1 + }, + { + "type": "tensor", + "shape": [1, 16, 320, 320], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 320, 320], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 16, 320, 320], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 160, 160], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 32, 160, 160], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 160, 160], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 160, 160], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 32, 160, 160], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 160, 160], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [16, 16] + }, + { + "type": "tensor", + "shape": [1, 16, 160, 160], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 16, 160, 160], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensorlist", + "items": [237, 238] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 8, 160, 160], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 8, 160, 160], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 8, 160, 160], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 160, 160], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 16, 160, 160], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 160, 160], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 16, 160, 160], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "list", + "items": [16, 16] + }, + { + "type": "tensor", + "shape": [1, 16, 160, 160], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 160, 160], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensorlist", + "items": [256, 257] + }, + { + "type": "tensorlist", + "items": [256, 238, 254] + }, + { + "type": "tensor", + "shape": [1, 48, 160, 160], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 160, 160], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 64, 160, 160], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 64, 160, 160], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 64, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 64, 80, 80], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 64, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 64, 80, 80], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [32, 32] + }, + { + "type": "tensor", + "shape": [1, 32, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 80, 80], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensorlist", + "items": [283, 284] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 80, 80], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 32, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 32, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 80, 80], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "list", + "items": [32, 32] + }, + { + "type": "tensor", + "shape": [1, 32, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 32, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensorlist", + "items": [302, 303] + }, + { + "type": "tensorlist", + "items": [302, 284, 300] + }, + { + "type": "tensor", + "shape": [1, 96, 80, 80], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 80, 80], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 128, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [64, 64] + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "tensorlist", + "items": [329, 330] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 7 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensorlist", + "items": [368, 375] + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "list", + "items": [64, 64] + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensorlist", + "items": [386, 387] + }, + { + "type": "tensorlist", + "items": [386, 330, 384] + }, + { + "type": "tensor", + "shape": [1, 192, 40, 40], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [128, 128] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "tensorlist", + "items": [413, 414] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 7 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 8 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 7 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensorlist", + "items": [452, 459] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "list", + "items": [128, 128] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensorlist", + "items": [470, 471] + }, + { + "type": "tensorlist", + "items": [470, 414, 468] + }, + { + "type": "tensor", + "shape": [1, 384, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 7 + }, + { + "type": "list", + "items": [5, 5] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "int32", + "memObjId": 0 + }, + { + "type": "tensorlist", + "items": [491, 492] + }, + { + "type": "list", + "items": [5, 5] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "int32", + "memObjId": 0 + }, + { + "type": "tensorlist", + "items": [498, 499] + }, + { + "type": "list", + "items": [5, 5] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "int32", + "memObjId": 0 + }, + { + "type": "tensorlist", + "items": [505, 506] + }, + { + "type": "tensorlist", + "items": [486, 491, 498, 505] + }, + { + "type": "tensor", + "shape": [1, 512, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "list", + "items": [128, 128] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 7 + }, + { + "type": "tensorlist", + "items": [526, 527] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 2, 128, 400] + }, + { + "type": "tensor", + "shape": [1, 2, 128, 400], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [32, 32, 64] + }, + { + "type": "scalar", + "value": 2 + }, + { + "type": "tensor", + "shape": [1, 2, 32, 400], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 2, 32, 400], + "dtype": "float16", + "memObjId": 8 + }, + { + "type": "tensor", + "shape": [1, 2, 64, 400], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "tensorlist", + "items": [540, 541, 542] + }, + { + "type": "tensor", + "shape": [1, 2, 32, 400], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [0, 1, 3, 2] + }, + { + "type": "tensor", + "shape": [1, 2, 400, 32], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [0] + }, + { + "type": "tensor", + "shape": [2, 400, 32], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "list", + "items": [0] + }, + { + "type": "tensor", + "shape": [2, 32, 400], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [2, 400, 400], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 2, 400, 400] + }, + { + "type": "tensor", + "shape": [1, 2, 400, 400], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "scalar", + "value": -1 + }, + { + "type": "tensor", + "shape": [1, 2, 400, 400], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [0, 1, 3, 2] + }, + { + "type": "tensor", + "shape": [1, 2, 400, 400], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [0] + }, + { + "type": "tensor", + "shape": [2, 64, 400], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "list", + "items": [0] + }, + { + "type": "tensor", + "shape": [2, 400, 400], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [2, 64, 400], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [2, 64, 400], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 128, 20, 20] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "list", + "items": [1, 128, 20, 20] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 128 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "list", + "items": [128, 128] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensorlist", + "items": [598, 599] + }, + { + "type": "tensorlist", + "items": [598, 596] + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 8 + }, + { + "type": "null" + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "tensor", + "shape": [1, 256, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensorlist", + "items": [612, 397] + }, + { + "type": "tensor", + "shape": [1, 384, 40, 40], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [64, 64] + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "tensorlist", + "items": [623, 624] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 7 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensorlist", + "items": [662, 669] + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "list", + "items": [64, 64] + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensorlist", + "items": [680, 681] + }, + { + "type": "tensorlist", + "items": [680, 624, 678] + }, + { + "type": "tensor", + "shape": [1, 192, 40, 40], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "null" + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "tensor", + "shape": [1, 128, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensorlist", + "items": [694, 313] + }, + { + "type": "tensor", + "shape": [1, 256, 80, 80], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 64, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 64, 80, 80], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [32, 32] + }, + { + "type": "tensor", + "shape": [1, 32, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 80, 80], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensorlist", + "items": [705, 706] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 7 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensorlist", + "items": [744, 751] + }, + { + "type": "tensor", + "shape": [1, 32, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 32, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 32, 80, 80], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "list", + "items": [32, 32] + }, + { + "type": "tensor", + "shape": [1, 32, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 32, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensorlist", + "items": [762, 763] + }, + { + "type": "tensorlist", + "items": [762, 706, 760] + }, + { + "type": "tensor", + "shape": [1, 96, 80, 80], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 64, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 64, 80, 80], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensorlist", + "items": [780, 691] + }, + { + "type": "tensor", + "shape": [1, 192, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [64, 64] + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "tensorlist", + "items": [791, 792] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 7 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 32, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensorlist", + "items": [830, 837] + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "list", + "items": [64, 64] + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 64, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensorlist", + "items": [848, 849] + }, + { + "type": "tensorlist", + "items": [848, 792, 846] + }, + { + "type": "tensor", + "shape": [1, 192, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensorlist", + "items": [866, 609] + }, + { + "type": "tensor", + "shape": [1, 384, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "list", + "items": [128, 128] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 8 + }, + { + "type": "tensorlist", + "items": [877, 878] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 64, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 7 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 2, 128, 400] + }, + { + "type": "tensor", + "shape": [1, 2, 128, 400], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [32, 32, 64] + }, + { + "type": "tensor", + "shape": [1, 2, 32, 400], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 2, 32, 400], + "dtype": "float16", + "memObjId": 9 + }, + { + "type": "tensor", + "shape": [1, 2, 64, 400], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "tensorlist", + "items": [905, 906, 907] + }, + { + "type": "tensor", + "shape": [1, 2, 32, 400], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [0, 1, 3, 2] + }, + { + "type": "tensor", + "shape": [1, 2, 400, 32], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [0] + }, + { + "type": "tensor", + "shape": [2, 400, 32], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "list", + "items": [0] + }, + { + "type": "tensor", + "shape": [2, 32, 400], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [2, 400, 400], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 2, 400, 400] + }, + { + "type": "tensor", + "shape": [1, 2, 400, 400], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 2, 400, 400], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [0, 1, 3, 2] + }, + { + "type": "tensor", + "shape": [1, 2, 400, 400], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [0] + }, + { + "type": "tensor", + "shape": [2, 64, 400], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "list", + "items": [0] + }, + { + "type": "tensor", + "shape": [2, 400, 400], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [2, 64, 400], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [2, 64, 400], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 128, 20, 20] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "list", + "items": [1, 128, 20, 20] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "list", + "items": [128, 128] + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 128, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensorlist", + "items": [961, 962] + }, + { + "type": "tensorlist", + "items": [961, 878, 959] + }, + { + "type": "tensor", + "shape": [1, 384, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 5 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 4, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 4, -1] + }, + { + "type": "tensor", + "shape": [1, 4, 6400], + "dtype": "float16", + "memObjId": 7 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 16, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 16, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 4, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 4, -1] + }, + { + "type": "tensor", + "shape": [1, 4, 1600], + "dtype": "float16", + "memObjId": 8 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 20, 20], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 16, 20, 20], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 16, 20, 20], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 16, 20, 20], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 16, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 4, 20, 20], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 4, -1] + }, + { + "type": "tensor", + "shape": [1, 4, 400], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensorlist", + "items": [993, 1014, 1035] + }, + { + "type": "tensor", + "shape": [1, 4, 8400], + "dtype": "float16", + "memObjId": 6 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 64 + }, + { + "type": "tensor", + "shape": [1, 64, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 64, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 64, 80, 80], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 80, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 80, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 80, 80, 80], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 80 + }, + { + "type": "tensor", + "shape": [1, 80, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 80, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 80, 80, 80], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 80, 80, 80], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 80, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 80, 80, 80], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 80, 80, 80], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 80, -1] + }, + { + "type": "tensor", + "shape": [1, 80, 6400], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 128, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 80, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 80, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 80, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 80, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 80, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 80, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 80, 40, 40], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 80, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 80, 40, 40], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 80, 40, 40], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 80, -1] + }, + { + "type": "tensor", + "shape": [1, 80, 1600], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "scalar", + "value": 256 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 256, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 80, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 80, 20, 20], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 80, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 80, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 80, 20, 20], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 80, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 80, 20, 20], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "tensor", + "shape": [1, 80, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensor", + "shape": [1, 80, 20, 20], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "list", + "items": [1, 1] + }, + { + "type": "list", + "items": [0, 0] + }, + { + "type": "tensor", + "shape": [1, 80, 20, 20], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "list", + "items": [1, 80, -1] + }, + { + "type": "tensor", + "shape": [1, 80, 400], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensorlist", + "items": [1074, 1109, 1145] + }, + { + "type": "tensor", + "shape": [1, 80, 8400], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "scalar", + "value": 0 + }, + { + "type": "tensor", + "shape": [1, 2, 8400], + "dtype": "float16", + "memObjId": 4 + }, + { + "type": "list", + "items": [2, 2] + }, + { + "type": "tensor", + "shape": [1, 2, 8400], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "tensor", + "shape": [1, 2, 8400], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensorlist", + "items": [1151, 1152] + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 2, 8400], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 2, 8400], + "dtype": "float16", + "memObjId": 3 + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 2, 8400], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 2, 8400], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensorlist", + "items": [1156, 1159] + }, + { + "type": "tensor", + "shape": [1, 4, 8400], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "tensor", + "shape": [1, 4, 8400], + "dtype": "float16", + "memObjId": 2 + }, + { + "type": "tensor", + "shape": [1, 80, 8400], + "dtype": "float16", + "memObjId": 0 + }, + { + "type": "null" + }, + { + "type": "tensor", + "shape": [1, 80, 8400], + "dtype": "float16", + "memObjId": 1 + }, + { + "type": "tensorlist", + "items": [1162, 1165] + }, + { + "type": "tensor", + "shape": [1, 84, 8400], + "dtype": "float16", + "memObjId": 0 + } + ], + "inputIds": [212], + "outputIds": [1167], + "memoryPlan": { + "pools": [ + { + "id": 0, + "size": 3276800 + }, + { + "id": 1, + "size": 3276800 + }, + { + "id": 2, + "size": 3276800 + }, + { + "id": 3, + "size": 819200 + }, + { + "id": 4, + "size": 819200 + }, + { + "id": 5, + "size": 409600 + }, + { + "id": 6, + "size": 204800 + }, + { + "id": 7, + "size": 204800 + }, + { + "id": 8, + "size": 204800 + }, + { + "id": 9, + "size": 51200 + } + ] + } +} diff --git a/example/public/assets/yolo26n/weights-fp16.bin b/example/public/assets/yolo26n/weights-fp16.bin new file mode 100644 index 0000000..c4cc3ca Binary files /dev/null and b/example/public/assets/yolo26n/weights-fp16.bin differ diff --git a/example/src/main.ts b/example/src/main.ts index f4f4d10..50e7b88 100644 --- a/example/src/main.ts +++ b/example/src/main.ts @@ -3,7 +3,7 @@ * The UI is plain DOM — the interesting code lives in each task's GPU loop. */ -import type { Task } from './tasks/types'; +import type { Precision, Task } from './tasks/types'; import { tasks } from './tasks'; /** Metrics tooltip lines shared by every task: `[token, description]` pairs. */ @@ -46,6 +46,24 @@ function renderTask(task: Task): void { const node = cloneTemplate('task-template'); fill(node, '.task-title', task.title); + const precisionToggle = node.querySelector('.precision-toggle') as HTMLInputElement; + let supportsF16 = false; + let pipelineStarting = false; + if (navigator.gpu) { + void navigator.gpu + .requestAdapter() + .then((adapter) => { + supportsF16 = adapter?.features.has('shader-f16') ?? false; + precisionToggle.disabled = !supportsF16 || pipelineStarting; + precisionToggle.closest('label')!.title = supportsF16 + ? 'Switch between FP32 and FP16' + : 'FP16 is unsupported on this device'; + }) + .catch(() => { + precisionToggle.closest('label')!.title = 'WebGPU is unavailable'; + }); + } + const tip = node.querySelector('.metrics-tip') as HTMLElement; for (const [token, description] of METRICS_TIP) { const line = document.createElement('p'); @@ -63,6 +81,8 @@ function renderTask(task: Task): void { const startBtn = node.querySelector('.action') as HTMLButtonElement; const setStatus = (text: string, state: 'idle' | 'starting' | 'live' | 'error') => { + pipelineStarting = state === 'starting'; + precisionToggle.disabled = !supportsF16 || pipelineStarting; statusText.textContent = text; statusDot.classList.toggle('live', state === 'live'); statusDot.classList.toggle('error', state === 'error'); @@ -70,27 +90,53 @@ function renderTask(task: Task): void { startBtn.setAttribute('aria-label', state === 'live' ? 'Camera running' : 'Start camera'); }; + let stopLoop: (() => void) | null = null; + let stream: MediaStream | null = null; + + const runPipeline = async (precision: Precision) => { + if (stopLoop) { + stopLoop(); + stopLoop = null; + } + metricsValues.textContent = ''; + stopLoop = await task.startFrameLoop( + video, + canvas, + (fps, ms) => { + const fpsStr = String(fps).padStart(2, ' '); + const msStr = ms.toFixed(1).padStart(4, ' '); + const modelFps = Math.round(1000 / ms); + metricsValues.textContent = `${fpsStr} FPS · ${msStr} MS (${modelFps} FPS)`; + }, + precision + ); + }; + startBtn.addEventListener('click', () => { void (async () => { setStatus('requesting camera', 'starting'); try { - const stream = await navigator.mediaDevices.getUserMedia({ + stream = await navigator.mediaDevices.getUserMedia({ video: { width: 1280, height: 720 }, audio: false, }); video.srcObject = stream; await video.play(); + setStatus('loading model', 'starting'); + const precision: Precision = precisionToggle.checked ? 'float16' : 'float32'; + await runPipeline(precision); + setStatus('camera live', 'live'); - const stopLoop = await task.startFrameLoop(video, canvas, (fps, ms) => { - const fpsStr = String(fps).padStart(2, ' '); - const msStr = ms.toFixed(1).padStart(4, ' '); - const modelFps = Math.round(1000 / ms); - metricsValues.textContent = `${fpsStr} FPS · ${msStr} MS (${modelFps} FPS)`; - }); cleanup = () => { - stopLoop(); - stream.getTracks().forEach((track) => track.stop()); + if (stopLoop) { + stopLoop(); + stopLoop = null; + } + if (stream) { + stream.getTracks().forEach((track) => track.stop()); + stream = null; + } }; } catch (error) { // eslint-disable-next-line no-console @@ -101,6 +147,24 @@ function renderTask(task: Task): void { })(); }); + precisionToggle.addEventListener('change', () => { + if (stream && stopLoop) { + void (async () => { + try { + setStatus('switching model', 'starting'); + const precision: Precision = precisionToggle.checked ? 'float16' : 'float32'; + await runPipeline(precision); + setStatus('camera live', 'live'); + } catch (error) { + // eslint-disable-next-line no-console + console.error(error); + metricsValues.textContent = ''; + setStatus('pipeline error', 'error'); + } + })(); + } + }); + app.replaceChildren(node); } diff --git a/example/src/tasks/depth-estimation/loop.ts b/example/src/tasks/depth-estimation/loop.ts index 4567de9..266d48b 100644 --- a/example/src/tasks/depth-estimation/loop.ts +++ b/example/src/tasks/depth-estimation/loop.ts @@ -5,7 +5,16 @@ import type { Program } from '../../../../src'; import { WgpuExecutionContext } from '../../../../src'; -import { clearStats, depthStats, drawDepth, preprocess } from './shaders'; +import type { Precision } from '../types'; +import { + clearStats, + depthStats, + depthStatsF16, + drawDepth, + drawDepthF16, + preprocess, + preprocessF16, +} from './shaders'; const size = [448, 448]; @@ -23,16 +32,31 @@ function onNextVideoFrame(video: HTMLVideoElement, callback: (now: number) => vo * @param video The playing camera video to source frames from. * @param canvas The canvas element to render the grayscale depth output to. * @param onMetrics Optional callback invoked with smoothed FPS and inference time in ms. + * @param precision Selected floating-point precision ('float32' or 'float16'). * @returns A stop function that halts the frame loop and cleans up resources. */ export async function startFrameLoop( video: HTMLVideoElement, canvas: HTMLCanvasElement, - onMetrics?: (fps: number, ms: number) => void + onMetrics?: (fps: number, ms: number) => void, + precision: Precision = 'float32' ): Promise<() => void> { + const isF16 = precision === 'float16'; + // Acquire WebGPU device const adapter = await navigator.gpu.requestAdapter(); - const device = await adapter?.requestDevice(); + if (adapter === null) { + throw new Error('no WebGPU adapter'); + } + + const hasF16 = adapter.features.has('shader-f16'); + if (isF16 && !hasF16) { + throw new Error('shader-f16 feature is not supported on this device'); + } + + const device = await adapter.requestDevice({ + requiredFeatures: isF16 ? ['shader-f16'] : [], + }); if (device === undefined) { throw new Error('no WebGPU device'); @@ -78,9 +102,11 @@ export async function startFrameLoop( // Load model program descriptor and weights const base = import.meta.env.BASE_URL; + const programFile = isF16 ? 'program-fp16.json' : 'program.json'; + const weightsFile = isF16 ? 'weights-fp16.bin' : 'weights.bin'; const [programRes, weightsRes] = await Promise.all([ - fetch(`${base}assets/depthart/program.json`), - fetch(`${base}assets/depthart/weights.bin`), + fetch(`${base}assets/depthart/${programFile}`), + fetch(`${base}assets/depthart/${weightsFile}`), ]); if (!programRes.ok || !weightsRes.ok) { @@ -92,15 +118,16 @@ export async function startFrameLoop( // Record the compute pipeline const ctx = new WgpuExecutionContext(device); - const inp = ctx.tensor('float32', [1, 3, ...size]); - const out = ctx.tensor('float32', [1, 1, ...size]); + const dtype = isF16 ? 'float16' : 'float32'; + const inp = ctx.tensor(dtype, [1, 3, ...size]); + const out = ctx.tensor(dtype, [1, 1, ...size]); const statsBuffer = ctx.storageBuffer(16); ctx - .recordShader(preprocess, [vidTexture, inp]) + .recordShader(isF16 ? preprocessF16 : preprocess, [vidTexture, inp]) .recordProgram(program, weights, [inp, out]) .recordShader(clearStats, [statsBuffer]) - .recordShader(depthStats, [out, statsBuffer]) - .recordShader(drawDepth, [out, statsBuffer, resTexture]); + .recordShader(isF16 ? depthStatsF16 : depthStats, [out, statsBuffer]) + .recordShader(isF16 ? drawDepthF16 : drawDepth, [out, statsBuffer, resTexture]); let running = true; let smoothedMs = 0; diff --git a/example/src/tasks/depth-estimation/shaders.ts b/example/src/tasks/depth-estimation/shaders.ts index f6b031d..599a6c8 100644 --- a/example/src/tasks/depth-estimation/shaders.ts +++ b/example/src/tasks/depth-estimation/shaders.ts @@ -208,3 +208,175 @@ export const drawDepth: Shader<[Tensor, GPUBuffer, GPUTexture]> = { }); }, }; + +/** + * Preprocess compute shader for FP16: reads normalized [0, 1] RGB texels from + * the video texture and writes ImageNet-normalized planar NCHW f16 elements. + */ +export const preprocessF16: Shader<[GPUTexture, Tensor]> = { + name: 'depth.preprocess.f16', + code: /* wgsl */ ` + enable f16; + + @group(0) @binding(0) var vid_tex: texture_2d; + @group(0) @binding(1) var output: array; + + @compute @workgroup_size(16, 16) + fn main(@builtin(global_invocation_id) gid: vec3) { + let dims = textureDimensions(vid_tex); + if (gid.x >= dims.x || gid.y >= dims.y) { + return; + } + + // Convert interleaved RGBA texel to planar NCHW channels + let pixel = textureLoad(vid_tex, vec2(gid.xy), 0); + let hw = dims.x * dims.y; + let offset = gid.y * dims.x + gid.x; + + output[offset] = f16((pixel.r - 0.485) / 0.229); + output[hw + offset] = f16((pixel.g - 0.456) / 0.224); + output[2u * hw + offset] = f16((pixel.b - 0.406) / 0.225); + } + `, + recordIn(ctx, [vidTexture, inp]) { + const bundle = createComputeBundle(ctx.device, preprocessF16.code as string, [ + { binding: 0, resource: vidTexture.createView() }, + { binding: 1, buffer: inp.buffer }, + ]); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX: Math.ceil(vidTexture.width / 16), + workgroupCountY: Math.ceil(vidTexture.height / 16), + }); + }, +}; + +/** + * Stats compute shader for FP16: min/max reduction over the FP16 depth map. + */ +export const depthStatsF16: Shader<[Tensor, GPUBuffer]> = { + name: 'depth.stats.f16', + code: /* wgsl */ ` + enable f16; + + @group(0) @binding(0) var depth: array; + @group(0) @binding(1) var stats: array>; + + override wg_size: u32 = 256; + + var tile_min: array; + var tile_max: array; + + // Total-order float keys: larger key <=> larger value + fn key(v: f32) -> u32 { + let bits = bitcast(v); + return select(bits ^ 0x80000000u, ~bits, (bits & 0x80000000u) != 0u); + } + + @compute @workgroup_size(wg_size) + fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(local_invocation_id) lid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + var k_min = 0xffffffffu; + var k_max = 0x00000000u; + if (idx < arrayLength(&depth)) { + let k = key(f32(depth[idx])); + k_min = k; + k_max = k; + } + + tile_min[lid.x] = k_min; + tile_max[lid.x] = k_max; + workgroupBarrier(); + + for (var stride = wg_size / 2u; stride > 0u; stride = stride / 2u) { + if (lid.x < stride) { + tile_min[lid.x] = min(tile_min[lid.x], tile_min[lid.x + stride]); + tile_max[lid.x] = max(tile_max[lid.x], tile_max[lid.x + stride]); + } + workgroupBarrier(); + } + + if (lid.x == 0u) { + atomicMin(&stats[0], tile_min[0]); + atomicMax(&stats[1], tile_max[0]); + } + } + `, + recordIn(ctx, [out, statsBuffer]) { + const wgSize = 256; + const totalWorkgroups = Math.ceil(out.numel / wgSize); + + const bundle = createComputeBundle( + ctx.device, + depthStatsF16.code as string, + [ + { binding: 0, buffer: out.buffer }, + { binding: 1, buffer: statsBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize } + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX: Math.min(totalWorkgroups, 65535), + workgroupCountY: Math.ceil(totalWorkgroups / 65535), + }); + }, +}; + +/** + * Draw compute shader for FP16: normalizes the FP16 depth map with min/max and renders grayscale. + */ +export const drawDepthF16: Shader<[Tensor, GPUBuffer, GPUTexture]> = { + name: 'depth.draw.f16', + code: /* wgsl */ ` + enable f16; + + @group(0) @binding(0) var depth: array; + @group(0) @binding(1) var stats: array; + @group(0) @binding(2) var res_tex: texture_storage_2d; + + // Inverse of the total-order key mapping in depth.stats + fn unkey(k: u32) -> f32 { + let bits = select(k ^ 0x80000000u, ~k, (k & 0x80000000u) == 0u); + return bitcast(bits); + } + + @compute @workgroup_size(16, 16) + fn main(@builtin(global_invocation_id) gid: vec3) { + let dims = textureDimensions(res_tex); + if (gid.x >= dims.x || gid.y >= dims.y) { + return; + } + + let lo = unkey(stats[0]); + let hi = unkey(stats[1]); + let offset = gid.y * dims.x + gid.x; + let t = clamp((f32(depth[offset]) - lo) / max(hi - lo, 1e-6), 0.0, 1.0); + + let gray = 1.0 - t; + textureStore(res_tex, vec2(gid.xy), vec4(gray, gray, gray, 1.0)); + } + `, + recordIn(ctx, [out, statsBuffer, resTexture]) { + const bundle = createComputeBundle(ctx.device, drawDepthF16.code as string, [ + { binding: 0, buffer: out.buffer }, + { binding: 1, buffer: statsBuffer }, + { binding: 2, resource: resTexture.createView() }, + ]); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX: Math.ceil(resTexture.width / 16), + workgroupCountY: Math.ceil(resTexture.height / 16), + }); + }, +}; diff --git a/example/src/tasks/selfie-segmentation/loop.ts b/example/src/tasks/selfie-segmentation/loop.ts index d166d54..40fbacb 100644 --- a/example/src/tasks/selfie-segmentation/loop.ts +++ b/example/src/tasks/selfie-segmentation/loop.ts @@ -5,7 +5,8 @@ import type { Program } from '../../../../src'; import { WgpuExecutionContext } from '../../../../src'; -import { postprocess, preprocess } from './shaders'; +import type { Precision } from '../types'; +import { postprocess, postprocessF16, preprocess, preprocessF16 } from './shaders'; const size = [256, 256]; @@ -23,16 +24,31 @@ function onNextVideoFrame(video: HTMLVideoElement, callback: (now: number) => vo * @param video The playing camera video to source frames from. * @param canvas The canvas element to render the segmented video output to. * @param onMetrics Optional callback invoked with smoothed FPS and inference time in ms. + * @param precision Selected floating-point precision ('float32' or 'float16'). * @returns A stop function that halts the frame loop and cleans up resources. */ export async function startFrameLoop( video: HTMLVideoElement, canvas: HTMLCanvasElement, - onMetrics?: (fps: number, ms: number) => void + onMetrics?: (fps: number, ms: number) => void, + precision: Precision = 'float32' ): Promise<() => void> { + const isF16 = precision === 'float16'; + // Acquire WebGPU device const adapter = await navigator.gpu.requestAdapter(); - const device = await adapter?.requestDevice(); + if (adapter === null) { + throw new Error('no WebGPU adapter'); + } + + const hasF16 = adapter.features.has('shader-f16'); + if (isF16 && !hasF16) { + throw new Error('shader-f16 feature is not supported on this device'); + } + + const device = await adapter.requestDevice({ + requiredFeatures: isF16 ? ['shader-f16'] : [], + }); if (device === undefined) { throw new Error('no WebGPU device'); @@ -78,9 +94,11 @@ export async function startFrameLoop( // Load model program descriptor and weights const base = import.meta.env.BASE_URL; + const programFile = isF16 ? 'program-fp16.json' : 'program.json'; + const weightsFile = isF16 ? 'weights-fp16.bin' : 'weights.bin'; const [programRes, weightsRes] = await Promise.all([ - fetch(`${base}assets/selfie-segmentation/program.json`), - fetch(`${base}assets/selfie-segmentation/weights.bin`), + fetch(`${base}assets/selfie-segmentation/${programFile}`), + fetch(`${base}assets/selfie-segmentation/${weightsFile}`), ]); if (!programRes.ok || !weightsRes.ok) { @@ -92,12 +110,13 @@ export async function startFrameLoop( // Record the compute pipeline const ctx = new WgpuExecutionContext(device); - const inp = ctx.tensor('float32', [1, 3, ...size]); - const out = ctx.tensor('float32', [1, 1, ...size]); + const dtype = isF16 ? 'float16' : 'float32'; + const inp = ctx.tensor(dtype, [1, 3, ...size]); + const out = ctx.tensor(dtype, [1, 1, ...size]); ctx - .recordShader(preprocess, [vidTexture, inp]) + .recordShader(isF16 ? preprocessF16 : preprocess, [vidTexture, inp]) .recordProgram(program, weights, [inp, out]) - .recordShader(postprocess, [vidTexture, out, resTexture]); + .recordShader(isF16 ? postprocessF16 : postprocess, [vidTexture, out, resTexture]); let running = true; let smoothedMs = 0; diff --git a/example/src/tasks/selfie-segmentation/shaders.ts b/example/src/tasks/selfie-segmentation/shaders.ts index 1d25ce0..58c3d9d 100644 --- a/example/src/tasks/selfie-segmentation/shaders.ts +++ b/example/src/tasks/selfie-segmentation/shaders.ts @@ -99,3 +99,97 @@ export const postprocess: Shader<[GPUTexture, Tensor, GPUTexture]> = { }); }, }; + +/** + * Preprocess compute shader for FP16: reads normalized [0, 1] RGB texels from the video + * texture and unpacks them into planar NCHW f16 elements. + */ +export const preprocessF16: Shader<[GPUTexture, Tensor]> = { + name: 'selfie.preprocess.f16', + code: /* wgsl */ ` + enable f16; + + @group(0) @binding(0) var vid_tex: texture_2d; + @group(0) @binding(1) var output: array; + + @compute @workgroup_size(16, 16) + fn main(@builtin(global_invocation_id) gid: vec3) { + let dims = textureDimensions(vid_tex); + if (gid.x >= dims.x || gid.y >= dims.y) { + return; + } + + // Convert interleaved RGBA texel to planar NCHW channels + let pixel = textureLoad(vid_tex, vec2(gid.xy), 0); + let hw = dims.x * dims.y; + let offset = gid.y * dims.x + gid.x; + + output[offset] = f16(pixel.r); + output[hw + offset] = f16(pixel.g); + output[2u * hw + offset] = f16(pixel.b); + } + `, + recordIn(ctx, [vidTexture, inp]) { + const bundle = createComputeBundle(ctx.device, preprocessF16.code as string, [ + { binding: 0, resource: vidTexture.createView() }, + { binding: 1, buffer: inp.buffer }, + ]); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX: Math.ceil(vidTexture.width / 16), + workgroupCountY: Math.ceil(vidTexture.height / 16), + }); + }, +}; + +/** + * Postprocess compute shader for FP16 mask: blends segmented mask with darkened background. + */ +export const postprocessF16: Shader<[GPUTexture, Tensor, GPUTexture]> = { + name: 'selfie.postprocess.f16', + code: /* wgsl */ ` + enable f16; + + @group(0) @binding(0) var vid_tex: texture_2d; + @group(0) @binding(1) var mask: array; + @group(0) @binding(2) var res_tex: texture_storage_2d; + + @compute @workgroup_size(16, 16) + fn main(@builtin(global_invocation_id) gid: vec3) { + let dims = textureDimensions(vid_tex); + if (gid.x >= dims.x || gid.y >= dims.y) { + return; + } + + let offset = gid.y * dims.x + gid.x; + let m = clamp(f32(mask[offset]), 0.0, 1.0); + + // Very dark background with slight transparency (premultiplied alpha) + let bg_alpha = 0.8; + let bg_rgb = vec3(0.02, 0.02, 0.03) * bg_alpha; + let bg = vec4(bg_rgb, bg_alpha); + + // Solid white fill for the person + let fg = vec4(1.0, 1.0, 1.0, 1.0); + + let color = mix(bg, fg, m); + textureStore(res_tex, vec2(gid.xy), color); + } + `, + recordIn(ctx, [vidTexture, out, resTexture]) { + const bundle = createComputeBundle(ctx.device, postprocessF16.code as string, [ + { binding: 0, resource: vidTexture.createView() }, + { binding: 1, buffer: out.buffer }, + { binding: 2, resource: resTexture.createView() }, + ]); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX: Math.ceil(vidTexture.width / 16), + workgroupCountY: Math.ceil(vidTexture.height / 16), + }); + }, +}; diff --git a/example/src/tasks/types.ts b/example/src/tasks/types.ts index 4eacd17..35cbafd 100644 --- a/example/src/tasks/types.ts +++ b/example/src/tasks/types.ts @@ -7,6 +7,9 @@ /** Per-frame metrics callback: smoothed FPS and inference time in ms. */ export type MetricsFn = (fps: number, ms: number) => void; +/** Floating point precision for execution: FP32 (default) or native FP16. */ +export type Precision = 'float32' | 'float16'; + /** * Starts a task's per-frame pipeline for the given camera video. * Resolves with a stop function that halts the loop and releases resources. @@ -14,7 +17,8 @@ export type MetricsFn = (fps: number, ms: number) => void; export type StartFrameLoop = ( video: HTMLVideoElement, canvas: HTMLCanvasElement, - onMetrics?: MetricsFn + onMetrics?: MetricsFn, + precision?: Precision ) => Promise<() => void>; /** Gallery tile + task page configuration. */ diff --git a/example/src/tasks/yolo26-detection/loop.ts b/example/src/tasks/yolo26-detection/loop.ts index 3dcbf2a..db9d840 100644 --- a/example/src/tasks/yolo26-detection/loop.ts +++ b/example/src/tasks/yolo26-detection/loop.ts @@ -5,12 +5,15 @@ */ import { WgpuExecutionContext, type Program } from '../../../../src'; +import type { Precision } from '../types'; import { clearDetections, compactDetections, + compactDetectionsF16, drawDetections, MAX_DETECTIONS, preprocess, + preprocessF16, } from './shaders'; const size = [640, 640]; @@ -29,16 +32,31 @@ function onNextVideoFrame(video: HTMLVideoElement, callback: (now: number) => vo * @param video The playing camera video to source frames from. * @param canvas The canvas element to render the graded video output to. * @param onMetrics Optional callback invoked with smoothed FPS and inference time in ms. + * @param precision Selected floating-point precision ('float32' or 'float16'). * @returns A stop function that halts the frame loop and cleans up resources. */ export async function startFrameLoop( video: HTMLVideoElement, canvas: HTMLCanvasElement, - onMetrics?: (fps: number, ms: number) => void + onMetrics?: (fps: number, ms: number) => void, + precision: Precision = 'float32' ): Promise<() => void> { + const isF16 = precision === 'float16'; + // Acquire WebGPU device const adapter = await navigator.gpu.requestAdapter(); - const device = await adapter?.requestDevice(); + if (adapter === null) { + throw new Error('no WebGPU adapter'); + } + + const hasF16 = adapter.features.has('shader-f16'); + if (isF16 && !hasF16) { + throw new Error('shader-f16 feature is not supported on this device'); + } + + const device = await adapter.requestDevice({ + requiredFeatures: isF16 ? ['shader-f16'] : [], + }); if (device === undefined) { throw new Error('no WebGPU device'); @@ -84,9 +102,11 @@ export async function startFrameLoop( // Load model program descriptor and weights const base = import.meta.env.BASE_URL; + const programFile = isF16 ? 'program-fp16.json' : 'program.json'; + const weightsFile = isF16 ? 'weights-fp16.bin' : 'weights.bin'; const [programRes, weightsRes] = await Promise.all([ - fetch(`${base}assets/yolo26n/program.json`), - fetch(`${base}assets/yolo26n/weights.bin`), + fetch(`${base}assets/yolo26n/${programFile}`), + fetch(`${base}assets/yolo26n/${weightsFile}`), ]); if (!programRes.ok || !weightsRes.ok) { @@ -98,15 +118,20 @@ export async function startFrameLoop( // Record the full pipeline once: model -> compaction -> draw. const ctx = new WgpuExecutionContext(device); - const inp = ctx.tensor('float32', [1, 3, ...size]); - const out = ctx.tensor('float32', [1, 84, 8400]); + const dtype = isF16 ? 'float16' : 'float32'; + const inp = ctx.tensor(dtype, [1, 3, ...size]); + const out = ctx.tensor(dtype, [1, 84, 8400]); const counterBuffer = ctx.storageBuffer(16); const detsBuffer = ctx.storageBuffer(MAX_DETECTIONS * 6 * 4); ctx - .recordShader(preprocess, [vidTexture, inp]) + .recordShader(isF16 ? preprocessF16 : preprocess, [vidTexture, inp]) .recordProgram(program, weights, [inp, out]) .recordShader(clearDetections, [counterBuffer]) - .recordShader(compactDetections, [out, counterBuffer, detsBuffer]) + .recordShader(isF16 ? compactDetectionsF16 : compactDetections, [ + out, + counterBuffer, + detsBuffer, + ]) .recordShader(drawDetections, [vidTexture, resTexture, counterBuffer, detsBuffer]); let running = true; diff --git a/example/src/tasks/yolo26-detection/shaders.ts b/example/src/tasks/yolo26-detection/shaders.ts index fe1a65d..fc4a297 100644 --- a/example/src/tasks/yolo26-detection/shaders.ts +++ b/example/src/tasks/yolo26-detection/shaders.ts @@ -280,3 +280,142 @@ export const drawDetections: Shader<[GPUTexture, GPUTexture, GPUBuffer, GPUBuffe }); }, }; + +/** + * Preprocess compute shader for FP16: reads normalized [0, 1] RGB texels from the video + * texture and unpacks them into planar NCHW f16 elements. + */ +export const preprocessF16: Shader<[GPUTexture, Tensor]> = { + name: 'yolo.preprocess.f16', + code: /* wgsl */ ` + enable f16; + + @group(0) @binding(0) var vid_tex: texture_2d; + @group(0) @binding(1) var output: array; + + @compute @workgroup_size(16, 16) + fn main(@builtin(global_invocation_id) gid: vec3) { + let dims = textureDimensions(vid_tex); + if (gid.x >= dims.x || gid.y >= dims.y) { + return; + } + + // Convert interleaved RGBA texel to planar NCHW channels + let pixel = textureLoad(vid_tex, vec2(gid.xy), 0); + let hw = dims.x * dims.y; + let offset = gid.y * dims.x + gid.x; + + output[offset] = f16(pixel.r); + output[hw + offset] = f16(pixel.g); + output[2u * hw + offset] = f16(pixel.b); + } + `, + recordIn(ctx, [vidTexture, inp]) { + const bundle = createComputeBundle(ctx.device, preprocessF16.code as string, [ + { binding: 0, resource: vidTexture.createView() }, + { binding: 1, buffer: inp.buffer }, + ]); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX: Math.ceil(vidTexture.width / 16), + workgroupCountY: Math.ceil(vidTexture.height / 16), + }); + }, +}; + +/** + * Detection compaction compute shader for FP16 predictions. + */ +export const compactDetectionsF16: Shader<[Tensor, GPUBuffer, GPUBuffer]> = { + name: 'yolo.compact_detections.f16', + code: /* wgsl */ ` + enable f16; + + @group(0) @binding(0) var preds: array; + @group(0) @binding(1) var counter: array>; + @group(0) @binding(2) var dets: array; + + override wg_size: u32 = 256; + override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + override conf: f32 = 0.5; + override max_det: u32 = 64; + override num_anchors: u32 = 8400; + override num_classes: u32 = 80; + + @compute @workgroup_size(wg_size) + fn main(@builtin(global_invocation_id) gid: vec3) { + let a = gid.y * stride_x + gid.x; + if (a >= num_anchors) { + return; + } + + // Best class over the anchor's score channel slice + var best: f32 = 0.0; + var best_class: u32 = 0u; + for (var c: u32 = 0u; c < num_classes; c = c + 1u) { + let s = f32(preds[(4u + c) * num_anchors + a]); + if (s > best) { + best = s; + best_class = c; + } + } + if (best < conf) { + return; + } + + let slot = atomicAdd(&counter[0], 1u); + if (slot >= max_det) { + return; + } + let b = slot * 6u; + dets[b] = f32(preds[a]); + dets[b + 1u] = f32(preds[num_anchors + a]); + dets[b + 2u] = f32(preds[2u * num_anchors + a]); + dets[b + 3u] = f32(preds[3u * num_anchors + a]); + dets[b + 4u] = best; + dets[b + 5u] = f32(best_class); + } + `, + recordIn(ctx, [preds, counterBuffer, detsBuffer]) { + if (preds.dtype !== 'float16') { + throw new Error('yolo.compact_detections.f16: predictions must be float16'); + } + + const wgSize = 256; + const numAnchors = preds.shape[2]; + const numClasses = preds.shape[1] - 4; + const totalWorkgroups = Math.ceil(numAnchors / wgSize); + const workgroupCountX = Math.min(totalWorkgroups, 65535); + + /* eslint-disable camelcase, @typescript-eslint/naming-convention -- WGSL override names */ + const constants = { + wg_size: wgSize, + stride_x: workgroupCountX * wgSize, + conf: CONF_THRESHOLD, + max_det: MAX_DETECTIONS, + num_anchors: numAnchors, + num_classes: numClasses, + }; + /* eslint-enable camelcase, @typescript-eslint/naming-convention */ + + const bundle = createComputeBundle( + ctx.device, + compactDetectionsF16.code as string, + [ + { binding: 0, buffer: preds.buffer }, + { binding: 1, buffer: counterBuffer }, + { binding: 2, buffer: detsBuffer }, + ], + constants + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY: Math.ceil(totalWorkgroups / workgroupCountX), + }); + }, +}; diff --git a/example/style.css b/example/style.css index 9954bfe..26a04d7 100644 --- a/example/style.css +++ b/example/style.css @@ -316,8 +316,9 @@ main { } .topbar { - width: min(800px, 100%, calc(100dvh - 16rem)); + width: min(800px, 100%); display: flex; + flex-wrap: wrap; align-items: center; gap: 1.25rem; } @@ -354,3 +355,101 @@ main { background: rgba(240, 246, 252, 0.1); border-color: var(--text); } + +.precision-control { + flex-shrink: 0; + margin-left: auto; + display: flex; + align-items: center; + gap: 0.5rem; +} + +.precision-label { + font-family: var(--mono); + font-size: 0.7rem; + letter-spacing: 0.12em; + text-transform: uppercase; + color: var(--muted); +} + +.precision-switch { + position: relative; + display: inline-block; +} + +.precision-toggle { + position: absolute; + inset: 0; + width: 100%; + height: 100%; + margin: 0; + opacity: 0; + z-index: 1; + cursor: pointer; +} + +.precision-track { + position: relative; + display: grid; + grid-template-columns: repeat(2, 1fr); + padding: 3px; + font-family: var(--mono); + font-size: 0.75rem; + letter-spacing: 0.06em; + background: var(--panel); + border: 1px solid var(--line-strong); + border-radius: 999px; +} + +.precision-track::before { + content: ''; + position: absolute; + top: 3px; + bottom: 3px; + left: 3px; + width: calc((100% - 6px) / 2); + background: var(--text); + border-radius: inherit; + transition: transform 0.18s ease; +} + +.precision-track > span { + position: relative; + padding: 0.45rem 0.7rem; + text-align: center; + color: var(--muted); + transition: color 0.18s ease; +} + +.precision-toggle:not(:checked) + .precision-track > span:first-child, +.precision-toggle:checked + .precision-track > span:last-child { + color: var(--panel); +} + +.precision-toggle:checked + .precision-track::before { + transform: translateX(100%); +} + +.precision-toggle:focus-visible + .precision-track { + outline: 2px solid var(--text); + outline-offset: 3px; +} + +.precision-switch:hover .precision-track { + border-color: var(--text); +} + +.precision-toggle:disabled + .precision-track { + opacity: 0.5; +} + +.precision-toggle:disabled { + cursor: not-allowed; +} + +@media (prefers-reduced-motion: reduce) { + .precision-track::before, + .precision-track > span { + transition: none; + } +} diff --git a/package-lock.json b/package-lock.json index 0134909..5c45d96 100644 --- a/package-lock.json +++ b/package-lock.json @@ -9,6 +9,7 @@ "version": "0.1.0", "license": "MIT", "devDependencies": { + "@petamoriken/float16": "^3.9.3", "@types/node": "^26.6.3", "@typescript-eslint/eslint-plugin": "^8.70.0", "@typescript-eslint/parser": "^8.70.0", @@ -298,6 +299,13 @@ "url": "https://github.com/sponsors/oxc-project" } }, + "node_modules/@petamoriken/float16": { + "version": "3.9.3", + "resolved": "https://registry.npmjs.org/@petamoriken/float16/-/float16-3.9.3.tgz", + "integrity": "sha512-8awtpHXCx/bNpFt4mt2xdkgtgVvKqty8VbjHI/WWWQuEw+KLzFot3f4+LkQY9YmOtq7A5GdOnqoIC8Pdygjk2g==", + "dev": true, + "license": "MIT" + }, "node_modules/@pkgr/core": { "version": "0.3.6", "resolved": "https://registry.npmjs.org/@pkgr/core/-/core-0.3.6.tgz", diff --git a/package.json b/package.json index 6860771..17943e9 100644 --- a/package.json +++ b/package.json @@ -8,6 +8,7 @@ "test": "vitest run", "example": "vite example", "build:example": "vite build example", + "benchmark:web": "vite utils/benchmark", "format": "prettier --write .", "lint": "eslint .", "lint:fix": "eslint . --fix", @@ -32,6 +33,7 @@ "url": "https://github.com/barhanc/wgpu-executorch/issues" }, "devDependencies": { + "@petamoriken/float16": "^3.9.3", "@types/node": "^26.6.3", "@typescript-eslint/eslint-plugin": "^8.70.0", "@typescript-eslint/parser": "^8.70.0", diff --git a/src/context.ts b/src/context.ts index 26e172f..5a84f75 100644 --- a/src/context.ts +++ b/src/context.ts @@ -182,7 +182,11 @@ export class WgpuExecutionContext { } this.#commands.push({ kind: 'copy', - copy: { src: src.buffer, dst: dst.buffer, size: src.byteLength }, + copy: { + src: src.buffer, + dst: dst.buffer, + size: Math.ceil(src.byteLength / 4) * 4, + }, }); return this; } @@ -324,6 +328,7 @@ export class WgpuExecutionContext { const encoder = this.#device.createCommandEncoder(); let pass: GPUComputePassEncoder | undefined; + let currentPipeline: GPUComputePipeline | undefined; for (const cmd of this.#commands) { switch (cmd.kind) { @@ -331,6 +336,7 @@ export class WgpuExecutionContext { if (pass) { pass.end(); pass = undefined; + currentPipeline = undefined; } encoder.copyBufferToBuffer(cmd.copy.src, 0, cmd.copy.dst, 0, cmd.copy.size); break; @@ -341,7 +347,10 @@ export class WgpuExecutionContext { pass = encoder.beginComputePass(); } const d = cmd.dispatch; - pass.setPipeline(d.pipeline); + if (d.pipeline !== currentPipeline) { + pass.setPipeline(d.pipeline); + currentPipeline = d.pipeline; + } pass.setBindGroup(0, d.bindGroup); pass.dispatchWorkgroups(d.workgroupCountX, d.workgroupCountY, d.workgroupCountZ ?? 1); break; diff --git a/src/shaders/add/add-broadcast-f16.wgsl.ts b/src/shaders/add/add-broadcast-f16.wgsl.ts new file mode 100644 index 0000000..fe13499 --- /dev/null +++ b/src/shaders/add/add-broadcast-f16.wgsl.ts @@ -0,0 +1,60 @@ +import { TENSOR_META_WGSL } from '../utils/meta'; + +/** + * WGSL compute shader for FP16 binary addition with broadcasting support. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_ADD_BROADCAST_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +${TENSOR_META_WGSL} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in1_meta: TensorMeta; +@group(0) @binding(5) var in2_meta: TensorMeta; + +override wg_size: u32 = 256u; +override alpha: f32 = 1.0; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3 +) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= out_meta.numel) { + return; + } + let a_h = f16(alpha); + + var same = true; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + if ( + in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] + ) { + same = false; + } + } + if (same) { + output[idx] = input1[idx] + a_h * input2[idx]; + return; + } + + var rem = idx; + var l1: u32 = 0u; + var l2: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; + } + output[idx] = input1[l1] + a_h * input2[l2]; +} +`; diff --git a/src/shaders/add/add-broadcast.wgsl.ts b/src/shaders/add/add-broadcast.wgsl.ts new file mode 100644 index 0000000..fd333b5 --- /dev/null +++ b/src/shaders/add/add-broadcast.wgsl.ts @@ -0,0 +1,57 @@ +import { TENSOR_META_WGSL } from '../utils/meta'; + +/** + * WGSL compute shader for binary addition with broadcasting support. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_ADD_BROADCAST = /* wgsl */ ` +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +${TENSOR_META_WGSL} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in1_meta: TensorMeta; +@group(0) @binding(5) var in2_meta: TensorMeta; + +override wg_size: u32 = 256u; +override alpha: f32 = 1.0; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3 +) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= out_meta.numel) { + return; + } + + var same = true; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + if ( + in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] + ) { + same = false; + } + } + if (same) { + output[idx] = input1[idx] + alpha * input2[idx]; + return; + } + + var rem = idx; + var l1: u32 = 0u; + var l2: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; + } + output[idx] = input1[l1] + alpha * input2[l2]; +} +`; diff --git a/src/shaders/add/add-elementwise-f16.wgsl.ts b/src/shaders/add/add-elementwise-f16.wgsl.ts new file mode 100644 index 0000000..fccb98d --- /dev/null +++ b/src/shaders/add/add-elementwise-f16.wgsl.ts @@ -0,0 +1,32 @@ +/** + * WGSL compute shader for contiguous 1D FP16 elementwise addition. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_ADD_ELEMENTWISE_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 256u; +override alpha: f32 = 1.0; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = input1[idx] + f16(alpha) * input2[idx]; +} +`; diff --git a/src/shaders/add/add-elementwise-vec4-f16.wgsl.ts b/src/shaders/add/add-elementwise-vec4-f16.wgsl.ts new file mode 100644 index 0000000..70e5994 --- /dev/null +++ b/src/shaders/add/add-elementwise-vec4-f16.wgsl.ts @@ -0,0 +1,32 @@ +/** + * WGSL compute shader for contiguous 64-bit/128-bit vectorized (vec4) elementwise addition. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_ADD_ELEMENTWISE_VEC4_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input1: array>; +@group(0) @binding(1) var input2: array>; +@group(0) @binding(2) var output: array>; + +struct Params { + num_elements: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 256u; +override alpha: f32 = 1.0; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = input1[idx] + f16(alpha) * input2[idx]; +} +`; diff --git a/src/shaders/add/add-elementwise-vec4.wgsl.ts b/src/shaders/add/add-elementwise-vec4.wgsl.ts new file mode 100644 index 0000000..c3ae8f1 --- /dev/null +++ b/src/shaders/add/add-elementwise-vec4.wgsl.ts @@ -0,0 +1,30 @@ +/** + * WGSL compute shader for contiguous 128-bit vectorized (vec4) elementwise addition. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_ADD_ELEMENTWISE_VEC4 = /* wgsl */ ` +@group(0) @binding(0) var input1: array>; +@group(0) @binding(1) var input2: array>; +@group(0) @binding(2) var output: array>; + +struct Params { + num_elements: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 256u; +override alpha: f32 = 1.0; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = input1[idx] + alpha * input2[idx]; +} +`; diff --git a/src/shaders/add/add-elementwise.wgsl.ts b/src/shaders/add/add-elementwise.wgsl.ts new file mode 100644 index 0000000..59a2dcc --- /dev/null +++ b/src/shaders/add/add-elementwise.wgsl.ts @@ -0,0 +1,30 @@ +/** + * WGSL compute shader for contiguous 1D elementwise addition. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_ADD_ELEMENTWISE = /* wgsl */ ` +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 256u; +override alpha: f32 = 1.0; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = input1[idx] + alpha * input2[idx]; +} +`; diff --git a/src/shaders/add/index.ts b/src/shaders/add/index.ts new file mode 100644 index 0000000..6684df5 --- /dev/null +++ b/src/shaders/add/index.ts @@ -0,0 +1,149 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { isBroadcastable } from '../utils/broadcasting'; +import { encodeTensorMeta } from '../utils/meta'; + +import { SHADER_ADD_BROADCAST } from './add-broadcast.wgsl'; +import { SHADER_ADD_BROADCAST_F16 } from './add-broadcast-f16.wgsl'; +import { SHADER_ADD_ELEMENTWISE } from './add-elementwise.wgsl'; +import { SHADER_ADD_ELEMENTWISE_F16 } from './add-elementwise-f16.wgsl'; +import { SHADER_ADD_ELEMENTWISE_VEC4 } from './add-elementwise-vec4.wgsl'; +import { SHADER_ADD_ELEMENTWISE_VEC4_F16 } from './add-elementwise-vec4-f16.wgsl'; + +const name = 'aten.add.Tensor'; + +/** + * Positional argument tuple for `aten::add.Tensor`: + * Standard ExecuTorch schema: `[in1, in2, alpha, out]` + * Or 3-arg variant: `[in1, in2, out]` + */ +export type AddArgs = + | readonly [in1: Tensor, in2: Tensor, out: Tensor] + | readonly [in1: Tensor, in2: Tensor, alpha: number, out: Tensor]; + +/** + * Internal dispatch builder for `aten::add.Tensor`. + * + * @param ctx The execution context. + * @param args Positional argument tuple for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: AddArgs): void { + const [in1, in2] = args; + const [alpha, out] = args.length === 4 ? [args[2], args[3]] : [1.0, args[2]]; + + if ( + (in1.dtype !== 'float32' && in1.dtype !== 'float16') || + in2.dtype !== in1.dtype || + out.dtype !== in1.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + if (!isBroadcastable(in1.shape, out.shape) || !isBroadcastable(in2.shape, out.shape)) { + throw new Error(`${name}: Input shapes are not broadcastable to output shape`); + } + + const device = ctx.device; + const isF16 = in1.dtype === 'float16'; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + const sameShape = + in1.shape.length === out.shape.length && + in2.shape.length === out.shape.length && + in1.shape.every((s, i) => s === out.shape[i]) && + in2.shape.every((s, i) => s === out.shape[i]); + + const wgSize = 256; + + if (sameShape) { + const numel = out.numel; + const useVec4 = numel % 4 === 0; + let code = useVec4 ? SHADER_ADD_ELEMENTWISE_VEC4 : SHADER_ADD_ELEMENTWISE; + if (isF16) { + code = useVec4 ? SHADER_ADD_ELEMENTWISE_VEC4_F16 : SHADER_ADD_ELEMENTWISE_F16; + } + const effectiveNumel = useVec4 ? numel / 4 : numel; + + const totalWorkgroups = Math.ceil(effectiveNumel / wgSize); + const workgroupCountX = Math.min(totalWorkgroups, 65535); + const workgroupCountY = Math.ceil(totalWorkgroups / 65535); + + const paramBuffer = ctx.uniformBuffer(new Uint32Array([effectiveNumel])); + + const bundle = createComputeBundle( + device, + code, + [ + { binding: 0, buffer: in1.buffer }, + { binding: 1, buffer: in2.buffer }, + { binding: 2, buffer: out.buffer }, + { binding: 3, buffer: paramBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize, alpha } + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); + return; + } + + const outRank = out.shape.length; + const outMetaBuffer = ctx.uniformBuffer(encodeTensorMeta(out)); + const in1MetaBuffer = ctx.uniformBuffer(encodeTensorMeta(in1, outRank)); + const in2MetaBuffer = ctx.uniformBuffer(encodeTensorMeta(in2, outRank)); + + const totalWorkgroups = Math.ceil(out.numel / wgSize); + const workgroupCountX = Math.min(totalWorkgroups, 65535); + const workgroupCountY = Math.ceil(totalWorkgroups / 65535); + + const bundle = createComputeBundle( + device, + isF16 ? SHADER_ADD_BROADCAST_F16 : SHADER_ADD_BROADCAST, + [ + { binding: 0, buffer: in1.buffer }, + { binding: 1, buffer: in2.buffer }, + { binding: 2, buffer: out.buffer }, + { binding: 3, buffer: outMetaBuffer }, + { binding: 4, buffer: in1MetaBuffer }, + { binding: 5, buffer: in2MetaBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize, alpha } + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::add.Tensor`. + * + * Computes elementwise `output = input1 + alpha * input2` with NumPy/PyTorch-style + * broadcasting (right-aligned, dimensions equal or 1, up to rank 8). Features an + * automatic fast path when input and output shapes match. + */ +export const add: Shader = { + name, + code: { + broadcast: SHADER_ADD_BROADCAST, + elementwise: SHADER_ADD_ELEMENTWISE, + elementwiseVec4: SHADER_ADD_ELEMENTWISE_VEC4, + broadcastF16: SHADER_ADD_BROADCAST_F16, + elementwiseF16: SHADER_ADD_ELEMENTWISE_F16, + elementwiseVec4F16: SHADER_ADD_ELEMENTWISE_VEC4_F16, + }, + recordIn, +}; diff --git a/src/shaders/bmm/bmm-64x64-f16.wgsl.ts b/src/shaders/bmm/bmm-64x64-f16.wgsl.ts new file mode 100644 index 0000000..b71c1a3 --- /dev/null +++ b/src/shaders/bmm/bmm-64x64-f16.wgsl.ts @@ -0,0 +1,121 @@ +/** + * WGSL compute shader for 64x64 register-tiled batched matrix multiplication in float16 (4x4 per thread across 256 threads). + * High occupancy, coalesced 128-bit memory transactions, minimal register pressure for all GPUs. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_BMM_64X64_F16 = /* wgsl */ ` +enable f16; + +struct Params { + B: u32, + M: u32, + N: u32, + K: u32, +}; + +@group(0) @binding(0) var a: array>; +@group(0) @binding(1) var b: array>; +@group(0) @binding(2) var out: array>; +@group(0) @binding(3) var params: Params; + +// 64 rows x 16 K elements (4 vec4s) = 256 vec4s (2 KB) +var a_sub: array, 4>, 64>; +// 16 K rows x 64 N elements (16 vec4s) = 256 vec4s (2 KB) +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let bn = wg_id.z; + let tile_row0 = wg_id.y * 64u; + let tile_col0_4 = wg_id.x * 16u; + let tid = local_id.y * 16u + local_id.x; // 0..255 + + let K_vecs = params.K / 4u; + let N_vecs = params.N / 4u; + + let a_base4 = bn * params.M * K_vecs; + let b_base4 = bn * params.K * N_vecs; + let out_base4 = bn * params.M * N_vecs; + + var acc0 = vec4(0.0h); + var acc1 = vec4(0.0h); + var acc2 = vec4(0.0h); + var acc3 = vec4(0.0h); + + // Load coordinates for 256 threads (each thread loads 1 vec4 of A and 1 vec4 of B) + let a_r = tid >> 2u; // 0..63 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < params.M; + + let b_r = tid >> 4u; // 0..15 + let b_c = tid & 15u; // 0..15 + let b_n4 = tile_col0_4 + b_c; + let b_c_valid = b_n4 < N_vecs; + + let num_tiles = (params.K + 15u) >> 4u; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t << 2u; + let k_start = t << 4u; + + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = a[a_base4 + a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0h); + } + + let bk = k_start + b_r; + if (bk < params.K && b_c_valid) { + b_sub[b_r][b_c] = b[b_base4 + bk * N_vecs + b_n4]; + } else { + b_sub[b_r][b_c] = vec4(0.0h); + } + workgroupBarrier(); + + let thread_r = local_id.y << 2u; // 0..60 (step 4) + let thread_c4 = local_id.x; // 0..15 + + for (var k4: u32 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + let thread_r = local_id.y << 2u; + let out_r0 = tile_row0 + thread_r; + let out_c4 = tile_col0_4 + local_id.x; + + if (out_c4 < N_vecs) { + if (out_r0 + 0u < params.M) { + out[out_base4 + (out_r0 + 0u) * N_vecs + out_c4] = acc0; + } + if (out_r0 + 1u < params.M) { + out[out_base4 + (out_r0 + 1u) * N_vecs + out_c4] = acc1; + } + if (out_r0 + 2u < params.M) { + out[out_base4 + (out_r0 + 2u) * N_vecs + out_c4] = acc2; + } + if (out_r0 + 3u < params.M) { + out[out_base4 + (out_r0 + 3u) * N_vecs + out_c4] = acc3; + } + } +} +`; diff --git a/src/shaders/bmm/bmm-64x64.wgsl.ts b/src/shaders/bmm/bmm-64x64.wgsl.ts new file mode 100644 index 0000000..fa6ff4f --- /dev/null +++ b/src/shaders/bmm/bmm-64x64.wgsl.ts @@ -0,0 +1,119 @@ +/** + * WGSL compute shader for 64x64 register-tiled batched matrix multiplication (4x4 per thread across 256 threads). + * High occupancy, coalesced 128-bit memory transactions, minimal register pressure for all GPUs. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_BMM_64X64 = /* wgsl */ ` +struct Params { + B: u32, + M: u32, + N: u32, + K: u32, +}; + +@group(0) @binding(0) var a: array>; +@group(0) @binding(1) var b: array>; +@group(0) @binding(2) var out: array>; +@group(0) @binding(3) var params: Params; + +// 64 rows x 16 K elements (4 vec4s) = 256 vec4s (4 KB) +var a_sub: array, 4>, 64>; +// 16 K rows x 64 N elements (16 vec4s) = 256 vec4s (4 KB) +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let bn = wg_id.z; + let tile_row0 = wg_id.y * 64u; + let tile_col0_4 = wg_id.x * 16u; + let tid = local_id.y * 16u + local_id.x; // 0..255 + + let K_vecs = params.K / 4u; + let N_vecs = params.N / 4u; + + let a_base4 = bn * params.M * K_vecs; + let b_base4 = bn * params.K * N_vecs; + let out_base4 = bn * params.M * N_vecs; + + var acc0 = vec4(0.0); + var acc1 = vec4(0.0); + var acc2 = vec4(0.0); + var acc3 = vec4(0.0); + + // Load coordinates for 256 threads (each thread loads 1 vec4 of A and 1 vec4 of B) + let a_r = tid >> 2u; // 0..63 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < params.M; + + let b_r = tid >> 4u; // 0..15 + let b_c = tid & 15u; // 0..15 + let b_n4 = tile_col0_4 + b_c; + let b_c_valid = b_n4 < N_vecs; + + let num_tiles = (params.K + 15u) >> 4u; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t << 2u; + let k_start = t << 4u; + + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = a[a_base4 + a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0); + } + + let bk = k_start + b_r; + if (bk < params.K && b_c_valid) { + b_sub[b_r][b_c] = b[b_base4 + bk * N_vecs + b_n4]; + } else { + b_sub[b_r][b_c] = vec4(0.0); + } + workgroupBarrier(); + + let thread_r = local_id.y << 2u; // 0..60 (step 4) + let thread_c4 = local_id.x; // 0..15 + + for (var k4: u32 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + let thread_r = local_id.y << 2u; + let out_r0 = tile_row0 + thread_r; + let out_c4 = tile_col0_4 + local_id.x; + + if (out_c4 < N_vecs) { + if (out_r0 + 0u < params.M) { + out[out_base4 + (out_r0 + 0u) * N_vecs + out_c4] = acc0; + } + if (out_r0 + 1u < params.M) { + out[out_base4 + (out_r0 + 1u) * N_vecs + out_c4] = acc1; + } + if (out_r0 + 2u < params.M) { + out[out_base4 + (out_r0 + 2u) * N_vecs + out_c4] = acc2; + } + if (out_r0 + 3u < params.M) { + out[out_base4 + (out_r0 + 3u) * N_vecs + out_c4] = acc3; + } + } +} +`; diff --git a/src/shaders/bmm/bmm-tiled-f16.wgsl.ts b/src/shaders/bmm/bmm-tiled-f16.wgsl.ts new file mode 100644 index 0000000..22a8b13 --- /dev/null +++ b/src/shaders/bmm/bmm-tiled-f16.wgsl.ts @@ -0,0 +1,105 @@ +/** + * WGSL compute shader for float16 batched matrix multiplication (tiled). + * Requires WebGPU shader-f16 feature. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_BMM_TILED_F16 = /* wgsl */ ` +enable f16; + +struct Params { + B: u32, + M: u32, + N: u32, + K: u32, +}; + +@group(0) @binding(0) var a: array; +@group(0) @binding(1) var b: array; +@group(0) @binding(2) var out: array; +@group(0) @binding(3) var params: Params; + +const TILE: u32 = 32u; +const RPT: u32 = 4u; + +var a_sub: array, 32>; +var b_sub: array, 32>; + +fn read_a(base: u32, row: u32, col: u32) -> f16 { + if (row < params.M && col < params.K) { + return a[base + row * params.K + col]; + } + return 0.0h; +} + +fn read_b(base: u32, krow: u32, col: u32) -> f16 { + if (krow < params.K && col < params.N) { + return b[base + krow * params.N + col]; + } + return 0.0h; +} + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3) { + let tiles_per_m = (params.M + TILE - 1u) / TILE; + let bn = wg_id.y / tiles_per_m; + let row_tile = wg_id.y % tiles_per_m; + let tile_row0 = row_tile * TILE; + let tile_col0 = wg_id.x * TILE; + let tile_row = local_id.y * RPT; + let tile_col = local_id.x * RPT; + + let a_base = bn * params.M * params.K; + let b_base = bn * params.K * params.N; + let out_base = bn * params.M * params.N; + + var acc: array, 4>; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + acc[ir][ic] = 0.0h; + } + } + + let num_tiles = (params.K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k_start = t * TILE; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) { + let col = local_id.x * RPT + kk; + a_sub[arow][col] = read_a(a_base, tile_row0 + arow, k_start + col); + b_sub[arow][col] = read_b(b_base, k_start + arow, tile_col0 + col); + } + } + workgroupBarrier(); + + for (var k: u32 = 0u; k < TILE; k = k + 1u) { + let b0 = b_sub[k][tile_col + 0u]; + let b1 = b_sub[k][tile_col + 1u]; + let b2 = b_sub[k][tile_col + 2u]; + let b3 = b_sub[k][tile_col + 3u]; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let aval = a_sub[tile_row + ir][k]; + acc[ir][0] = acc[ir][0] + aval * b0; + acc[ir][1] = acc[ir][1] + aval * b1; + acc[ir][2] = acc[ir][2] + aval * b2; + acc[ir][3] = acc[ir][3] + aval * b3; + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + let r = tile_row0 + tile_row + ir; + let c = tile_col0 + tile_col + ic; + if (r < params.M && c < params.N) { + out[out_base + r * params.N + c] = acc[ir][ic]; + } + } + } +} +`; diff --git a/src/shaders/bmm/bmm-tiled.wgsl.ts b/src/shaders/bmm/bmm-tiled.wgsl.ts new file mode 100644 index 0000000..ff0e7fb --- /dev/null +++ b/src/shaders/bmm/bmm-tiled.wgsl.ts @@ -0,0 +1,103 @@ +/** + * WGSL compute shader for batched matrix multiplication (tiled). + * + * Source: ExecuTorch (backends/webgpu/runtime/ops/bmm/bmm_tiled.wgsl) + * GitHub: https://github.com/pytorch/executorch/blob/main/backends/webgpu/runtime/ops/bmm/bmm_tiled.wgsl + * License: BSD-3-Clause (Copyright (c) Meta Platforms, Inc. and affiliates) + */ +export const SHADER_BMM_TILED = /* wgsl */ ` +struct Params { + B: u32, + M: u32, + N: u32, + K: u32, +}; + +@group(0) @binding(0) var a: array; +@group(0) @binding(1) var b: array; +@group(0) @binding(2) var out: array; +@group(0) @binding(3) var params: Params; + +const TILE: u32 = 32u; +const RPT: u32 = 4u; + +var a_sub: array, 32>; +var b_sub: array, 32>; + +fn read_a(base: u32, row: u32, col: u32) -> f32 { + if (row < params.M && col < params.K) { + return a[base + row * params.K + col]; + } + return 0.0; +} + +fn read_b(base: u32, krow: u32, col: u32) -> f32 { + if (krow < params.K && col < params.N) { + return b[base + krow * params.N + col]; + } + return 0.0; +} + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3) { + let tiles_per_m = (params.M + TILE - 1u) / TILE; + let bn = wg_id.y / tiles_per_m; + let row_tile = wg_id.y % tiles_per_m; + let tile_row0 = row_tile * TILE; + let tile_col0 = wg_id.x * TILE; + let tile_row = local_id.y * RPT; + let tile_col = local_id.x * RPT; + + let a_base = bn * params.M * params.K; + let b_base = bn * params.K * params.N; + let out_base = bn * params.M * params.N; + + var acc: array, 4>; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + acc[ir][ic] = 0.0; + } + } + + let num_tiles = (params.K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k_start = t * TILE; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) { + let col = local_id.x * RPT + kk; + a_sub[arow][col] = read_a(a_base, tile_row0 + arow, k_start + col); + b_sub[arow][col] = read_b(b_base, k_start + arow, tile_col0 + col); + } + } + workgroupBarrier(); + + for (var k: u32 = 0u; k < TILE; k = k + 1u) { + let b0 = b_sub[k][tile_col + 0u]; + let b1 = b_sub[k][tile_col + 1u]; + let b2 = b_sub[k][tile_col + 2u]; + let b3 = b_sub[k][tile_col + 3u]; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let aval = a_sub[tile_row + ir][k]; + acc[ir][0] = acc[ir][0] + aval * b0; + acc[ir][1] = acc[ir][1] + aval * b1; + acc[ir][2] = acc[ir][2] + aval * b2; + acc[ir][3] = acc[ir][3] + aval * b3; + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + let r = tile_row0 + tile_row + ir; + let c = tile_col0 + tile_col + ic; + if (r < params.M && c < params.N) { + out[out_base + r * params.N + c] = acc[ir][ic]; + } + } + } +} +`; diff --git a/src/shaders/bmm/bmm-vec4-f16.wgsl.ts b/src/shaders/bmm/bmm-vec4-f16.wgsl.ts new file mode 100644 index 0000000..dfe61bf --- /dev/null +++ b/src/shaders/bmm/bmm-vec4-f16.wgsl.ts @@ -0,0 +1,95 @@ +/** + * WGSL compute shader for float16 batched matrix multiplication vectorized with vec4. + * Requires WebGPU shader-f16 feature. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_BMM_VEC4_F16 = /* wgsl */ ` +enable f16; + +struct Params { + B: u32, + M: u32, + N: u32, + K: u32, +}; + +@group(0) @binding(0) var a: array>; +@group(0) @binding(1) var b: array>; +@group(0) @binding(2) var out: array>; +@group(0) @binding(3) var params: Params; + +const TILE: u32 = 32u; +const RPT: u32 = 4u; +const TILE4: u32 = 8u; + +var a_sub: array, 8>, 32>; +var b_sub: array, 8>, 32>; + +fn read_a4(base4: u32, row: u32, k4: u32) -> vec4 { + if (row < params.M && k4 * 4u < params.K) { + return a[base4 + row * (params.K / 4u) + k4]; + } + return vec4(0.0h); +} + +fn read_b4(base4: u32, krow: u32, n4: u32) -> vec4 { + if (krow < params.K && n4 * 4u < params.N) { + return b[base4 + krow * (params.N / 4u) + n4]; + } + return vec4(0.0h); +} + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let bn = wg_id.z; + let tile_row0 = wg_id.y * TILE; + let tile_col0_4 = wg_id.x * TILE4; + let tile_row = local_id.y * RPT; + let tile_col4 = local_id.x; + + let a_base4 = bn * params.M * (params.K / 4u); + let b_base4 = bn * params.K * (params.N / 4u); + let out_base4 = bn * params.M * (params.N / 4u); + + var acc: array, 4>; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + acc[ir] = vec4(0.0h); + } + + let num_tiles = (params.K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t * TILE4; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + a_sub[arow][local_id.x] = read_a4(a_base4, tile_row0 + arow, k4_start + local_id.x); + b_sub[arow][local_id.x] = read_b4(b_base4, t * TILE + arow, tile_col0_4 + local_id.x); + } + workgroupBarrier(); + + for (var k4: u32 = 0u; k4 < TILE4; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][tile_col4]; + let b1 = b_sub[k4 * 4u + 1u][tile_col4]; + let b2 = b_sub[k4 * 4u + 2u][tile_col4]; + let b3 = b_sub[k4 * 4u + 3u][tile_col4]; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let ac = a_sub[tile_row + ir][k4]; + acc[ir] = acc[ir] + b0 * ac.x + b1 * ac.y + b2 * ac.z + b3 * ac.w; + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let r = tile_row0 + tile_row + ir; + let c4 = tile_col0_4 + tile_col4; + if (r < params.M && c4 * 4u < params.N) { + out[out_base4 + r * (params.N / 4u) + c4] = acc[ir]; + } + } +} +`; diff --git a/src/shaders/bmm/bmm-vec4.wgsl.ts b/src/shaders/bmm/bmm-vec4.wgsl.ts new file mode 100644 index 0000000..cd052e2 --- /dev/null +++ b/src/shaders/bmm/bmm-vec4.wgsl.ts @@ -0,0 +1,94 @@ +/** + * WGSL compute shader for vectorized 128-bit memory-bandwidth optimized batched + * matrix multiplication. Operates on vec4 for 4x wider memory transactions + * when K and N are multiples of 4. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_BMM_VEC4 = /* wgsl */ ` +struct Params { + B: u32, + M: u32, + N: u32, + K: u32, +}; + +@group(0) @binding(0) var a: array>; +@group(0) @binding(1) var b: array>; +@group(0) @binding(2) var out: array>; +@group(0) @binding(3) var params: Params; + +const TILE: u32 = 32u; +const RPT: u32 = 4u; +const TILE4: u32 = 8u; + +var a_sub: array, 8>, 32>; +var b_sub: array, 8>, 32>; + +fn read_a4(base4: u32, row: u32, k4: u32) -> vec4 { + if (row < params.M && k4 * 4u < params.K) { + return a[base4 + row * (params.K / 4u) + k4]; + } + return vec4(0.0); +} + +fn read_b4(base4: u32, krow: u32, n4: u32) -> vec4 { + if (krow < params.K && n4 * 4u < params.N) { + return b[base4 + krow * (params.N / 4u) + n4]; + } + return vec4(0.0); +} + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let bn = wg_id.z; + let tile_row0 = wg_id.y * TILE; + let tile_col0_4 = wg_id.x * TILE4; + let tile_row = local_id.y * RPT; + let tile_col4 = local_id.x; + + let a_base4 = bn * params.M * (params.K / 4u); + let b_base4 = bn * params.K * (params.N / 4u); + let out_base4 = bn * params.M * (params.N / 4u); + + var acc: array, 4>; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + acc[ir] = vec4(0.0); + } + + let num_tiles = (params.K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t * TILE4; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + a_sub[arow][local_id.x] = read_a4(a_base4, tile_row0 + arow, k4_start + local_id.x); + b_sub[arow][local_id.x] = read_b4(b_base4, t * TILE + arow, tile_col0_4 + local_id.x); + } + workgroupBarrier(); + + for (var k4: u32 = 0u; k4 < TILE4; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][tile_col4]; + let b1 = b_sub[k4 * 4u + 1u][tile_col4]; + let b2 = b_sub[k4 * 4u + 2u][tile_col4]; + let b3 = b_sub[k4 * 4u + 3u][tile_col4]; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let ac = a_sub[tile_row + ir][k4]; + acc[ir] = acc[ir] + b0 * ac.x + b1 * ac.y + b2 * ac.z + b3 * ac.w; + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let r = tile_row0 + tile_row + ir; + let c4 = tile_col0_4 + tile_col4; + if (r < params.M && c4 * 4u < params.N) { + out[out_base4 + r * (params.N / 4u) + c4] = acc[ir]; + } + } +} +`; diff --git a/src/shaders/bmm/index.ts b/src/shaders/bmm/index.ts new file mode 100644 index 0000000..ffce2b0 --- /dev/null +++ b/src/shaders/bmm/index.ts @@ -0,0 +1,116 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { SHADER_BMM_TILED } from './bmm-tiled.wgsl'; +import { SHADER_BMM_VEC4 } from './bmm-vec4.wgsl'; +import { SHADER_BMM_64X64 } from './bmm-64x64.wgsl'; +import { SHADER_BMM_TILED_F16 } from './bmm-tiled-f16.wgsl'; +import { SHADER_BMM_VEC4_F16 } from './bmm-vec4-f16.wgsl'; +import { SHADER_BMM_64X64_F16 } from './bmm-64x64-f16.wgsl'; + +const name = 'aten.bmm.default'; + +/** + * Positional argument tuple for `aten::bmm.default`: + * `[a, b, out]` + */ +export type BmmArgs = readonly [a: Tensor, b: Tensor, out: Tensor]; + +/** + * Internal dispatch builder for `aten::bmm.default`. + * + * @param ctx The execution context. + * @param args Positional argument tuple for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: BmmArgs): void { + const [a, b, out] = args; + + if ( + (a.dtype !== 'float32' && a.dtype !== 'float16') || + b.dtype !== a.dtype || + out.dtype !== a.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + if (a.shape.length !== 3 || b.shape.length !== 3 || out.shape.length !== 3) { + throw new Error(`${name}: inputs and output must be 3D`); + } + + const batch = a.shape[0]; + const m = a.shape[1]; + const k = a.shape[2]; + const n = b.shape[2]; + if (b.shape[0] !== batch || b.shape[1] !== k) { + throw new Error(`${name}: batch/K mismatch between a and b`); + } + if (out.shape[0] !== batch || out.shape[1] !== m || out.shape[2] !== n) { + throw new Error(`${name}: output shape [${out.shape}] != [${batch}, ${m}, ${n}]`); + } + + const device = ctx.device; + const isF16 = a.dtype === 'float16'; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + const paramsBuffer = ctx.uniformBuffer(new Uint32Array([batch, m, n, k])); + + // Use 64x64 register-tiled GEMM when dimensions are sufficiently large and divisible by 4, + // falling back to 32x32 vec4 or 32x32 scalar tiled paths. + const useVec4 = k % 4 === 0 && n % 4 === 0; + const use64x64 = useVec4 && m >= 64 && n >= 64 && k >= 16; + + let code = isF16 ? SHADER_BMM_TILED_F16 : SHADER_BMM_TILED; + let workgroupCountX = Math.ceil(n / 32); + let workgroupCountY = batch * Math.ceil(m / 32); + let workgroupCountZ = 1; + + if (use64x64) { + code = isF16 ? SHADER_BMM_64X64_F16 : SHADER_BMM_64X64; + workgroupCountX = Math.ceil(n / 64); + workgroupCountY = Math.ceil(m / 64); + workgroupCountZ = batch; + } else if (useVec4) { + code = isF16 ? SHADER_BMM_VEC4_F16 : SHADER_BMM_VEC4; + workgroupCountX = Math.ceil(n / 32); + workgroupCountY = Math.ceil(m / 32); + workgroupCountZ = batch; + } + + const bundle = createComputeBundle(device, code, [ + { binding: 0, buffer: a.buffer }, + { binding: 1, buffer: b.buffer }, + { binding: 2, buffer: out.buffer }, + { binding: 3, buffer: paramsBuffer }, + ]); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + workgroupCountZ, + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::bmm.default`. + * + * Batched matrix multiplication via shared-memory tiling — + * Automatically selects 64x64 register-tiled or 32x32 vectorized vec4 path + * when dimensions permit. Supports float32 and float16 with `shader-f16`. + */ +export const bmm: Shader = { + name, + code: { + tiled: SHADER_BMM_TILED, + vec4: SHADER_BMM_VEC4, + tiled64x64: SHADER_BMM_64X64, + tiledF16: SHADER_BMM_TILED_F16, + vec4F16: SHADER_BMM_VEC4_F16, + tiled64x64F16: SHADER_BMM_64X64_F16, + }, + recordIn, +}; diff --git a/src/shaders/cat/cat-block-f16.wgsl.ts b/src/shaders/cat/cat-block-f16.wgsl.ts new file mode 100644 index 0000000..f939fc9 --- /dev/null +++ b/src/shaders/cat/cat-block-f16.wgsl.ts @@ -0,0 +1,44 @@ +/** + * WGSL compute shader for contiguous block-scatter tensor concatenation (f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CAT_BLOCK_F16 = /* wgsl */ ` +enable f16; + +struct BlockParams { + c_in: u32, + c_out: u32, + inner_size: u32, + off_k: u32, + total_elements: u32, + _p0: u32, + _p1: u32, + _p2: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var params: BlockParams; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3 +) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.total_elements) { + return; + } + let inner_idx = idx % params.inner_size; + let slice_idx = idx / params.inner_size; + let c = slice_idx % params.c_in; + let b = slice_idx / params.c_in; + + let out_idx = (b * params.c_out + params.off_k + c) * params.inner_size + inner_idx; + output[out_idx] = input[idx]; +} +`; diff --git a/src/shaders/cat/cat-block-vec4-f16.wgsl.ts b/src/shaders/cat/cat-block-vec4-f16.wgsl.ts new file mode 100644 index 0000000..b544d77 --- /dev/null +++ b/src/shaders/cat/cat-block-vec4-f16.wgsl.ts @@ -0,0 +1,44 @@ +/** + * WGSL compute shader for vectorized contiguous block-scatter tensor concatenation (vec4). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CAT_BLOCK_VEC4_F16 = /* wgsl */ ` +enable f16; + +struct BlockParams { + c_in: u32, + c_out: u32, + inner_size: u32, + off_k: u32, + total_elements: u32, + _p0: u32, + _p1: u32, + _p2: u32, +}; + +@group(0) @binding(0) var input: array>; +@group(0) @binding(1) var output: array>; +@group(0) @binding(2) var params: BlockParams; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3 +) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.total_elements) { + return; + } + let inner_idx = idx % params.inner_size; + let slice_idx = idx / params.inner_size; + let c = slice_idx % params.c_in; + let b = slice_idx / params.c_in; + + let out_idx = (b * params.c_out + params.off_k + c) * params.inner_size + inner_idx; + output[out_idx] = input[idx]; +} +`; diff --git a/src/shaders/cat/cat-block-vec4.wgsl.ts b/src/shaders/cat/cat-block-vec4.wgsl.ts new file mode 100644 index 0000000..a34c44c --- /dev/null +++ b/src/shaders/cat/cat-block-vec4.wgsl.ts @@ -0,0 +1,42 @@ +/** + * WGSL compute shader for 128-bit vectorized contiguous block-scatter tensor concatenation (vec4). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CAT_BLOCK_VEC4 = /* wgsl */ ` +struct BlockParams { + c_in: u32, + c_out: u32, + inner_size: u32, + off_k: u32, + total_elements: u32, + _p0: u32, + _p1: u32, + _p2: u32, +}; + +@group(0) @binding(0) var input: array>; +@group(0) @binding(1) var output: array>; +@group(0) @binding(2) var params: BlockParams; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3 +) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.total_elements) { + return; + } + let inner_idx = idx % params.inner_size; + let slice_idx = idx / params.inner_size; + let c = slice_idx % params.c_in; + let b = slice_idx / params.c_in; + + let out_idx = (b * params.c_out + params.off_k + c) * params.inner_size + inner_idx; + output[out_idx] = input[idx]; +} +`; diff --git a/src/shaders/cat/cat-block.wgsl.ts b/src/shaders/cat/cat-block.wgsl.ts new file mode 100644 index 0000000..a3ada86 --- /dev/null +++ b/src/shaders/cat/cat-block.wgsl.ts @@ -0,0 +1,43 @@ +/** + * WGSL compute shader for contiguous block-scatter tensor concatenation (f32). + * Eliminates coordinate-unraveling integer divisions across dimensions. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CAT_BLOCK = /* wgsl */ ` +struct BlockParams { + c_in: u32, + c_out: u32, + inner_size: u32, + off_k: u32, + total_elements: u32, + _p0: u32, + _p1: u32, + _p2: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var params: BlockParams; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3 +) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.total_elements) { + return; + } + let inner_idx = idx % params.inner_size; + let slice_idx = idx / params.inner_size; + let c = slice_idx % params.c_in; + let b = slice_idx / params.c_in; + + let out_idx = (b * params.c_out + params.off_k + c) * params.inner_size + inner_idx; + output[out_idx] = input[idx]; +} +`; diff --git a/src/shaders/cat/cat-f16.wgsl.ts b/src/shaders/cat/cat-f16.wgsl.ts new file mode 100644 index 0000000..77d2dcb --- /dev/null +++ b/src/shaders/cat/cat-f16.wgsl.ts @@ -0,0 +1,47 @@ +import { TENSOR_META_WGSL } from '../utils/meta'; + +/** + * WGSL compute shader for general tensor concatenation (f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CAT_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +${TENSOR_META_WGSL} +@group(0) @binding(2) var out_meta: TensorMeta; +@group(0) @binding(3) var in_meta: TensorMeta; + +struct Params { + concat_dim: u32, + off_k: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let in_bufi = gid.x; + if (in_bufi >= in_meta.numel) { + return; + } + + var rem = in_bufi; + var out_bufi: u32 = 0u; + for (var d: u32 = 0u; d < in_meta.ndim; d = d + 1u) { + let coord = rem / in_meta.strides[d >> 2u][d & 3u]; + rem = rem % in_meta.strides[d >> 2u][d & 3u]; + var out_coord = coord; + if (d == params.concat_dim) { + out_coord = coord + params.off_k; + } + out_bufi = out_bufi + out_coord * out_meta.strides[d >> 2u][d & 3u]; + } + output[out_bufi] = input[in_bufi]; +} +`; diff --git a/src/shaders/cat/cat.wgsl.ts b/src/shaders/cat/cat.wgsl.ts new file mode 100644 index 0000000..4214020 --- /dev/null +++ b/src/shaders/cat/cat.wgsl.ts @@ -0,0 +1,45 @@ +import { TENSOR_META_WGSL } from '../utils/meta'; + +/** + * WGSL compute shader for general tensor concatenation (f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CAT = /* wgsl */ ` +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +${TENSOR_META_WGSL} +@group(0) @binding(2) var out_meta: TensorMeta; +@group(0) @binding(3) var in_meta: TensorMeta; + +struct Params { + concat_dim: u32, + off_k: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let in_bufi = gid.x; + if (in_bufi >= in_meta.numel) { + return; + } + + var rem = in_bufi; + var out_bufi: u32 = 0u; + for (var d: u32 = 0u; d < in_meta.ndim; d = d + 1u) { + let coord = rem / in_meta.strides[d >> 2u][d & 3u]; + rem = rem % in_meta.strides[d >> 2u][d & 3u]; + var out_coord = coord; + if (d == params.concat_dim) { + out_coord = coord + params.off_k; + } + out_bufi = out_bufi + out_coord * out_meta.strides[d >> 2u][d & 3u]; + } + output[out_bufi] = input[in_bufi]; +} +`; diff --git a/src/shaders/cat/index.ts b/src/shaders/cat/index.ts new file mode 100644 index 0000000..b9887d7 --- /dev/null +++ b/src/shaders/cat/index.ts @@ -0,0 +1,135 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { SHADER_CAT } from './cat.wgsl'; +import { SHADER_CAT_BLOCK } from './cat-block.wgsl'; +import { SHADER_CAT_BLOCK_VEC4 } from './cat-block-vec4.wgsl'; +import { SHADER_CAT_F16 } from './cat-f16.wgsl'; +import { SHADER_CAT_BLOCK_F16 } from './cat-block-f16.wgsl'; +import { SHADER_CAT_BLOCK_VEC4_F16 } from './cat-block-vec4-f16.wgsl'; + +const name = 'aten.cat.default'; + +/** + * Positional argument tuple for `aten::cat.default`: + * `[tensors, dim, out]` + */ +export type CatArgs = readonly [tensors: readonly Tensor[], dim: number, out: Tensor]; + +/** + * Internal dispatch builder for `aten::cat.default`. + * + * @param ctx The execution context. + * @param args Positional argument tuple for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: CatArgs): void { + const [tensors, dimArg, out] = args; + + if (tensors.length === 0) { + throw new Error(`${name}: empty input list`); + } + + const ndim = out.shape.length; + let dim = dimArg; + if (dim < 0) { + dim += ndim; + } + if (dim < 0 || dim >= ndim) { + throw new Error(`${name}: dim out of range`); + } + + const isF16 = out.dtype === 'float16'; + const device = ctx.device; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + let innerSize = 1; + for (let d = dim + 1; d < ndim; d++) { + innerSize *= out.shape[d]; + } + + let offK = 0; + let concatSum = 0; + const wgSize = 256; + + for (const inp of tensors) { + if ((inp.dtype !== 'float32' && inp.dtype !== 'float16') || inp.dtype !== out.dtype) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + if (inp.shape.length !== ndim) { + throw new Error(`${name}: input rank != output rank`); + } + for (let d = 0; d < ndim; d++) { + if (d !== dim && inp.shape[d] !== out.shape[d]) { + throw new Error(`${name}: non-concat dim ${d} size mismatch`); + } + } + concatSum += inp.shape[dim]; + + const cIn = inp.shape[dim]; + const cOut = out.shape[dim]; + const isVec4 = innerSize % 4 === 0; + const effectiveInner = isVec4 ? innerSize / 4 : innerSize; + const totalElements = isVec4 ? inp.numel / 4 : inp.numel; + + const paramsBuffer = ctx.uniformBuffer( + new Uint32Array([cIn, cOut, effectiveInner, offK, totalElements, 0, 0, 0]) + ); + offK += cIn; + + let code = isVec4 ? SHADER_CAT_BLOCK_VEC4 : SHADER_CAT_BLOCK; + if (isF16) { + code = isVec4 ? SHADER_CAT_BLOCK_VEC4_F16 : SHADER_CAT_BLOCK_F16; + } + + const bundle = createComputeBundle( + device, + code, + [ + { binding: 0, buffer: inp.buffer }, + { binding: 1, buffer: out.buffer }, + { binding: 2, buffer: paramsBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize } + ); + + const totalWorkgroups = Math.ceil(totalElements / wgSize); + const workgroupCountX = Math.min(totalWorkgroups, 65535); + const workgroupCountY = Math.ceil(totalWorkgroups / 65535); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); + } + + if (concatSum !== out.shape[dim]) { + throw new Error(`${name}: concat dim sizes do not sum to output`); + } +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::cat.default`. + * + * Concatenates a list of tensors along one dimension. Uses fast block-scatter + * and 128-bit vectorization for contiguous slices. + * Supports float32 and float16 with the `shader-f16` WebGPU feature extension. + */ +export const cat: Shader = { + name, + code: { + default: SHADER_CAT, + block: SHADER_CAT_BLOCK, + blockVec4: SHADER_CAT_BLOCK_VEC4, + defaultF16: SHADER_CAT_F16, + blockF16: SHADER_CAT_BLOCK_F16, + blockVec4F16: SHADER_CAT_BLOCK_VEC4_F16, + }, + recordIn, +}; diff --git a/src/shaders/clone/index.ts b/src/shaders/clone/index.ts new file mode 100644 index 0000000..fde81ee --- /dev/null +++ b/src/shaders/clone/index.ts @@ -0,0 +1,36 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; + +import { recordFlatCopy } from '../utils/flat-copy'; + +const name = 'aten.clone.default'; + +/** + * Positional argument tuple for `aten::clone.default`: + * `[in, memory_format, out]` + */ +export type CloneArgs = readonly [inTensor: Tensor, memoryFormat: undefined, outTensor: Tensor]; + +/** + * Internal dispatch builder for `aten::clone.default`. + * + * @param ctx The execution context. + * @param args Positional argument tuple for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: CloneArgs): void { + const [inTensor, , outTensor] = args; + recordFlatCopy(ctx, inTensor, outTensor, name); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::clone.default`. + * + * Clone = flat byte copy (survives the serializer's redundant-op elimination). + * Supports any dtype (float32, float16, etc.) via buffer copy. + */ +export const clone: Shader = { + name, + code: '/* flat byte copy — no kernel dispatch */', + recordIn, +}; diff --git a/src/shaders/constant-pad-nd/constant-pad-nd-f16.wgsl.ts b/src/shaders/constant-pad-nd/constant-pad-nd-f16.wgsl.ts new file mode 100644 index 0000000..0b3cf17 --- /dev/null +++ b/src/shaders/constant-pad-nd/constant-pad-nd-f16.wgsl.ts @@ -0,0 +1,64 @@ +/** + * WGSL compute shader for constant padding on up to 4D tensors (f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONSTANT_PAD_ND_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +struct Params { + out_dims: vec4, + in_dims: vec4, + left: vec4, + out_numel: u32, + value: f32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3 +) { + let i = gid.x + gid.y * (num_workgroups.x * wg_size); + if (i >= params.out_numel) { + return; + } + + var rem = i; + let o3 = rem % params.out_dims.w; + rem = rem / params.out_dims.w; + let o2 = rem % params.out_dims.z; + rem = rem / params.out_dims.z; + let o1 = rem % params.out_dims.y; + rem = rem / params.out_dims.y; + let o0 = rem % params.out_dims.x; + + let c0 = o0 - params.left.x; + let c1 = o1 - params.left.y; + let c2 = o2 - params.left.z; + let c3 = o3 - params.left.w; + + let in0 = o0 >= params.left.x && c0 < params.in_dims.x; + let in1 = o1 >= params.left.y && c1 < params.in_dims.y; + let in2 = o2 >= params.left.z && c2 < params.in_dims.z; + let in3 = o3 >= params.left.w && c3 < params.in_dims.w; + + if (in0 && in1 && in2 && in3) { + let in_idx = + ((c0 * params.in_dims.y + c1) * params.in_dims.z + c2) * params.in_dims.w + + c3; + out[i] = inp[in_idx]; + } else { + out[i] = f16(params.value); + } +} +`; diff --git a/src/shaders/constant-pad-nd/constant-pad-nd-spatial-f16.wgsl.ts b/src/shaders/constant-pad-nd/constant-pad-nd-spatial-f16.wgsl.ts new file mode 100644 index 0000000..bcb7208 --- /dev/null +++ b/src/shaders/constant-pad-nd/constant-pad-nd-spatial-f16.wgsl.ts @@ -0,0 +1,52 @@ +/** + * WGSL compute shader for spatial constant padding using 2D grid (f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONSTANT_PAD_ND_SPATIAL_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +struct Params { + out_dims: vec4, + in_dims: vec4, + left: vec4, + out_numel: u32, + value: f32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(2) var params: Params; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(workgroup_id) wg_id: vec3 +) { + let ow = gid.x; + let oh = gid.y; + let slice = wg_id.z; + + if (ow >= params.out_dims.w || oh >= params.out_dims.z) { + return; + } + + let iw = ow - params.left.w; + let ih = oh - params.left.z; + + let out_idx = (slice * params.out_dims.z + oh) * params.out_dims.w + ow; + + let in_w = ow >= params.left.w && iw < params.in_dims.w; + let in_h = oh >= params.left.z && ih < params.in_dims.z; + + if (in_w && in_h) { + let in_idx = (slice * params.in_dims.z + ih) * params.in_dims.w + iw; + out[out_idx] = inp[in_idx]; + } else { + out[out_idx] = f16(params.value); + } +} +`; diff --git a/src/shaders/constant-pad-nd/constant-pad-nd-spatial.wgsl.ts b/src/shaders/constant-pad-nd/constant-pad-nd-spatial.wgsl.ts new file mode 100644 index 0000000..70aea70 --- /dev/null +++ b/src/shaders/constant-pad-nd/constant-pad-nd-spatial.wgsl.ts @@ -0,0 +1,50 @@ +/** + * WGSL compute shader for spatial constant padding using 2D grid (f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONSTANT_PAD_ND_SPATIAL = /* wgsl */ ` +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +struct Params { + out_dims: vec4, + in_dims: vec4, + left: vec4, + out_numel: u32, + value: f32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(2) var params: Params; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(workgroup_id) wg_id: vec3 +) { + let ow = gid.x; + let oh = gid.y; + let slice = wg_id.z; + + if (ow >= params.out_dims.w || oh >= params.out_dims.z) { + return; + } + + let iw = ow - params.left.w; + let ih = oh - params.left.z; + + let out_idx = (slice * params.out_dims.z + oh) * params.out_dims.w + ow; + + let in_w = ow >= params.left.w && iw < params.in_dims.w; + let in_h = oh >= params.left.z && ih < params.in_dims.z; + + if (in_w && in_h) { + let in_idx = (slice * params.in_dims.z + ih) * params.in_dims.w + iw; + out[out_idx] = inp[in_idx]; + } else { + out[out_idx] = params.value; + } +} +`; diff --git a/src/shaders/constant-pad-nd/constant-pad-nd.wgsl.ts b/src/shaders/constant-pad-nd/constant-pad-nd.wgsl.ts new file mode 100644 index 0000000..e57a71d --- /dev/null +++ b/src/shaders/constant-pad-nd/constant-pad-nd.wgsl.ts @@ -0,0 +1,62 @@ +/** + * WGSL compute shader for constant padding on up to 4D tensors (f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONSTANT_PAD_ND = /* wgsl */ ` +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +struct Params { + out_dims: vec4, + in_dims: vec4, + left: vec4, + out_numel: u32, + value: f32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3 +) { + let i = gid.x + gid.y * (num_workgroups.x * wg_size); + if (i >= params.out_numel) { + return; + } + + var rem = i; + let o3 = rem % params.out_dims.w; + rem = rem / params.out_dims.w; + let o2 = rem % params.out_dims.z; + rem = rem / params.out_dims.z; + let o1 = rem % params.out_dims.y; + rem = rem / params.out_dims.y; + let o0 = rem % params.out_dims.x; + + let c0 = o0 - params.left.x; + let c1 = o1 - params.left.y; + let c2 = o2 - params.left.z; + let c3 = o3 - params.left.w; + + let in0 = o0 >= params.left.x && c0 < params.in_dims.x; + let in1 = o1 >= params.left.y && c1 < params.in_dims.y; + let in2 = o2 >= params.left.z && c2 < params.in_dims.z; + let in3 = o3 >= params.left.w && c3 < params.in_dims.w; + + if (in0 && in1 && in2 && in3) { + let in_idx = + ((c0 * params.in_dims.y + c1) * params.in_dims.z + c2) * params.in_dims.w + + c3; + out[i] = inp[in_idx]; + } else { + out[i] = params.value; + } +} +`; diff --git a/src/shaders/constant-pad-nd/index.ts b/src/shaders/constant-pad-nd/index.ts new file mode 100644 index 0000000..af79923 --- /dev/null +++ b/src/shaders/constant-pad-nd/index.ts @@ -0,0 +1,161 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { SHADER_CONSTANT_PAD_ND } from './constant-pad-nd.wgsl'; +import { SHADER_CONSTANT_PAD_ND_SPATIAL } from './constant-pad-nd-spatial.wgsl'; +import { SHADER_CONSTANT_PAD_ND_F16 } from './constant-pad-nd-f16.wgsl'; +import { SHADER_CONSTANT_PAD_ND_SPATIAL_F16 } from './constant-pad-nd-spatial-f16.wgsl'; + +const name = 'aten.constant_pad_nd.default'; + +/** + * Positional argument tuple for `aten::constant_pad_nd.default`: + * - `[inTensor, pad, value, outTensor]` + * - `[inTensor, pad, outTensor]` (value defaults to 0.0) + */ +export type ConstantPadNdArgs = + | readonly [inTensor: Tensor, pad: readonly number[], outTensor: Tensor] + | readonly [inTensor: Tensor, pad: readonly number[], value: number, outTensor: Tensor]; + +/** + * Internal dispatch builder for `aten::constant_pad_nd.default`. + * + * @param ctx The execution context. + * @param args Positional arguments for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: ConstantPadNdArgs): void { + const [inTensor, pad] = args; + const [value, outTensor] = args.length === 4 ? [args[2], args[3]] : [0.0, args[2]]; + + if ( + (inTensor.dtype !== 'float32' && inTensor.dtype !== 'float16') || + outTensor.dtype !== inTensor.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + + const isF16 = inTensor.dtype === 'float16'; + const device = ctx.device; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + const nd = inTensor.shape.length; + if (nd === 0 || nd > 4) { + throw new Error(`${name}: rank must be 1..4 (got ${nd})`); + } + if (outTensor.shape.length !== nd) { + throw new Error(`${name}: in/out rank mismatch (${nd} !== ${outTensor.shape.length})`); + } + if (pad.length % 2 !== 0) { + throw new Error(`${name}: pad must be even-length`); + } + + const npad = pad.length / 2; + if (npad > nd) { + throw new Error(`${name}: pad longer than rank`); + } + + // Per-dim left/right pad (pad list is reversed-dim, from the LAST dim) + const left = [0, 0, 0, 0]; + const right = [0, 0, 0, 0]; + for (let k = 0; k < npad; k++) { + const d = nd - 1 - k; + left[d] = pad[2 * k]; + right[d] = pad[2 * k + 1]; + } + + for (let d = 0; d < nd; d++) { + if (left[d] < 0 || right[d] < 0) { + throw new Error(`${name}: negative pad (cropping) not supported`); + } + const expected = inTensor.shape[d] + left[d] + right[d]; + if (outTensor.shape[d] !== expected) { + throw new Error(`${name}: output dim ${d} mismatch (${expected} != ${outTensor.shape[d]})`); + } + } + + // Right-align dims into 4D tuples: leading unused dims = 1 (pad = 0) + const padLeft = (arr: readonly number[], fill: number): number[] => [ + ...Array(4 - arr.length).fill(fill), + ...arr, + ]; + + const left4 = padLeft(left.slice(0, nd), 0); + const right4 = padLeft(right.slice(0, nd), 0); + const outDims = padLeft(outTensor.shape, 1); + const inDims = padLeft(inTensor.shape, 1); + + const buffer = new ArrayBuffer(64); + new Uint32Array(buffer).set([...outDims, ...inDims, ...left4, outTensor.numel]); + new Float32Array(buffer, 52, 1)[0] = value; + + const uniformBuffer = ctx.uniformBuffer(buffer); + + const isSpatialOnly = left4[0] === 0 && left4[1] === 0 && right4[0] === 0 && right4[1] === 0; + const sliceCount = outDims[0] * outDims[1]; + + if (isSpatialOnly && sliceCount <= 65535) { + const bundle = createComputeBundle( + device, + isF16 ? SHADER_CONSTANT_PAD_ND_SPATIAL_F16 : SHADER_CONSTANT_PAD_ND_SPATIAL, + [ + { binding: 0, buffer: outTensor.buffer }, + { binding: 1, buffer: inTensor.buffer }, + { binding: 2, buffer: uniformBuffer }, + ] + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX: Math.ceil(outDims[3] / 16), + workgroupCountY: Math.ceil(outDims[2] / 16), + workgroupCountZ: sliceCount, + }); + return; + } + + const wgSize = 256; + const totalWorkgroups = Math.ceil(outTensor.numel / wgSize); + const workgroupCountX = Math.min(totalWorkgroups, 65535); + const workgroupCountY = Math.ceil(totalWorkgroups / 65535); + + const bundle = createComputeBundle( + device, + isF16 ? SHADER_CONSTANT_PAD_ND_F16 : SHADER_CONSTANT_PAD_ND, + [ + { binding: 0, buffer: outTensor.buffer }, + { binding: 1, buffer: inTensor.buffer }, + { binding: 2, buffer: uniformBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize } + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::constant_pad_nd.default`. + * + * Copies input tensor into output tensor with constant padding on up to 4D tensors. + * Supports float32 and float16 with the `shader-f16` WebGPU feature extension. + */ +export const constantPadNd: Shader = { + name, + code: { + default: SHADER_CONSTANT_PAD_ND, + spatial: SHADER_CONSTANT_PAD_ND_SPATIAL, + defaultF16: SHADER_CONSTANT_PAD_ND_F16, + spatialF16: SHADER_CONSTANT_PAD_ND_SPATIAL_F16, + }, + recordIn, +}; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-1x1-single-channel-f16.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-1x1-single-channel-f16.wgsl.ts new file mode 100644 index 0000000..e5de28c --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-1x1-single-channel-f16.wgsl.ts @@ -0,0 +1,59 @@ +/** + * FP16 pointwise convolution for a single output channel, with FP32 accumulation. + * Keeps packed FP16 storage and widens operands before multiplication. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_1X1_SINGLE_CHANNEL_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_x: array>; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) gid: vec3) { + let plane = (params.H_out * params.W_out) >> 2u; + let pos = gid.x; + let batch = gid.y; + if (pos >= plane) { return; } + + var acc = vec4(0.0); + let base = batch * params.IC * plane + pos; + for (var ic = 0u; ic < params.IC; ic = ic + 1u) { + acc += vec4(t_x[base + ic * plane]) * f32(t_weight[ic]); + } + if (params.has_bias != 0u) { acc += vec4(f32(t_bias[0])); } + t_out[batch * plane + pos] = vec4(clamp(acc, vec4(params.output_min), vec4(params.output_max))); +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-1x1-single-channel.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-1x1-single-channel.wgsl.ts new file mode 100644 index 0000000..8b53302 --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-1x1-single-channel.wgsl.ts @@ -0,0 +1,57 @@ +/** + * Pointwise convolution for a single output channel, with four spatial values per lane. + * Uses coalesced vector loads without staging mostly empty GEMM tiles. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_1X1_SINGLE_CHANNEL = /* wgsl */ ` +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_x: array>; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) gid: vec3) { + let plane = (params.H_out * params.W_out) >> 2u; + let pos = gid.x; + let batch = gid.y; + if (pos >= plane) { return; } + + var acc = vec4(0.0); + let base = batch * params.IC * plane + pos; + for (var ic = 0u; ic < params.IC; ic = ic + 1u) { + acc += t_x[base + ic * plane] * t_weight[ic]; + } + if (params.has_bias != 0u) { acc += vec4(t_bias[0]); } + t_out[batch * plane + pos] = clamp(acc, vec4(params.output_min), vec4(params.output_max)); +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-1x1-tile32-f16.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-1x1-tile32-f16.wgsl.ts new file mode 100644 index 0000000..8087ef9 --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-1x1-tile32-f16.wgsl.ts @@ -0,0 +1,160 @@ +/** + * WGSL compute shader for 32x64 register-tiled 1x1 convolution with clamp bounds in FP16 (N = 1, 4x4 per thread across 128 threads). + * Optimized for small channel counts (OC <= 32) to maximize wave occupancy and eliminate thread divergence. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_1X1_TILE32_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_x: array>; +@group(0) @binding(2) var t_weight: array>; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +var a_sub: array, 4>, 32>; +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(local_invocation_index) tid: u32 +) { + let M = params.OC; + let N = params.H_out * params.W_out; + let K = params.IC; + let K_vecs = K >> 2u; + let N_vecs = N >> 2u; + + let tile_row0 = wg_id.y * 32u; + let tile_col0_4 = wg_id.x * 16u; + + var acc0 = vec4(0.0h); + var acc1 = vec4(0.0h); + var acc2 = vec4(0.0h); + var acc3 = vec4(0.0h); + + let a_r = tid >> 2u; // 0..31 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < M; + + let b_r0 = tid >> 4u; // 0..7 + let b_c0 = tid & 15u; // 0..15 + let b_n4_0 = tile_col0_4 + b_c0; + let b_c0_valid = b_n4_0 < N_vecs; + + let b_r1 = b_r0 + 8u; // 8..15 + let b_c1 = b_c0; + let b_n4_1 = b_n4_0; + let b_c1_valid = b_c0_valid; + + let num_tiles = (K + 15u) >> 4u; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t << 2u; + let k_start = t << 4u; + + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = t_weight[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0h); + } + + let bk0 = k_start + b_r0; + if (bk0 < K && b_c0_valid) { + b_sub[b_r0][b_c0] = t_x[bk0 * N_vecs + b_n4_0]; + } else { + b_sub[b_r0][b_c0] = vec4(0.0h); + } + + let bk1 = k_start + b_r1; + if (bk1 < K && b_c1_valid) { + b_sub[b_r1][b_c1] = t_x[bk1 * N_vecs + b_n4_1]; + } else { + b_sub[b_r1][b_c1] = vec4(0.0h); + } + + workgroupBarrier(); + + let thread_r = local_id.y << 2u; + let thread_c4 = local_id.x; + + for (var k4: u32 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + let thread_r = local_id.y << 2u; + let out_r0 = tile_row0 + thread_r; + let out_c4 = tile_col0_4 + local_id.x; + let min_v = vec4(f16(params.output_min)); + let max_v = vec4(f16(params.output_max)); + + if (out_c4 < N_vecs) { + if (out_r0 + 0u < M) { + var v0 = acc0; + if (params.has_bias != 0u) { v0 = v0 + vec4(t_bias[out_r0 + 0u]); } + t_out[(out_r0 + 0u) * N_vecs + out_c4] = clamp(v0, min_v, max_v); + } + if (out_r0 + 1u < M) { + var v1 = acc1; + if (params.has_bias != 0u) { v1 = v1 + vec4(t_bias[out_r0 + 1u]); } + t_out[(out_r0 + 1u) * N_vecs + out_c4] = clamp(v1, min_v, max_v); + } + if (out_r0 + 2u < M) { + var v2 = acc2; + if (params.has_bias != 0u) { v2 = v2 + vec4(t_bias[out_r0 + 2u]); } + t_out[(out_r0 + 2u) * N_vecs + out_c4] = clamp(v2, min_v, max_v); + } + if (out_r0 + 3u < M) { + var v3 = acc3; + if (params.has_bias != 0u) { v3 = v3 + vec4(t_bias[out_r0 + 3u]); } + t_out[(out_r0 + 3u) * N_vecs + out_c4] = clamp(v3, min_v, max_v); + } + } +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-1x1-tile32.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-1x1-tile32.wgsl.ts new file mode 100644 index 0000000..419ab22 --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-1x1-tile32.wgsl.ts @@ -0,0 +1,160 @@ +/** + * WGSL compute shader for 32x64 register-tiled 1x1 convolution with clamp bounds (N = 1, 4x4 per thread across 128 threads). + * Optimized for small channel counts (OC <= 32) to maximize wave occupancy and eliminate thread divergence. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_1X1_TILE32 = /* wgsl */ ` +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_x: array>; +@group(0) @binding(2) var t_weight: array>; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +// 32 rows x 16 K elements (4 vec4s) = 128 vec4s (2 KB) +var a_sub: array, 4>, 32>; +// 16 K rows x 64 N elements (16 vec4s) = 256 vec4s (4 KB) +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(local_invocation_index) tid: u32 +) { + let M = params.OC; + let N = params.H_out * params.W_out; + let K = params.IC; + let K_vecs = K >> 2u; + let N_vecs = N >> 2u; + + let tile_row0 = wg_id.y * 32u; + let tile_col0_4 = wg_id.x * 16u; + + var acc0 = vec4(0.0); + var acc1 = vec4(0.0); + var acc2 = vec4(0.0); + var acc3 = vec4(0.0); + + let a_r = tid >> 2u; // 0..31 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < M; + + let b_r0 = tid >> 4u; // 0..7 + let b_c0 = tid & 15u; // 0..15 + let b_n4_0 = tile_col0_4 + b_c0; + let b_c0_valid = b_n4_0 < N_vecs; + + let b_r1 = b_r0 + 8u; // 8..15 + let b_c1 = b_c0; + let b_n4_1 = b_n4_0; + let b_c1_valid = b_c0_valid; + + let num_tiles = (K + 15u) >> 4u; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t << 2u; + let k_start = t << 4u; + + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = t_weight[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0); + } + + let bk0 = k_start + b_r0; + if (bk0 < K && b_c0_valid) { + b_sub[b_r0][b_c0] = t_x[bk0 * N_vecs + b_n4_0]; + } else { + b_sub[b_r0][b_c0] = vec4(0.0); + } + + let bk1 = k_start + b_r1; + if (bk1 < K && b_c1_valid) { + b_sub[b_r1][b_c1] = t_x[bk1 * N_vecs + b_n4_1]; + } else { + b_sub[b_r1][b_c1] = vec4(0.0); + } + + workgroupBarrier(); + + let thread_r = local_id.y << 2u; + let thread_c4 = local_id.x; + + for (var k4: u32 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + let thread_r = local_id.y << 2u; + let out_r0 = tile_row0 + thread_r; + let out_c4 = tile_col0_4 + local_id.x; + let min_v = vec4(params.output_min); + let max_v = vec4(params.output_max); + + if (out_c4 < N_vecs) { + if (out_r0 + 0u < M) { + var v0 = acc0; + if (params.has_bias != 0u) { v0 = v0 + vec4(t_bias[out_r0 + 0u]); } + t_out[(out_r0 + 0u) * N_vecs + out_c4] = clamp(v0, min_v, max_v); + } + if (out_r0 + 1u < M) { + var v1 = acc1; + if (params.has_bias != 0u) { v1 = v1 + vec4(t_bias[out_r0 + 1u]); } + t_out[(out_r0 + 1u) * N_vecs + out_c4] = clamp(v1, min_v, max_v); + } + if (out_r0 + 2u < M) { + var v2 = acc2; + if (params.has_bias != 0u) { v2 = v2 + vec4(t_bias[out_r0 + 2u]); } + t_out[(out_r0 + 2u) * N_vecs + out_c4] = clamp(v2, min_v, max_v); + } + if (out_r0 + 3u < M) { + var v3 = acc3; + if (params.has_bias != 0u) { v3 = v3 + vec4(t_bias[out_r0 + 3u]); } + t_out[(out_r0 + 3u) * N_vecs + out_c4] = clamp(v3, min_v, max_v); + } + } +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-1x1-tile64-f16.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-1x1-tile64-f16.wgsl.ts new file mode 100644 index 0000000..722c94b --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-1x1-tile64-f16.wgsl.ts @@ -0,0 +1,150 @@ +/** + * WGSL compute shader for 64x64 register-tiled 1x1 convolution with clamp bounds in float16 (N = 1, 4x4 per thread across 256 threads). + * High occupancy, coalesced 128-bit memory transactions, minimal register pressure for all GPUs. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_1X1_64X64_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_x: array>; +@group(0) @binding(2) var t_weight: array>; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +// 64 rows x 16 K elements (4 vec4s) = 256 vec4s (2 KB) +var a_sub: array, 4>, 64>; +// 16 K rows x 64 N elements (16 vec4s) = 256 vec4s (2 KB) +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let M = params.OC; + let N = params.H_out * params.W_out; + let K = params.IC; + let K_vecs = K >> 2u; + let N_vecs = N >> 2u; + + let tile_row0 = wg_id.y * 64u; + let tile_col0_4 = wg_id.x * 16u; + let tid = local_id.y * 16u + local_id.x; // 0..255 + + var acc0 = vec4(0.0h); + var acc1 = vec4(0.0h); + var acc2 = vec4(0.0h); + var acc3 = vec4(0.0h); + + // Load coordinates for 256 threads (each thread loads 1 vec4 of A/weight and 1 vec4 of B/input) + let a_r = tid >> 2u; // 0..63 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < M; + + let b_r = tid >> 4u; // 0..15 + let b_c = tid & 15u; // 0..15 + let b_n4 = tile_col0_4 + b_c; + let b_c_valid = b_n4 < N_vecs; + + let num_tiles = (K + 15u) >> 4u; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t << 2u; + let k_start = t << 4u; + + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = t_weight[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0h); + } + + let bk = k_start + b_r; + if (bk < K && b_c_valid) { + b_sub[b_r][b_c] = t_x[bk * N_vecs + b_n4]; + } else { + b_sub[b_r][b_c] = vec4(0.0h); + } + workgroupBarrier(); + + let thread_r = local_id.y << 2u; // 0..60 (step 4) + let thread_c4 = local_id.x; // 0..15 + + for (var k4: u32 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + let thread_r = local_id.y << 2u; + let out_r0 = tile_row0 + thread_r; + let out_c4 = tile_col0_4 + local_id.x; + let min_v = vec4(f16(params.output_min)); + let max_v = vec4(f16(params.output_max)); + + if (out_c4 < N_vecs) { + if (out_r0 + 0u < M) { + var v0 = acc0; + if (params.has_bias != 0u) { v0 = v0 + vec4(t_bias[out_r0 + 0u]); } + t_out[(out_r0 + 0u) * N_vecs + out_c4] = clamp(v0, min_v, max_v); + } + if (out_r0 + 1u < M) { + var v1 = acc1; + if (params.has_bias != 0u) { v1 = v1 + vec4(t_bias[out_r0 + 1u]); } + t_out[(out_r0 + 1u) * N_vecs + out_c4] = clamp(v1, min_v, max_v); + } + if (out_r0 + 2u < M) { + var v2 = acc2; + if (params.has_bias != 0u) { v2 = v2 + vec4(t_bias[out_r0 + 2u]); } + t_out[(out_r0 + 2u) * N_vecs + out_c4] = clamp(v2, min_v, max_v); + } + if (out_r0 + 3u < M) { + var v3 = acc3; + if (params.has_bias != 0u) { v3 = v3 + vec4(t_bias[out_r0 + 3u]); } + t_out[(out_r0 + 3u) * N_vecs + out_c4] = clamp(v3, min_v, max_v); + } + } +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-1x1-tile64.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-1x1-tile64.wgsl.ts new file mode 100644 index 0000000..a941391 --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-1x1-tile64.wgsl.ts @@ -0,0 +1,148 @@ +/** + * WGSL compute shader for 64x64 register-tiled 1x1 convolution with clamp bounds (N = 1, 4x4 per thread across 256 threads). + * High occupancy, coalesced 128-bit memory transactions, minimal register pressure for all GPUs. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_1X1_64X64 = /* wgsl */ ` +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_x: array>; +@group(0) @binding(2) var t_weight: array>; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +// 64 rows x 16 K elements (4 vec4s) = 256 vec4s (4 KB) +var a_sub: array, 4>, 64>; +// 16 K rows x 64 N elements (16 vec4s) = 256 vec4s (4 KB) +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let M = params.OC; + let N = params.H_out * params.W_out; + let K = params.IC; + let K_vecs = K >> 2u; + let N_vecs = N >> 2u; + + let tile_row0 = wg_id.y * 64u; + let tile_col0_4 = wg_id.x * 16u; + let tid = local_id.y * 16u + local_id.x; // 0..255 + + var acc0 = vec4(0.0); + var acc1 = vec4(0.0); + var acc2 = vec4(0.0); + var acc3 = vec4(0.0); + + // Load coordinates for 256 threads (each thread loads 1 vec4 of A/weight and 1 vec4 of B/input) + let a_r = tid >> 2u; // 0..63 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < M; + + let b_r = tid >> 4u; // 0..15 + let b_c = tid & 15u; // 0..15 + let b_n4 = tile_col0_4 + b_c; + let b_c_valid = b_n4 < N_vecs; + + let num_tiles = (K + 15u) >> 4u; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t << 2u; + let k_start = t << 4u; + + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = t_weight[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0); + } + + let bk = k_start + b_r; + if (bk < K && b_c_valid) { + b_sub[b_r][b_c] = t_x[bk * N_vecs + b_n4]; + } else { + b_sub[b_r][b_c] = vec4(0.0); + } + workgroupBarrier(); + + let thread_r = local_id.y << 2u; // 0..60 (step 4) + let thread_c4 = local_id.x; // 0..15 + + for (var k4: u32 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + let thread_r = local_id.y << 2u; + let out_r0 = tile_row0 + thread_r; + let out_c4 = tile_col0_4 + local_id.x; + let min_v = vec4(params.output_min); + let max_v = vec4(params.output_max); + + if (out_c4 < N_vecs) { + if (out_r0 + 0u < M) { + var v0 = acc0; + if (params.has_bias != 0u) { v0 = v0 + vec4(t_bias[out_r0 + 0u]); } + t_out[(out_r0 + 0u) * N_vecs + out_c4] = clamp(v0, min_v, max_v); + } + if (out_r0 + 1u < M) { + var v1 = acc1; + if (params.has_bias != 0u) { v1 = v1 + vec4(t_bias[out_r0 + 1u]); } + t_out[(out_r0 + 1u) * N_vecs + out_c4] = clamp(v1, min_v, max_v); + } + if (out_r0 + 2u < M) { + var v2 = acc2; + if (params.has_bias != 0u) { v2 = v2 + vec4(t_bias[out_r0 + 2u]); } + t_out[(out_r0 + 2u) * N_vecs + out_c4] = clamp(v2, min_v, max_v); + } + if (out_r0 + 3u < M) { + var v3 = acc3; + if (params.has_bias != 0u) { v3 = v3 + vec4(t_bias[out_r0 + 3u]); } + t_out[(out_r0 + 3u) * N_vecs + out_c4] = clamp(v3, min_v, max_v); + } + } +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-1x1-vec4.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-1x1-vec4.wgsl.ts new file mode 100644 index 0000000..19202d0 --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-1x1-vec4.wgsl.ts @@ -0,0 +1,127 @@ +/** + * WGSL compute shader for fused 1x1 convolution with clamp bounds (vec4 vectorized GEMM). + * Processes 128-bit vec4 memory transactions along channels (K) and spatial dimensions (N). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_1X1_VEC4 = /* wgsl */ ` +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_x: array>; +@group(0) @binding(2) var t_weight: array>; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +const TILE: u32 = 32u; +const RPT: u32 = 4u; +const TILE4: u32 = 8u; + +var a_sub: array, 8>, 32>; +var b_sub: array, 8>, 32>; + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let M = params.OC; + let spatial = params.H_out * params.W_out; + let sp_vecs = spatial / 4u; + let N_vecs = params.N * sp_vecs; + let K = params.IC; + let K_vecs = K / 4u; + + let tile_row0 = wg_id.y * TILE; + let tile_col0_4 = wg_id.x * TILE4; + let tile_row = local_id.y * RPT; + let tile_col4 = local_id.x; + + var acc: array, 4>; + acc[0] = vec4(0.0); + acc[1] = vec4(0.0); + acc[2] = vec4(0.0); + acc[3] = vec4(0.0); + + let num_tiles = (K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t * TILE4; + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + let m = tile_row0 + arow; + let k4 = k4_start + local_id.x; + if (m < M && k4 < K_vecs) { + a_sub[arow][local_id.x] = t_weight[m * K_vecs + k4]; + } else { + a_sub[arow][local_id.x] = vec4(0.0); + } + + let b_k = t * TILE + arow; + let n4 = tile_col0_4 + local_id.x; + if (b_k < K && n4 < N_vecs) { + let b_idx = n4 / sp_vecs; + let sp4 = n4 % sp_vecs; + b_sub[arow][local_id.x] = t_x[(b_idx * params.IC + b_k) * sp_vecs + sp4]; + } else { + b_sub[arow][local_id.x] = vec4(0.0); + } + } + workgroupBarrier(); + + for (var k4: u32 = 0u; k4 < TILE4; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][tile_col4]; + let b1 = b_sub[k4 * 4u + 1u][tile_col4]; + let b2 = b_sub[k4 * 4u + 2u][tile_col4]; + let b3 = b_sub[k4 * 4u + 3u][tile_col4]; + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let ac = a_sub[tile_row + ir][k4]; + acc[ir] = acc[ir] + b0 * ac.x + b1 * ac.y + b2 * ac.z + b3 * ac.w; + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let m = tile_row0 + tile_row + ir; + let n4 = tile_col0_4 + tile_col4; + if (m < M && n4 < N_vecs) { + var val = acc[ir]; + if (params.has_bias != 0u) { + val = val + vec4(t_bias[m]); + } + val = clamp(val, vec4(params.output_min), vec4(params.output_max)); + let b_idx = n4 / sp_vecs; + let sp4 = n4 % sp_vecs; + t_out[(b_idx * params.OC + m) * sp_vecs + sp4] = val; + } + } +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-1x1.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-1x1.wgsl.ts new file mode 100644 index 0000000..fb5fe6c --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-1x1.wgsl.ts @@ -0,0 +1,134 @@ +/** + * WGSL compute shader for fused 1x1 convolution with clamp bounds (scalar GEMM). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_1X1 = /* wgsl */ ` +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +const TILE: u32 = 32u; +const RPT: u32 = 4u; + +var a_sub: array, 32>; +var b_sub: array, 32>; + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let M = params.OC; + let spatial = params.H_out * params.W_out; + let N_total = params.N * spatial; + let K = params.IC; + + let tile_row0 = wg_id.y * TILE; + let tile_col0 = wg_id.x * TILE; + let tile_row = local_id.y * RPT; + let tile_col = local_id.x * RPT; + + var acc: array, 4>; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + acc[ir][ic] = 0.0; + } + } + + let num_tiles = (K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k_start = t * TILE; + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) { + let col = local_id.x * RPT + kk; + let m = tile_row0 + arow; + let k = k_start + col; + if (m < M && k < K) { + a_sub[arow][col] = t_weight[m * K + k]; + } else { + a_sub[arow][col] = 0.0; + } + + let b_k = k_start + arow; + let n = tile_col0 + col; + if (b_k < K && n < N_total) { + let b_idx = n / spatial; + let sp = n % spatial; + b_sub[arow][col] = t_x[(b_idx * params.IC + b_k) * spatial + sp]; + } else { + b_sub[arow][col] = 0.0; + } + } + } + workgroupBarrier(); + + for (var k: u32 = 0u; k < TILE; k = k + 1u) { + let b0 = b_sub[k][tile_col + 0u]; + let b1 = b_sub[k][tile_col + 1u]; + let b2 = b_sub[k][tile_col + 2u]; + let b3 = b_sub[k][tile_col + 3u]; + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let aval = a_sub[tile_row + ir][k]; + acc[ir][0] = acc[ir][0] + aval * b0; + acc[ir][1] = acc[ir][1] + aval * b1; + acc[ir][2] = acc[ir][2] + aval * b2; + acc[ir][3] = acc[ir][3] + aval * b3; + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let m = tile_row0 + tile_row + ir; + if (m < M) { + var bias_val: f32 = 0.0; + if (params.has_bias != 0u) { + bias_val = t_bias[m]; + } + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + let n = tile_col0 + tile_col + ic; + if (n < N_total) { + let val = acc[ir][ic] + bias_val; + let b_idx = n / spatial; + let sp = n % spatial; + t_out[(b_idx * params.OC + m) * spatial + sp] = clamp(val, params.output_min, params.output_max); + } + } + } + } +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-depthwise-3x3-f16.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-depthwise-3x3-f16.wgsl.ts new file mode 100644 index 0000000..69cfb37 --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-depthwise-3x3-f16.wgsl.ts @@ -0,0 +1,170 @@ +/** + * WGSL compute shader for depthwise 3x3 convolution with clamp in FP16 and shared-memory halo tile caching. + * Loads a 34x34 halo tile cooperatively across 256 threads, eliminating redundant global memory reads. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_DEPTHWISE_3X3_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +var in_tile: array, 34>; +var w_tile: array; +var bias_val: f16; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(local_invocation_id) local_id: vec3, + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_index) tid: u32 +) { + let oc = wg_id.z % params.OC; + let b = wg_id.z / params.OC; + + if (tid < 9u) { + w_tile[tid] = t_weight[oc * 9u + tid]; + } + if (tid == 0u) { + if (params.has_bias != 0u) { + bias_val = t_bias[oc]; + } else { + bias_val = 0.0h; + } + } + + let in_hw = params.H_in * params.W_in; + let in_w = params.W_in; + let in_base = (b * params.IC + oc) * in_hw; + let iH = i32(params.H_in); + let iW = i32(params.W_in); + + let tile_ih_start = i32(wg_id.y * 32u) - 1; + let tile_iw_start = i32(wg_id.x * 32u) - 1; + + for (var idx = tid; idx < 1156u; idx = idx + 256u) { + let tr = idx / 34u; + let tc = idx % 34u; + let ih = tile_ih_start + i32(tr); + let iw = tile_iw_start + i32(tc); + if (ih >= 0 && ih < iH && iw >= 0 && iw < iW) { + in_tile[tr][tc] = t_x[in_base + u32(ih) * in_w + u32(iw)]; + } else { + in_tile[tr][tc] = 0.0h; + } + } + workgroupBarrier(); + + let ow0 = wg_id.x * 32u + local_id.x * 2u; + let oh0 = wg_id.y * 32u + local_id.y * 2u; + + if (ow0 >= params.W_out || oh0 >= params.H_out || b >= params.N) { + return; + } + + var acc00 = bias_val; + var acc01 = bias_val; + var acc10 = bias_val; + var acc11 = bias_val; + + let tr = local_id.y * 2u; + let tc = local_id.x * 2u; + + let w0 = w_tile[0]; let w1 = w_tile[1]; let w2 = w_tile[2]; + let w3 = w_tile[3]; let w4 = w_tile[4]; let w5 = w_tile[5]; + let w6 = w_tile[6]; let w7 = w_tile[7]; let w8 = w_tile[8]; + + // Row 0 + let r0_c0 = in_tile[tr + 0u][tc + 0u]; + let r0_c1 = in_tile[tr + 0u][tc + 1u]; + let r0_c2 = in_tile[tr + 0u][tc + 2u]; + let r0_c3 = in_tile[tr + 0u][tc + 3u]; + + acc00 = acc00 + r0_c0 * w0 + r0_c1 * w1 + r0_c2 * w2; + acc01 = acc01 + r0_c1 * w0 + r0_c2 * w1 + r0_c3 * w2; + + // Row 1 + let r1_c0 = in_tile[tr + 1u][tc + 0u]; + let r1_c1 = in_tile[tr + 1u][tc + 1u]; + let r1_c2 = in_tile[tr + 1u][tc + 2u]; + let r1_c3 = in_tile[tr + 1u][tc + 3u]; + + acc00 = acc00 + r1_c0 * w3 + r1_c1 * w4 + r1_c2 * w5; + acc01 = acc01 + r1_c1 * w3 + r1_c2 * w4 + r1_c3 * w5; + acc10 = acc10 + r1_c0 * w0 + r1_c1 * w1 + r1_c2 * w2; + acc11 = acc11 + r1_c1 * w0 + r1_c2 * w1 + r1_c3 * w2; + + // Row 2 + let r2_c0 = in_tile[tr + 2u][tc + 0u]; + let r2_c1 = in_tile[tr + 2u][tc + 1u]; + let r2_c2 = in_tile[tr + 2u][tc + 2u]; + let r2_c3 = in_tile[tr + 2u][tc + 3u]; + + acc00 = acc00 + r2_c0 * w6 + r2_c1 * w7 + r2_c2 * w8; + acc01 = acc01 + r2_c1 * w6 + r2_c2 * w7 + r2_c3 * w8; + acc10 = acc10 + r2_c0 * w3 + r2_c1 * w4 + r2_c2 * w5; + acc11 = acc11 + r2_c1 * w3 + r2_c2 * w4 + r2_c3 * w5; + + // Row 3 + let r3_c0 = in_tile[tr + 3u][tc + 0u]; + let r3_c1 = in_tile[tr + 3u][tc + 1u]; + let r3_c2 = in_tile[tr + 3u][tc + 2u]; + let r3_c3 = in_tile[tr + 3u][tc + 3u]; + + acc10 = acc10 + r3_c0 * w6 + r3_c1 * w7 + r3_c2 * w8; + acc11 = acc11 + r3_c1 * w6 + r3_c2 * w7 + r3_c3 * w8; + + let out_hw = params.H_out * params.W_out; + let out_w = params.W_out; + let out_base = (b * params.OC + oc) * out_hw; + + let min_v = f16(params.output_min); + let max_v = f16(params.output_max); + + let out_idx00 = out_base + oh0 * out_w + ow0; + t_out[out_idx00] = clamp(acc00, min_v, max_v); + + if (ow0 + 1u < params.W_out) { + t_out[out_idx00 + 1u] = clamp(acc01, min_v, max_v); + } + if (oh0 + 1u < params.H_out) { + let out_idx10 = out_base + (oh0 + 1u) * out_w + ow0; + t_out[out_idx10] = clamp(acc10, min_v, max_v); + if (ow0 + 1u < params.W_out) { + t_out[out_idx10 + 1u] = clamp(acc11, min_v, max_v); + } + } +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-depthwise-3x3.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-depthwise-3x3.wgsl.ts new file mode 100644 index 0000000..d3cde79 --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-depthwise-3x3.wgsl.ts @@ -0,0 +1,168 @@ +/** + * WGSL compute shader for depthwise 3x3 convolution with clamp and shared-memory halo tile caching. + * Loads a 34x34 halo tile cooperatively across 256 threads, eliminating redundant global memory reads. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_DEPTHWISE_3X3 = /* wgsl */ ` +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +var in_tile: array, 34>; +var w_tile: array; +var bias_val: f32; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(local_invocation_id) local_id: vec3, + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_index) tid: u32 +) { + let oc = wg_id.z % params.OC; + let b = wg_id.z / params.OC; + + if (tid < 9u) { + w_tile[tid] = t_weight[oc * 9u + tid]; + } + if (tid == 0u) { + if (params.has_bias != 0u) { + bias_val = t_bias[oc]; + } else { + bias_val = 0.0; + } + } + + let in_hw = params.H_in * params.W_in; + let in_w = params.W_in; + let in_base = (b * params.IC + oc) * in_hw; + let iH = i32(params.H_in); + let iW = i32(params.W_in); + + let tile_ih_start = i32(wg_id.y * 32u) - 1; + let tile_iw_start = i32(wg_id.x * 32u) - 1; + + for (var idx = tid; idx < 1156u; idx = idx + 256u) { + let tr = idx / 34u; + let tc = idx % 34u; + let ih = tile_ih_start + i32(tr); + let iw = tile_iw_start + i32(tc); + if (ih >= 0 && ih < iH && iw >= 0 && iw < iW) { + in_tile[tr][tc] = t_x[in_base + u32(ih) * in_w + u32(iw)]; + } else { + in_tile[tr][tc] = 0.0; + } + } + workgroupBarrier(); + + let ow0 = wg_id.x * 32u + local_id.x * 2u; + let oh0 = wg_id.y * 32u + local_id.y * 2u; + + if (ow0 >= params.W_out || oh0 >= params.H_out || b >= params.N) { + return; + } + + var acc00 = bias_val; + var acc01 = bias_val; + var acc10 = bias_val; + var acc11 = bias_val; + + let tr = local_id.y * 2u; + let tc = local_id.x * 2u; + + let w0 = w_tile[0]; let w1 = w_tile[1]; let w2 = w_tile[2]; + let w3 = w_tile[3]; let w4 = w_tile[4]; let w5 = w_tile[5]; + let w6 = w_tile[6]; let w7 = w_tile[7]; let w8 = w_tile[8]; + + // Row 0 + let r0_c0 = in_tile[tr + 0u][tc + 0u]; + let r0_c1 = in_tile[tr + 0u][tc + 1u]; + let r0_c2 = in_tile[tr + 0u][tc + 2u]; + let r0_c3 = in_tile[tr + 0u][tc + 3u]; + + acc00 = acc00 + r0_c0 * w0 + r0_c1 * w1 + r0_c2 * w2; + acc01 = acc01 + r0_c1 * w0 + r0_c2 * w1 + r0_c3 * w2; + + // Row 1 + let r1_c0 = in_tile[tr + 1u][tc + 0u]; + let r1_c1 = in_tile[tr + 1u][tc + 1u]; + let r1_c2 = in_tile[tr + 1u][tc + 2u]; + let r1_c3 = in_tile[tr + 1u][tc + 3u]; + + acc00 = acc00 + r1_c0 * w3 + r1_c1 * w4 + r1_c2 * w5; + acc01 = acc01 + r1_c1 * w3 + r1_c2 * w4 + r1_c3 * w5; + acc10 = acc10 + r1_c0 * w0 + r1_c1 * w1 + r1_c2 * w2; + acc11 = acc11 + r1_c1 * w0 + r1_c2 * w1 + r1_c3 * w2; + + // Row 2 + let r2_c0 = in_tile[tr + 2u][tc + 0u]; + let r2_c1 = in_tile[tr + 2u][tc + 1u]; + let r2_c2 = in_tile[tr + 2u][tc + 2u]; + let r2_c3 = in_tile[tr + 2u][tc + 3u]; + + acc00 = acc00 + r2_c0 * w6 + r2_c1 * w7 + r2_c2 * w8; + acc01 = acc01 + r2_c1 * w6 + r2_c2 * w7 + r2_c3 * w8; + acc10 = acc10 + r2_c0 * w3 + r2_c1 * w4 + r2_c2 * w5; + acc11 = acc11 + r2_c1 * w3 + r2_c2 * w4 + r2_c3 * w5; + + // Row 3 + let r3_c0 = in_tile[tr + 3u][tc + 0u]; + let r3_c1 = in_tile[tr + 3u][tc + 1u]; + let r3_c2 = in_tile[tr + 3u][tc + 2u]; + let r3_c3 = in_tile[tr + 3u][tc + 3u]; + + acc10 = acc10 + r3_c0 * w6 + r3_c1 * w7 + r3_c2 * w8; + acc11 = acc11 + r3_c1 * w6 + r3_c2 * w7 + r3_c3 * w8; + + let out_hw = params.H_out * params.W_out; + let out_w = params.W_out; + let out_base = (b * params.OC + oc) * out_hw; + + let min_v = params.output_min; + let max_v = params.output_max; + + let out_idx00 = out_base + oh0 * out_w + ow0; + t_out[out_idx00] = clamp(acc00, min_v, max_v); + + if (ow0 + 1u < params.W_out) { + t_out[out_idx00 + 1u] = clamp(acc01, min_v, max_v); + } + if (oh0 + 1u < params.H_out) { + let out_idx10 = out_base + (oh0 + 1u) * out_w + ow0; + t_out[out_idx10] = clamp(acc10, min_v, max_v); + if (ow0 + 1u < params.W_out) { + t_out[out_idx10 + 1u] = clamp(acc11, min_v, max_v); + } + } +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-depthwise-f16.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-depthwise-f16.wgsl.ts new file mode 100644 index 0000000..1c32f07 --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-depthwise-f16.wgsl.ts @@ -0,0 +1,163 @@ +/** + * WGSL compute shader for float16 depthwise convolution with clamp bounds (up to 7x7 kernels). + * Uses 2D spatial workgroup layout (16x16) with 2x2 register tiling per thread (32x32 per workgroup). + * Weights and bias are cached in shared memory then private registers to avoid bank conflicts. + * Requires WebGPU shader-f16 feature. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_DEPTHWISE_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +var w_tile: array; +var bias_val: f16; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(workgroup_id) wg_id: vec3 +) { + let oc = wg_id.z % params.OC; + let b = wg_id.z / params.OC; + + let tid = local_id.y * 16u + local_id.x; + let khw = params.Kh * params.Kw; + if (tid < khw) { + w_tile[tid] = t_weight[oc * khw + tid]; + } + if (tid == 0u) { + if (params.has_bias != 0u) { + bias_val = t_bias[oc]; + } else { + bias_val = 0.0h; + } + } + workgroupBarrier(); + + let ow0 = gid.x * 2u; + let oh0 = gid.y * 2u; + + if (ow0 >= params.W_out || oh0 >= params.H_out || b >= params.N) { + return; + } + + var acc00 = bias_val; + var acc01 = bias_val; + var acc10 = bias_val; + var acc11 = bias_val; + + let sH = i32(params.stride_h); + let sW = i32(params.stride_w); + let dH = i32(params.dil_h); + let dW = i32(params.dil_w); + let iH = i32(params.H_in); + let iW = i32(params.W_in); + + let ih0_base = i32(oh0) * sH - i32(params.pad_h); + let iw0_base = i32(ow0) * sW - i32(params.pad_w); + let in_base = (b * params.IC + oc) * params.H_in; + + var w_reg: array; + for (var i = 0u; i < khw; i = i + 1u) { + w_reg[i] = w_tile[i]; + } + + for (var kh = 0u; kh < params.Kh; kh = kh + 1u) { + let ih0 = ih0_base + i32(kh) * dH; + let ih1 = ih0 + sH; + let valid_h0 = ih0 >= 0 && ih0 < iH; + let valid_h1 = ih1 >= 0 && ih1 < iH; + + var row_offset0: u32 = 0u; + if (valid_h0) { + row_offset0 = (in_base + u32(ih0)) * params.W_in; + } + var row_offset1: u32 = 0u; + if (valid_h1) { + row_offset1 = (in_base + u32(ih1)) * params.W_in; + } + let w_offset = kh * params.Kw; + + for (var kw = 0u; kw < params.Kw; kw = kw + 1u) { + let iw0 = iw0_base + i32(kw) * dW; + let iw1 = iw0 + sW; + let valid_w0 = iw0 >= 0 && iw0 < iW; + let valid_w1 = iw1 >= 0 && iw1 < iW; + let w_val = w_reg[w_offset + kw]; + + if (valid_h0) { + if (valid_w0) { + acc00 = acc00 + t_x[row_offset0 + u32(iw0)] * w_val; + } + if (valid_w1) { + acc01 = acc01 + t_x[row_offset0 + u32(iw1)] * w_val; + } + } + if (valid_h1) { + if (valid_w0) { + acc10 = acc10 + t_x[row_offset1 + u32(iw0)] * w_val; + } + if (valid_w1) { + acc11 = acc11 + t_x[row_offset1 + u32(iw1)] * w_val; + } + } + } + } + + let min_clamp = f16(params.output_min); + let max_clamp = f16(params.output_max); + acc00 = clamp(acc00, min_clamp, max_clamp); + acc01 = clamp(acc01, min_clamp, max_clamp); + acc10 = clamp(acc10, min_clamp, max_clamp); + acc11 = clamp(acc11, min_clamp, max_clamp); + + let out_ch_base = (b * params.OC + oc) * params.H_out; + let out_row0 = (out_ch_base + oh0) * params.W_out; + t_out[out_row0 + ow0] = acc00; + if (ow0 + 1u < params.W_out) { + t_out[out_row0 + ow0 + 1u] = acc01; + } + if (oh0 + 1u < params.H_out) { + let out_row1 = (out_ch_base + oh0 + 1u) * params.W_out; + t_out[out_row1 + ow0] = acc10; + if (ow0 + 1u < params.W_out) { + t_out[out_row1 + ow0 + 1u] = acc11; + } + } +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-depthwise.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-depthwise.wgsl.ts new file mode 100644 index 0000000..b6da632 --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-depthwise.wgsl.ts @@ -0,0 +1,158 @@ +/** + * WGSL compute shader for depthwise convolution with clamp bounds (up to 7x7 kernels). + * Uses 2D spatial workgroup layout (16x16) with 2x2 register tiling per thread (32x32 per workgroup). + * Weights and bias are cached in shared memory then private registers to avoid bank conflicts. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_DEPTHWISE = /* wgsl */ ` +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +var w_tile: array; +var bias_val: f32; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(workgroup_id) wg_id: vec3 +) { + let oc = wg_id.z % params.OC; + let b = wg_id.z / params.OC; + + let tid = local_id.y * 16u + local_id.x; + let khw = params.Kh * params.Kw; + if (tid < khw) { + w_tile[tid] = t_weight[oc * khw + tid]; + } + if (tid == 0u) { + if (params.has_bias != 0u) { + bias_val = t_bias[oc]; + } else { + bias_val = 0.0; + } + } + workgroupBarrier(); + + let ow0 = gid.x * 2u; + let oh0 = gid.y * 2u; + + if (ow0 >= params.W_out || oh0 >= params.H_out || b >= params.N) { + return; + } + + var acc00 = bias_val; + var acc01 = bias_val; + var acc10 = bias_val; + var acc11 = bias_val; + + let sH = i32(params.stride_h); + let sW = i32(params.stride_w); + let dH = i32(params.dil_h); + let dW = i32(params.dil_w); + let iH = i32(params.H_in); + let iW = i32(params.W_in); + + let ih0_base = i32(oh0) * sH - i32(params.pad_h); + let iw0_base = i32(ow0) * sW - i32(params.pad_w); + let in_base = (b * params.IC + oc) * params.H_in; + + var w_reg: array; + for (var i = 0u; i < khw; i = i + 1u) { + w_reg[i] = w_tile[i]; + } + + for (var kh = 0u; kh < params.Kh; kh = kh + 1u) { + let ih0 = ih0_base + i32(kh) * dH; + let ih1 = ih0 + sH; + let valid_h0 = ih0 >= 0 && ih0 < iH; + let valid_h1 = ih1 >= 0 && ih1 < iH; + + var row_offset0: u32 = 0u; + if (valid_h0) { + row_offset0 = (in_base + u32(ih0)) * params.W_in; + } + var row_offset1: u32 = 0u; + if (valid_h1) { + row_offset1 = (in_base + u32(ih1)) * params.W_in; + } + let w_offset = kh * params.Kw; + + for (var kw = 0u; kw < params.Kw; kw = kw + 1u) { + let iw0 = iw0_base + i32(kw) * dW; + let iw1 = iw0 + sW; + let valid_w0 = iw0 >= 0 && iw0 < iW; + let valid_w1 = iw1 >= 0 && iw1 < iW; + let w_val = w_reg[w_offset + kw]; + + if (valid_h0) { + if (valid_w0) { + acc00 = acc00 + t_x[row_offset0 + u32(iw0)] * w_val; + } + if (valid_w1) { + acc01 = acc01 + t_x[row_offset0 + u32(iw1)] * w_val; + } + } + if (valid_h1) { + if (valid_w0) { + acc10 = acc10 + t_x[row_offset1 + u32(iw0)] * w_val; + } + if (valid_w1) { + acc11 = acc11 + t_x[row_offset1 + u32(iw1)] * w_val; + } + } + } + } + + acc00 = clamp(acc00, params.output_min, params.output_max); + acc01 = clamp(acc01, params.output_min, params.output_max); + acc10 = clamp(acc10, params.output_min, params.output_max); + acc11 = clamp(acc11, params.output_min, params.output_max); + + let out_ch_base = (b * params.OC + oc) * params.H_out; + let out_row0 = (out_ch_base + oh0) * params.W_out; + t_out[out_row0 + ow0] = acc00; + if (ow0 + 1u < params.W_out) { + t_out[out_row0 + ow0 + 1u] = acc01; + } + if (oh0 + 1u < params.H_out) { + let out_row1 = (out_ch_base + oh0 + 1u) * params.W_out; + t_out[out_row1 + ow0] = acc10; + if (ow0 + 1u < params.W_out) { + t_out[out_row1 + ow0 + 1u] = acc11; + } + } +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-f16.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-f16.wgsl.ts new file mode 100644 index 0000000..caa30cb --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-f16.wgsl.ts @@ -0,0 +1,91 @@ +/** + * Direct 2D convolution with clamp compute shader (fallback for general grouped cases) in float16. + * Requires WebGPU shader-f16 feature. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.numel) { + return; + } + let ow = idx % params.W_out; + var r = idx / params.W_out; + let oh = r % params.H_out; + r = r / params.H_out; + let oc = r % params.OC; + let n = r / params.OC; + + var acc: f16 = 0.0h; + if (params.has_bias != 0u) { + acc = t_bias[oc]; + } + let oc_per_group = params.OC / params.groups; + let g = oc / oc_per_group; + let ic_base = g * params.ic_per_group; + for (var ic_local: u32 = 0u; ic_local < params.ic_per_group; ic_local = ic_local + 1u) { + let ic = ic_base + ic_local; + for (var kh: u32 = 0u; kh < params.Kh; kh = kh + 1u) { + let ih = i32(oh) * i32(params.stride_h) - i32(params.pad_h) + + i32(kh) * i32(params.dil_h); + if (ih < 0 || ih >= i32(params.H_in)) { + continue; + } + let in_row = ((n * params.IC + ic) * params.H_in + u32(ih)) * params.W_in; + let w_row = + ((oc * params.ic_per_group + ic_local) * params.Kh + kh) * params.Kw; + for (var kw: u32 = 0u; kw < params.Kw; kw = kw + 1u) { + let iw = i32(ow) * i32(params.stride_w) - i32(params.pad_w) + + i32(kw) * i32(params.dil_w); + if (iw < 0 || iw >= i32(params.W_in)) { + continue; + } + acc = acc + t_x[in_row + u32(iw)] * t_weight[w_row + kw]; + } + } + } + t_out[idx] = clamp(acc, f16(params.output_min), f16(params.output_max)); +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-gemm-f16.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-gemm-f16.wgsl.ts new file mode 100644 index 0000000..40d3024 --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-gemm-f16.wgsl.ts @@ -0,0 +1,194 @@ +/** + * WGSL compute shader for float16 2D convolution with clamp bounds as implicit-im2col tiled GEMM. + * Requires WebGPU shader-f16 feature. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_GEMM_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +const TILE: u32 = 32u; +const RPT: u32 = 4u; + +var a_sub: array, 32>; +var b_sub: array, 8>, 32>; + +fn gK() -> u32 { return params.IC * params.Kh * params.Kw; } +fn gN() -> u32 { return params.N * params.H_out * params.W_out; } + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3) { + let M = params.OC; + let N = gN(); + let K = gK(); + let tile_row0 = wg_id.y * TILE; + let tile_col0 = wg_id.x * TILE; + let tile_row = local_id.y * RPT; + let tile_col = local_id.x * RPT; + + var acc: array, 4>; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + acc[ir] = vec4(0.0h); + } + + let ohw = params.H_out * params.W_out; + let khw = params.Kh * params.Kw; + let iH = i32(params.H_in); + let iW = i32(params.W_in); + let sH = i32(params.stride_h); + let sW = i32(params.stride_w); + let pH = i32(params.pad_h); + let pW = i32(params.pad_w); + let dH = i32(params.dil_h); + let dW = i32(params.dil_w); + + var col_b: array; + var col_valid: array; + var col_oh_arr: array; + var col_ow_arr: array; + var col_ih_base: array; + var col_iw_base: array; + for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) { + let n = tile_col0 + local_id.x * RPT + kk; + if (n < N) { + col_valid[kk] = true; + let b_val = n / ohw; + let sp = n % ohw; + col_b[kk] = b_val; + let oh_val = i32(sp / params.W_out); + let ow_val = i32(sp % params.W_out); + col_oh_arr[kk] = oh_val; + col_ow_arr[kk] = ow_val; + col_ih_base[kk] = oh_val * sH - pH; + col_iw_base[kk] = ow_val * sW - pW; + } else { + col_valid[kk] = false; + col_b[kk] = 0u; + col_oh_arr[kk] = 0; + col_ow_arr[kk] = 0; + col_ih_base[kk] = 0; + col_iw_base[kk] = 0; + } + } + + let in_hw = params.H_in * params.W_in; + let in_w = params.W_in; + let num_tiles = (K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k_start = t * TILE; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + let m = tile_row0 + arow; + let kk_idx = k_start + arow; + + var k_valid = false; + var k_ic = 0u; + var kh_dH = 0; + var kw_dW = 0; + if (kk_idx < K) { + k_valid = true; + k_ic = kk_idx / khw; + let rem = kk_idx % khw; + kh_dH = i32(rem / params.Kw) * dH; + kw_dW = i32(rem % params.Kw) * dW; + } + + let m_valid = m < M; + let w_row_base = m * K + k_start + local_id.x * RPT; + + var b_row: vec4; + for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) { + let col = local_id.x * RPT + kk; + let k_col = k_start + col; + if (m_valid && k_col < K) { + a_sub[arow][col] = t_weight[w_row_base + kk]; + } else { + a_sub[arow][col] = 0.0h; + } + + if (k_valid && col_valid[kk]) { + let ih = col_ih_base[kk] + kh_dH; + let iw = col_iw_base[kk] + kw_dW; + if (ih >= 0 && ih < iH && iw >= 0 && iw < iW) { + let ch_idx = col_b[kk] * params.IC + k_ic; + b_row[kk] = t_x[ch_idx * in_hw + u32(ih) * in_w + u32(iw)]; + } else { + b_row[kk] = 0.0h; + } + } else { + b_row[kk] = 0.0h; + } + } + b_sub[arow][local_id.x] = b_row; + } + workgroupBarrier(); + for (var k: u32 = 0u; k < TILE; k = k + 1u) { + let b_vec = b_sub[k][local_id.x]; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let aval = a_sub[tile_row + ir][k]; + acc[ir] = acc[ir] + aval * b_vec; + } + } + workgroupBarrier(); + } + + let clamp_min = vec4(f16(params.output_min)); + let clamp_max = vec4(f16(params.output_max)); + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let m = tile_row0 + tile_row + ir; + if (m < M) { + var bias_val = 0.0h; + if (params.has_bias != 0u) { + bias_val = t_bias[m]; + } + let v = clamp(acc[ir] + vec4(bias_val), clamp_min, clamp_max); + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + let n = tile_col0 + tile_col + ic; + if (n < N) { + let b = col_b[ic]; + let oh = u32(col_oh_arr[ic]); + let ow = u32(col_ow_arr[ic]); + t_out[((b * params.OC + m) * params.H_out + oh) * params.W_out + ow] = v[ic]; + } + } + } + } +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-gemm.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-gemm.wgsl.ts new file mode 100644 index 0000000..074d339 --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-gemm.wgsl.ts @@ -0,0 +1,198 @@ +/** + * WGSL compute shader for fused convolution with clamp bounds as + * implicit-im2col tiled GEMM (groups == 1). + * + * Source: ExecuTorch (backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm.wgsl) + * GitHub: https://github.com/pytorch/executorch/blob/main/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm.wgsl + * License: BSD-3-Clause (Copyright (c) Meta Platforms, Inc. and affiliates) + */ +export const SHADER_CONV_WITH_CLAMP_GEMM = /* wgsl */ ` +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +// Tiled GEMM over implicit im2col (groups == 1): M=OC, N=N*H_out*W_out, +// K=IC*Kh*Kw. A=weight [OC,K] (contiguous OIHW); B=input sampled on the fly, +// out-of-range -> 0.0 implements padding. bias per-ROW (OC). Output written +// NCHW with the clamp bounds applied. +const TILE: u32 = 32u; +const RPT: u32 = 4u; + +var a_sub: array, 32>; +var b_sub: array, 8>, 32>; + +fn gK() -> u32 { return params.IC * params.Kh * params.Kw; } +fn gN() -> u32 { return params.N * params.H_out * params.W_out; } + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3) { + let M = params.OC; + let N = gN(); + let K = gK(); + let tile_row0 = wg_id.y * TILE; + let tile_col0 = wg_id.x * TILE; + let tile_row = local_id.y * RPT; + let tile_col = local_id.x * RPT; + + var acc: array, 4>; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + acc[ir] = vec4(0.0); + } + + let ohw = params.H_out * params.W_out; + let khw = params.Kh * params.Kw; + let iH = i32(params.H_in); + let iW = i32(params.W_in); + let sH = i32(params.stride_h); + let sW = i32(params.stride_w); + let pH = i32(params.pad_h); + let pW = i32(params.pad_w); + let dH = i32(params.dil_h); + let dW = i32(params.dil_w); + + // Precompute spatial coords for this thread's 4 columns in B (invariant across K tiles) + var col_b: array; + var col_b_ic: array; + var col_valid: array; + var col_oh_arr: array; + var col_ow_arr: array; + var col_ih_base: array; + var col_iw_base: array; + for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) { + let n = tile_col0 + local_id.x * RPT + kk; + if (n < N) { + col_valid[kk] = true; + let b_val = n / ohw; + let sp = n % ohw; + col_b[kk] = b_val; + col_b_ic[kk] = b_val * params.IC; + let oh_val = i32(sp / params.W_out); + let ow_val = i32(sp % params.W_out); + col_oh_arr[kk] = oh_val; + col_ow_arr[kk] = ow_val; + col_ih_base[kk] = oh_val * sH - pH; + col_iw_base[kk] = ow_val * sW - pW; + } else { + col_valid[kk] = false; + col_b[kk] = 0u; + col_b_ic[kk] = 0u; + col_oh_arr[kk] = 0; + col_ow_arr[kk] = 0; + col_ih_base[kk] = 0; + col_iw_base[kk] = 0; + } + } + + let in_hw = params.H_in * params.W_in; + let in_w = params.W_in; + let num_tiles = (K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k_start = t * TILE; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + let m = tile_row0 + arow; + let kk_idx = k_start + arow; + + var k_valid = false; + var k_ic = 0u; + var kh_dH = 0; + var kw_dW = 0; + if (kk_idx < K) { + k_valid = true; + k_ic = kk_idx / khw; + let rem = kk_idx % khw; + kh_dH = i32(rem / params.Kw) * dH; + kw_dW = i32(rem % params.Kw) * dW; + } + + let m_valid = m < M; + let w_row_base = m * K + k_start + local_id.x * RPT; + + var b_row: vec4; + for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) { + let col = local_id.x * RPT + kk; + let k_col = k_start + col; + if (m_valid && k_col < K) { + a_sub[arow][col] = t_weight[w_row_base + kk]; + } else { + a_sub[arow][col] = 0.0; + } + + if (k_valid && col_valid[kk]) { + let ih = col_ih_base[kk] + kh_dH; + let iw = col_iw_base[kk] + kw_dW; + if (ih >= 0 && ih < iH && iw >= 0 && iw < iW) { + let ch_idx = col_b_ic[kk] + k_ic; + b_row[kk] = t_x[ch_idx * in_hw + u32(ih) * in_w + u32(iw)]; + } else { + b_row[kk] = 0.0; + } + } else { + b_row[kk] = 0.0; + } + } + b_sub[arow][local_id.x] = b_row; + } + workgroupBarrier(); + for (var k: u32 = 0u; k < TILE; k = k + 1u) { + let b_vec = b_sub[k][local_id.x]; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let aval = a_sub[tile_row + ir][k]; + acc[ir] = acc[ir] + aval * b_vec; + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let m = tile_row0 + tile_row + ir; + if (m < M) { + var bias_val = 0.0; + if (params.has_bias != 0u) { + bias_val = t_bias[m]; + } + let v = clamp(acc[ir] + vec4(bias_val), vec4(params.output_min), vec4(params.output_max)); + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + let n = tile_col0 + tile_col + ic; + if (n < N) { + let b = col_b[ic]; + let oh = u32(col_oh_arr[ic]); + let ow = u32(col_ow_arr[ic]); + t_out[((b * params.OC + m) * params.H_out + oh) * params.W_out + ow] = v[ic]; + } + } + } + } +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-tile32-f16.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-tile32-f16.wgsl.ts new file mode 100644 index 0000000..40e184c --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-tile32-f16.wgsl.ts @@ -0,0 +1,224 @@ +/** + * WGSL compute shader for 32x64 register-tiled 2D convolution with clamped activation in FP16 (4x4 per thread across 128 threads). + * Optimized for small channel counts (OC <= 32) to maximize wave occupancy and eliminate thread divergence. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_TILE32_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_in: array; +@group(0) @binding(2) var t_weight: array>; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +var a_sub: array, 4>, 32>; +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(local_invocation_index) tid: u32 +) { + let tile_row0 = wg_id.y * 32u; + let tile_col0 = wg_id.x * 64u; + let M = params.OC; + let N = params.N * params.H_out * params.W_out; + let K = params.IC * params.Kh * params.Kw; + let K_vecs = K >> 2u; + let khw = params.Kh * params.Kw; + let ohw = params.H_out * params.W_out; + + let in_hw = params.H_in * params.W_in; + let in_w = params.W_in; + let sH = i32(params.stride_h); + let sW = i32(params.stride_w); + let pH = i32(params.pad_h); + let pW = i32(params.pad_w); + let dH = i32(params.dil_h); + let dW = i32(params.dil_w); + + var acc0 = vec4(0.0h); + var acc1 = vec4(0.0h); + var acc2 = vec4(0.0h); + var acc3 = vec4(0.0h); + + let a_r = tid >> 2u; // 0..31 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < M; + + let b_r0 = tid >> 4u; // 0..7 + let b_c4_0 = tid & 15u; // 0..15 + let b_r1 = b_r0 + 8u; // 8..15 + let b_c4_1 = b_c4_0; + + let n_base0 = tile_col0 + (b_c4_0 << 2u); + + var col_b0: array; + var col_ih_base0: array; + var col_iw_base0: array; + var col_valid0: array; + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + let n = n_base0 + kk; + if (n < N) { + col_valid0[kk] = true; + let b_val = n / ohw; + let sp = n % ohw; + col_b0[kk] = b_val; + let oh = sp / params.W_out; + let ow = sp % params.W_out; + col_ih_base0[kk] = i32(oh) * sH - pH; + col_iw_base0[kk] = i32(ow) * sW - pW; + } else { + col_valid0[kk] = false; + col_b0[kk] = 0u; + col_ih_base0[kk] = 0; + col_iw_base0[kk] = 0; + } + } + + let num_tiles = (K + 15u) >> 4u; + for (var t = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t << 2u; + let k_start = t << 4u; + + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = t_weight[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0h); + } + + let bk0 = k_start + b_r0; + if (bk0 < K) { + let ic0 = bk0 / khw; + let rem0 = bk0 % khw; + let kh0 = i32(rem0 / params.Kw) * dH; + let kw0 = i32(rem0 % params.Kw) * dW; + + var b_row0 = vec4(0.0h); + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + if (col_valid0[kk]) { + let ih0 = col_ih_base0[kk] + kh0; + let iw0 = col_iw_base0[kk] + kw0; + if (ih0 >= 0 && ih0 < i32(params.H_in) && iw0 >= 0 && iw0 < i32(params.W_in)) { + let in_idx0 = (col_b0[kk] * params.IC + ic0) * in_hw + u32(ih0) * in_w + u32(iw0); + b_row0[kk] = t_in[in_idx0]; + } + } + } + b_sub[b_r0][b_c4_0] = b_row0; + } else { + b_sub[b_r0][b_c4_0] = vec4(0.0h); + } + + let bk1 = k_start + b_r1; + if (bk1 < K) { + let ic1 = bk1 / khw; + let rem1 = bk1 % khw; + let kh1 = i32(rem1 / params.Kw) * dH; + let kw1 = i32(rem1 % params.Kw) * dW; + + var b_row1 = vec4(0.0h); + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + if (col_valid0[kk]) { + let ih1 = col_ih_base0[kk] + kh1; + let iw1 = col_iw_base0[kk] + kw1; + if (ih1 >= 0 && ih1 < i32(params.H_in) && iw1 >= 0 && iw1 < i32(params.W_in)) { + let in_idx1 = (col_b0[kk] * params.IC + ic1) * in_hw + u32(ih1) * in_w + u32(iw1); + b_row1[kk] = t_in[in_idx1]; + } + } + } + b_sub[b_r1][b_c4_1] = b_row1; + } else { + b_sub[b_r1][b_c4_1] = vec4(0.0h); + } + + workgroupBarrier(); + + let thread_r = local_id.y << 2u; + let thread_c4 = local_id.x; + + for (var k4 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + let thread_r = local_id.y << 2u; + let out_c_base = tile_col0 + (local_id.x << 2u); + let min_v = f16(params.output_min); + let max_v = f16(params.output_max); + + for (var ir = 0u; ir < 4u; ir = ir + 1u) { + let m = tile_row0 + thread_r + ir; + if (m < M) { + var bias_val = 0.0h; + if (params.has_bias != 0u) { + bias_val = t_bias[m]; + } + var acc_v = acc0; + if (ir == 1u) { acc_v = acc1; } + else if (ir == 2u) { acc_v = acc2; } + else if (ir == 3u) { acc_v = acc3; } + acc_v = acc_v + vec4(bias_val); + + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + let n = out_c_base + kk; + if (n < N) { + let b_val = n / ohw; + let sp = n % ohw; + let oh = sp / params.W_out; + let ow = sp % params.W_out; + let val = clamp(acc_v[kk], min_v, max_v); + t_out[((b_val * params.OC + m) * params.H_out + oh) * params.W_out + ow] = val; + } + } + } + } +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-tile32.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-tile32.wgsl.ts new file mode 100644 index 0000000..5283f4e --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-tile32.wgsl.ts @@ -0,0 +1,226 @@ +/** + * WGSL compute shader for 32x64 register-tiled 2D convolution with clamped activation (4x4 per thread across 128 threads). + * Optimized for small channel counts (OC <= 32) to maximize wave occupancy and eliminate thread divergence. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_TILE32 = /* wgsl */ ` +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_in: array; +@group(0) @binding(2) var t_weight: array>; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +// a_sub: 32 output channels x 16 K elements (4 vec4s) = 128 vec4s (2 KB) +var a_sub: array, 4>, 32>; +// b_sub: 16 K elements x 64 spatial elements (16 vec4s) = 256 vec4s (4 KB) +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(local_invocation_index) tid: u32 +) { + let tile_row0 = wg_id.y * 32u; + let tile_col0 = wg_id.x * 64u; + let M = params.OC; + let N = params.N * params.H_out * params.W_out; + let K = params.IC * params.Kh * params.Kw; + let K_vecs = K >> 2u; + let khw = params.Kh * params.Kw; + let ohw = params.H_out * params.W_out; + + let in_hw = params.H_in * params.W_in; + let in_w = params.W_in; + let sH = i32(params.stride_h); + let sW = i32(params.stride_w); + let pH = i32(params.pad_h); + let pW = i32(params.pad_w); + let dH = i32(params.dil_h); + let dW = i32(params.dil_w); + + var acc0 = vec4(0.0); + var acc1 = vec4(0.0); + var acc2 = vec4(0.0); + var acc3 = vec4(0.0); + + let a_r = tid >> 2u; // 0..31 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < M; + + let b_r0 = tid >> 4u; // 0..7 + let b_c4_0 = tid & 15u; // 0..15 + let b_r1 = b_r0 + 8u; // 8..15 + let b_c4_1 = b_c4_0; + + let n_base0 = tile_col0 + (b_c4_0 << 2u); + + var col_b0: array; + var col_ih_base0: array; + var col_iw_base0: array; + var col_valid0: array; + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + let n = n_base0 + kk; + if (n < N) { + col_valid0[kk] = true; + let b_val = n / ohw; + let sp = n % ohw; + col_b0[kk] = b_val; + let oh = sp / params.W_out; + let ow = sp % params.W_out; + col_ih_base0[kk] = i32(oh) * sH - pH; + col_iw_base0[kk] = i32(ow) * sW - pW; + } else { + col_valid0[kk] = false; + col_b0[kk] = 0u; + col_ih_base0[kk] = 0; + col_iw_base0[kk] = 0; + } + } + + let num_tiles = (K + 15u) >> 4u; + for (var t = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t << 2u; + let k_start = t << 4u; + + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = t_weight[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0); + } + + // Load first B slice (rows 0..7) + let bk0 = k_start + b_r0; + if (bk0 < K) { + let ic0 = bk0 / khw; + let rem0 = bk0 % khw; + let kh0 = i32(rem0 / params.Kw) * dH; + let kw0 = i32(rem0 % params.Kw) * dW; + + var b_row0 = vec4(0.0); + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + if (col_valid0[kk]) { + let ih0 = col_ih_base0[kk] + kh0; + let iw0 = col_iw_base0[kk] + kw0; + if (ih0 >= 0 && ih0 < i32(params.H_in) && iw0 >= 0 && iw0 < i32(params.W_in)) { + let in_idx0 = (col_b0[kk] * params.IC + ic0) * in_hw + u32(ih0) * in_w + u32(iw0); + b_row0[kk] = t_in[in_idx0]; + } + } + } + b_sub[b_r0][b_c4_0] = b_row0; + } else { + b_sub[b_r0][b_c4_0] = vec4(0.0); + } + + // Load second B slice (rows 8..15) + let bk1 = k_start + b_r1; + if (bk1 < K) { + let ic1 = bk1 / khw; + let rem1 = bk1 % khw; + let kh1 = i32(rem1 / params.Kw) * dH; + let kw1 = i32(rem1 % params.Kw) * dW; + + var b_row1 = vec4(0.0); + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + if (col_valid0[kk]) { + let ih1 = col_ih_base0[kk] + kh1; + let iw1 = col_iw_base0[kk] + kw1; + if (ih1 >= 0 && ih1 < i32(params.H_in) && iw1 >= 0 && iw1 < i32(params.W_in)) { + let in_idx1 = (col_b0[kk] * params.IC + ic1) * in_hw + u32(ih1) * in_w + u32(iw1); + b_row1[kk] = t_in[in_idx1]; + } + } + } + b_sub[b_r1][b_c4_1] = b_row1; + } else { + b_sub[b_r1][b_c4_1] = vec4(0.0); + } + + workgroupBarrier(); + + let thread_r = local_id.y << 2u; + let thread_c4 = local_id.x; + + for (var k4 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + let thread_r = local_id.y << 2u; + let out_c_base = tile_col0 + (local_id.x << 2u); + let min_v = params.output_min; + let max_v = params.output_max; + + for (var ir = 0u; ir < 4u; ir = ir + 1u) { + let m = tile_row0 + thread_r + ir; + if (m < M) { + var bias_val = 0.0; + if (params.has_bias != 0u) { + bias_val = t_bias[m]; + } + var acc_v = acc0; + if (ir == 1u) { acc_v = acc1; } + else if (ir == 2u) { acc_v = acc2; } + else if (ir == 3u) { acc_v = acc3; } + acc_v = acc_v + vec4(bias_val); + + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + let n = out_c_base + kk; + if (n < N) { + let b_val = n / ohw; + let sp = n % ohw; + let oh = sp / params.W_out; + let ow = sp % params.W_out; + let val = clamp(acc_v[kk], min_v, max_v); + t_out[((b_val * params.OC + m) * params.H_out + oh) * params.W_out + ow] = val; + } + } + } + } +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-tile64-f16.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-tile64-f16.wgsl.ts new file mode 100644 index 0000000..a95957d --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-tile64-f16.wgsl.ts @@ -0,0 +1,205 @@ +/** + * WGSL compute shader for 64x64 register-tiled 2D convolution with clamp bounds in float16 (4x4 per thread across 256 threads). + * High occupancy, coalesced 128-bit memory transactions, minimal register pressure for all GPUs. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_TILE64_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array>; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +// a_sub: 64 output channels x 16 K elements (4 vec4s) = 256 vec4s (2 KB) +var a_sub: array, 4>, 64>; +// b_sub: 16 K elements x 64 spatial elements (16 vec4s) = 256 vec4s (2 KB) +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(local_invocation_index) tid: u32 +) { + let tile_row0 = wg_id.y * 64u; + let tile_col0 = wg_id.x * 64u; + let M = params.OC; + let N = params.N * params.H_out * params.W_out; + let K = params.IC * params.Kh * params.Kw; + let K_vecs = K >> 2u; + let khw = params.Kh * params.Kw; + let ohw = params.H_out * params.W_out; + + let iH = i32(params.H_in); + let iW = i32(params.W_in); + let in_hw = params.H_in * params.W_in; + let in_w = params.W_in; + let sH = i32(params.stride_h); + let sW = i32(params.stride_w); + let pH = i32(params.pad_h); + let pW = i32(params.pad_w); + let dH = i32(params.dil_h); + let dW = i32(params.dil_w); + + var acc0 = vec4(0.0h); + var acc1 = vec4(0.0h); + var acc2 = vec4(0.0h); + var acc3 = vec4(0.0h); + + // Cooperative load coordinates for 256 threads (1 vec4 weight + 1 vec4 input per thread) + let a_r = tid >> 2u; // 0..63 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < M; + + let b_r = tid >> 4u; // 0..15 (row in K tile) + let b_c4 = tid & 15u; // 0..15 (col in spatial tile) + let n_base = tile_col0 + (b_c4 << 2u); + + // Precompute invariant spatial coordinates for this thread's 4 input columns once + var col_b: array; + var col_ih_base: array; + var col_iw_base: array; + var col_valid: array; + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + let n = n_base + kk; + if (n < N) { + col_valid[kk] = true; + let b_val = n / ohw; + let sp = n % ohw; + col_b[kk] = b_val; + let oh = sp / params.W_out; + let ow = sp % params.W_out; + col_ih_base[kk] = i32(oh) * sH - pH; + col_iw_base[kk] = i32(ow) * sW - pW; + } else { + col_valid[kk] = false; + col_b[kk] = 0u; + col_ih_base[kk] = 0; + col_iw_base[kk] = 0; + } + } + + let num_tiles = (K + 15u) >> 4u; + for (var t = 0u; t < num_tiles; t = t + 1u) { + let k_start = t << 4u; + let k4_start = t << 2u; + + // 1. Coalesced 128-bit load of weights (256 threads load 256 vec4s) + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = t_weight[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0h); + } + + // 2. Coalesced load of input using precomputed spatial bases + let kk_idx = k_start + b_r; + var b_row = vec4(0.0h); + if (kk_idx < K) { + let k_ic = kk_idx / khw; + let rem = kk_idx % khw; + let kh = i32(rem / params.Kw) * dH; + let kw = i32(rem % params.Kw) * dW; + + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + if (col_valid[kk]) { + let ih = col_ih_base[kk] + kh; + let iw = col_iw_base[kk] + kw; + if (ih >= 0 && ih < iH && iw >= 0 && iw < iW) { + let ch_idx = col_b[kk] * params.IC + k_ic; + b_row[kk] = t_x[ch_idx * in_hw + u32(ih) * in_w + u32(iw)]; + } + } + } + } + b_sub[b_r][b_c4] = b_row; + workgroupBarrier(); + + // 3. Compute phase: 4x4 register accumulation per thread + let thread_r = local_id.y << 2u; + let thread_c4 = local_id.x; + + for (var k4 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + // 4. Output write: 4 rows x 4 cols per thread with clamp + let thread_r = local_id.y << 2u; + let out_c_base = tile_col0 + (local_id.x << 2u); + let min_v = f16(params.output_min); + let max_v = f16(params.output_max); + + for (var ir = 0u; ir < 4u; ir = ir + 1u) { + let m = tile_row0 + thread_r + ir; + if (m < M) { + var bias_val = 0.0h; + if (params.has_bias != 0u) { + bias_val = t_bias[m]; + } + var acc_v = acc0; + if (ir == 1u) { acc_v = acc1; } + else if (ir == 2u) { acc_v = acc2; } + else if (ir == 3u) { acc_v = acc3; } + acc_v = acc_v + vec4(bias_val); + + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + let n = out_c_base + kk; + if (n < N) { + let b_val = n / ohw; + let sp = n % ohw; + let oh = sp / params.W_out; + let ow = sp % params.W_out; + let val = clamp(acc_v[kk], min_v, max_v); + t_out[((b_val * params.OC + m) * params.H_out + oh) * params.W_out + ow] = val; + } + } + } + } +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp-tile64.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp-tile64.wgsl.ts new file mode 100644 index 0000000..c16abea --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp-tile64.wgsl.ts @@ -0,0 +1,203 @@ +/** + * WGSL compute shader for 64x64 register-tiled 2D convolution with clamp bounds (4x4 per thread across 256 threads). + * High occupancy, coalesced 128-bit memory transactions, minimal register pressure for all GPUs. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV_WITH_CLAMP_TILE64 = /* wgsl */ ` +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array>; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +// a_sub: 64 output channels x 16 K elements (4 vec4s) = 256 vec4s (4 KB) +var a_sub: array, 4>, 64>; +// b_sub: 16 K elements x 64 spatial elements (16 vec4s) = 256 vec4s (4 KB) +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(local_invocation_index) tid: u32 +) { + let tile_row0 = wg_id.y * 64u; + let tile_col0 = wg_id.x * 64u; + let M = params.OC; + let N = params.N * params.H_out * params.W_out; + let K = params.IC * params.Kh * params.Kw; + let K_vecs = K >> 2u; + let khw = params.Kh * params.Kw; + let ohw = params.H_out * params.W_out; + + let iH = i32(params.H_in); + let iW = i32(params.W_in); + let in_hw = params.H_in * params.W_in; + let in_w = params.W_in; + let sH = i32(params.stride_h); + let sW = i32(params.stride_w); + let pH = i32(params.pad_h); + let pW = i32(params.pad_w); + let dH = i32(params.dil_h); + let dW = i32(params.dil_w); + + var acc0 = vec4(0.0); + var acc1 = vec4(0.0); + var acc2 = vec4(0.0); + var acc3 = vec4(0.0); + + // Cooperative load coordinates for 256 threads (1 vec4 weight + 1 vec4 input per thread) + let a_r = tid >> 2u; // 0..63 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < M; + + let b_r = tid >> 4u; // 0..15 (row in K tile) + let b_c4 = tid & 15u; // 0..15 (col in spatial tile) + let n_base = tile_col0 + (b_c4 << 2u); + + // Precompute invariant spatial coordinates for this thread's 4 input columns once + var col_b: array; + var col_ih_base: array; + var col_iw_base: array; + var col_valid: array; + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + let n = n_base + kk; + if (n < N) { + col_valid[kk] = true; + let b_val = n / ohw; + let sp = n % ohw; + col_b[kk] = b_val; + let oh = sp / params.W_out; + let ow = sp % params.W_out; + col_ih_base[kk] = i32(oh) * sH - pH; + col_iw_base[kk] = i32(ow) * sW - pW; + } else { + col_valid[kk] = false; + col_b[kk] = 0u; + col_ih_base[kk] = 0; + col_iw_base[kk] = 0; + } + } + + let num_tiles = (K + 15u) >> 4u; + for (var t = 0u; t < num_tiles; t = t + 1u) { + let k_start = t << 4u; + let k4_start = t << 2u; + + // 1. Coalesced 128-bit load of weights (256 threads load 256 vec4s) + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = t_weight[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0); + } + + // 2. Coalesced load of input using precomputed spatial bases + let kk_idx = k_start + b_r; + var b_row = vec4(0.0); + if (kk_idx < K) { + let k_ic = kk_idx / khw; + let rem = kk_idx % khw; + let kh = i32(rem / params.Kw) * dH; + let kw = i32(rem % params.Kw) * dW; + + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + if (col_valid[kk]) { + let ih = col_ih_base[kk] + kh; + let iw = col_iw_base[kk] + kw; + if (ih >= 0 && ih < iH && iw >= 0 && iw < iW) { + let ch_idx = col_b[kk] * params.IC + k_ic; + b_row[kk] = t_x[ch_idx * in_hw + u32(ih) * in_w + u32(iw)]; + } + } + } + } + b_sub[b_r][b_c4] = b_row; + workgroupBarrier(); + + // 3. Compute phase: 4x4 register accumulation per thread + let thread_r = local_id.y << 2u; + let thread_c4 = local_id.x; + + for (var k4 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + // 4. Output write: 4 rows x 4 cols per thread with clamp + let thread_r = local_id.y << 2u; + let out_c_base = tile_col0 + (local_id.x << 2u); + let min_v = params.output_min; + let max_v = params.output_max; + + for (var ir = 0u; ir < 4u; ir = ir + 1u) { + let m = tile_row0 + thread_r + ir; + if (m < M) { + var bias_val = 0.0; + if (params.has_bias != 0u) { + bias_val = t_bias[m]; + } + var acc_v = acc0; + if (ir == 1u) { acc_v = acc1; } + else if (ir == 2u) { acc_v = acc2; } + else if (ir == 3u) { acc_v = acc3; } + acc_v = acc_v + vec4(bias_val); + + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + let n = out_c_base + kk; + if (n < N) { + let b_val = n / ohw; + let sp = n % ohw; + let oh = sp / params.W_out; + let ow = sp % params.W_out; + let val = clamp(acc_v[kk], min_v, max_v); + t_out[((b_val * params.OC + m) * params.H_out + oh) * params.W_out + ow] = val; + } + } + } + } +} +`; diff --git a/src/shaders/conv-with-clamp/conv-with-clamp.wgsl.ts b/src/shaders/conv-with-clamp/conv-with-clamp.wgsl.ts new file mode 100644 index 0000000..f5846ab --- /dev/null +++ b/src/shaders/conv-with-clamp/conv-with-clamp.wgsl.ts @@ -0,0 +1,94 @@ +/** + * Direct 2D convolution compute shader with clamp (fallback for general grouped cases). + * + * Source: ExecuTorch (backends/webgpu/runtime/ops/et_vk_conv2d/conv2d.wgsl) + * GitHub: https://github.com/pytorch/executorch/blob/main/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d.wgsl + * License: BSD-3-Clause (Copyright (c) Meta Platforms, Inc. and affiliates) + */ +export const SHADER_CONV_WITH_CLAMP = /* wgsl */ ` +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_x: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_bias: array; + +struct Params { + N: u32, + IC: u32, + H_in: u32, + W_in: u32, + OC: u32, + H_out: u32, + W_out: u32, + Kh: u32, + Kw: u32, + stride_h: u32, + stride_w: u32, + pad_h: u32, + pad_w: u32, + dil_h: u32, + dil_w: u32, + has_bias: u32, + numel: u32, + groups: u32, + ic_per_group: u32, + pad0: u32, + pad1: u32, + pad2: u32, + output_min: f32, + output_max: f32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.numel) { + return; + } + // Unravel the NCHW output index -> (n, oc, oh, ow). + let ow = idx % params.W_out; + var r = idx / params.W_out; + let oh = r % params.H_out; + r = r / params.H_out; + let oc = r % params.OC; + let n = r / params.OC; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = t_bias[oc]; + } + // Grouped conv: output channel oc belongs to group g and dots only that + // group's ic_per_group input channels. weight is [OC, IC/groups, Kh, Kw]; + // groups==1 (ic_per_group==IC, g==0) is the general dense case. + let oc_per_group = params.OC / params.groups; + let g = oc / oc_per_group; + let ic_base = g * params.ic_per_group; + for (var ic_local: u32 = 0u; ic_local < params.ic_per_group; ic_local = ic_local + 1u) { + let ic = ic_base + ic_local; + for (var kh: u32 = 0u; kh < params.Kh; kh = kh + 1u) { + let ih = i32(oh) * i32(params.stride_h) - i32(params.pad_h) + + i32(kh) * i32(params.dil_h); + if (ih < 0 || ih >= i32(params.H_in)) { + continue; + } + let in_row = ((n * params.IC + ic) * params.H_in + u32(ih)) * params.W_in; + let w_row = + ((oc * params.ic_per_group + ic_local) * params.Kh + kh) * params.Kw; + for (var kw: u32 = 0u; kw < params.Kw; kw = kw + 1u) { + let iw = i32(ow) * i32(params.stride_w) - i32(params.pad_w) + + i32(kw) * i32(params.dil_w); + if (iw < 0 || iw >= i32(params.W_in)) { + continue; + } + acc = acc + t_x[in_row + u32(iw)] * t_weight[w_row + kw]; + } + } + } + t_out[idx] = clamp(acc, params.output_min, params.output_max); +} +`; diff --git a/src/shaders/conv-with-clamp/index.ts b/src/shaders/conv-with-clamp/index.ts new file mode 100644 index 0000000..c449ec2 --- /dev/null +++ b/src/shaders/conv-with-clamp/index.ts @@ -0,0 +1,344 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { SHADER_CONV_WITH_CLAMP } from './conv-with-clamp.wgsl'; +import { SHADER_CONV_WITH_CLAMP_GEMM } from './conv-with-clamp-gemm.wgsl'; +import { SHADER_CONV_WITH_CLAMP_1X1 } from './conv-with-clamp-1x1.wgsl'; +import { SHADER_CONV_WITH_CLAMP_1X1_SINGLE_CHANNEL } from './conv-with-clamp-1x1-single-channel.wgsl'; +import { SHADER_CONV_WITH_CLAMP_1X1_SINGLE_CHANNEL_F16 } from './conv-with-clamp-1x1-single-channel-f16.wgsl'; +import { SHADER_CONV_WITH_CLAMP_1X1_VEC4 } from './conv-with-clamp-1x1-vec4.wgsl'; +import { SHADER_CONV_WITH_CLAMP_1X1_64X64 } from './conv-with-clamp-1x1-tile64.wgsl'; +import { SHADER_CONV_WITH_CLAMP_1X1_64X64_F16 } from './conv-with-clamp-1x1-tile64-f16.wgsl'; +import { SHADER_CONV_WITH_CLAMP_TILE64 } from './conv-with-clamp-tile64.wgsl'; +import { SHADER_CONV_WITH_CLAMP_TILE32 } from './conv-with-clamp-tile32.wgsl'; +import { SHADER_CONV_WITH_CLAMP_1X1_TILE32 } from './conv-with-clamp-1x1-tile32.wgsl'; +import { SHADER_CONV_WITH_CLAMP_1X1_TILE32_F16 } from './conv-with-clamp-1x1-tile32-f16.wgsl'; +import { SHADER_CONV_WITH_CLAMP_DEPTHWISE } from './conv-with-clamp-depthwise.wgsl'; +import { SHADER_CONV_WITH_CLAMP_DEPTHWISE_3X3 } from './conv-with-clamp-depthwise-3x3.wgsl'; +import { SHADER_CONV_WITH_CLAMP_F16 } from './conv-with-clamp-f16.wgsl'; +import { SHADER_CONV_WITH_CLAMP_GEMM_F16 } from './conv-with-clamp-gemm-f16.wgsl'; +import { SHADER_CONV_WITH_CLAMP_TILE64_F16 } from './conv-with-clamp-tile64-f16.wgsl'; +import { SHADER_CONV_WITH_CLAMP_TILE32_F16 } from './conv-with-clamp-tile32-f16.wgsl'; +import { SHADER_CONV_WITH_CLAMP_DEPTHWISE_F16 } from './conv-with-clamp-depthwise-f16.wgsl'; +import { SHADER_CONV_WITH_CLAMP_DEPTHWISE_3X3_F16 } from './conv-with-clamp-depthwise-3x3-f16.wgsl'; + +const name = 'et_vk.conv_with_clamp.default'; + +/** + * Positional argument tuple for `et_vk::conv_with_clamp.default`: + * `[in, weight, bias, stride, padding, dilation, transposed, output_padding, groups, min, max, out]` + */ +export type ConvWithClampArgs = readonly [ + inTensor: Tensor, + weight: Tensor, + bias: Tensor | undefined, + stride: readonly number[], + padding: readonly number[], + dilation: readonly number[], + transposed: boolean, + outputPadding: readonly number[], + groups: number, + outputMin: number | undefined, + outputMax: number | undefined, + outTensor: Tensor, +]; + +function parseHw(v: readonly number[], argName: string): [number, number] { + if (v.length === 1) return [v[0], v[0]]; + if (v.length === 2) return [v[0], v[1]]; + throw new Error(`${name}: ${argName} must be 1 or 2 elements`); +} + +/** + * Internal dispatch builder for `et_vk::conv_with_clamp.default`. + * + * @param ctx The execution context. + * @param args Positional arguments for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: ConvWithClampArgs): void { + const [ + inTensor, + weight, + bias, + stride, + padding, + dilation, + transposed, + outputPadding, + groups, + outputMin, + outputMax, + outTensor, + ] = args; + + if ( + (inTensor.dtype !== 'float32' && inTensor.dtype !== 'float16') || + weight.dtype !== inTensor.dtype || + outTensor.dtype !== inTensor.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + if (inTensor.shape.length !== 4 || weight.shape.length !== 4 || outTensor.shape.length !== 4) { + throw new Error(`${name}: expected 4D tensors for input, weight, and output`); + } + if (bias !== undefined && bias.dtype !== inTensor.dtype) { + throw new Error(`${name}: bias dtype must match input dtype`); + } + if (transposed) { + throw new Error(`${name}: transposed conv is unsupported in conv_with_clamp`); + } + if (groups < 1) { + throw new Error(`${name}: groups must be >= 1`); + } + + const [sH, sW] = parseHw(stride, 'stride'); + const [pH, pW] = parseHw(padding, 'padding'); + const [dH, dW] = parseHw(dilation, 'dilation'); + + if (sH <= 0 || sW <= 0) { + throw new Error(`${name}: stride must be positive`); + } + + const [n, ic, hIn, wIn] = inTensor.shape as [number, number, number, number]; + const oc = weight.shape[0]; + const kh = weight.shape[2]; + const kw = weight.shape[3]; + + if (outputPadding.some((v) => v !== 0)) { + throw new Error(`${name}: non-zero output_padding unsupported for non-transposed conv`); + } + + const icPerGroup = weight.shape[1]; + if (ic % groups !== 0 || oc % groups !== 0 || icPerGroup * groups !== ic) { + throw new Error(`${name}: bad shape (IC/groups mismatch)`); + } + + const hEff = hIn + 2 * pH - dH * (kh - 1) - 1; + const wEff = wIn + 2 * pW - dW * (kw - 1) - 1; + if (hEff < 0 || wEff < 0) { + throw new Error(`${name}: invalid geometry (kernel larger than input)`); + } + + const expectedHout = Math.floor(hEff / sH) + 1; + const expectedWout = Math.floor(wEff / sW) + 1; + if (outTensor.shape[2] !== expectedHout || outTensor.shape[3] !== expectedWout) { + throw new Error(`${name}: output dims inconsistent with conv2d formula`); + } + + const hasBias = bias !== undefined; + if (hasBias && (bias.buffer === null || bias.shape[0] !== oc)) { + throw new Error(`${name}: bias must be fp32 [OC]`); + } + + const minVal = outputMin !== undefined ? outputMin : -Infinity; + const maxVal = outputMax !== undefined ? outputMax : Infinity; + + const numel = outTensor.numel; + const device = ctx.device; + const isF16 = inTensor.dtype === 'float16'; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + const wgSize = 64; + + const buffer = new ArrayBuffer(96); + const u32View = new Uint32Array(buffer); + const f32View = new Float32Array(buffer); + + // ConvWithClampParams (96 B): geometry, kernel, stride, pad, dilation, flags + // prettier-ignore + u32View.set([ + n, ic, hIn, wIn, + oc, expectedHout, expectedWout, + kh, kw, + sH, sW, + pH, pW, + dH, dW, + hasBias ? 1 : 0, numel, + groups, icPerGroup, + 0, 0, 0, + ]); + f32View[22] = minVal; // output_min + f32View[23] = maxVal; // output_max + + const uniformBuffer = ctx.uniformBuffer(buffer); + const biasBuffer = hasBias ? bias.buffer : weight.buffer; + + // Variant selection: 1x1 convolutions use the pointwise GEMM kernel (no im2col), + // with 128-bit vec4 memory optimization when dimensions are divisible by 4. + // General dense convolutions use implicit-im2col tiled GEMM, and + // grouped/depthwise convs use the direct kernel. + const is1x1 = + kh === 1 && kw === 1 && sH === 1 && sW === 1 && pH === 0 && pW === 0 && dH === 1 && dW === 1; + // A single output channel leaves almost all rows of a GEMM tile unused. + const isSingleChannel = + is1x1 && + groups === 1 && + oc === 1 && + ic <= 32 && + (expectedHout * expectedWout) % 4 === 0 && + Math.ceil((expectedHout * expectedWout) / 256) <= 65535 && + n <= 65535; + const is1x1Vec4 = + is1x1 && ic % 4 === 0 && (expectedHout * expectedWout) % 4 === 0 && oc % 4 === 0; + const is1x1Tile32 = + is1x1Vec4 && n === 1 && oc <= 32 && expectedHout * expectedWout >= 64 && ic >= 16; + const is1x1Tile64 = + is1x1Vec4 && n === 1 && oc > 32 && expectedHout * expectedWout >= 64 && ic >= 16; + + const isDenseTile32 = + groups === 1 && + oc <= 32 && + oc >= 8 && + n * expectedHout * expectedWout >= 64 && + (ic * kh * kw) % 4 === 0; + const isDenseTile64 = + groups === 1 && oc > 32 && n * expectedHout * expectedWout >= 64 && (ic * kh * kw) % 4 === 0; + + const isDepthwise = + groups === oc && ic === oc && icPerGroup === 1 && kh * kw <= 49 && n * oc <= 65535; + + const isDepthwise3x3 = + isDepthwise && + kh === 3 && + kw === 3 && + sH === 1 && + sW === 1 && + pH === 1 && + pW === 1 && + dH === 1 && + dW === 1; + + let code = isF16 ? SHADER_CONV_WITH_CLAMP_F16 : SHADER_CONV_WITH_CLAMP; + if (isSingleChannel) { + code = isF16 + ? SHADER_CONV_WITH_CLAMP_1X1_SINGLE_CHANNEL_F16 + : SHADER_CONV_WITH_CLAMP_1X1_SINGLE_CHANNEL; + } else if (isDepthwise) { + if (isDepthwise3x3) { + code = isF16 + ? SHADER_CONV_WITH_CLAMP_DEPTHWISE_3X3_F16 + : SHADER_CONV_WITH_CLAMP_DEPTHWISE_3X3; + } else { + code = isF16 ? SHADER_CONV_WITH_CLAMP_DEPTHWISE_F16 : SHADER_CONV_WITH_CLAMP_DEPTHWISE; + } + } else if (groups === 1) { + if (is1x1Tile32) { + code = isF16 ? SHADER_CONV_WITH_CLAMP_1X1_TILE32_F16 : SHADER_CONV_WITH_CLAMP_1X1_TILE32; + } else if (is1x1Tile64) { + code = isF16 ? SHADER_CONV_WITH_CLAMP_1X1_64X64_F16 : SHADER_CONV_WITH_CLAMP_1X1_64X64; + } else if (isDenseTile32) { + code = isF16 ? SHADER_CONV_WITH_CLAMP_TILE32_F16 : SHADER_CONV_WITH_CLAMP_TILE32; + } else if (isDenseTile64) { + code = isF16 ? SHADER_CONV_WITH_CLAMP_TILE64_F16 : SHADER_CONV_WITH_CLAMP_TILE64; + } else if (is1x1Vec4 && !isF16) { + code = SHADER_CONV_WITH_CLAMP_1X1_VEC4; + } else if (is1x1 && !isF16) { + code = SHADER_CONV_WITH_CLAMP_1X1; + } else { + code = isF16 ? SHADER_CONV_WITH_CLAMP_GEMM_F16 : SHADER_CONV_WITH_CLAMP_GEMM; + } + } + + let workgroupCountX: number; + let workgroupCountY: number; + let workgroupCountZ = 1; + let constants: Record; + + if (isSingleChannel) { + workgroupCountX = Math.ceil((expectedHout * expectedWout) / 256); + workgroupCountY = n; + constants = {}; + } else if (isDepthwise) { + // Spatial 16x16 workgroup per (batch, oc) channel slice with 2x2 register tiling (32x32 per workgroup) + workgroupCountX = Math.ceil(expectedWout / 32); + workgroupCountY = Math.ceil(expectedHout / 32); + workgroupCountZ = n * oc; + constants = {}; + } else if (groups !== 1) { + // Direct kernel: one thread per output element, 2D spill + const totalWorkgroups = Math.ceil(numel / wgSize); + workgroupCountX = Math.min(totalWorkgroups, 65535); + workgroupCountY = Math.ceil(totalWorkgroups / 65535); + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + constants = { wg_size: wgSize }; + } else if (is1x1Tile32) { + // 32x64 register-tiled GEMM for pointwise 1x1 conv with clamp + workgroupCountX = Math.ceil((expectedHout * expectedWout) / 64); + workgroupCountY = Math.ceil(oc / 32); + constants = {}; + } else if (is1x1Tile64) { + // 64x64 register-tiled GEMM for pointwise 1x1 conv with clamp + workgroupCountX = Math.ceil((expectedHout * expectedWout) / 64); + workgroupCountY = Math.ceil(oc / 64); + constants = {}; + } else if (isDenseTile32) { + // 32x64 register-tiled GEMM for general dense convolution with clamp + workgroupCountX = Math.ceil((n * expectedHout * expectedWout) / 64); + workgroupCountY = Math.ceil(oc / 32); + constants = {}; + } else if (isDenseTile64) { + // 64x64 register-tiled GEMM for general dense convolution with clamp + workgroupCountX = Math.ceil((n * expectedHout * expectedWout) / 64); + workgroupCountY = Math.ceil(oc / 64); + constants = {}; + } else { + // GEMM kernel: fixed 8x8 workgroup covering a 32x32 output tile; + // grid over (N = batch*H_out*W_out cols, M = OC rows) + const tile = 32; + workgroupCountX = Math.ceil((n * expectedHout * expectedWout) / tile); + workgroupCountY = Math.ceil(oc / tile); + if (workgroupCountX > 65535 || workgroupCountY > 65535) { + throw new Error(`${name}: GEMM tile grid exceeds dispatch limit`); + } + constants = {}; + } + + const bundle = createComputeBundle( + device, + code, + [ + { binding: 0, buffer: outTensor.buffer }, + { binding: 1, buffer: inTensor.buffer }, + { binding: 2, buffer: weight.buffer }, + { binding: 3, buffer: biasBuffer }, + { binding: 4, buffer: uniformBuffer }, + ], + constants + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + workgroupCountZ, + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `et_vk::conv_with_clamp.default`. + * + * Implements 2D direct and grouped convolutions with clamping. Dense + * convolutions are evaluated as implicit-im2col tiled GEMM, grouped and + * depthwise convolutions via the direct kernel. + */ +export const convWithClamp: Shader = { + name, + code: { + convWithClamp: SHADER_CONV_WITH_CLAMP, + convWithClampGemm: SHADER_CONV_WITH_CLAMP_GEMM, + convWithClamp1x1: SHADER_CONV_WITH_CLAMP_1X1, + convWithClamp1x1Vec4: SHADER_CONV_WITH_CLAMP_1X1_VEC4, + convWithClamp1x1Tile64: SHADER_CONV_WITH_CLAMP_1X1_64X64, + convWithClampTile64: SHADER_CONV_WITH_CLAMP_TILE64, + convWithClampDepthwise3x3: SHADER_CONV_WITH_CLAMP_DEPTHWISE, + convWithClampF16: SHADER_CONV_WITH_CLAMP_F16, + convWithClampGemmF16: SHADER_CONV_WITH_CLAMP_GEMM_F16, + convWithClampTile64F16: SHADER_CONV_WITH_CLAMP_TILE64_F16, + convWithClampDepthwiseF16: SHADER_CONV_WITH_CLAMP_DEPTHWISE_F16, + }, + recordIn, +}; diff --git a/src/shaders/convolution/conv-transpose2d-f16.wgsl.ts b/src/shaders/convolution/conv-transpose2d-f16.wgsl.ts new file mode 100644 index 0000000..e32f30c --- /dev/null +++ b/src/shaders/convolution/conv-transpose2d-f16.wgsl.ts @@ -0,0 +1,97 @@ +/** + * WGSL compute shader for float16 transposed 2D convolution (gather form). + * + * Source: ExecuTorch (backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d.wgsl) + * Adapted for WebGPU float16 (shader-f16). + * License: BSD-3-Clause (Copyright (c) Meta Platforms, Inc. and affiliates) + * Project: https://github.com/barhanc/wgpu-executorch + */ +export const SHADER_CONV_TRANSPOSE2D_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, + IC: u32, + IH: u32, + IW: u32, + OC: u32, + OH: u32, + OW: u32, + KH: u32, + KW: u32, + sH: u32, + sW: u32, + pH: u32, + pW: u32, + dH: u32, + dW: u32, + groups: u32, + has_bias: u32, + _p0: u32, + _p1: u32, + _p2: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.B * params.OC * params.OH * params.OW; + let i = gid.y * stride_x + gid.x; + if (i >= total) { + return; + } + let ow = i % params.OW; + let oh = (i / params.OW) % params.OH; + let oc = (i / (params.OW * params.OH)) % params.OC; + let b = i / (params.OW * params.OH * params.OC); + + let icpg = params.IC / params.groups; // input channels per group + let ocpg = params.OC / params.groups; // output channels per group + let g = oc / ocpg; + let ic0 = g * icpg; + let oc_in_g = oc % ocpg; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = f32(bias[oc]); + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + for (var kh: u32 = 0u; kh < params.KH; kh = kh + 1u) { + let num_h = i32(oh) + i32(params.pH) - i32(kh) * i32(params.dH); + if (num_h % i32(params.sH) != 0) { + continue; + } + let ih = num_h / i32(params.sH); + if (ih < 0 || ih >= iH) { + continue; + } + for (var kw: u32 = 0u; kw < params.KW; kw = kw + 1u) { + let num_w = i32(ow) + i32(params.pW) - i32(kw) * i32(params.dW); + if (num_w % i32(params.sW) != 0) { + continue; + } + let iw = num_w / i32(params.sW); + if (iw < 0 || iw >= iW) { + continue; + } + for (var icg: u32 = 0u; icg < icpg; icg = icg + 1u) { + let ic = ic0 + icg; + let in_idx = ((b * params.IC + ic) * params.IH + u32(ih)) * params.IW + u32(iw); + let w_idx = ((ic * ocpg + oc_in_g) * params.KH + kh) * params.KW + kw; + acc = acc + f32(input[in_idx]) * f32(weight[w_idx]); + } + } + } + out[i] = f16(acc); +} +`; diff --git a/src/shaders/convolution/conv-transpose2d.wgsl.ts b/src/shaders/convolution/conv-transpose2d.wgsl.ts new file mode 100644 index 0000000..45443a7 --- /dev/null +++ b/src/shaders/convolution/conv-transpose2d.wgsl.ts @@ -0,0 +1,101 @@ +/** + * WGSL compute shader for transposed 2D convolution (gather form). + * + * Source: ExecuTorch (backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d.wgsl) + * GitHub: https://github.com/pytorch/executorch/blob/main/backends/webgpu/runtime/ops/et_vk_conv2d/conv_transpose2d.wgsl + * License: BSD-3-Clause (Copyright (c) Meta Platforms, Inc. and affiliates) + */ +export const SHADER_CONV_TRANSPOSE2D = /* wgsl */ ` +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, + IC: u32, + IH: u32, + IW: u32, + OC: u32, + OH: u32, + OW: u32, + KH: u32, + KW: u32, + sH: u32, + sW: u32, + pH: u32, + pW: u32, + dH: u32, + dW: u32, + groups: u32, + has_bias: u32, + _p0: u32, + _p1: u32, + _p2: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// Transposed 2D convolution (gather form), NCHW row-major, fp32. ONE thread per +// (b, oc, oh, ow) output element. weight layout = torch convT [IC, OC/groups, +// KH, KW] (NOT flipped). For each kernel tap (kh,kw): an input row ih +// contributes iff (oh + pH - kh*dH) is divisible by sH and ih in range (the +// scatter-inversion). CPU-derisked vs torch.conv_transpose2d to fp64 round-off +// (/tmp/convtr_derisk.py), incl. non-square spatial + non-square kernel. +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.B * params.OC * params.OH * params.OW; + let i = gid.y * stride_x + gid.x; + if (i >= total) { + return; + } + let ow = i % params.OW; + let oh = (i / params.OW) % params.OH; + let oc = (i / (params.OW * params.OH)) % params.OC; + let b = i / (params.OW * params.OH * params.OC); + + let icpg = params.IC / params.groups; // input channels per group + let ocpg = params.OC / params.groups; // output channels per group + let g = oc / ocpg; + let ic0 = g * icpg; + let oc_in_g = oc % ocpg; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = bias[oc]; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + for (var kh: u32 = 0u; kh < params.KH; kh = kh + 1u) { + let num_h = i32(oh) + i32(params.pH) - i32(kh) * i32(params.dH); + if (num_h % i32(params.sH) != 0) { + continue; + } + let ih = num_h / i32(params.sH); + if (ih < 0 || ih >= iH) { + continue; + } + for (var kw: u32 = 0u; kw < params.KW; kw = kw + 1u) { + let num_w = i32(ow) + i32(params.pW) - i32(kw) * i32(params.dW); + if (num_w % i32(params.sW) != 0) { + continue; + } + let iw = num_w / i32(params.sW); + if (iw < 0 || iw >= iW) { + continue; + } + for (var icg: u32 = 0u; icg < icpg; icg = icg + 1u) { + let ic = ic0 + icg; + let in_idx = ((b * params.IC + ic) * params.IH + u32(ih)) * params.IW + u32(iw); + // weight [IC, OC/groups, KH, KW]: index (ic, oc_in_g, kh, kw) + let w_idx = ((ic * ocpg + oc_in_g) * params.KH + kh) * params.KW + kw; + acc = acc + input[in_idx] * weight[w_idx]; + } + } + } + out[i] = acc; +} +`; diff --git a/src/shaders/convolution/conv2d-1x1-tile32-f16.wgsl.ts b/src/shaders/convolution/conv2d-1x1-tile32-f16.wgsl.ts new file mode 100644 index 0000000..3b4fa15 --- /dev/null +++ b/src/shaders/convolution/conv2d-1x1-tile32-f16.wgsl.ts @@ -0,0 +1,130 @@ +/** + * WGSL compute shader for 32x64 register-tiled 1x1 convolution in FP16 (B = 1, 4x4 per thread across 128 threads). + * Optimized for small channel counts (OC <= 32) to maximize wave occupancy and eliminate thread divergence. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV2D_1X1_TILE32_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var out: array>; +@group(0) @binding(1) var input: array>; +@group(0) @binding(2) var weight: array>; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(4) var params: Params; + +var a_sub: array, 4>, 32>; +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(local_invocation_index) tid: u32 +) { + let M = params.OC; + let N = params.OH * params.OW; + let K = params.IC; + let K_vecs = K >> 2u; + let N_vecs = N >> 2u; + + let tile_row0 = wg_id.y * 32u; + let tile_col0_4 = wg_id.x * 16u; + + var acc0 = vec4(0.0h); + var acc1 = vec4(0.0h); + var acc2 = vec4(0.0h); + var acc3 = vec4(0.0h); + + let a_r = tid >> 2u; // 0..31 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < M; + + let b_r0 = tid >> 4u; // 0..7 + let b_c0 = tid & 15u; // 0..15 + let b_n4_0 = tile_col0_4 + b_c0; + let b_c0_valid = b_n4_0 < N_vecs; + + let b_r1 = b_r0 + 8u; // 8..15 + let b_c1 = b_c0; + let b_n4_1 = b_n4_0; + let b_c1_valid = b_c0_valid; + + let num_tiles = (K + 15u) >> 4u; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t << 2u; + let k_start = t << 4u; + + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = weight[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0h); + } + + let bk0 = k_start + b_r0; + if (bk0 < K && b_c0_valid) { + b_sub[b_r0][b_c0] = input[bk0 * N_vecs + b_n4_0]; + } else { + b_sub[b_r0][b_c0] = vec4(0.0h); + } + + let bk1 = k_start + b_r1; + if (bk1 < K && b_c1_valid) { + b_sub[b_r1][b_c1] = input[bk1 * N_vecs + b_n4_1]; + } else { + b_sub[b_r1][b_c1] = vec4(0.0h); + } + + workgroupBarrier(); + + let thread_r = local_id.y << 2u; + let thread_c4 = local_id.x; + + for (var k4: u32 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + let thread_r = local_id.y << 2u; + let out_r0 = tile_row0 + thread_r; + let out_c4 = tile_col0_4 + local_id.x; + + if (out_c4 < N_vecs) { + var b_val = vec4(0.0h); + if (params.has_bias != 0u) { + if (out_r0 + 0u < M) { b_val.x = bias[out_r0 + 0u]; } + if (out_r0 + 1u < M) { b_val.y = bias[out_r0 + 1u]; } + if (out_r0 + 2u < M) { b_val.z = bias[out_r0 + 2u]; } + if (out_r0 + 3u < M) { b_val.w = bias[out_r0 + 3u]; } + } + + if (out_r0 + 0u < M) { out[(out_r0 + 0u) * N_vecs + out_c4] = acc0 + vec4(b_val.x); } + if (out_r0 + 1u < M) { out[(out_r0 + 1u) * N_vecs + out_c4] = acc1 + vec4(b_val.y); } + if (out_r0 + 2u < M) { out[(out_r0 + 2u) * N_vecs + out_c4] = acc2 + vec4(b_val.z); } + if (out_r0 + 3u < M) { out[(out_r0 + 3u) * N_vecs + out_c4] = acc3 + vec4(b_val.w); } + } +} +`; diff --git a/src/shaders/convolution/conv2d-1x1-tile32.wgsl.ts b/src/shaders/convolution/conv2d-1x1-tile32.wgsl.ts new file mode 100644 index 0000000..0c82392 --- /dev/null +++ b/src/shaders/convolution/conv2d-1x1-tile32.wgsl.ts @@ -0,0 +1,132 @@ +/** + * WGSL compute shader for 32x64 register-tiled 1x1 convolution (B = 1, 4x4 per thread across 128 threads). + * Optimized for small channel counts (OC <= 32) to maximize wave occupancy and eliminate thread divergence. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV2D_1X1_TILE32 = /* wgsl */ ` +@group(0) @binding(0) var out: array>; +@group(0) @binding(1) var input: array>; +@group(0) @binding(2) var weight: array>; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(4) var params: Params; + +// 32 rows x 16 K elements (4 vec4s) = 128 vec4s (2 KB) +var a_sub: array, 4>, 32>; +// 16 K rows x 64 N elements (16 vec4s) = 256 vec4s (4 KB) +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(local_invocation_index) tid: u32 +) { + let M = params.OC; + let N = params.OH * params.OW; + let K = params.IC; + let K_vecs = K >> 2u; + let N_vecs = N >> 2u; + + let tile_row0 = wg_id.y * 32u; + let tile_col0_4 = wg_id.x * 16u; + + var acc0 = vec4(0.0); + var acc1 = vec4(0.0); + var acc2 = vec4(0.0); + var acc3 = vec4(0.0); + + // Load A coordinates for 128 threads (each thread loads 1 vec4 of weight) + let a_r = tid >> 2u; // 0..31 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < M; + + // Load B coordinates for 128 threads (each thread loads 2 vec4s of input) + let b_r0 = tid >> 4u; // 0..7 + let b_c0 = tid & 15u; // 0..15 + let b_n4_0 = tile_col0_4 + b_c0; + let b_c0_valid = b_n4_0 < N_vecs; + + let b_r1 = b_r0 + 8u; // 8..15 + let b_c1 = b_c0; + let b_n4_1 = b_n4_0; + let b_c1_valid = b_c0_valid; + + let num_tiles = (K + 15u) >> 4u; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t << 2u; + let k_start = t << 4u; + + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = weight[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0); + } + + let bk0 = k_start + b_r0; + if (bk0 < K && b_c0_valid) { + b_sub[b_r0][b_c0] = input[bk0 * N_vecs + b_n4_0]; + } else { + b_sub[b_r0][b_c0] = vec4(0.0); + } + + let bk1 = k_start + b_r1; + if (bk1 < K && b_c1_valid) { + b_sub[b_r1][b_c1] = input[bk1 * N_vecs + b_n4_1]; + } else { + b_sub[b_r1][b_c1] = vec4(0.0); + } + + workgroupBarrier(); + + let thread_r = local_id.y << 2u; // 0..28 (step 4) + let thread_c4 = local_id.x; // 0..15 + + for (var k4: u32 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + let thread_r = local_id.y << 2u; + let out_r0 = tile_row0 + thread_r; + let out_c4 = tile_col0_4 + local_id.x; + + if (out_c4 < N_vecs) { + var b_val = vec4(0.0); + if (params.has_bias != 0u) { + if (out_r0 + 0u < M) { b_val.x = bias[out_r0 + 0u]; } + if (out_r0 + 1u < M) { b_val.y = bias[out_r0 + 1u]; } + if (out_r0 + 2u < M) { b_val.z = bias[out_r0 + 2u]; } + if (out_r0 + 3u < M) { b_val.w = bias[out_r0 + 3u]; } + } + + if (out_r0 + 0u < M) { out[(out_r0 + 0u) * N_vecs + out_c4] = acc0 + vec4(b_val.x); } + if (out_r0 + 1u < M) { out[(out_r0 + 1u) * N_vecs + out_c4] = acc1 + vec4(b_val.y); } + if (out_r0 + 2u < M) { out[(out_r0 + 2u) * N_vecs + out_c4] = acc2 + vec4(b_val.z); } + if (out_r0 + 3u < M) { out[(out_r0 + 3u) * N_vecs + out_c4] = acc3 + vec4(b_val.w); } + } +} +`; diff --git a/src/shaders/convolution/conv2d-1x1-tile64-f16.wgsl.ts b/src/shaders/convolution/conv2d-1x1-tile64-f16.wgsl.ts new file mode 100644 index 0000000..d14e57f --- /dev/null +++ b/src/shaders/convolution/conv2d-1x1-tile64-f16.wgsl.ts @@ -0,0 +1,128 @@ +/** + * WGSL compute shader for 64x64 register-tiled 1x1 convolution in float16 (B = 1, 4x4 per thread across 256 threads). + * High occupancy, coalesced 128-bit memory transactions, minimal register pressure for all GPUs. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV2D_1X1_64X64_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var out: array>; +@group(0) @binding(1) var input: array>; +@group(0) @binding(2) var weight: array>; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(4) var params: Params; + +// 64 rows x 16 K elements (4 vec4s) = 256 vec4s (2 KB) +var a_sub: array, 4>, 64>; +// 16 K rows x 64 N elements (16 vec4s) = 256 vec4s (2 KB) +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let M = params.OC; + let N = params.OH * params.OW; + let K = params.IC; + let K_vecs = K >> 2u; + let N_vecs = N >> 2u; + + let tile_row0 = wg_id.y * 64u; + let tile_col0_4 = wg_id.x * 16u; + let tid = local_id.y * 16u + local_id.x; // 0..255 + + var acc0 = vec4(0.0h); + var acc1 = vec4(0.0h); + var acc2 = vec4(0.0h); + var acc3 = vec4(0.0h); + + // Load coordinates for 256 threads (each thread loads 1 vec4 of A/weight and 1 vec4 of B/input) + let a_r = tid >> 2u; // 0..63 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < M; + + let b_r = tid >> 4u; // 0..15 + let b_c = tid & 15u; // 0..15 + let b_n4 = tile_col0_4 + b_c; + let b_c_valid = b_n4 < N_vecs; + + let num_tiles = (K + 15u) >> 4u; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t << 2u; + let k_start = t << 4u; + + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = weight[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0h); + } + + let bk = k_start + b_r; + if (bk < K && b_c_valid) { + b_sub[b_r][b_c] = input[bk * N_vecs + b_n4]; + } else { + b_sub[b_r][b_c] = vec4(0.0h); + } + workgroupBarrier(); + + let thread_r = local_id.y << 2u; // 0..60 (step 4) + let thread_c4 = local_id.x; // 0..15 + + for (var k4: u32 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + let thread_r = local_id.y << 2u; + let out_r0 = tile_row0 + thread_r; + let out_c4 = tile_col0_4 + local_id.x; + + if (out_c4 < N_vecs) { + if (out_r0 + 0u < M) { + var v0 = acc0; + if (params.has_bias != 0u) { v0 = v0 + vec4(bias[out_r0 + 0u]); } + out[(out_r0 + 0u) * N_vecs + out_c4] = v0; + } + if (out_r0 + 1u < M) { + var v1 = acc1; + if (params.has_bias != 0u) { v1 = v1 + vec4(bias[out_r0 + 1u]); } + out[(out_r0 + 1u) * N_vecs + out_c4] = v1; + } + if (out_r0 + 2u < M) { + var v2 = acc2; + if (params.has_bias != 0u) { v2 = v2 + vec4(bias[out_r0 + 2u]); } + out[(out_r0 + 2u) * N_vecs + out_c4] = v2; + } + if (out_r0 + 3u < M) { + var v3 = acc3; + if (params.has_bias != 0u) { v3 = v3 + vec4(bias[out_r0 + 3u]); } + out[(out_r0 + 3u) * N_vecs + out_c4] = v3; + } + } +} +`; diff --git a/src/shaders/convolution/conv2d-1x1-tile64.wgsl.ts b/src/shaders/convolution/conv2d-1x1-tile64.wgsl.ts new file mode 100644 index 0000000..a807d53 --- /dev/null +++ b/src/shaders/convolution/conv2d-1x1-tile64.wgsl.ts @@ -0,0 +1,126 @@ +/** + * WGSL compute shader for 64x64 register-tiled 1x1 convolution (B = 1, 4x4 per thread across 256 threads). + * High occupancy, coalesced 128-bit memory transactions, minimal register pressure for all GPUs. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV2D_1X1_64X64 = /* wgsl */ ` +@group(0) @binding(0) var out: array>; +@group(0) @binding(1) var input: array>; +@group(0) @binding(2) var weight: array>; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(4) var params: Params; + +// 64 rows x 16 K elements (4 vec4s) = 256 vec4s (4 KB) +var a_sub: array, 4>, 64>; +// 16 K rows x 64 N elements (16 vec4s) = 256 vec4s (4 KB) +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let M = params.OC; + let N = params.OH * params.OW; + let K = params.IC; + let K_vecs = K >> 2u; + let N_vecs = N >> 2u; + + let tile_row0 = wg_id.y * 64u; + let tile_col0_4 = wg_id.x * 16u; + let tid = local_id.y * 16u + local_id.x; // 0..255 + + var acc0 = vec4(0.0); + var acc1 = vec4(0.0); + var acc2 = vec4(0.0); + var acc3 = vec4(0.0); + + // Load coordinates for 256 threads (each thread loads 1 vec4 of A/weight and 1 vec4 of B/input) + let a_r = tid >> 2u; // 0..63 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < M; + + let b_r = tid >> 4u; // 0..15 + let b_c = tid & 15u; // 0..15 + let b_n4 = tile_col0_4 + b_c; + let b_c_valid = b_n4 < N_vecs; + + let num_tiles = (K + 15u) >> 4u; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t << 2u; + let k_start = t << 4u; + + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = weight[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0); + } + + let bk = k_start + b_r; + if (bk < K && b_c_valid) { + b_sub[b_r][b_c] = input[bk * N_vecs + b_n4]; + } else { + b_sub[b_r][b_c] = vec4(0.0); + } + workgroupBarrier(); + + let thread_r = local_id.y << 2u; // 0..60 (step 4) + let thread_c4 = local_id.x; // 0..15 + + for (var k4: u32 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + let thread_r = local_id.y << 2u; + let out_r0 = tile_row0 + thread_r; + let out_c4 = tile_col0_4 + local_id.x; + + if (out_c4 < N_vecs) { + if (out_r0 + 0u < M) { + var v0 = acc0; + if (params.has_bias != 0u) { v0 = v0 + vec4(bias[out_r0 + 0u]); } + out[(out_r0 + 0u) * N_vecs + out_c4] = v0; + } + if (out_r0 + 1u < M) { + var v1 = acc1; + if (params.has_bias != 0u) { v1 = v1 + vec4(bias[out_r0 + 1u]); } + out[(out_r0 + 1u) * N_vecs + out_c4] = v1; + } + if (out_r0 + 2u < M) { + var v2 = acc2; + if (params.has_bias != 0u) { v2 = v2 + vec4(bias[out_r0 + 2u]); } + out[(out_r0 + 2u) * N_vecs + out_c4] = v2; + } + if (out_r0 + 3u < M) { + var v3 = acc3; + if (params.has_bias != 0u) { v3 = v3 + vec4(bias[out_r0 + 3u]); } + out[(out_r0 + 3u) * N_vecs + out_c4] = v3; + } + } +} +`; diff --git a/src/shaders/convolution/conv2d-1x1-vec4.wgsl.ts b/src/shaders/convolution/conv2d-1x1-vec4.wgsl.ts new file mode 100644 index 0000000..2eb5d53 --- /dev/null +++ b/src/shaders/convolution/conv2d-1x1-vec4.wgsl.ts @@ -0,0 +1,106 @@ +/** + * WGSL compute shader for direct 2D 1x1 convolution (vec4 vectorized GEMM). + * Processes 128-bit vec4 memory transactions along channels (K) and spatial dimensions (N). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV2D_1X1_VEC4 = /* wgsl */ ` +@group(0) @binding(0) var out: array>; +@group(0) @binding(1) var input: array>; +@group(0) @binding(2) var weight: array>; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(4) var params: Params; + +const TILE: u32 = 32u; +const RPT: u32 = 4u; +const TILE4: u32 = 8u; + +var a_sub: array, 8>, 32>; +var b_sub: array, 8>, 32>; + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let M = params.OC; + let spatial = params.OH * params.OW; + let sp_vecs = spatial / 4u; + let N_vecs = params.B * sp_vecs; + let K = params.IC; + let K_vecs = K / 4u; + + let tile_row0 = wg_id.y * TILE; + let tile_col0_4 = wg_id.x * TILE4; + let tile_row = local_id.y * RPT; + let tile_col4 = local_id.x; + + var acc: array, 4>; + acc[0] = vec4(0.0); + acc[1] = vec4(0.0); + acc[2] = vec4(0.0); + acc[3] = vec4(0.0); + + let num_tiles = (K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t * TILE4; + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + let m = tile_row0 + arow; + let k4 = k4_start + local_id.x; + if (m < M && k4 < K_vecs) { + a_sub[arow][local_id.x] = weight[m * K_vecs + k4]; + } else { + a_sub[arow][local_id.x] = vec4(0.0); + } + + let b_k = t * TILE + arow; + let n4 = tile_col0_4 + local_id.x; + if (b_k < K && n4 < N_vecs) { + let b_idx = n4 / sp_vecs; + let sp4 = n4 % sp_vecs; + b_sub[arow][local_id.x] = input[(b_idx * params.IC + b_k) * sp_vecs + sp4]; + } else { + b_sub[arow][local_id.x] = vec4(0.0); + } + } + workgroupBarrier(); + + for (var k4: u32 = 0u; k4 < TILE4; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][tile_col4]; + let b1 = b_sub[k4 * 4u + 1u][tile_col4]; + let b2 = b_sub[k4 * 4u + 2u][tile_col4]; + let b3 = b_sub[k4 * 4u + 3u][tile_col4]; + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let ac = a_sub[tile_row + ir][k4]; + acc[ir] = acc[ir] + b0 * ac.x + b1 * ac.y + b2 * ac.z + b3 * ac.w; + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let m = tile_row0 + tile_row + ir; + let n4 = tile_col0_4 + tile_col4; + if (m < M && n4 < N_vecs) { + var val = acc[ir]; + if (params.has_bias != 0u) { + val = val + vec4(bias[m]); + } + let b_idx = n4 / sp_vecs; + let sp4 = n4 % sp_vecs; + out[(b_idx * params.OC + m) * sp_vecs + sp4] = val; + } + } +} +`; diff --git a/src/shaders/convolution/conv2d-1x1.wgsl.ts b/src/shaders/convolution/conv2d-1x1.wgsl.ts new file mode 100644 index 0000000..e53170f --- /dev/null +++ b/src/shaders/convolution/conv2d-1x1.wgsl.ts @@ -0,0 +1,115 @@ +/** + * WGSL compute shader for direct 2D 1x1 convolution (pointwise GEMM). + * Avoids im2col overhead since 1x1 conv with stride 1, padding 0 is a direct matrix multiply. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV2D_1X1 = /* wgsl */ ` +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(4) var params: Params; + +const TILE: u32 = 32u; +const RPT: u32 = 4u; + +var a_sub: array, 32>; +var b_sub: array, 32>; + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let M = params.OC; + let spatial = params.OH * params.OW; + let N_total = params.B * spatial; + let K = params.IC; + + let tile_row0 = wg_id.y * TILE; + let tile_col0 = wg_id.x * TILE; + let tile_row = local_id.y * RPT; + let tile_col = local_id.x * RPT; + + var acc: array, 4>; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + acc[ir][ic] = 0.0; + } + } + + let num_tiles = (K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k_start = t * TILE; + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) { + let col = local_id.x * RPT + kk; + let m = tile_row0 + arow; + let k = k_start + col; + if (m < M && k < K) { + a_sub[arow][col] = weight[m * K + k]; + } else { + a_sub[arow][col] = 0.0; + } + + let b_k = k_start + arow; + let n = tile_col0 + col; + if (b_k < K && n < N_total) { + let b_idx = n / spatial; + let sp = n % spatial; + b_sub[arow][col] = input[(b_idx * params.IC + b_k) * spatial + sp]; + } else { + b_sub[arow][col] = 0.0; + } + } + } + workgroupBarrier(); + + for (var k: u32 = 0u; k < TILE; k = k + 1u) { + let b0 = b_sub[k][tile_col + 0u]; + let b1 = b_sub[k][tile_col + 1u]; + let b2 = b_sub[k][tile_col + 2u]; + let b3 = b_sub[k][tile_col + 3u]; + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let aval = a_sub[tile_row + ir][k]; + acc[ir][0] = acc[ir][0] + aval * b0; + acc[ir][1] = acc[ir][1] + aval * b1; + acc[ir][2] = acc[ir][2] + aval * b2; + acc[ir][3] = acc[ir][3] + aval * b3; + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let m = tile_row0 + tile_row + ir; + if (m < M) { + var bias_val: f32 = 0.0; + if (params.has_bias != 0u) { + bias_val = bias[m]; + } + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + let n = tile_col0 + tile_col + ic; + if (n < N_total) { + let val = acc[ir][ic] + bias_val; + let b_idx = n / spatial; + let sp = n % spatial; + out[(b_idx * params.OC + m) * spatial + sp] = val; + } + } + } + } +} +`; diff --git a/src/shaders/convolution/conv2d-depthwise-3x3-f16.wgsl.ts b/src/shaders/convolution/conv2d-depthwise-3x3-f16.wgsl.ts new file mode 100644 index 0000000..d6f2ceb --- /dev/null +++ b/src/shaders/convolution/conv2d-depthwise-3x3-f16.wgsl.ts @@ -0,0 +1,148 @@ +/** + * WGSL compute shader for depthwise 3x3 convolution in FP16 with shared-memory halo tile caching. + * Loads a 34x34 halo tile cooperatively across 256 threads, eliminating redundant global memory reads. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV2D_DEPTHWISE_3X3_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(4) var params: Params; + +var in_tile: array, 34>; +var w_tile: array; +var bias_val: f16; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(local_invocation_id) local_id: vec3, + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_index) tid: u32 +) { + let oc = wg_id.z % params.OC; + let b = wg_id.z / params.OC; + + if (tid < 9u) { + w_tile[tid] = weight[oc * 9u + tid]; + } + if (tid == 0u) { + if (params.has_bias != 0u) { + bias_val = bias[oc]; + } else { + bias_val = 0.0h; + } + } + + let in_hw = params.IH * params.IW; + let in_w = params.IW; + let in_base = (b * params.IC + oc) * in_hw; + let iH = i32(params.IH); + let iW = i32(params.IW); + + let tile_ih_start = i32(wg_id.y * 32u) - 1; + let tile_iw_start = i32(wg_id.x * 32u) - 1; + + // Cooperatively load 34x34 halo tile (1156 f16s) into shared memory + for (var idx = tid; idx < 1156u; idx = idx + 256u) { + let tr = idx / 34u; + let tc = idx % 34u; + let ih = tile_ih_start + i32(tr); + let iw = tile_iw_start + i32(tc); + if (ih >= 0 && ih < iH && iw >= 0 && iw < iW) { + in_tile[tr][tc] = input[in_base + u32(ih) * in_w + u32(iw)]; + } else { + in_tile[tr][tc] = 0.0h; + } + } + workgroupBarrier(); + + let ow0 = wg_id.x * 32u + local_id.x * 2u; + let oh0 = wg_id.y * 32u + local_id.y * 2u; + + if (ow0 >= params.OW || oh0 >= params.OH || b >= params.B) { + return; + } + + var acc00 = bias_val; + var acc01 = bias_val; + var acc10 = bias_val; + var acc11 = bias_val; + + let tr = local_id.y * 2u; + let tc = local_id.x * 2u; + + let w0 = w_tile[0]; let w1 = w_tile[1]; let w2 = w_tile[2]; + let w3 = w_tile[3]; let w4 = w_tile[4]; let w5 = w_tile[5]; + let w6 = w_tile[6]; let w7 = w_tile[7]; let w8 = w_tile[8]; + + // Row 0 + let r0_c0 = in_tile[tr + 0u][tc + 0u]; + let r0_c1 = in_tile[tr + 0u][tc + 1u]; + let r0_c2 = in_tile[tr + 0u][tc + 2u]; + let r0_c3 = in_tile[tr + 0u][tc + 3u]; + + acc00 = acc00 + r0_c0 * w0 + r0_c1 * w1 + r0_c2 * w2; + acc01 = acc01 + r0_c1 * w0 + r0_c2 * w1 + r0_c3 * w2; + + // Row 1 + let r1_c0 = in_tile[tr + 1u][tc + 0u]; + let r1_c1 = in_tile[tr + 1u][tc + 1u]; + let r1_c2 = in_tile[tr + 1u][tc + 2u]; + let r1_c3 = in_tile[tr + 1u][tc + 3u]; + + acc00 = acc00 + r1_c0 * w3 + r1_c1 * w4 + r1_c2 * w5; + acc01 = acc01 + r1_c1 * w3 + r1_c2 * w4 + r1_c3 * w5; + acc10 = acc10 + r1_c0 * w0 + r1_c1 * w1 + r1_c2 * w2; + acc11 = acc11 + r1_c1 * w0 + r1_c2 * w1 + r1_c3 * w2; + + // Row 2 + let r2_c0 = in_tile[tr + 2u][tc + 0u]; + let r2_c1 = in_tile[tr + 2u][tc + 1u]; + let r2_c2 = in_tile[tr + 2u][tc + 2u]; + let r2_c3 = in_tile[tr + 2u][tc + 3u]; + + acc00 = acc00 + r2_c0 * w6 + r2_c1 * w7 + r2_c2 * w8; + acc01 = acc01 + r2_c1 * w6 + r2_c2 * w7 + r2_c3 * w8; + acc10 = acc10 + r2_c0 * w3 + r2_c1 * w4 + r2_c2 * w5; + acc11 = acc11 + r2_c1 * w3 + r2_c2 * w4 + r2_c3 * w5; + + // Row 3 + let r3_c0 = in_tile[tr + 3u][tc + 0u]; + let r3_c1 = in_tile[tr + 3u][tc + 1u]; + let r3_c2 = in_tile[tr + 3u][tc + 2u]; + let r3_c3 = in_tile[tr + 3u][tc + 3u]; + + acc10 = acc10 + r3_c0 * w6 + r3_c1 * w7 + r3_c2 * w8; + acc11 = acc11 + r3_c1 * w6 + r3_c2 * w7 + r3_c3 * w8; + + let out_hw = params.OH * params.OW; + let out_w = params.OW; + let out_base = (b * params.OC + oc) * out_hw; + + let out_idx00 = out_base + oh0 * out_w + ow0; + out[out_idx00] = acc00; + + if (ow0 + 1u < params.OW) { + out[out_idx00 + 1u] = acc01; + } + if (oh0 + 1u < params.OH) { + let out_idx10 = out_base + (oh0 + 1u) * out_w + ow0; + out[out_idx10] = acc10; + if (ow0 + 1u < params.OW) { + out[out_idx10 + 1u] = acc11; + } + } +} +`; diff --git a/src/shaders/convolution/conv2d-depthwise-3x3.wgsl.ts b/src/shaders/convolution/conv2d-depthwise-3x3.wgsl.ts new file mode 100644 index 0000000..fda2573 --- /dev/null +++ b/src/shaders/convolution/conv2d-depthwise-3x3.wgsl.ts @@ -0,0 +1,147 @@ +/** + * WGSL compute shader for depthwise 3x3 convolution with shared-memory halo tile caching. + * Loads a 34x34 halo tile cooperatively across 256 threads, eliminating redundant global memory reads. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV2D_DEPTHWISE_3X3 = /* wgsl */ ` +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(4) var params: Params; + +var in_tile: array, 34>; +var w_tile: array; +var bias_val: f32; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(local_invocation_id) local_id: vec3, + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_index) tid: u32 +) { + let oc = wg_id.z % params.OC; + let b = wg_id.z / params.OC; + + if (tid < 9u) { + w_tile[tid] = weight[oc * 9u + tid]; + } + if (tid == 0u) { + if (params.has_bias != 0u) { + bias_val = bias[oc]; + } else { + bias_val = 0.0; + } + } + + let in_hw = params.IH * params.IW; + let in_w = params.IW; + let in_base = (b * params.IC + oc) * in_hw; + let iH = i32(params.IH); + let iW = i32(params.IW); + + let tile_ih_start = i32(wg_id.y * 32u) - 1; + let tile_iw_start = i32(wg_id.x * 32u) - 1; + + // Cooperatively load 34x34 halo tile (1156 floats) into shared memory + for (var idx = tid; idx < 1156u; idx = idx + 256u) { + let tr = idx / 34u; + let tc = idx % 34u; + let ih = tile_ih_start + i32(tr); + let iw = tile_iw_start + i32(tc); + if (ih >= 0 && ih < iH && iw >= 0 && iw < iW) { + in_tile[tr][tc] = input[in_base + u32(ih) * in_w + u32(iw)]; + } else { + in_tile[tr][tc] = 0.0; + } + } + workgroupBarrier(); + + let ow0 = wg_id.x * 32u + local_id.x * 2u; + let oh0 = wg_id.y * 32u + local_id.y * 2u; + + if (ow0 >= params.OW || oh0 >= params.OH || b >= params.B) { + return; + } + + var acc00 = bias_val; + var acc01 = bias_val; + var acc10 = bias_val; + var acc11 = bias_val; + + let tr = local_id.y * 2u; + let tc = local_id.x * 2u; + + // Load weights from shared memory into private registers + let w0 = w_tile[0]; let w1 = w_tile[1]; let w2 = w_tile[2]; + let w3 = w_tile[3]; let w4 = w_tile[4]; let w5 = w_tile[5]; + let w6 = w_tile[6]; let w7 = w_tile[7]; let w8 = w_tile[8]; + + // Row 0 + let r0_c0 = in_tile[tr + 0u][tc + 0u]; + let r0_c1 = in_tile[tr + 0u][tc + 1u]; + let r0_c2 = in_tile[tr + 0u][tc + 2u]; + let r0_c3 = in_tile[tr + 0u][tc + 3u]; + + acc00 = acc00 + r0_c0 * w0 + r0_c1 * w1 + r0_c2 * w2; + acc01 = acc01 + r0_c1 * w0 + r0_c2 * w1 + r0_c3 * w2; + + // Row 1 + let r1_c0 = in_tile[tr + 1u][tc + 0u]; + let r1_c1 = in_tile[tr + 1u][tc + 1u]; + let r1_c2 = in_tile[tr + 1u][tc + 2u]; + let r1_c3 = in_tile[tr + 1u][tc + 3u]; + + acc00 = acc00 + r1_c0 * w3 + r1_c1 * w4 + r1_c2 * w5; + acc01 = acc01 + r1_c1 * w3 + r1_c2 * w4 + r1_c3 * w5; + acc10 = acc10 + r1_c0 * w0 + r1_c1 * w1 + r1_c2 * w2; + acc11 = acc11 + r1_c1 * w0 + r1_c2 * w1 + r1_c3 * w2; + + // Row 2 + let r2_c0 = in_tile[tr + 2u][tc + 0u]; + let r2_c1 = in_tile[tr + 2u][tc + 1u]; + let r2_c2 = in_tile[tr + 2u][tc + 2u]; + let r2_c3 = in_tile[tr + 2u][tc + 3u]; + + acc00 = acc00 + r2_c0 * w6 + r2_c1 * w7 + r2_c2 * w8; + acc01 = acc01 + r2_c1 * w6 + r2_c2 * w7 + r2_c3 * w8; + acc10 = acc10 + r2_c0 * w3 + r2_c1 * w4 + r2_c2 * w5; + acc11 = acc11 + r2_c1 * w3 + r2_c2 * w4 + r2_c3 * w5; + + // Row 3 + let r3_c0 = in_tile[tr + 3u][tc + 0u]; + let r3_c1 = in_tile[tr + 3u][tc + 1u]; + let r3_c2 = in_tile[tr + 3u][tc + 2u]; + let r3_c3 = in_tile[tr + 3u][tc + 3u]; + + acc10 = acc10 + r3_c0 * w6 + r3_c1 * w7 + r3_c2 * w8; + acc11 = acc11 + r3_c1 * w6 + r3_c2 * w7 + r3_c3 * w8; + + let out_hw = params.OH * params.OW; + let out_w = params.OW; + let out_base = (b * params.OC + oc) * out_hw; + + let out_idx00 = out_base + oh0 * out_w + ow0; + out[out_idx00] = acc00; + + if (ow0 + 1u < params.OW) { + out[out_idx00 + 1u] = acc01; + } + if (oh0 + 1u < params.OH) { + let out_idx10 = out_base + (oh0 + 1u) * out_w + ow0; + out[out_idx10] = acc10; + if (ow0 + 1u < params.OW) { + out[out_idx10 + 1u] = acc11; + } + } +} +`; diff --git a/src/shaders/convolution/conv2d-depthwise-f16.wgsl.ts b/src/shaders/convolution/conv2d-depthwise-f16.wgsl.ts new file mode 100644 index 0000000..0e80d68 --- /dev/null +++ b/src/shaders/convolution/conv2d-depthwise-f16.wgsl.ts @@ -0,0 +1,136 @@ +/** + * WGSL compute shader for float16 depthwise convolution (up to 7x7 kernels). + * Uses 2D spatial workgroup layout (16x16) with 2x2 register tiling per thread (32x32 per workgroup). + * Weights and bias are cached in shared memory then private registers to avoid bank conflicts. + * Requires WebGPU shader-f16 feature. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV2D_DEPTHWISE_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(4) var params: Params; + +var w_tile: array; +var bias_val: f16; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(workgroup_id) wg_id: vec3 +) { + let oc = wg_id.z % params.OC; + let b = wg_id.z / params.OC; + + let tid = local_id.y * 16u + local_id.x; + let khw = params.KH * params.KW; + if (tid < khw) { + w_tile[tid] = weight[oc * khw + tid]; + } + if (tid == 0u) { + if (params.has_bias != 0u) { + bias_val = bias[oc]; + } else { + bias_val = 0.0h; + } + } + workgroupBarrier(); + + let ow0 = gid.x * 2u; + let oh0 = gid.y * 2u; + + if (ow0 >= params.OW || oh0 >= params.OH || b >= params.B) { + return; + } + + var acc00 = bias_val; + var acc01 = bias_val; + var acc10 = bias_val; + var acc11 = bias_val; + + let sH = i32(params.sH); + let sW = i32(params.sW); + let dH = i32(params.dH); + let dW = i32(params.dW); + let iH = i32(params.IH); + let iW = i32(params.IW); + + let ih0_base = i32(oh0) * sH - i32(params.pH); + let iw0_base = i32(ow0) * sW - i32(params.pW); + let in_base = (b * params.IC + oc) * params.IH; + + var w_reg: array; + for (var i = 0u; i < khw; i = i + 1u) { + w_reg[i] = w_tile[i]; + } + + for (var kh = 0u; kh < params.KH; kh = kh + 1u) { + let ih0 = ih0_base + i32(kh) * dH; + let ih1 = ih0 + sH; + let valid_h0 = ih0 >= 0 && ih0 < iH; + let valid_h1 = ih1 >= 0 && ih1 < iH; + + var row_offset0: u32 = 0u; + if (valid_h0) { + row_offset0 = (in_base + u32(ih0)) * params.IW; + } + var row_offset1: u32 = 0u; + if (valid_h1) { + row_offset1 = (in_base + u32(ih1)) * params.IW; + } + let w_offset = kh * params.KW; + + for (var kw = 0u; kw < params.KW; kw = kw + 1u) { + let iw0 = iw0_base + i32(kw) * dW; + let iw1 = iw0 + sW; + let valid_w0 = iw0 >= 0 && iw0 < iW; + let valid_w1 = iw1 >= 0 && iw1 < iW; + let w_val = w_reg[w_offset + kw]; + + if (valid_h0) { + if (valid_w0) { + acc00 = acc00 + input[row_offset0 + u32(iw0)] * w_val; + } + if (valid_w1) { + acc01 = acc01 + input[row_offset0 + u32(iw1)] * w_val; + } + } + if (valid_h1) { + if (valid_w0) { + acc10 = acc10 + input[row_offset1 + u32(iw0)] * w_val; + } + if (valid_w1) { + acc11 = acc11 + input[row_offset1 + u32(iw1)] * w_val; + } + } + } + } + + let out_ch_base = (b * params.OC + oc) * params.OH; + let out_row0 = (out_ch_base + oh0) * params.OW; + out[out_row0 + ow0] = acc00; + if (ow0 + 1u < params.OW) { + out[out_row0 + ow0 + 1u] = acc01; + } + if (oh0 + 1u < params.OH) { + let out_row1 = (out_ch_base + oh0 + 1u) * params.OW; + out[out_row1 + ow0] = acc10; + if (ow0 + 1u < params.OW) { + out[out_row1 + ow0 + 1u] = acc11; + } + } +} +`; diff --git a/src/shaders/convolution/conv2d-depthwise.wgsl.ts b/src/shaders/convolution/conv2d-depthwise.wgsl.ts new file mode 100644 index 0000000..4c364d5 --- /dev/null +++ b/src/shaders/convolution/conv2d-depthwise.wgsl.ts @@ -0,0 +1,133 @@ +/** + * WGSL compute shader for depthwise convolution (up to 7x7 kernels). + * Uses 2D spatial workgroup layout (16x16) with 2x2 register tiling per thread (32x32 per workgroup). + * Weights and bias are cached in shared memory then private registers to avoid bank conflicts. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV2D_DEPTHWISE = /* wgsl */ ` +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(4) var params: Params; + +var w_tile: array; +var bias_val: f32; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(workgroup_id) wg_id: vec3 +) { + let oc = wg_id.z % params.OC; + let b = wg_id.z / params.OC; + + let tid = local_id.y * 16u + local_id.x; + let khw = params.KH * params.KW; + if (tid < khw) { + w_tile[tid] = weight[oc * khw + tid]; + } + if (tid == 0u) { + if (params.has_bias != 0u) { + bias_val = bias[oc]; + } else { + bias_val = 0.0; + } + } + workgroupBarrier(); + + let ow0 = gid.x * 2u; + let oh0 = gid.y * 2u; + + if (ow0 >= params.OW || oh0 >= params.OH || b >= params.B) { + return; + } + + var acc00 = bias_val; + var acc01 = bias_val; + var acc10 = bias_val; + var acc11 = bias_val; + + let sH = i32(params.sH); + let sW = i32(params.sW); + let dH = i32(params.dH); + let dW = i32(params.dW); + let iH = i32(params.IH); + let iW = i32(params.IW); + + let ih0_base = i32(oh0) * sH - i32(params.pH); + let iw0_base = i32(ow0) * sW - i32(params.pW); + let in_base = (b * params.IC + oc) * params.IH; + + var w_reg: array; + for (var i = 0u; i < khw; i = i + 1u) { + w_reg[i] = w_tile[i]; + } + + for (var kh = 0u; kh < params.KH; kh = kh + 1u) { + let ih0 = ih0_base + i32(kh) * dH; + let ih1 = ih0 + sH; + let valid_h0 = ih0 >= 0 && ih0 < iH; + let valid_h1 = ih1 >= 0 && ih1 < iH; + + var row_offset0: u32 = 0u; + if (valid_h0) { + row_offset0 = (in_base + u32(ih0)) * params.IW; + } + var row_offset1: u32 = 0u; + if (valid_h1) { + row_offset1 = (in_base + u32(ih1)) * params.IW; + } + let w_offset = kh * params.KW; + + for (var kw = 0u; kw < params.KW; kw = kw + 1u) { + let iw0 = iw0_base + i32(kw) * dW; + let iw1 = iw0 + sW; + let valid_w0 = iw0 >= 0 && iw0 < iW; + let valid_w1 = iw1 >= 0 && iw1 < iW; + let w_val = w_reg[w_offset + kw]; + + if (valid_h0) { + if (valid_w0) { + acc00 = acc00 + input[row_offset0 + u32(iw0)] * w_val; + } + if (valid_w1) { + acc01 = acc01 + input[row_offset0 + u32(iw1)] * w_val; + } + } + if (valid_h1) { + if (valid_w0) { + acc10 = acc10 + input[row_offset1 + u32(iw0)] * w_val; + } + if (valid_w1) { + acc11 = acc11 + input[row_offset1 + u32(iw1)] * w_val; + } + } + } + } + + let out_ch_base = (b * params.OC + oc) * params.OH; + let out_row0 = (out_ch_base + oh0) * params.OW; + out[out_row0 + ow0] = acc00; + if (ow0 + 1u < params.OW) { + out[out_row0 + ow0 + 1u] = acc01; + } + if (oh0 + 1u < params.OH) { + let out_row1 = (out_ch_base + oh0 + 1u) * params.OW; + out[out_row1 + ow0] = acc10; + if (ow0 + 1u < params.OW) { + out[out_row1 + ow0 + 1u] = acc11; + } + } +} +`; diff --git a/src/shaders/convolution/conv2d-f16.wgsl.ts b/src/shaders/convolution/conv2d-f16.wgsl.ts new file mode 100644 index 0000000..1f9f272 --- /dev/null +++ b/src/shaders/convolution/conv2d-f16.wgsl.ts @@ -0,0 +1,89 @@ +/** + * Direct 2D convolution compute shader (fallback for general grouped cases) in float16. + * Requires WebGPU shader-f16 feature. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV2D_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, + IC: u32, + IH: u32, + IW: u32, + OC: u32, + OH: u32, + OW: u32, + KH: u32, + KW: u32, + sH: u32, + sW: u32, + pH: u32, + pW: u32, + dH: u32, + dW: u32, + groups: u32, + has_bias: u32, + _p0: u32, + _p1: u32, + _p2: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.B * params.OC * params.OH * params.OW; + let i = gid.y * stride_x + gid.x; + if (i >= total) { + return; + } + let ow = i % params.OW; + let oh = (i / params.OW) % params.OH; + let oc = (i / (params.OW * params.OH)) % params.OC; + let b = i / (params.OW * params.OH * params.OC); + + let icpg = params.IC / params.groups; + let ocpg = params.OC / params.groups; + let g = oc / ocpg; + let ic0 = g * icpg; + + var acc: f16 = 0.0h; + if (params.has_bias != 0u) { + acc = bias[oc]; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + for (var icg: u32 = 0u; icg < icpg; icg = icg + 1u) { + let ic = ic0 + icg; + let in_c_base = (b * params.IC + ic) * params.IH; + let w_c_base = (oc * icpg + icg) * params.KH; + for (var kh: u32 = 0u; kh < params.KH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + let in_row = (in_c_base + u32(ih)) * params.IW; + let w_row = (w_c_base + kh) * params.KW; + for (var kw: u32 = 0u; kw < params.KW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + acc = acc + input[in_row + u32(iw)] * weight[w_row + kw]; + } + } + } + out[i] = acc; +} +`; diff --git a/src/shaders/convolution/conv2d-gemm-f16.wgsl.ts b/src/shaders/convolution/conv2d-gemm-f16.wgsl.ts new file mode 100644 index 0000000..bb4da84 --- /dev/null +++ b/src/shaders/convolution/conv2d-gemm-f16.wgsl.ts @@ -0,0 +1,171 @@ +/** + * WGSL compute shader for float16 2D convolution as implicit-im2col tiled GEMM. + * Requires WebGPU shader-f16 feature. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV2D_GEMM_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(4) var params: Params; + +const TILE: u32 = 32u; +const RPT: u32 = 4u; + +var a_sub: array, 32>; +var b_sub: array, 8>, 32>; + +fn gK() -> u32 { return params.IC * params.KH * params.KW; } +fn gN() -> u32 { return params.B * params.OH * params.OW; } + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3) { + let M = params.OC; + let N = gN(); + let K = gK(); + let tile_row0 = wg_id.y * TILE; + let tile_col0 = wg_id.x * TILE; + let tile_row = local_id.y * RPT; + let tile_col = local_id.x * RPT; + + var acc: array, 4>; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + acc[ir] = vec4(0.0h); + } + + let ohw = params.OH * params.OW; + let khw = params.KH * params.KW; + let iH = i32(params.IH); + let iW = i32(params.IW); + let sH = i32(params.sH); + let sW = i32(params.sW); + let pH = i32(params.pH); + let pW = i32(params.pW); + let dH = i32(params.dH); + let dW = i32(params.dW); + + var col_b: array; + var col_valid: array; + var col_oh_arr: array; + var col_ow_arr: array; + var col_ih_base: array; + var col_iw_base: array; + for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) { + let n = tile_col0 + local_id.x * RPT + kk; + if (n < N) { + col_valid[kk] = true; + let b_val = n / ohw; + let sp = n % ohw; + col_b[kk] = b_val; + let oh_val = i32(sp / params.OW); + let ow_val = i32(sp % params.OW); + col_oh_arr[kk] = oh_val; + col_ow_arr[kk] = ow_val; + col_ih_base[kk] = oh_val * sH - pH; + col_iw_base[kk] = ow_val * sW - pW; + } else { + col_valid[kk] = false; + col_b[kk] = 0u; + col_oh_arr[kk] = 0; + col_ow_arr[kk] = 0; + col_ih_base[kk] = 0; + col_iw_base[kk] = 0; + } + } + + let in_hw = params.IH * params.IW; + let in_w = params.IW; + let num_tiles = (K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k_start = t * TILE; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + let m = tile_row0 + arow; + let kk_idx = k_start + arow; + + var k_valid = false; + var k_ic = 0u; + var kh_dH = 0; + var kw_dW = 0; + if (kk_idx < K) { + k_valid = true; + k_ic = kk_idx / khw; + let rem = kk_idx % khw; + kh_dH = i32(rem / params.KW) * dH; + kw_dW = i32(rem % params.KW) * dW; + } + + let m_valid = m < M; + let w_row_base = m * K + k_start + local_id.x * RPT; + + var b_row: vec4; + for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) { + let col = local_id.x * RPT + kk; + let k_col = k_start + col; + if (m_valid && k_col < K) { + a_sub[arow][col] = weight[w_row_base + kk]; + } else { + a_sub[arow][col] = 0.0h; + } + + if (k_valid && col_valid[kk]) { + let ih = col_ih_base[kk] + kh_dH; + let iw = col_iw_base[kk] + kw_dW; + if (ih >= 0 && ih < iH && iw >= 0 && iw < iW) { + let ch_idx = col_b[kk] * params.IC + k_ic; + b_row[kk] = input[ch_idx * in_hw + u32(ih) * in_w + u32(iw)]; + } else { + b_row[kk] = 0.0h; + } + } else { + b_row[kk] = 0.0h; + } + } + b_sub[arow][local_id.x] = b_row; + } + workgroupBarrier(); + for (var k: u32 = 0u; k < TILE; k = k + 1u) { + let b_vec = b_sub[k][local_id.x]; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let aval = a_sub[tile_row + ir][k]; + acc[ir] = acc[ir] + aval * b_vec; + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let m = tile_row0 + tile_row + ir; + if (m < M) { + var bias_val = 0.0h; + if (params.has_bias != 0u) { + bias_val = bias[m]; + } + let v = acc[ir] + vec4(bias_val); + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + let n = tile_col0 + tile_col + ic; + if (n < N) { + let b = col_b[ic]; + let oh = u32(col_oh_arr[ic]); + let ow = u32(col_ow_arr[ic]); + out[((b * params.OC + m) * params.OH + oh) * params.OW + ow] = v[ic]; + } + } + } + } +} +`; diff --git a/src/shaders/convolution/conv2d-gemm.wgsl.ts b/src/shaders/convolution/conv2d-gemm.wgsl.ts new file mode 100644 index 0000000..e92699d --- /dev/null +++ b/src/shaders/convolution/conv2d-gemm.wgsl.ts @@ -0,0 +1,177 @@ +/** + * WGSL compute shader for direct 2D convolution as implicit-im2col tiled GEMM. + * + * Source: ExecuTorch (backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm.wgsl) + * GitHub: https://github.com/pytorch/executorch/blob/main/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d_gemm.wgsl + * License: BSD-3-Clause (Copyright (c) Meta Platforms, Inc. and affiliates) + */ +export const SHADER_CONV2D_GEMM = /* wgsl */ ` +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(4) var params: Params; + +// Standard (groups==1) conv2d as an implicit-im2col tiled GEMM, reusing the +// linear_fp32_tiled skeleton (TILE=32, RPT=4, shared-mem tiles). M=OC, N=B*OH*OW, +// K=IC*KH*KW. A=weight [OC,K] (contiguous OIHW); B=input sampled on the fly (im2col), +// out-of-range -> 0.0 implements padding. bias is per-ROW (OC). Output written NCHW. +const TILE: u32 = 32u; +const RPT: u32 = 4u; + +var a_sub: array, 32>; +var b_sub: array, 8>, 32>; + +fn gK() -> u32 { return params.IC * params.KH * params.KW; } +fn gN() -> u32 { return params.B * params.OH * params.OW; } + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3) { + let M = params.OC; + let N = gN(); + let K = gK(); + let tile_row0 = wg_id.y * TILE; + let tile_col0 = wg_id.x * TILE; + let tile_row = local_id.y * RPT; + let tile_col = local_id.x * RPT; + + var acc: array, 4>; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + acc[ir] = vec4(0.0); + } + + let ohw = params.OH * params.OW; + let khw = params.KH * params.KW; + let iH = i32(params.IH); + let iW = i32(params.IW); + let sH = i32(params.sH); + let sW = i32(params.sW); + let pH = i32(params.pH); + let pW = i32(params.pW); + let dH = i32(params.dH); + let dW = i32(params.dW); + + // Precompute spatial coords for this thread's 4 columns in B (invariant across K tiles) + var col_b: array; + var col_b_ic: array; + var col_valid: array; + var col_oh_arr: array; + var col_ow_arr: array; + var col_ih_base: array; + var col_iw_base: array; + for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) { + let n = tile_col0 + local_id.x * RPT + kk; + if (n < N) { + col_valid[kk] = true; + let b_val = n / ohw; + let sp = n % ohw; + col_b[kk] = b_val; + col_b_ic[kk] = b_val * params.IC; + let oh_val = i32(sp / params.OW); + let ow_val = i32(sp % params.OW); + col_oh_arr[kk] = oh_val; + col_ow_arr[kk] = ow_val; + col_ih_base[kk] = oh_val * sH - pH; + col_iw_base[kk] = ow_val * sW - pW; + } else { + col_valid[kk] = false; + col_b[kk] = 0u; + col_b_ic[kk] = 0u; + col_oh_arr[kk] = 0; + col_ow_arr[kk] = 0; + col_ih_base[kk] = 0; + col_iw_base[kk] = 0; + } + } + + let in_hw = params.IH * params.IW; + let in_w = params.IW; + let num_tiles = (K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k_start = t * TILE; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + let m = tile_row0 + arow; + let kk_idx = k_start + arow; + + var k_valid = false; + var k_ic = 0u; + var kh_dH = 0; + var kw_dW = 0; + if (kk_idx < K) { + k_valid = true; + k_ic = kk_idx / khw; + let rem = kk_idx % khw; + kh_dH = i32(rem / params.KW) * dH; + kw_dW = i32(rem % params.KW) * dW; + } + + let m_valid = m < M; + let w_row_base = m * K + k_start + local_id.x * RPT; + + var b_row: vec4; + for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) { + let col = local_id.x * RPT + kk; + let k_col = k_start + col; + if (m_valid && k_col < K) { + a_sub[arow][col] = weight[w_row_base + kk]; + } else { + a_sub[arow][col] = 0.0; + } + + if (k_valid && col_valid[kk]) { + let ih = col_ih_base[kk] + kh_dH; + let iw = col_iw_base[kk] + kw_dW; + if (ih >= 0 && ih < iH && iw >= 0 && iw < iW) { + let ch_idx = col_b[kk] + k_ic; + b_row[kk] = input[ch_idx * in_hw + u32(ih) * in_w + u32(iw)]; + } else { + b_row[kk] = 0.0; + } + } else { + b_row[kk] = 0.0; + } + } + b_sub[arow][local_id.x] = b_row; + } + workgroupBarrier(); + for (var k: u32 = 0u; k < TILE; k = k + 1u) { + let b_vec = b_sub[k][local_id.x]; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let aval = a_sub[tile_row + ir][k]; + acc[ir] = acc[ir] + aval * b_vec; + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let m = tile_row0 + tile_row + ir; + if (m < M) { + var bias_val = 0.0; + if (params.has_bias != 0u) { + bias_val = bias[m]; + } + let v = acc[ir] + vec4(bias_val); + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + let n = tile_col0 + tile_col + ic; + if (n < N) { + let b = col_b[ic]; + let oh = u32(col_oh_arr[ic]); + let ow = u32(col_ow_arr[ic]); + out[((b * params.OC + m) * params.OH + oh) * params.OW + ow] = v[ic]; + } + } + } + } +} +`; diff --git a/src/shaders/convolution/conv2d-tile32-f16.wgsl.ts b/src/shaders/convolution/conv2d-tile32-f16.wgsl.ts new file mode 100644 index 0000000..33e380a --- /dev/null +++ b/src/shaders/convolution/conv2d-tile32-f16.wgsl.ts @@ -0,0 +1,200 @@ +/** + * WGSL compute shader for 32x64 register-tiled 2D convolution GEMM in FP16 (4x4 per thread across 128 threads). + * Optimized for small channel counts (OC <= 32) to maximize wave occupancy and eliminate thread divergence. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV2D_TILE32_F16 = /* wgsl */ ` +enable f16; + +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array>; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var params: Params; + +var a_sub: array, 4>, 32>; +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(local_invocation_index) tid: u32 +) { + let tile_row0 = wg_id.y * 32u; + let tile_col0 = wg_id.x * 64u; + let M = params.OC; + let N = params.B * params.OH * params.OW; + let K = params.IC * params.KH * params.KW; + let K_vecs = K >> 2u; + let khw = params.KH * params.KW; + let ohw = params.OH * params.OW; + + let iH = i32(params.IH); + let iW = i32(params.IW); + let in_hw = params.IH * params.IW; + let in_w = params.IW; + let sH = i32(params.sH); + let sW = i32(params.sW); + let pH = i32(params.pH); + let pW = i32(params.pW); + let dH = i32(params.dH); + let dW = i32(params.dW); + + var acc0 = vec4(0.0h); + var acc1 = vec4(0.0h); + var acc2 = vec4(0.0h); + var acc3 = vec4(0.0h); + + let a_r = tid >> 2u; // 0..31 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < M; + + let b_r0 = tid >> 4u; // 0..7 + let b_c4_0 = tid & 15u; // 0..15 + let b_r1 = b_r0 + 8u; // 8..15 + let b_c4_1 = b_c4_0; + + let n_base0 = tile_col0 + (b_c4_0 << 2u); + + var col_b0: array; + var col_ih_base0: array; + var col_iw_base0: array; + var col_valid0: array; + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + let n = n_base0 + kk; + if (n < N) { + col_valid0[kk] = true; + let b_val = n / ohw; + let sp = n % ohw; + col_b0[kk] = b_val; + let oh = sp / params.OW; + let ow = sp % params.OW; + col_ih_base0[kk] = i32(oh) * sH - pH; + col_iw_base0[kk] = i32(ow) * sW - pW; + } else { + col_valid0[kk] = false; + col_b0[kk] = 0u; + col_ih_base0[kk] = 0; + col_iw_base0[kk] = 0; + } + } + + let num_tiles = (K + 15u) >> 4u; + for (var t = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t << 2u; + let k_start = t << 4u; + + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = weight[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0h); + } + + let bk0 = k_start + b_r0; + if (bk0 < K) { + let ic0 = bk0 / khw; + let rem0 = bk0 % khw; + let kh0 = i32(rem0 / params.KW) * dH; + let kw0 = i32(rem0 % params.KW) * dW; + + var b_vec0 = vec4(0.0h); + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + if (col_valid0[kk]) { + let ih0 = col_ih_base0[kk] + kh0; + let iw0 = col_iw_base0[kk] + kw0; + if (ih0 >= 0 && ih0 < iH && iw0 >= 0 && iw0 < iW) { + let in_idx0 = (col_b0[kk] * params.IC + ic0) * in_hw + u32(ih0) * in_w + u32(iw0); + b_vec0[kk] = input[in_idx0]; + } + } + } + b_sub[b_r0][b_c4_0] = b_vec0; + } else { + b_sub[b_r0][b_c4_0] = vec4(0.0h); + } + + let bk1 = k_start + b_r1; + if (bk1 < K) { + let ic1 = bk1 / khw; + let rem1 = bk1 % khw; + let kh1 = i32(rem1 / params.KW) * dH; + let kw1 = i32(rem1 % params.KW) * dW; + + var b_vec1 = vec4(0.0h); + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + if (col_valid0[kk]) { + let ih1 = col_ih_base0[kk] + kh1; + let iw1 = col_iw_base0[kk] + kw1; + if (ih1 >= 0 && ih1 < iH && iw1 >= 0 && iw1 < iW) { + let in_idx1 = (col_b0[kk] * params.IC + ic1) * in_hw + u32(ih1) * in_w + u32(iw1); + b_vec1[kk] = input[in_idx1]; + } + } + } + b_sub[b_r1][b_c4_1] = b_vec1; + } else { + b_sub[b_r1][b_c4_1] = vec4(0.0h); + } + + workgroupBarrier(); + + let thread_r = local_id.y << 2u; + let thread_c4 = local_id.x; + + for (var k4 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + let thread_r = local_id.y << 2u; + let out_r0 = tile_row0 + thread_r; + let out_c_base = tile_col0 + (local_id.x << 2u); + + var bias_val = vec4(0.0h); + if (params.has_bias != 0u) { + if (out_r0 + 0u < M) { bias_val.x = bias[out_r0 + 0u]; } + if (out_r0 + 1u < M) { bias_val.y = bias[out_r0 + 1u]; } + if (out_r0 + 2u < M) { bias_val.z = bias[out_r0 + 2u]; } + if (out_r0 + 3u < M) { bias_val.w = bias[out_r0 + 3u]; } + } + + acc0 = acc0 + vec4(bias_val.x); + acc1 = acc1 + vec4(bias_val.y); + acc2 = acc2 + vec4(bias_val.z); + acc3 = acc3 + vec4(bias_val.w); + + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + let out_c = out_c_base + kk; + if (out_c < N) { + if (out_r0 + 0u < M) { out[(out_r0 + 0u) * N + out_c] = acc0[kk]; } + if (out_r0 + 1u < M) { out[(out_r0 + 1u) * N + out_c] = acc1[kk]; } + if (out_r0 + 2u < M) { out[(out_r0 + 2u) * N + out_c] = acc2[kk]; } + if (out_r0 + 3u < M) { out[(out_r0 + 3u) * N + out_c] = acc3[kk]; } + } + } +} +`; diff --git a/src/shaders/convolution/conv2d-tile32.wgsl.ts b/src/shaders/convolution/conv2d-tile32.wgsl.ts new file mode 100644 index 0000000..79bb9da --- /dev/null +++ b/src/shaders/convolution/conv2d-tile32.wgsl.ts @@ -0,0 +1,204 @@ +/** + * WGSL compute shader for 32x64 register-tiled 2D convolution GEMM (4x4 per thread across 128 threads). + * Optimized for small channel counts (OC <= 32) to maximize wave occupancy and eliminate thread divergence. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV2D_TILE32 = /* wgsl */ ` +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array>; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var params: Params; + +// a_sub: 32 output channels x 16 K elements (4 vec4s) = 128 vec4s (2 KB) +var a_sub: array, 4>, 32>; +// b_sub: 16 K elements x 64 spatial elements (16 vec4s) = 256 vec4s (4 KB) +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(local_invocation_index) tid: u32 +) { + let tile_row0 = wg_id.y * 32u; + let tile_col0 = wg_id.x * 64u; + let M = params.OC; + let N = params.B * params.OH * params.OW; + let K = params.IC * params.KH * params.KW; + let K_vecs = K >> 2u; + let khw = params.KH * params.KW; + let ohw = params.OH * params.OW; + + let iH = i32(params.IH); + let iW = i32(params.IW); + let in_hw = params.IH * params.IW; + let in_w = params.IW; + let sH = i32(params.sH); + let sW = i32(params.sW); + let pH = i32(params.pH); + let pW = i32(params.pW); + let dH = i32(params.dH); + let dW = i32(params.dW); + + var acc0 = vec4(0.0); + var acc1 = vec4(0.0); + var acc2 = vec4(0.0); + var acc3 = vec4(0.0); + + // 128 threads cooperatively load 128 vec4s of A (32 rows x 4 vec4s) + let a_r = tid >> 2u; // 0..31 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < M; + + // 128 threads load 256 vec4s of B (16 rows x 16 vec4s) -> 2 passes per thread + let b_r0 = tid >> 4u; // 0..7 + let b_c4_0 = tid & 15u; // 0..15 + let b_r1 = b_r0 + 8u; // 8..15 + let b_c4_1 = b_c4_0; + + let n_base0 = tile_col0 + (b_c4_0 << 2u); + + var col_b0: array; + var col_ih_base0: array; + var col_iw_base0: array; + var col_valid0: array; + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + let n = n_base0 + kk; + if (n < N) { + col_valid0[kk] = true; + let b_val = n / ohw; + let sp = n % ohw; + col_b0[kk] = b_val; + let oh = sp / params.OW; + let ow = sp % params.OW; + col_ih_base0[kk] = i32(oh) * sH - pH; + col_iw_base0[kk] = i32(ow) * sW - pW; + } else { + col_valid0[kk] = false; + col_b0[kk] = 0u; + col_ih_base0[kk] = 0; + col_iw_base0[kk] = 0; + } + } + + let num_tiles = (K + 15u) >> 4u; + for (var t = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t << 2u; + let k_start = t << 4u; + + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = weight[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0); + } + + // Load first B slice (rows 0..7) + let bk0 = k_start + b_r0; + if (bk0 < K) { + let ic0 = bk0 / khw; + let rem0 = bk0 % khw; + let kh0 = i32(rem0 / params.KW) * dH; + let kw0 = i32(rem0 % params.KW) * dW; + + var b_vec0 = vec4(0.0); + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + if (col_valid0[kk]) { + let ih0 = col_ih_base0[kk] + kh0; + let iw0 = col_iw_base0[kk] + kw0; + if (ih0 >= 0 && ih0 < iH && iw0 >= 0 && iw0 < iW) { + let in_idx0 = (col_b0[kk] * params.IC + ic0) * in_hw + u32(ih0) * in_w + u32(iw0); + b_vec0[kk] = input[in_idx0]; + } + } + } + b_sub[b_r0][b_c4_0] = b_vec0; + } else { + b_sub[b_r0][b_c4_0] = vec4(0.0); + } + + // Load second B slice (rows 8..15) + let bk1 = k_start + b_r1; + if (bk1 < K) { + let ic1 = bk1 / khw; + let rem1 = bk1 % khw; + let kh1 = i32(rem1 / params.KW) * dH; + let kw1 = i32(rem1 % params.KW) * dW; + + var b_vec1 = vec4(0.0); + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + if (col_valid0[kk]) { + let ih1 = col_ih_base0[kk] + kh1; + let iw1 = col_iw_base0[kk] + kw1; + if (ih1 >= 0 && ih1 < iH && iw1 >= 0 && iw1 < iW) { + let in_idx1 = (col_b0[kk] * params.IC + ic1) * in_hw + u32(ih1) * in_w + u32(iw1); + b_vec1[kk] = input[in_idx1]; + } + } + } + b_sub[b_r1][b_c4_1] = b_vec1; + } else { + b_sub[b_r1][b_c4_1] = vec4(0.0); + } + + workgroupBarrier(); + + let thread_r = local_id.y << 2u; // 0..28 (step 4) + let thread_c4 = local_id.x; // 0..15 + + for (var k4 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + let thread_r = local_id.y << 2u; + let out_r0 = tile_row0 + thread_r; + let out_c_base = tile_col0 + (local_id.x << 2u); + + var bias_val = vec4(0.0); + if (params.has_bias != 0u) { + if (out_r0 + 0u < M) { bias_val.x = bias[out_r0 + 0u]; } + if (out_r0 + 1u < M) { bias_val.y = bias[out_r0 + 1u]; } + if (out_r0 + 2u < M) { bias_val.z = bias[out_r0 + 2u]; } + if (out_r0 + 3u < M) { bias_val.w = bias[out_r0 + 3u]; } + } + + acc0 = acc0 + vec4(bias_val.x); + acc1 = acc1 + vec4(bias_val.y); + acc2 = acc2 + vec4(bias_val.z); + acc3 = acc3 + vec4(bias_val.w); + + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + let out_c = out_c_base + kk; + if (out_c < N) { + if (out_r0 + 0u < M) { out[(out_r0 + 0u) * N + out_c] = acc0[kk]; } + if (out_r0 + 1u < M) { out[(out_r0 + 1u) * N + out_c] = acc1[kk]; } + if (out_r0 + 2u < M) { out[(out_r0 + 2u) * N + out_c] = acc2[kk]; } + if (out_r0 + 3u < M) { out[(out_r0 + 3u) * N + out_c] = acc3[kk]; } + } + } +} +`; diff --git a/src/shaders/convolution/conv2d-tile64-f16.wgsl.ts b/src/shaders/convolution/conv2d-tile64-f16.wgsl.ts new file mode 100644 index 0000000..3a3c956 --- /dev/null +++ b/src/shaders/convolution/conv2d-tile64-f16.wgsl.ts @@ -0,0 +1,181 @@ +/** + * WGSL compute shader for 64x64 register-tiled 2D convolution GEMM in float16 (4x4 per thread across 256 threads). + * High occupancy, coalesced 128-bit memory transactions, minimal register pressure for all GPUs. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV2D_TILE64_F16 = /* wgsl */ ` +enable f16; + +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array>; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var params: Params; + +// a_sub: 64 output channels x 16 K elements (4 vec4s) = 256 vec4s (2 KB) +var a_sub: array, 4>, 64>; +// b_sub: 16 K elements x 64 spatial elements (16 vec4s) = 256 vec4s (2 KB) +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(local_invocation_index) tid: u32 +) { + let tile_row0 = wg_id.y * 64u; + let tile_col0 = wg_id.x * 64u; + let M = params.OC; + let N = params.B * params.OH * params.OW; + let K = params.IC * params.KH * params.KW; + let K_vecs = K >> 2u; + let khw = params.KH * params.KW; + let ohw = params.OH * params.OW; + + let iH = i32(params.IH); + let iW = i32(params.IW); + let in_hw = params.IH * params.IW; + let in_w = params.IW; + let sH = i32(params.sH); + let sW = i32(params.sW); + let pH = i32(params.pH); + let pW = i32(params.pW); + let dH = i32(params.dH); + let dW = i32(params.dW); + + var acc0 = vec4(0.0h); + var acc1 = vec4(0.0h); + var acc2 = vec4(0.0h); + var acc3 = vec4(0.0h); + + // Cooperative load coordinates for 256 threads (1 vec4 weight + 1 vec4 input per thread) + let a_r = tid >> 2u; // 0..63 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < M; + + let b_r = tid >> 4u; // 0..15 (row in K tile) + let b_c4 = tid & 15u; // 0..15 (col in spatial tile) + let n_base = tile_col0 + (b_c4 << 2u); + + // Precompute invariant spatial coordinates for this thread's 4 input columns once + var col_b: array; + var col_ih_base: array; + var col_iw_base: array; + var col_valid: array; + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + let n = n_base + kk; + if (n < N) { + col_valid[kk] = true; + let b_val = n / ohw; + let sp = n % ohw; + col_b[kk] = b_val; + let oh = sp / params.OW; + let ow = sp % params.OW; + col_ih_base[kk] = i32(oh) * sH - pH; + col_iw_base[kk] = i32(ow) * sW - pW; + } else { + col_valid[kk] = false; + col_b[kk] = 0u; + col_ih_base[kk] = 0; + col_iw_base[kk] = 0; + } + } + + let num_tiles = (K + 15u) >> 4u; + for (var t = 0u; t < num_tiles; t = t + 1u) { + let k_start = t << 4u; + let k4_start = t << 2u; + + // 1. Coalesced 128-bit load of weights (256 threads load 256 vec4s) + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = weight[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0h); + } + + // 2. Coalesced load of input using precomputed spatial bases + let kk_idx = k_start + b_r; + var b_row = vec4(0.0h); + if (kk_idx < K) { + let k_ic = kk_idx / khw; + let rem = kk_idx % khw; + let kh = i32(rem / params.KW) * dH; + let kw = i32(rem % params.KW) * dW; + + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + if (col_valid[kk]) { + let ih = col_ih_base[kk] + kh; + let iw = col_iw_base[kk] + kw; + if (ih >= 0 && ih < iH && iw >= 0 && iw < iW) { + let ch_idx = col_b[kk] * params.IC + k_ic; + b_row[kk] = input[ch_idx * in_hw + u32(ih) * in_w + u32(iw)]; + } + } + } + } + b_sub[b_r][b_c4] = b_row; + workgroupBarrier(); + + // 3. Compute phase: 4x4 register accumulation per thread + let thread_r = local_id.y << 2u; + let thread_c4 = local_id.x; + + for (var k4 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + // 4. Output write: 4 rows x 4 cols per thread + let thread_r = local_id.y << 2u; + let out_c_base = tile_col0 + (local_id.x << 2u); + + for (var ir = 0u; ir < 4u; ir = ir + 1u) { + let m = tile_row0 + thread_r + ir; + if (m < M) { + var bias_val = 0.0h; + if (params.has_bias != 0u) { + bias_val = bias[m]; + } + var acc_v = acc0; + if (ir == 1u) { acc_v = acc1; } + else if (ir == 2u) { acc_v = acc2; } + else if (ir == 3u) { acc_v = acc3; } + acc_v = acc_v + vec4(bias_val); + + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + let n = out_c_base + kk; + if (n < N) { + let b_val = n / ohw; + let sp = n % ohw; + let oh = sp / params.OW; + let ow = sp % params.OW; + out[((b_val * params.OC + m) * params.OH + oh) * params.OW + ow] = acc_v[kk]; + } + } + } + } +} +`; diff --git a/src/shaders/convolution/conv2d-tile64.wgsl.ts b/src/shaders/convolution/conv2d-tile64.wgsl.ts new file mode 100644 index 0000000..f8ac116 --- /dev/null +++ b/src/shaders/convolution/conv2d-tile64.wgsl.ts @@ -0,0 +1,179 @@ +/** + * WGSL compute shader for 64x64 register-tiled 2D convolution GEMM (4x4 per thread across 256 threads). + * High occupancy, coalesced 128-bit memory transactions, minimal register pressure for all GPUs. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_CONV2D_TILE64 = /* wgsl */ ` +struct Params { + B: u32, IC: u32, IH: u32, IW: u32, + OC: u32, OH: u32, OW: u32, KH: u32, KW: u32, + sH: u32, sW: u32, pH: u32, pW: u32, dH: u32, dW: u32, + groups: u32, has_bias: u32, _p0: u32, _p1: u32, _p2: u32, +} +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array>; +@group(0) @binding(3) var bias: array; +@group(0) @binding(4) var params: Params; + +// a_sub: 64 output channels x 16 K elements (4 vec4s) = 256 vec4s (4 KB) +var a_sub: array, 4>, 64>; +// b_sub: 16 K elements x 64 spatial elements (16 vec4s) = 256 vec4s (4 KB) +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3, + @builtin(local_invocation_index) tid: u32 +) { + let tile_row0 = wg_id.y * 64u; + let tile_col0 = wg_id.x * 64u; + let M = params.OC; + let N = params.B * params.OH * params.OW; + let K = params.IC * params.KH * params.KW; + let K_vecs = K >> 2u; + let khw = params.KH * params.KW; + let ohw = params.OH * params.OW; + + let iH = i32(params.IH); + let iW = i32(params.IW); + let in_hw = params.IH * params.IW; + let in_w = params.IW; + let sH = i32(params.sH); + let sW = i32(params.sW); + let pH = i32(params.pH); + let pW = i32(params.pW); + let dH = i32(params.dH); + let dW = i32(params.dW); + + var acc0 = vec4(0.0); + var acc1 = vec4(0.0); + var acc2 = vec4(0.0); + var acc3 = vec4(0.0); + + // Cooperative load coordinates for 256 threads (1 vec4 weight + 1 vec4 input per thread) + let a_r = tid >> 2u; // 0..63 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < M; + + let b_r = tid >> 4u; // 0..15 (row in K tile) + let b_c4 = tid & 15u; // 0..15 (col in spatial tile) + let n_base = tile_col0 + (b_c4 << 2u); + + // Precompute invariant spatial coordinates for this thread's 4 input columns once + var col_b: array; + var col_ih_base: array; + var col_iw_base: array; + var col_valid: array; + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + let n = n_base + kk; + if (n < N) { + col_valid[kk] = true; + let b_val = n / ohw; + let sp = n % ohw; + col_b[kk] = b_val; + let oh = sp / params.OW; + let ow = sp % params.OW; + col_ih_base[kk] = i32(oh) * sH - pH; + col_iw_base[kk] = i32(ow) * sW - pW; + } else { + col_valid[kk] = false; + col_b[kk] = 0u; + col_ih_base[kk] = 0; + col_iw_base[kk] = 0; + } + } + + let num_tiles = (K + 15u) >> 4u; + for (var t = 0u; t < num_tiles; t = t + 1u) { + let k_start = t << 4u; + let k4_start = t << 2u; + + // 1. Coalesced 128-bit load of weights (256 threads load 256 vec4s) + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = weight[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0); + } + + // 2. Coalesced load of input using precomputed spatial bases + let kk_idx = k_start + b_r; + var b_row = vec4(0.0); + if (kk_idx < K) { + let k_ic = kk_idx / khw; + let rem = kk_idx % khw; + let kh = i32(rem / params.KW) * dH; + let kw = i32(rem % params.KW) * dW; + + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + if (col_valid[kk]) { + let ih = col_ih_base[kk] + kh; + let iw = col_iw_base[kk] + kw; + if (ih >= 0 && ih < iH && iw >= 0 && iw < iW) { + let ch_idx = col_b[kk] * params.IC + k_ic; + b_row[kk] = input[ch_idx * in_hw + u32(ih) * in_w + u32(iw)]; + } + } + } + } + b_sub[b_r][b_c4] = b_row; + workgroupBarrier(); + + // 3. Compute phase: 4x4 register accumulation per thread + let thread_r = local_id.y << 2u; + let thread_c4 = local_id.x; + + for (var k4 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + // 4. Output write: 4 rows x 4 cols per thread + let thread_r = local_id.y << 2u; + let out_c_base = tile_col0 + (local_id.x << 2u); + + for (var ir = 0u; ir < 4u; ir = ir + 1u) { + let m = tile_row0 + thread_r + ir; + if (m < M) { + var bias_val = 0.0; + if (params.has_bias != 0u) { + bias_val = bias[m]; + } + var acc_v = acc0; + if (ir == 1u) { acc_v = acc1; } + else if (ir == 2u) { acc_v = acc2; } + else if (ir == 3u) { acc_v = acc3; } + acc_v = acc_v + vec4(bias_val); + + for (var kk = 0u; kk < 4u; kk = kk + 1u) { + let n = out_c_base + kk; + if (n < N) { + let b_val = n / ohw; + let sp = n % ohw; + let oh = sp / params.OW; + let ow = sp % params.OW; + out[((b_val * params.OC + m) * params.OH + oh) * params.OW + ow] = acc_v[kk]; + } + } + } + } +} +`; diff --git a/src/shaders/convolution/conv2d.wgsl.ts b/src/shaders/convolution/conv2d.wgsl.ts new file mode 100644 index 0000000..842363b --- /dev/null +++ b/src/shaders/convolution/conv2d.wgsl.ts @@ -0,0 +1,90 @@ +/** + * Direct 2D convolution compute shader (fallback for general grouped cases). + * + * Source: ExecuTorch (backends/webgpu/runtime/ops/et_vk_conv2d/conv2d.wgsl) + * GitHub: https://github.com/pytorch/executorch/blob/main/backends/webgpu/runtime/ops/et_vk_conv2d/conv2d.wgsl + * License: BSD-3-Clause (Copyright (c) Meta Platforms, Inc. and affiliates) + */ +export const SHADER_CONV2D = /* wgsl */ ` +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var input: array; +@group(0) @binding(2) var weight: array; +@group(0) @binding(3) var bias: array; + +struct Params { + B: u32, + IC: u32, + IH: u32, + IW: u32, + OC: u32, + OH: u32, + OW: u32, + KH: u32, + KW: u32, + sH: u32, + sW: u32, + pH: u32, + pW: u32, + dH: u32, + dW: u32, + groups: u32, + has_bias: u32, + _p0: u32, + _p1: u32, + _p2: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; // = count_x * wg_size; set by host for 2D-spill + +// Direct 2D convolution (non-transposed), NCHW row-major, fp32. ONE thread per +// (b, oc, oh, ow) output element. Supports general stride/padding/dilation and +// groups. input [B,IC,IH,IW]; weight [OC, IC/groups, KH, KW]; bias [OC] (gated). +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.B * params.OC * params.OH * params.OW; + let i = gid.y * stride_x + gid.x; + if (i >= total) { + return; + } + let ow = i % params.OW; + let oh = (i / params.OW) % params.OH; + let oc = (i / (params.OW * params.OH)) % params.OC; + let b = i / (params.OW * params.OH * params.OC); + + let icpg = params.IC / params.groups; // input channels per group + let ocpg = params.OC / params.groups; // output channels per group + let g = oc / ocpg; + let ic0 = g * icpg; + + var acc: f32 = 0.0; + if (params.has_bias != 0u) { + acc = bias[oc]; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + for (var icg: u32 = 0u; icg < icpg; icg = icg + 1u) { + let ic = ic0 + icg; + let in_c_base = (b * params.IC + ic) * params.IH; // *IW added per-row below + let w_c_base = (oc * icpg + icg) * params.KH; // *KW added per-row below + for (var kh: u32 = 0u; kh < params.KH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + let in_row = (in_c_base + u32(ih)) * params.IW; + let w_row = (w_c_base + kh) * params.KW; + for (var kw: u32 = 0u; kw < params.KW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + acc = acc + input[in_row + u32(iw)] * weight[w_row + kw]; + } + } + } + out[i] = acc; +} +`; diff --git a/src/shaders/convolution/index.ts b/src/shaders/convolution/index.ts new file mode 100644 index 0000000..f4e75a5 --- /dev/null +++ b/src/shaders/convolution/index.ts @@ -0,0 +1,337 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { SHADER_CONV2D } from './conv2d.wgsl'; +import { SHADER_CONV2D_GEMM } from './conv2d-gemm.wgsl'; +import { SHADER_CONV2D_1X1 } from './conv2d-1x1.wgsl'; +import { SHADER_CONV2D_1X1_VEC4 } from './conv2d-1x1-vec4.wgsl'; +import { SHADER_CONV2D_1X1_64X64 } from './conv2d-1x1-tile64.wgsl'; +import { SHADER_CONV2D_1X1_64X64_F16 } from './conv2d-1x1-tile64-f16.wgsl'; +import { SHADER_CONV2D_TILE64 } from './conv2d-tile64.wgsl'; +import { SHADER_CONV2D_TILE32 } from './conv2d-tile32.wgsl'; +import { SHADER_CONV2D_1X1_TILE32 } from './conv2d-1x1-tile32.wgsl'; +import { SHADER_CONV2D_1X1_TILE32_F16 } from './conv2d-1x1-tile32-f16.wgsl'; +import { SHADER_CONV2D_DEPTHWISE } from './conv2d-depthwise.wgsl'; +import { SHADER_CONV2D_DEPTHWISE_3X3 } from './conv2d-depthwise-3x3.wgsl'; +import { SHADER_CONV_TRANSPOSE2D } from './conv-transpose2d.wgsl'; +import { SHADER_CONV_TRANSPOSE2D_F16 } from './conv-transpose2d-f16.wgsl'; +import { SHADER_CONV2D_F16 } from './conv2d-f16.wgsl'; +import { SHADER_CONV2D_GEMM_F16 } from './conv2d-gemm-f16.wgsl'; +import { SHADER_CONV2D_TILE64_F16 } from './conv2d-tile64-f16.wgsl'; +import { SHADER_CONV2D_TILE32_F16 } from './conv2d-tile32-f16.wgsl'; +import { SHADER_CONV2D_DEPTHWISE_F16 } from './conv2d-depthwise-f16.wgsl'; +import { SHADER_CONV2D_DEPTHWISE_3X3_F16 } from './conv2d-depthwise-3x3-f16.wgsl'; + +const name = 'aten.convolution.default'; + +/** + * Positional argument tuple for `aten::convolution.default`: + * `[in, weight, bias, stride, padding, dilation, transposed, output_padding, groups, out]` + */ +export type ConvolutionArgs = readonly [ + inTensor: Tensor, + weight: Tensor, + bias: Tensor | undefined, + stride: readonly number[], + padding: readonly number[], + dilation: readonly number[], + transposed: boolean, + outputPadding: readonly number[], + groups: number, + outTensor: Tensor, +]; + +function parseHw(v: readonly number[], argName: string): [number, number] { + if (v.length === 1) return [v[0], v[0]]; + if (v.length === 2) return [v[0], v[1]]; + throw new Error(`${name}: ${argName} must be 1 or 2 elements`); +} + +/** + * Internal dispatch builder for `aten::convolution.default`. + * + * @param ctx The execution context. + * @param args Positional arguments for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: ConvolutionArgs): void { + const [ + inTensor, + weight, + bias, + stride, + padding, + dilation, + transposed, + outputPadding, + groups, + outTensor, + ] = args; + + if ( + (inTensor.dtype !== 'float32' && inTensor.dtype !== 'float16') || + weight.dtype !== inTensor.dtype || + outTensor.dtype !== inTensor.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + if (bias !== undefined && bias.dtype !== inTensor.dtype) { + throw new Error(`${name}: bias dtype must match input dtype`); + } + if (inTensor.shape.length !== 4 || weight.shape.length !== 4 || outTensor.shape.length !== 4) { + throw new Error(`${name}: expected 4D tensors for input, weight, and output`); + } + + const [sH, sW] = parseHw(stride, 'stride'); + const [pH, pW] = parseHw(padding, 'padding'); + const [dH, dW] = parseHw(dilation, 'dilation'); + const [opH, opW] = parseHw(outputPadding, 'output_padding'); + + const [b, ic, ih, iw] = inTensor.shape as [number, number, number, number]; + const hasBias = bias !== undefined; + + let oc: number; + let kh: number; + let kw: number; + + if (transposed) { + // weight layout for transposed conv: [IC, OC / groups, KH, KW] + kh = weight.shape[2]; + kw = weight.shape[3]; + const ocpg = weight.shape[1]; + oc = ocpg * groups; + + if (weight.shape[0] !== ic) { + throw new Error(`${name}: weight dim0 !== IC (${weight.shape[0]} != ${ic})`); + } + if (opH >= sH || opW >= sW) { + throw new Error(`${name}: output_padding >= stride`); + } + + const expectedOh = (ih - 1) * sH - 2 * pH + dH * (kh - 1) + opH + 1; + const expectedOw = (iw - 1) * sW - 2 * pW + dW * (kw - 1) + opW + 1; + if (outTensor.shape[2] !== expectedOh || outTensor.shape[3] !== expectedOw) { + throw new Error(`${name}: expected output spatial [${expectedOh}, ${expectedOw}]`); + } + } else { + // weight layout for direct conv: [OC, IC / groups, KH, KW] + oc = weight.shape[0]; + kh = weight.shape[2]; + kw = weight.shape[3]; + + if (outputPadding.some((v) => v !== 0)) { + throw new Error(`${name}: non-zero output_padding unsupported for non-transposed conv`); + } + if (weight.shape[1] !== ic / groups) { + throw new Error(`${name}: weight in-ch (${weight.shape[1]}) != IC/groups (${ic / groups})`); + } + + const ohNum = ih + 2 * pH - dH * (kh - 1) - 1; + const owNum = iw + 2 * pW - dW * (kw - 1) - 1; + if (ohNum < 0 || owNum < 0) { + throw new Error(`${name}: invalid geometry (kernel larger than input)`); + } + const expectedOh = Math.floor(ohNum / sH) + 1; + const expectedOw = Math.floor(owNum / sW) + 1; + if (outTensor.shape[2] !== expectedOh || outTensor.shape[3] !== expectedOw) { + throw new Error(`${name}: expected output spatial [${expectedOh}, ${expectedOw}]`); + } + } + + const [, , oh, ow] = outTensor.shape as [number, number, number, number]; + + const device = ctx.device; + const isF16 = inTensor.dtype === 'float16'; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + const wgSize = 256; + + // ConvParams (80 B): input/output geometry, kernel, stride, pad, dilation, flags + const buffer = new ArrayBuffer(80); + // prettier-ignore + new Uint32Array(buffer).set([ + b, ic, ih, iw, + oc, oh, ow, + kh, kw, + sH, sW, + pH, pW, + dH, dW, + groups, hasBias ? 1 : 0, + ]); + + const uniformBuffer = ctx.uniformBuffer(buffer); + const biasBuffer = hasBias ? bias.buffer : weight.buffer; + + // Variant selection: 1x1 convolutions use the pointwise GEMM kernel (no im2col), + // with 128-bit vec4 memory optimization when dimensions are divisible by 4. + // General dense convolutions use implicit-im2col tiled GEMM, grouped/depthwise + // convs through the direct kernel, and transposed convs through the gather kernel. + const is1x1 = + !transposed && + kh === 1 && + kw === 1 && + sH === 1 && + sW === 1 && + pH === 0 && + pW === 0 && + dH === 1 && + dW === 1; + const is1x1Vec4 = is1x1 && ic % 4 === 0 && (oh * ow) % 4 === 0 && oc % 4 === 0; + const is1x1Tile32 = is1x1Vec4 && b === 1 && oc <= 32 && oh * ow >= 64 && ic >= 16; + const is1x1Tile64 = is1x1Vec4 && b === 1 && oc > 32 && oh * ow >= 64 && ic >= 16; + const isDenseTile32 = + !transposed && + groups === 1 && + oc <= 32 && + oc >= 8 && + b * oh * ow >= 64 && + (ic * kh * kw) % 4 === 0; + const isDenseTile64 = + !transposed && groups === 1 && oc > 32 && b * oh * ow >= 64 && (ic * kh * kw) % 4 === 0; + + const isDepthwise = + !transposed && + groups === oc && + ic === oc && + weight.shape[1] === 1 && + kh * kw <= 49 && + b * oc <= 65535; + + const isDepthwise3x3 = + isDepthwise && + kh === 3 && + kw === 3 && + sH === 1 && + sW === 1 && + pH === 1 && + pW === 1 && + dH === 1 && + dW === 1; + + let code = isF16 ? SHADER_CONV2D_F16 : SHADER_CONV2D; + if (transposed) { + code = isF16 ? SHADER_CONV_TRANSPOSE2D_F16 : SHADER_CONV_TRANSPOSE2D; + } else if (isDepthwise) { + if (isDepthwise3x3) { + code = isF16 ? SHADER_CONV2D_DEPTHWISE_3X3_F16 : SHADER_CONV2D_DEPTHWISE_3X3; + } else { + code = isF16 ? SHADER_CONV2D_DEPTHWISE_F16 : SHADER_CONV2D_DEPTHWISE; + } + } else if (groups === 1) { + if (is1x1Tile32) { + code = isF16 ? SHADER_CONV2D_1X1_TILE32_F16 : SHADER_CONV2D_1X1_TILE32; + } else if (is1x1Tile64) { + code = isF16 ? SHADER_CONV2D_1X1_64X64_F16 : SHADER_CONV2D_1X1_64X64; + } else if (isDenseTile32) { + code = isF16 ? SHADER_CONV2D_TILE32_F16 : SHADER_CONV2D_TILE32; + } else if (isDenseTile64) { + code = isF16 ? SHADER_CONV2D_TILE64_F16 : SHADER_CONV2D_TILE64; + } else if (is1x1Vec4 && !isF16) { + code = SHADER_CONV2D_1X1_VEC4; + } else if (is1x1 && !isF16) { + code = SHADER_CONV2D_1X1; + } else { + code = isF16 ? SHADER_CONV2D_GEMM_F16 : SHADER_CONV2D_GEMM; + } + } + + let workgroupCountX: number; + let workgroupCountY: number; + let workgroupCountZ = 1; + let constants: Record; + + if (isDepthwise) { + // Spatial 16x16 workgroup per (batch, oc) channel slice with 2x2 register tiling (32x32 per workgroup) + workgroupCountX = Math.ceil(ow / 32); + workgroupCountY = Math.ceil(oh / 32); + workgroupCountZ = b * oc; + constants = {}; + } else if (transposed || groups !== 1) { + // Direct / transposed kernels: one thread per output element, 2D spill + const totalWorkgroups = Math.ceil(outTensor.numel / wgSize); + workgroupCountX = Math.min(totalWorkgroups, 65535); + workgroupCountY = Math.ceil(totalWorkgroups / 65535); + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + constants = { wg_size: wgSize, stride_x: workgroupCountX * wgSize }; + } else if (is1x1Tile32) { + // 32x64 register-tiled GEMM for pointwise 1x1 conv + workgroupCountX = Math.ceil((oh * ow) / 64); + workgroupCountY = Math.ceil(oc / 32); + constants = {}; + } else if (is1x1Tile64) { + // 64x64 register-tiled GEMM for pointwise 1x1 conv + workgroupCountX = Math.ceil((oh * ow) / 64); + workgroupCountY = Math.ceil(oc / 64); + constants = {}; + } else if (isDenseTile32) { + // 32x64 register-tiled GEMM for general dense convolution + workgroupCountX = Math.ceil((b * oh * ow) / 64); + workgroupCountY = Math.ceil(oc / 32); + constants = {}; + } else if (isDenseTile64) { + // 64x64 register-tiled GEMM for general dense convolution + workgroupCountX = Math.ceil((b * oh * ow) / 64); + workgroupCountY = Math.ceil(oc / 64); + constants = {}; + } else { + // GEMM kernel: fixed 8x8 workgroup covering a 32x32 output tile; + // grid over (N = B*OH*OW cols, M = OC rows) + const tile = 32; + workgroupCountX = Math.ceil((b * oh * ow) / tile); + workgroupCountY = Math.ceil(oc / tile); + if (workgroupCountX > 65535 || workgroupCountY > 65535) { + throw new Error(`${name}: GEMM tile grid exceeds dispatch limit`); + } + constants = {}; + } + + const bundle = createComputeBundle( + device, + code, + [ + { binding: 0, buffer: outTensor.buffer }, + { binding: 1, buffer: inTensor.buffer }, + { binding: 2, buffer: weight.buffer }, + { binding: 3, buffer: biasBuffer }, + { binding: 4, buffer: uniformBuffer }, + ], + constants + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + workgroupCountZ, + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::convolution.default`. + * + * Implements 2D dense and grouped convolutions as well as 2D transposed + * convolutions. Dense convolutions are evaluated as implicit-im2col tiled GEMM + * (1.1-2.4x over the direct kernel in the ExecuTorch reference), grouped and + * depthwise convolutions via the direct kernel. + */ +export const convolution: Shader = { + name, + code: { + conv2d: SHADER_CONV2D, + conv2dGemm: SHADER_CONV2D_GEMM, + conv2d1x1: SHADER_CONV2D_1X1, + conv2d1x1Vec4: SHADER_CONV2D_1X1_VEC4, + conv2d1x1Tile64: SHADER_CONV2D_1X1_64X64, + conv2dTile64: SHADER_CONV2D_TILE64, + conv2dDepthwise3x3: SHADER_CONV2D_DEPTHWISE, + convTranspose2d: SHADER_CONV_TRANSPOSE2D, + conv2dF16: SHADER_CONV2D_F16, + conv2dGemmF16: SHADER_CONV2D_GEMM_F16, + conv2dTile64F16: SHADER_CONV2D_TILE64_F16, + conv2dDepthwiseF16: SHADER_CONV2D_DEPTHWISE_F16, + }, + recordIn, +}; diff --git a/src/shaders/hardswish/hardswish-f16.wgsl.ts b/src/shaders/hardswish/hardswish-f16.wgsl.ts new file mode 100644 index 0000000..6427f49 --- /dev/null +++ b/src/shaders/hardswish/hardswish-f16.wgsl.ts @@ -0,0 +1,31 @@ +/** + * WGSL compute shader for FP16 Hardswish activation (scalar f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_HARDSWISH_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + let x = input[idx]; + output[idx] = clamp(x + 3.0h, 0.0h, 6.0h) * x / 6.0h; +} +`; diff --git a/src/shaders/hardswish/hardswish-vec4-f16.wgsl.ts b/src/shaders/hardswish/hardswish-vec4-f16.wgsl.ts new file mode 100644 index 0000000..92c8a3c --- /dev/null +++ b/src/shaders/hardswish/hardswish-vec4-f16.wgsl.ts @@ -0,0 +1,31 @@ +/** + * WGSL compute shader for FP16 Hardswish activation (vec4 64-bit/128-bit vectorized). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_HARDSWISH_VEC4_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input: array>; +@group(0) @binding(1) var output: array>; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + let x = input[idx]; + output[idx] = clamp(x + 3.0h, vec4(0.0h), vec4(6.0h)) * x / 6.0h; +} +`; diff --git a/src/shaders/hardswish/hardswish-vec4.wgsl.ts b/src/shaders/hardswish/hardswish-vec4.wgsl.ts new file mode 100644 index 0000000..e8d51a9 --- /dev/null +++ b/src/shaders/hardswish/hardswish-vec4.wgsl.ts @@ -0,0 +1,29 @@ +/** + * WGSL compute shader for Hardswish activation (vec4 128-bit vectorized). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_HARDSWISH_VEC4 = /* wgsl */ ` +@group(0) @binding(0) var input: array>; +@group(0) @binding(1) var output: array>; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + let x = input[idx]; + output[idx] = clamp(x + 3.0, vec4(0.0), vec4(6.0)) * x / 6.0; +} +`; diff --git a/src/shaders/hardswish/hardswish.wgsl.ts b/src/shaders/hardswish/hardswish.wgsl.ts new file mode 100644 index 0000000..46924f3 --- /dev/null +++ b/src/shaders/hardswish/hardswish.wgsl.ts @@ -0,0 +1,29 @@ +/** + * WGSL compute shader for Hardswish activation (scalar f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_HARDSWISH = /* wgsl */ ` +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + let x = input[idx]; + output[idx] = clamp(x + 3.0, 0.0, 6.0) * x / 6.0; +} +`; diff --git a/src/shaders/hardswish/index.ts b/src/shaders/hardswish/index.ts new file mode 100644 index 0000000..9916bf9 --- /dev/null +++ b/src/shaders/hardswish/index.ts @@ -0,0 +1,95 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { SHADER_HARDSWISH } from './hardswish.wgsl'; +import { SHADER_HARDSWISH_VEC4 } from './hardswish-vec4.wgsl'; +import { SHADER_HARDSWISH_F16 } from './hardswish-f16.wgsl'; +import { SHADER_HARDSWISH_VEC4_F16 } from './hardswish-vec4-f16.wgsl'; + +const name = 'aten.hardswish.default'; + +/** + * Positional argument tuple for `aten::hardswish.default`: + * `[inTensor, outTensor]` + */ +export type HardswishArgs = readonly [inTensor: Tensor, outTensor: Tensor]; + +/** + * Internal dispatch builder for `aten::hardswish.default`. + * + * @param ctx The execution context. + * @param args Positional arguments for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, [inTensor, outTensor]: HardswishArgs): void { + if ( + (inTensor.dtype !== 'float32' && inTensor.dtype !== 'float16') || + outTensor.dtype !== inTensor.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + if (inTensor.numel !== outTensor.numel) { + throw new Error(`${name}: input numel != output numel`); + } + + const device = ctx.device; + const isF16 = inTensor.dtype === 'float16'; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + const wgSize = 256; + const isVec4 = outTensor.numel % 4 === 0; + const elements = isVec4 ? outTensor.numel / 4 : outTensor.numel; + + let code = isVec4 ? SHADER_HARDSWISH_VEC4 : SHADER_HARDSWISH; + if (isF16) { + code = isVec4 ? SHADER_HARDSWISH_VEC4_F16 : SHADER_HARDSWISH_F16; + } + + const totalWorkgroups = Math.ceil(elements / wgSize); + const workgroupCountX = Math.min(totalWorkgroups, 65535); + const workgroupCountY = Math.ceil(totalWorkgroups / 65535); + + const buffer = new ArrayBuffer(16); + new Uint32Array(buffer)[0] = elements; + const uniformBuffer = ctx.uniformBuffer(buffer); + + const bundle = createComputeBundle( + device, + code, + [ + { binding: 0, buffer: inTensor.buffer }, + { binding: 1, buffer: outTensor.buffer }, + { binding: 2, buffer: uniformBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize } + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::hardswish.default`. + * + * Computes elementwise `out = x * clamp(x + 3, 0, 6) / 6`. + * Automatically vectorizes using 128-bit vec4 or vec4 when tensor size is divisible by 4. + * Supports float32 and float16 with the `shader-f16` WebGPU feature extension. + */ +export const hardswish: Shader = { + name, + code: { + scalar: SHADER_HARDSWISH, + vec4: SHADER_HARDSWISH_VEC4, + scalarF16: SHADER_HARDSWISH_F16, + vec4F16: SHADER_HARDSWISH_VEC4_F16, + }, + recordIn, +}; diff --git a/src/shaders/max-pool-2d/index.ts b/src/shaders/max-pool-2d/index.ts new file mode 100644 index 0000000..8ec6230 --- /dev/null +++ b/src/shaders/max-pool-2d/index.ts @@ -0,0 +1,177 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { SHADER_MAX_POOL2D } from './max-pool-2d.wgsl'; +import { SHADER_MAX_POOL2D_SPATIAL } from './max-pool-2d-spatial.wgsl'; +import { SHADER_MAX_POOL2D_F16 } from './max-pool-2d-f16.wgsl'; +import { SHADER_MAX_POOL2D_SPATIAL_F16 } from './max-pool-2d-spatial-f16.wgsl'; + +const name = 'aten.max_pool2d_with_indices.default'; + +/** + * Positional argument tuple for `aten::max_pool2d_with_indices.default`: + * `[in, kernel_size, stride, padding, dilation, ceil_mode, outs]` + * where `outs` is `[values, indices]`. + */ +export type MaxPool2dArgs = readonly [ + inTensor: Tensor, + kernelSize: readonly number[], + stride: readonly number[], + padding: readonly number[], + dilation: readonly number[], + ceilMode: boolean, + outs: readonly Tensor[], +]; + +function parseHw(v: readonly number[], argName: string): [number, number] { + if (v.length === 1) return [v[0], v[0]]; + if (v.length === 2) return [v[0], v[1]]; + throw new Error(`${name}: ${argName} must be 1 or 2 elements`); +} + +/** + * Internal dispatch builder for `aten::max_pool2d_with_indices.default`. + * + * @param ctx The execution context. + * @param args Positional argument tuple for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: MaxPool2dArgs): void { + const [inTensor, kernelSize, stride, padding, dilation, ceilMode, outs] = args; + + if (outs.length !== 2) { + throw new Error(`${name}: expected 2 outputs [values, indices], got ${outs.length}`); + } + const [outVals, outIdx] = outs; + + if ( + (inTensor.dtype !== 'float32' && inTensor.dtype !== 'float16') || + outVals.dtype !== inTensor.dtype + ) { + throw new Error(`${name}: values tensors must have matching float32 or float16 dtype`); + } + if (outIdx.dtype !== 'int32') { + throw new Error(`${name}: indices output must be int32`); + } + if (inTensor.shape.length !== 4) { + throw new Error(`${name}: input must be 4D NCHW`); + } + + const device = ctx.device; + const isF16 = inTensor.dtype === 'float16'; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + const [b, c, ih, iw] = inTensor.shape as [number, number, number, number]; + const [kH, kW] = parseHw(kernelSize, 'kernel_size'); + const [sH, sW] = parseHw(stride, 'stride'); + const [pH, pW] = parseHw(padding, 'padding'); + const [dH, dW] = parseHw(dilation, 'dilation'); + + const ohNum = ih + 2 * pH - (dH * (kH - 1) + 1); + const owNum = iw + 2 * pW - (dW * (kW - 1) + 1); + const oh = Math.max(0, (ceilMode ? Math.ceil(ohNum / sH) : Math.floor(ohNum / sH)) + 1); + const ow = Math.max(0, (ceilMode ? Math.ceil(owNum / sW) : Math.floor(owNum / sW)) + 1); + + if ( + outVals.shape.length !== 4 || + outVals.shape[0] !== b || + outVals.shape[1] !== c || + outVals.shape[2] !== oh || + outVals.shape[3] !== ow + ) { + throw new Error(`${name}: values output shape ${outVals.shape} != [${b}, ${c}, ${oh}, ${ow}]`); + } + if ( + outIdx.shape.length !== 4 || + outIdx.shape[0] !== b || + outIdx.shape[1] !== c || + outIdx.shape[2] !== oh || + outIdx.shape[3] !== ow + ) { + throw new Error(`${name}: indices output shape ${outIdx.shape} != [${b}, ${c}, ${oh}, ${ow}]`); + } + + // PoolParams: input/output geometry, kernel, stride, pad, dilation + // prettier-ignore + const paramsBuffer = ctx.uniformBuffer( + new Uint32Array([ + b, c, ih, iw, + oh, ow, + kH, kW, + sH, sW, + pH, pW, + dH, dW, + 0, 0, + ]) + ); + + // Use 2D spatial grid when workgroup dimensions fit within limits (eliminates unraveling divisions) + if (b * c <= 65535 && Math.ceil(ow / 16) <= 65535 && Math.ceil(oh / 16) <= 65535) { + const bundle = createComputeBundle( + device, + isF16 ? SHADER_MAX_POOL2D_SPATIAL_F16 : SHADER_MAX_POOL2D_SPATIAL, + [ + { binding: 0, buffer: outVals.buffer }, + { binding: 1, buffer: inTensor.buffer }, + { binding: 2, buffer: paramsBuffer }, + { binding: 3, buffer: outIdx.buffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { write_indices: 1 } + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX: Math.ceil(ow / 16), + workgroupCountY: Math.ceil(oh / 16), + workgroupCountZ: b * c, + }); + return; + } + + const wgSize = 256; + const total = b * c * oh * ow; + const workgroupCountX = Math.min(Math.ceil(total / wgSize), 65535); + + const bundle = createComputeBundle( + device, + isF16 ? SHADER_MAX_POOL2D_F16 : SHADER_MAX_POOL2D, + [ + { binding: 0, buffer: outVals.buffer }, + { binding: 1, buffer: inTensor.buffer }, + { binding: 2, buffer: paramsBuffer }, + { binding: 3, buffer: outIdx.buffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize, stride_x: workgroupCountX * wgSize, write_indices: 1 } + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY: Math.ceil(total / wgSize / 65535), + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::max_pool2d_with_indices.default`. + * + * Computes max pooling over a 2D NCHW input — windowed gather max per output + * element, plus flat in-window argmax indices as an int32 tensor. + * Supports float32 and float16 with the `shader-f16` WebGPU feature extension. + */ +export const maxPool2d: Shader = { + name, + code: { + default: SHADER_MAX_POOL2D, + spatial: SHADER_MAX_POOL2D_SPATIAL, + defaultF16: SHADER_MAX_POOL2D_F16, + spatialF16: SHADER_MAX_POOL2D_SPATIAL_F16, + }, + recordIn, +}; diff --git a/src/shaders/max-pool-2d/max-pool-2d-f16.wgsl.ts b/src/shaders/max-pool-2d/max-pool-2d-f16.wgsl.ts new file mode 100644 index 0000000..e648bd7 --- /dev/null +++ b/src/shaders/max-pool-2d/max-pool-2d-f16.wgsl.ts @@ -0,0 +1,79 @@ +/** + * WGSL compute shader for 2D max pooling with 1D-flattened grid (f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_MAX_POOL2D_F16 = /* wgsl */ ` +enable f16; + +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + kH: u32, + kW: u32, + sH: u32, + sW: u32, + pH: u32, + pW: u32, + dH: u32, + dW: u32, + _p0: u32, + _p1: u32, +} + +@group(0) @binding(0) var out_vals: array; +@group(0) @binding(1) var inp: array; +@group(0) @binding(2) var params: Params; +@group(0) @binding(3) var out_idx: array; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; +override write_indices: u32 = 0u; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.N * params.C * params.OH * params.OW; + let i = gid.y * stride_x + gid.x; + if (i >= total) { + return; + } + let ow = i % params.OW; + let oh = (i / params.OW) % params.OH; + let c = (i / (params.OW * params.OH)) % params.C; + let n = i / (params.OW * params.OH * params.C); + + let iH = i32(params.IH); + let iW = i32(params.IW); + let in_c_base = (n * params.C + c) * params.IH; + + var best: f16 = -65504.0h; + var best_idx: i32 = 0; + for (var kh: u32 = 0u; kh < params.kH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + let in_row = (in_c_base + u32(ih)) * params.IW; + for (var kw: u32 = 0u; kw < params.kW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + let v = inp[in_row + u32(iw)]; + if (v > best) { + best = v; + best_idx = ih * iW + iw; + } + } + } + out_vals[i] = best; + if (write_indices != 0u) { + out_idx[i] = best_idx; + } +} +`; diff --git a/src/shaders/max-pool-2d/max-pool-2d-spatial-f16.wgsl.ts b/src/shaders/max-pool-2d/max-pool-2d-spatial-f16.wgsl.ts new file mode 100644 index 0000000..71968c8 --- /dev/null +++ b/src/shaders/max-pool-2d/max-pool-2d-spatial-f16.wgsl.ts @@ -0,0 +1,81 @@ +/** + * WGSL compute shader for 2D max pooling with 2D spatial grid (f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_MAX_POOL2D_SPATIAL_F16 = /* wgsl */ ` +enable f16; + +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + kH: u32, + kW: u32, + sH: u32, + sW: u32, + pH: u32, + pW: u32, + dH: u32, + dW: u32, + _p0: u32, + _p1: u32, +} + +@group(0) @binding(0) var out_vals: array; +@group(0) @binding(1) var inp: array; +@group(0) @binding(2) var params: Params; +@group(0) @binding(3) var out_idx: array; + +override write_indices: u32 = 0u; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(workgroup_id) wg_id: vec3 +) { + let ow = gid.x; + let oh = gid.y; + let c = wg_id.z % params.C; + let n = wg_id.z / params.C; + + if (ow >= params.OW || oh >= params.OH || n >= params.N) { + return; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + let in_c_base = (n * params.C + c) * params.IH; + + var best: f16 = -65504.0h; + var best_idx: i32 = 0; + for (var kh: u32 = 0u; kh < params.kH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + let in_row = (in_c_base + u32(ih)) * params.IW; + for (var kw: u32 = 0u; kw < params.kW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + let v = inp[in_row + u32(iw)]; + if (v > best) { + best = v; + best_idx = ih * iW + iw; + } + } + } + + let out_pos = ((n * params.C + c) * params.OH + oh) * params.OW + ow; + out_vals[out_pos] = best; + if (write_indices != 0u) { + out_idx[out_pos] = best_idx; + } +} +`; diff --git a/src/shaders/max-pool-2d/max-pool-2d-spatial.wgsl.ts b/src/shaders/max-pool-2d/max-pool-2d-spatial.wgsl.ts new file mode 100644 index 0000000..327eb17 --- /dev/null +++ b/src/shaders/max-pool-2d/max-pool-2d-spatial.wgsl.ts @@ -0,0 +1,79 @@ +/** + * WGSL compute shader for 2D max pooling with 2D spatial grid (f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_MAX_POOL2D_SPATIAL = /* wgsl */ ` +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + kH: u32, + kW: u32, + sH: u32, + sW: u32, + pH: u32, + pW: u32, + dH: u32, + dW: u32, + _p0: u32, + _p1: u32, +} + +@group(0) @binding(0) var out_vals: array; +@group(0) @binding(1) var inp: array; +@group(0) @binding(2) var params: Params; +@group(0) @binding(3) var out_idx: array; + +override write_indices: u32 = 0u; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(workgroup_id) wg_id: vec3 +) { + let ow = gid.x; + let oh = gid.y; + let c = wg_id.z % params.C; + let n = wg_id.z / params.C; + + if (ow >= params.OW || oh >= params.OH || n >= params.N) { + return; + } + + let iH = i32(params.IH); + let iW = i32(params.IW); + let in_c_base = (n * params.C + c) * params.IH; + + var best: f32 = -3.4e38; + var best_idx: i32 = 0; + for (var kh: u32 = 0u; kh < params.kH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + let in_row = (in_c_base + u32(ih)) * params.IW; + for (var kw: u32 = 0u; kw < params.kW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + let v = inp[in_row + u32(iw)]; + if (v > best) { + best = v; + best_idx = ih * iW + iw; + } + } + } + + let out_pos = ((n * params.C + c) * params.OH + oh) * params.OW + ow; + out_vals[out_pos] = best; + if (write_indices != 0u) { + out_idx[out_pos] = best_idx; + } +} +`; diff --git a/src/shaders/max-pool-2d/max-pool-2d.wgsl.ts b/src/shaders/max-pool-2d/max-pool-2d.wgsl.ts new file mode 100644 index 0000000..9062f7f --- /dev/null +++ b/src/shaders/max-pool-2d/max-pool-2d.wgsl.ts @@ -0,0 +1,77 @@ +/** + * WGSL compute shader for 2D max pooling with 1D-flattened grid (f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_MAX_POOL2D = /* wgsl */ ` +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + kH: u32, + kW: u32, + sH: u32, + sW: u32, + pH: u32, + pW: u32, + dH: u32, + dW: u32, + _p0: u32, + _p1: u32, +} + +@group(0) @binding(0) var out_vals: array; +@group(0) @binding(1) var inp: array; +@group(0) @binding(2) var params: Params; +@group(0) @binding(3) var out_idx: array; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; +override write_indices: u32 = 0u; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.N * params.C * params.OH * params.OW; + let i = gid.y * stride_x + gid.x; + if (i >= total) { + return; + } + let ow = i % params.OW; + let oh = (i / params.OW) % params.OH; + let c = (i / (params.OW * params.OH)) % params.C; + let n = i / (params.OW * params.OH * params.C); + + let iH = i32(params.IH); + let iW = i32(params.IW); + let in_c_base = (n * params.C + c) * params.IH; + + var best: f32 = -3.4e38; + var best_idx: i32 = 0; + for (var kh: u32 = 0u; kh < params.kH; kh = kh + 1u) { + let ih = i32(oh) * i32(params.sH) - i32(params.pH) + i32(kh) * i32(params.dH); + if (ih < 0 || ih >= iH) { + continue; + } + let in_row = (in_c_base + u32(ih)) * params.IW; + for (var kw: u32 = 0u; kw < params.kW; kw = kw + 1u) { + let iw = i32(ow) * i32(params.sW) - i32(params.pW) + i32(kw) * i32(params.dW); + if (iw < 0 || iw >= iW) { + continue; + } + let v = inp[in_row + u32(iw)]; + if (v > best) { + best = v; + best_idx = ih * iW + iw; + } + } + } + out_vals[i] = best; + if (write_indices != 0u) { + out_idx[i] = best_idx; + } +} +`; diff --git a/src/shaders/mean-dim/index.ts b/src/shaders/mean-dim/index.ts new file mode 100644 index 0000000..89fadd7 --- /dev/null +++ b/src/shaders/mean-dim/index.ts @@ -0,0 +1,158 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { SHADER_MEAN_DIM } from './mean-dim.wgsl'; +import { SHADER_MEAN_DIM_F16 } from './mean-dim-f16.wgsl'; + +const name = 'aten.mean.dim'; + +/** + * Positional argument tuple for `aten::mean.dim`: + * `[in, dim, keepdim, dtype, out]` + */ +export type MeanDimArgs = readonly [ + inTensor: Tensor, + dim: readonly number[] | undefined, + keepdim: boolean, + dtype: unknown, + outTensor: Tensor, +]; + +/** + * Normalizes and validates contiguous reduced dims, folding into [outer, r, inner]. + * + * @param shape Input tensor shape. + * @param dimList Dimensions to reduce over. + * @returns Flattened outer, r, and inner sizes. + */ +function decomposeDims( + shape: readonly number[], + dimList: readonly number[] +): { outer: number; r: number; inner: number } { + const ndim = shape.length; + if (ndim === 0 || dimList.length === 0) { + throw new Error(`${name}: dim out of range`); + } + + const normalized = dimList.map((d) => (d < 0 ? d + ndim : d)); + for (const d of normalized) { + if (d < 0 || d >= ndim) { + throw new Error(`${name}: dim out of range`); + } + } + + normalized.sort((a, b) => a - b); + for (let i = 1; i < normalized.length; i++) { + if (normalized[i] === normalized[i - 1]) { + throw new Error(`${name}: duplicate reduced dim`); + } + if (normalized[i] !== normalized[i - 1] + 1) { + throw new Error(`${name}: only contiguous reduced dims supported`); + } + } + + const firstRd = normalized[0]; + const lastRd = normalized[normalized.length - 1]; + + let outer = 1; + let r = 1; + let inner = 1; + + for (let d = 0; d < firstRd; d++) { + outer *= shape[d]; + } + for (let d = firstRd; d <= lastRd; d++) { + r *= shape[d]; + } + for (let d = lastRd + 1; d < ndim; d++) { + inner *= shape[d]; + } + + return { outer, r, inner }; +} + +/** + * Internal dispatch builder for `aten::mean.dim`. + * + * @param ctx The execution context. + * @param args Positional arguments for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: MeanDimArgs): void { + const inTensor = args[0]; + const dim = args[1]; + const outTensor = args[4]; + + if ( + (inTensor.dtype !== 'float32' && inTensor.dtype !== 'float16') || + outTensor.dtype !== inTensor.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + if (inTensor.shape.length === 0) { + throw new Error(`${name}: scalar input unsupported`); + } + + const device = ctx.device; + const isF16 = inTensor.dtype === 'float16'; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + const reduceDims = dim?.length ? dim : [...inTensor.shape.keys()]; + + const { outer, r, inner } = decomposeDims(inTensor.shape, reduceDims); + if (outer === 0 || r === 0 || inner === 0) { + throw new Error(`${name}: zero-sized reduction`); + } + + const outputs = outer * inner; + if (outTensor.numel !== outputs) { + throw new Error(`${name}: output numel mismatch`); + } + + const wgSize = 256; + const workgroupCountX = Math.min(outputs, 65535); + const workgroupCountY = Math.ceil(outputs / 65535); + + const buffer = new ArrayBuffer(16); + const u32View = new Uint32Array(buffer); + u32View[0] = outer; + u32View[1] = r; + u32View[2] = inner; + u32View[3] = 1; // is_mean = 1 + + const uniformBuffer = ctx.uniformBuffer(buffer); + + const bundle = createComputeBundle( + device, + isF16 ? SHADER_MEAN_DIM_F16 : SHADER_MEAN_DIM, + [ + { binding: 0, buffer: inTensor.buffer }, + { binding: 1, buffer: outTensor.buffer }, + { binding: 2, buffer: uniformBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize } + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::mean.dim`. + * + * Implements reduction across contiguous tensor dimensions with cooperative workgroups. + * Supports float32 and float16 with the `shader-f16` WebGPU feature extension. + */ +export const meanDim: Shader = { + name, + code: { default: SHADER_MEAN_DIM, defaultF16: SHADER_MEAN_DIM_F16 }, + recordIn, +}; diff --git a/src/shaders/mean-dim/mean-dim-f16.wgsl.ts b/src/shaders/mean-dim/mean-dim-f16.wgsl.ts new file mode 100644 index 0000000..deae72c --- /dev/null +++ b/src/shaders/mean-dim/mean-dim-f16.wgsl.ts @@ -0,0 +1,63 @@ +/** + * WGSL compute shader for dimension reduction mean/sum (f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_MEAN_DIM_F16 = /* wgsl */ ` +enable f16; + +struct Params { + outer_: u32, + r_: u32, + inner_: u32, + is_mean: u32, +}; + +@group(0) @binding(0) var inp: array; +@group(0) @binding(1) var out: array; +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +var partials: array; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let t = wid.x + wid.y * num_workgroups.x; + let outs = params.outer_ * params.inner_; + if (t >= outs) { + return; + } + let oo = t / params.inner_; + let ii = t % params.inner_; + let base = oo * params.r_ * params.inner_ + ii; + + var acc: f16 = 0.0h; + var k: u32 = lid.x; + while (k < params.r_) { + acc = acc + inp[base + k * params.inner_]; + k = k + wg_size; + } + partials[lid.x] = acc; + + for (var stride = 128u; stride > 0u; stride = stride >> 1u) { + workgroupBarrier(); + if (lid.x < stride && lid.x + stride < wg_size) { + partials[lid.x] = partials[lid.x] + partials[lid.x + stride]; + } + } + workgroupBarrier(); + + if (lid.x == 0u) { + var s: f16 = partials[0]; + if (params.is_mean == 1u) { + s = s / f16(params.r_); + } + out[t] = s; + } +} +`; diff --git a/src/shaders/mean-dim/mean-dim.wgsl.ts b/src/shaders/mean-dim/mean-dim.wgsl.ts new file mode 100644 index 0000000..8831bf4 --- /dev/null +++ b/src/shaders/mean-dim/mean-dim.wgsl.ts @@ -0,0 +1,65 @@ +/** + * WGSL compute shader for dimension reduction mean/sum (f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_MEAN_DIM = /* wgsl */ ` +struct Params { + outer_: u32, + r_: u32, + inner_: u32, + is_mean: u32, +}; + +@group(0) @binding(0) var inp: array; +@group(0) @binding(1) var out: array; +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +// Cooperative shared-memory reduction, one workgroup per output element: each +// thread sums a strided slice of the reduced dim into a shared partial, then +// thread 0 folds the partials. Fixed 256 upper bound >= any clamped wg_size; +// only [0, wg_size) is used. +var partials: array; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let t = wid.x + wid.y * num_workgroups.x; + let outs = params.outer_ * params.inner_; + if (t >= outs) { + return; + } + let oo = t / params.inner_; + let ii = t % params.inner_; + let base = oo * params.r_ * params.inner_ + ii; + + var acc: f32 = 0.0; + var k: u32 = lid.x; + while (k < params.r_) { + acc = acc + inp[base + k * params.inner_]; + k = k + wg_size; + } + partials[lid.x] = acc; + + for (var stride = 128u; stride > 0u; stride = stride >> 1u) { + workgroupBarrier(); + if (lid.x < stride && lid.x + stride < wg_size) { + partials[lid.x] = partials[lid.x] + partials[lid.x + stride]; + } + } + workgroupBarrier(); + + if (lid.x == 0u) { + var s: f32 = partials[0]; + if (params.is_mean == 1u) { + s = s / f32(params.r_); + } + out[t] = s; + } +} +`; diff --git a/src/shaders/mm/index.ts b/src/shaders/mm/index.ts new file mode 100644 index 0000000..1a78395 --- /dev/null +++ b/src/shaders/mm/index.ts @@ -0,0 +1,114 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { SHADER_MM_TILED } from './mm-tiled.wgsl'; +import { SHADER_MM_VEC4 } from './mm-vec4.wgsl'; +import { SHADER_MM_64X64 } from './mm-64x64.wgsl'; +import { SHADER_MM_TILED_F16 } from './mm-tiled-f16.wgsl'; +import { SHADER_MM_VEC4_F16 } from './mm-vec4-f16.wgsl'; +import { SHADER_MM_64X64_F16 } from './mm-64x64-f16.wgsl'; + +const name = 'aten.mm.default'; + +/** + * Output tile dimension per workgroup (32x32). + */ +const TILE = 32; + +/** + * Positional argument tuple for `aten::mm.default`: + * `[in1, in2, out]` + */ +export type MmArgs = readonly [in1: Tensor, in2: Tensor, out: Tensor]; + +/** + * Internal dispatch builder for `aten::mm.default`. + * + * @param ctx The execution context. + * @param args Positional argument tuple for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, [in1, in2, out]: MmArgs): void { + if ( + (in1.dtype !== 'float32' && in1.dtype !== 'float16') || + in2.dtype !== in1.dtype || + out.dtype !== in1.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + if (in1.shape.length !== 2 || in2.shape.length !== 2 || out.shape.length !== 2) { + throw new Error(`${name}: Tensors must be 2D matrices`); + } + + const [m, k1] = in1.shape as [number, number]; + const [k2, n] = in2.shape as [number, number]; + const [outM, outN] = out.shape as [number, number]; + + if (k1 !== k2) { + throw new Error(`${name}: Matrix inner dimensions must match (${k1} !== ${k2})`); + } + if (outM !== m || outN !== n) { + throw new Error(`${name}: Output shape [${outM}, ${outN}] !== expected [${m}, ${n}]`); + } + + const device = ctx.device; + const isF16 = in1.dtype === 'float16'; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + const paramsBuffer = ctx.uniformBuffer(new Uint32Array([m, n, k1, 0])); + + // Use 64x64 register-tiled GEMM when dimensions are sufficiently large and divisible by 4, + // falling back to 32x32 vec4 or 32x32 scalar tiled paths. + const useVec4 = k1 % 4 === 0 && n % 4 === 0; + const use64x64 = useVec4 && m >= 64 && n >= 64 && k1 >= 16; + + let code = isF16 ? SHADER_MM_TILED_F16 : SHADER_MM_TILED; + let workgroupCountX = Math.ceil(n / TILE); + let workgroupCountY = Math.ceil(m / TILE); + + if (use64x64) { + code = isF16 ? SHADER_MM_64X64_F16 : SHADER_MM_64X64; + workgroupCountX = Math.ceil(n / 64); + workgroupCountY = Math.ceil(m / 64); + } else if (useVec4) { + code = isF16 ? SHADER_MM_VEC4_F16 : SHADER_MM_VEC4; + } + + const bundle = createComputeBundle(device, code, [ + { binding: 0, buffer: in1.buffer }, + { binding: 1, buffer: in2.buffer }, + { binding: 2, buffer: out.buffer }, + { binding: 3, buffer: paramsBuffer }, + ]); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::mm.default`. + * + * Computes matrix multiplication `out = A @ B` using a shared-memory tiled GEMM. + * Automatically selects 64x64 register-tiled or 32x32 128-bit vectorized paths + * when dimensions permit, falling back to the standard scalar tiled path. + * Supports both float32 and float16 with the `shader-f16` WebGPU feature extension. + */ +export const mm: Shader = { + name, + code: { + tiled: SHADER_MM_TILED, + vec4: SHADER_MM_VEC4, + tiled64x64: SHADER_MM_64X64, + tiledF16: SHADER_MM_TILED_F16, + vec4F16: SHADER_MM_VEC4_F16, + tiled64x64F16: SHADER_MM_64X64_F16, + }, + recordIn, +}; diff --git a/src/shaders/mm/mm-64x64-f16.wgsl.ts b/src/shaders/mm/mm-64x64-f16.wgsl.ts new file mode 100644 index 0000000..b0f7346 --- /dev/null +++ b/src/shaders/mm/mm-64x64-f16.wgsl.ts @@ -0,0 +1,117 @@ +/** + * WGSL compute shader for 64x64 tiled GEMM in float16 with 2D register tiling (4x4 per thread across 256 threads). + * High-occupancy, 128-bit coalesced memory transactions, minimal register pressure for all GPUs. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_MM_64X64_F16 = /* wgsl */ ` +enable f16; + +struct Params { + M: u32, + N: u32, + K: u32, + pad_: u32, +}; + +@group(0) @binding(0) var a: array>; +@group(0) @binding(1) var b: array>; +@group(0) @binding(2) var out: array>; +@group(0) @binding(3) var params: Params; + +// 64 rows x 16 K elements (4 vec4s) = 256 vec4s (2 KB) +var a_sub: array, 4>, 64>; +// 16 K rows x 64 N elements (16 vec4s) = 256 vec4s (2 KB) +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let tile_row0 = wg_id.y * 64u; + let tile_col0_4 = wg_id.x * 16u; + let tid = local_id.y * 16u + local_id.x; // 0..255 + let K_vecs = params.K / 4u; + let N_vecs = params.N / 4u; + + var acc0 = vec4(0.0h); + var acc1 = vec4(0.0h); + var acc2 = vec4(0.0h); + var acc3 = vec4(0.0h); + + // Load coordinates for 256 threads (each thread loads 1 vec4 of A and 1 vec4 of B) + let a_r = tid >> 2u; // 0..63 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < params.M; + + let b_r = tid >> 4u; // 0..15 + let b_c = tid & 15u; // 0..15 + let b_n4 = tile_col0_4 + b_c; + let b_c_valid = b_n4 < N_vecs; + + let num_tiles = (params.K + 15u) >> 4u; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t << 2u; + let k_start = t << 4u; + + // Coalesced 128-bit load of A (256 threads load 256 vec4s) + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = a[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0h); + } + + // Coalesced 128-bit load of B (256 threads load 256 vec4s) + let bk = k_start + b_r; + if (bk < params.K && b_c_valid) { + b_sub[b_r][b_c] = b[bk * N_vecs + b_n4]; + } else { + b_sub[b_r][b_c] = vec4(0.0h); + } + workgroupBarrier(); + + let thread_r = local_id.y << 2u; // 0..60 (step 4) + let thread_c4 = local_id.x; // 0..15 + + for (var k4: u32 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + let thread_r = local_id.y << 2u; + let out_r0 = tile_row0 + thread_r; + let out_c4 = tile_col0_4 + local_id.x; + + if (out_c4 < N_vecs) { + if (out_r0 + 0u < params.M) { + out[(out_r0 + 0u) * N_vecs + out_c4] = acc0; + } + if (out_r0 + 1u < params.M) { + out[(out_r0 + 1u) * N_vecs + out_c4] = acc1; + } + if (out_r0 + 2u < params.M) { + out[(out_r0 + 2u) * N_vecs + out_c4] = acc2; + } + if (out_r0 + 3u < params.M) { + out[(out_r0 + 3u) * N_vecs + out_c4] = acc3; + } + } +} +`; diff --git a/src/shaders/mm/mm-64x64.wgsl.ts b/src/shaders/mm/mm-64x64.wgsl.ts new file mode 100644 index 0000000..f08aa04 --- /dev/null +++ b/src/shaders/mm/mm-64x64.wgsl.ts @@ -0,0 +1,115 @@ +/** + * WGSL compute shader for 64x64 tiled GEMM with 2D register tiling (4x4 per thread across 256 threads). + * High-occupancy, 128-bit coalesced memory transactions, minimal register pressure for all GPUs. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_MM_64X64 = /* wgsl */ ` +struct Params { + M: u32, + N: u32, + K: u32, + pad_: u32, +}; + +@group(0) @binding(0) var a: array>; +@group(0) @binding(1) var b: array>; +@group(0) @binding(2) var out: array>; +@group(0) @binding(3) var params: Params; + +// 64 rows x 16 K elements (4 vec4s) = 256 vec4s (4 KB) +var a_sub: array, 4>, 64>; +// 16 K rows x 64 N elements (16 vec4s) = 256 vec4s (4 KB) +var b_sub: array, 16>, 16>; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let tile_row0 = wg_id.y * 64u; + let tile_col0_4 = wg_id.x * 16u; + let tid = local_id.y * 16u + local_id.x; // 0..255 + let K_vecs = params.K / 4u; + let N_vecs = params.N / 4u; + + var acc0 = vec4(0.0); + var acc1 = vec4(0.0); + var acc2 = vec4(0.0); + var acc3 = vec4(0.0); + + // Load coordinates for 256 threads (each thread loads 1 vec4 of A and 1 vec4 of B) + let a_r = tid >> 2u; // 0..63 + let a_c = tid & 3u; // 0..3 + let a_m = tile_row0 + a_r; + let a_valid = a_m < params.M; + + let b_r = tid >> 4u; // 0..15 + let b_c = tid & 15u; // 0..15 + let b_n4 = tile_col0_4 + b_c; + let b_c_valid = b_n4 < N_vecs; + + let num_tiles = (params.K + 15u) >> 4u; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t << 2u; + let k_start = t << 4u; + + // Coalesced 128-bit load of A (256 threads load 256 vec4s) + if (a_valid && (k4_start + a_c) < K_vecs) { + a_sub[a_r][a_c] = a[a_m * K_vecs + k4_start + a_c]; + } else { + a_sub[a_r][a_c] = vec4(0.0); + } + + // Coalesced 128-bit load of B (256 threads load 256 vec4s) + let bk = k_start + b_r; + if (bk < params.K && b_c_valid) { + b_sub[b_r][b_c] = b[bk * N_vecs + b_n4]; + } else { + b_sub[b_r][b_c] = vec4(0.0); + } + workgroupBarrier(); + + let thread_r = local_id.y << 2u; // 0..60 (step 4) + let thread_c4 = local_id.x; // 0..15 + + for (var k4: u32 = 0u; k4 < 4u; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][thread_c4]; + let b1 = b_sub[k4 * 4u + 1u][thread_c4]; + let b2 = b_sub[k4 * 4u + 2u][thread_c4]; + let b3 = b_sub[k4 * 4u + 3u][thread_c4]; + + let a0 = a_sub[thread_r + 0u][k4]; + let a1 = a_sub[thread_r + 1u][k4]; + let a2 = a_sub[thread_r + 2u][k4]; + let a3 = a_sub[thread_r + 3u][k4]; + + acc0 = acc0 + b0 * a0.x + b1 * a0.y + b2 * a0.z + b3 * a0.w; + acc1 = acc1 + b0 * a1.x + b1 * a1.y + b2 * a1.z + b3 * a1.w; + acc2 = acc2 + b0 * a2.x + b1 * a2.y + b2 * a2.z + b3 * a2.w; + acc3 = acc3 + b0 * a3.x + b1 * a3.y + b2 * a3.z + b3 * a3.w; + } + workgroupBarrier(); + } + + let thread_r = local_id.y << 2u; + let out_r0 = tile_row0 + thread_r; + let out_c4 = tile_col0_4 + local_id.x; + + if (out_c4 < N_vecs) { + if (out_r0 + 0u < params.M) { + out[(out_r0 + 0u) * N_vecs + out_c4] = acc0; + } + if (out_r0 + 1u < params.M) { + out[(out_r0 + 1u) * N_vecs + out_c4] = acc1; + } + if (out_r0 + 2u < params.M) { + out[(out_r0 + 2u) * N_vecs + out_c4] = acc2; + } + if (out_r0 + 3u < params.M) { + out[(out_r0 + 3u) * N_vecs + out_c4] = acc3; + } + } +} +`; diff --git a/src/shaders/mm/mm-tiled-f16.wgsl.ts b/src/shaders/mm/mm-tiled-f16.wgsl.ts new file mode 100644 index 0000000..863caa1 --- /dev/null +++ b/src/shaders/mm/mm-tiled-f16.wgsl.ts @@ -0,0 +1,94 @@ +/** + * WGSL compute shader for float16 matrix multiplication with shared-memory tiling + * (32x32 tile, 4x4 per thread). Requires WebGPU shader-f16 feature. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_MM_TILED_F16 = /* wgsl */ ` +enable f16; + +struct Params { + M: u32, + N: u32, + K: u32, + pad_: u32, +}; + +@group(0) @binding(0) var a: array; +@group(0) @binding(1) var b: array; +@group(0) @binding(2) var out: array; +@group(0) @binding(3) var params: Params; + +const TILE: u32 = 32u; +const RPT: u32 = 4u; + +var a_sub: array, 32>; +var b_sub: array, 32>; + +fn read_a(row: u32, col: u32) -> f16 { + if (row < params.M && col < params.K) { + return a[row * params.K + col]; + } + return 0.0h; +} + +fn read_b(row: u32, col: u32) -> f16 { + if (row < params.K && col < params.N) { + return b[row * params.N + col]; + } + return 0.0h; +} + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let tile_row0 = wg_id.y * TILE; + let tile_col0 = wg_id.x * TILE; + let tile_row = local_id.y * RPT; + let tile_col = local_id.x * RPT; + + var acc: array, 4>; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + acc[ir][ic] = 0.0h; + } + } + + let num_tiles = (params.K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k_start = t * TILE; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) { + let col = local_id.x * RPT + kk; + a_sub[arow][col] = read_a(tile_row0 + arow, k_start + col); + b_sub[arow][col] = read_b(k_start + arow, tile_col0 + col); + } + } + workgroupBarrier(); + + for (var k: u32 = 0u; k < TILE; k = k + 1u) { + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let aval = a_sub[tile_row + ir][k]; + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + acc[ir][ic] = acc[ir][ic] + aval * b_sub[k][tile_col + ic]; + } + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + let r = tile_row0 + tile_row + ir; + let c = tile_col0 + tile_col + ic; + if (r < params.M && c < params.N) { + out[r * params.N + c] = acc[ir][ic]; + } + } + } +} +`; diff --git a/src/shaders/mm/mm-tiled.wgsl.ts b/src/shaders/mm/mm-tiled.wgsl.ts new file mode 100644 index 0000000..069ec36 --- /dev/null +++ b/src/shaders/mm/mm-tiled.wgsl.ts @@ -0,0 +1,93 @@ +/** + * WGSL compute shader for matrix multiplication with shared-memory tiling + * (32x32 tile, 4x4 per thread). + * + * Source: ExecuTorch (backends/webgpu/runtime/ops/mm/mm_tiled.wgsl) + * GitHub: https://github.com/pytorch/executorch/blob/main/backends/webgpu/runtime/ops/mm/mm_tiled.wgsl + * License: BSD-3-Clause (Copyright (c) Meta Platforms, Inc. and affiliates) + */ +export const SHADER_MM_TILED = /* wgsl */ ` +struct Params { + M: u32, + N: u32, + K: u32, + pad_: u32, +}; + +@group(0) @binding(0) var a: array; +@group(0) @binding(1) var b: array; +@group(0) @binding(2) var out: array; +@group(0) @binding(3) var params: Params; + +const TILE: u32 = 32u; +const RPT: u32 = 4u; + +var a_sub: array, 32>; +var b_sub: array, 32>; + +fn read_a(row: u32, col: u32) -> f32 { + if (row < params.M && col < params.K) { + return a[row * params.K + col]; + } + return 0.0; +} + +fn read_b(row: u32, col: u32) -> f32 { + if (row < params.K && col < params.N) { + return b[row * params.N + col]; + } + return 0.0; +} + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let tile_row0 = wg_id.y * TILE; + let tile_col0 = wg_id.x * TILE; + let tile_row = local_id.y * RPT; + let tile_col = local_id.x * RPT; + + var acc: array, 4>; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + acc[ir][ic] = 0.0; + } + } + + let num_tiles = (params.K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k_start = t * TILE; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + for (var kk: u32 = 0u; kk < RPT; kk = kk + 1u) { + let col = local_id.x * RPT + kk; + a_sub[arow][col] = read_a(tile_row0 + arow, k_start + col); + b_sub[arow][col] = read_b(k_start + arow, tile_col0 + col); + } + } + workgroupBarrier(); + + for (var k: u32 = 0u; k < TILE; k = k + 1u) { + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let aval = a_sub[tile_row + ir][k]; + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + acc[ir][ic] = acc[ir][ic] + aval * b_sub[k][tile_col + ic]; + } + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + for (var ic: u32 = 0u; ic < RPT; ic = ic + 1u) { + let r = tile_row0 + tile_row + ir; + let c = tile_col0 + tile_col + ic; + if (r < params.M && c < params.N) { + out[r * params.N + c] = acc[ir][ic]; + } + } + } +} +`; diff --git a/src/shaders/mm/mm-vec4-f16.wgsl.ts b/src/shaders/mm/mm-vec4-f16.wgsl.ts new file mode 100644 index 0000000..7524043 --- /dev/null +++ b/src/shaders/mm/mm-vec4-f16.wgsl.ts @@ -0,0 +1,90 @@ +/** + * WGSL compute shader for float16 matrix multiplication vectorized with vec4. + * Requires WebGPU shader-f16 feature. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_MM_VEC4_F16 = /* wgsl */ ` +enable f16; + +struct Params { + M: u32, + N: u32, + K: u32, + pad_: u32, +}; + +@group(0) @binding(0) var a: array>; +@group(0) @binding(1) var b: array>; +@group(0) @binding(2) var out: array>; +@group(0) @binding(3) var params: Params; + +const TILE: u32 = 32u; +const RPT: u32 = 4u; +const TILE4: u32 = 8u; + +var a_sub: array, 8>, 32>; +var b_sub: array, 8>, 32>; + +fn read_a4(row: u32, k4: u32) -> vec4 { + if (row < params.M && k4 * 4u < params.K) { + return a[row * (params.K / 4u) + k4]; + } + return vec4(0.0h); +} + +fn read_b4(krow: u32, n4: u32) -> vec4 { + if (krow < params.K && n4 * 4u < params.N) { + return b[krow * (params.N / 4u) + n4]; + } + return vec4(0.0h); +} + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let tile_row0 = wg_id.y * TILE; + let tile_col0_4 = wg_id.x * TILE4; + let tile_row = local_id.y * RPT; + let tile_col4 = local_id.x; + + var acc: array, 4>; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + acc[ir] = vec4(0.0h); + } + + let num_tiles = (params.K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t * TILE4; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + a_sub[arow][local_id.x] = read_a4(tile_row0 + arow, k4_start + local_id.x); + b_sub[arow][local_id.x] = read_b4(t * TILE + arow, tile_col0_4 + local_id.x); + } + workgroupBarrier(); + + for (var k4: u32 = 0u; k4 < TILE4; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][tile_col4]; + let b1 = b_sub[k4 * 4u + 1u][tile_col4]; + let b2 = b_sub[k4 * 4u + 2u][tile_col4]; + let b3 = b_sub[k4 * 4u + 3u][tile_col4]; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let ac = a_sub[tile_row + ir][k4]; + acc[ir] = acc[ir] + b0 * ac.x + b1 * ac.y + b2 * ac.z + b3 * ac.w; + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let r = tile_row0 + tile_row + ir; + let c4 = tile_col0_4 + tile_col4; + if (r < params.M && c4 * 4u < params.N) { + out[r * (params.N / 4u) + c4] = acc[ir]; + } + } +} +`; diff --git a/src/shaders/mm/mm-vec4.wgsl.ts b/src/shaders/mm/mm-vec4.wgsl.ts new file mode 100644 index 0000000..e53f8ed --- /dev/null +++ b/src/shaders/mm/mm-vec4.wgsl.ts @@ -0,0 +1,90 @@ +/** + * WGSL compute shader for vectorized 128-bit memory-bandwidth optimized matrix + * multiplication. Operates on vec4 for 4x wider memory transactions when K + * and N are multiples of 4. + * + * Source: ExecuTorch (backends/webgpu/runtime/ops/mm/mm_vec4.wgsl) + * GitHub: https://github.com/pytorch/executorch/blob/main/backends/webgpu/runtime/ops/mm/mm_vec4.wgsl + * License: BSD-3-Clause (Copyright (c) Meta Platforms, Inc. and affiliates) + */ +export const SHADER_MM_VEC4 = /* wgsl */ ` +struct Params { + M: u32, + N: u32, + K: u32, + pad_: u32, +}; + +@group(0) @binding(0) var a: array>; +@group(0) @binding(1) var b: array>; +@group(0) @binding(2) var out: array>; +@group(0) @binding(3) var params: Params; + +const TILE: u32 = 32u; +const RPT: u32 = 4u; +const TILE4: u32 = 8u; + +var a_sub: array, 8>, 32>; +var b_sub: array, 8>, 32>; + +fn read_a4(row: u32, k4: u32) -> vec4 { + if (row < params.M && k4 * 4u < params.K) { + return a[row * (params.K / 4u) + k4]; + } + return vec4(0.0); +} + +fn read_b4(krow: u32, n4: u32) -> vec4 { + if (krow < params.K && n4 * 4u < params.N) { + return b[krow * (params.N / 4u) + n4]; + } + return vec4(0.0); +} + +@compute @workgroup_size(8, 8, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let tile_row0 = wg_id.y * TILE; + let tile_col0_4 = wg_id.x * TILE4; + let tile_row = local_id.y * RPT; + let tile_col4 = local_id.x; + + var acc: array, 4>; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + acc[ir] = vec4(0.0); + } + + let num_tiles = (params.K + TILE - 1u) / TILE; + for (var t: u32 = 0u; t < num_tiles; t = t + 1u) { + let k4_start = t * TILE4; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let arow = local_id.y * RPT + ir; + a_sub[arow][local_id.x] = read_a4(tile_row0 + arow, k4_start + local_id.x); + b_sub[arow][local_id.x] = read_b4(t * TILE + arow, tile_col0_4 + local_id.x); + } + workgroupBarrier(); + + for (var k4: u32 = 0u; k4 < TILE4; k4 = k4 + 1u) { + let b0 = b_sub[k4 * 4u + 0u][tile_col4]; + let b1 = b_sub[k4 * 4u + 1u][tile_col4]; + let b2 = b_sub[k4 * 4u + 2u][tile_col4]; + let b3 = b_sub[k4 * 4u + 3u][tile_col4]; + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let ac = a_sub[tile_row + ir][k4]; + acc[ir] = acc[ir] + b0 * ac.x + b1 * ac.y + b2 * ac.z + b3 * ac.w; + } + } + workgroupBarrier(); + } + + for (var ir: u32 = 0u; ir < RPT; ir = ir + 1u) { + let r = tile_row0 + tile_row + ir; + let c4 = tile_col0_4 + tile_col4; + if (r < params.M && c4 * 4u < params.N) { + out[r * (params.N / 4u) + c4] = acc[ir]; + } + } +} +`; diff --git a/src/shaders/mul/index.ts b/src/shaders/mul/index.ts new file mode 100644 index 0000000..ae00b1f --- /dev/null +++ b/src/shaders/mul/index.ts @@ -0,0 +1,144 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { isBroadcastable } from '../utils/broadcasting'; +import { encodeTensorMeta } from '../utils/meta'; + +import { SHADER_MUL_BROADCAST } from './mul-broadcast.wgsl'; +import { SHADER_MUL_BROADCAST_F16 } from './mul-broadcast-f16.wgsl'; +import { SHADER_MUL_ELEMENTWISE } from './mul-elementwise.wgsl'; +import { SHADER_MUL_ELEMENTWISE_F16 } from './mul-elementwise-f16.wgsl'; +import { SHADER_MUL_ELEMENTWISE_VEC4 } from './mul-elementwise-vec4.wgsl'; +import { SHADER_MUL_ELEMENTWISE_VEC4_F16 } from './mul-elementwise-vec4-f16.wgsl'; + +const name = 'aten.mul.Tensor'; + +/** + * Positional argument tuple for `aten::mul.Tensor`: + * `[in1, in2, out]` + */ +export type MulArgs = readonly [in1: Tensor, in2: Tensor, out: Tensor]; + +/** + * Internal dispatch builder for `aten::mul.Tensor`. + * + * @param ctx The execution context. + * @param args Positional argument tuple for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: MulArgs): void { + const [in1, in2, out] = args; + + if ( + (in1.dtype !== 'float32' && in1.dtype !== 'float16') || + in2.dtype !== in1.dtype || + out.dtype !== in1.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + if (!isBroadcastable(in1.shape, out.shape) || !isBroadcastable(in2.shape, out.shape)) { + throw new Error(`${name}: Input shapes are not broadcastable to output shape`); + } + + const device = ctx.device; + const isF16 = in1.dtype === 'float16'; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + const sameShape = + in1.shape.length === out.shape.length && + in2.shape.length === out.shape.length && + in1.shape.every((s, i) => s === out.shape[i]) && + in2.shape.every((s, i) => s === out.shape[i]); + + const wgSize = 256; + + if (sameShape) { + const numel = out.numel; + const useVec4 = numel % 4 === 0; + let code = useVec4 ? SHADER_MUL_ELEMENTWISE_VEC4 : SHADER_MUL_ELEMENTWISE; + if (isF16) { + code = useVec4 ? SHADER_MUL_ELEMENTWISE_VEC4_F16 : SHADER_MUL_ELEMENTWISE_F16; + } + const effectiveNumel = useVec4 ? numel / 4 : numel; + + const totalWorkgroups = Math.ceil(effectiveNumel / wgSize); + const workgroupCountX = Math.min(totalWorkgroups, 65535); + const workgroupCountY = Math.ceil(totalWorkgroups / 65535); + + const paramBuffer = ctx.uniformBuffer(new Uint32Array([effectiveNumel])); + + const bundle = createComputeBundle( + device, + code, + [ + { binding: 0, buffer: in1.buffer }, + { binding: 1, buffer: in2.buffer }, + { binding: 2, buffer: out.buffer }, + { binding: 3, buffer: paramBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize } + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); + return; + } + + const outRank = out.shape.length; + const outMetaBuffer = ctx.uniformBuffer(encodeTensorMeta(out)); + const in1MetaBuffer = ctx.uniformBuffer(encodeTensorMeta(in1, outRank)); + const in2MetaBuffer = ctx.uniformBuffer(encodeTensorMeta(in2, outRank)); + + const totalWorkgroups = Math.ceil(out.numel / wgSize); + const workgroupCountX = Math.min(totalWorkgroups, 65535); + const workgroupCountY = Math.ceil(totalWorkgroups / 65535); + + const bundle = createComputeBundle( + device, + isF16 ? SHADER_MUL_BROADCAST_F16 : SHADER_MUL_BROADCAST, + [ + { binding: 0, buffer: in1.buffer }, + { binding: 1, buffer: in2.buffer }, + { binding: 2, buffer: out.buffer }, + { binding: 3, buffer: outMetaBuffer }, + { binding: 4, buffer: in1MetaBuffer }, + { binding: 5, buffer: in2MetaBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize } + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::mul.Tensor`. + * + * Computes elementwise `output = input1 * input2` with NumPy/PyTorch-style + * broadcasting (right-aligned, dimensions equal or 1, up to rank 8). + */ +export const mul: Shader = { + name, + code: { + broadcast: SHADER_MUL_BROADCAST, + elementwise: SHADER_MUL_ELEMENTWISE, + elementwiseVec4: SHADER_MUL_ELEMENTWISE_VEC4, + broadcastF16: SHADER_MUL_BROADCAST_F16, + elementwiseF16: SHADER_MUL_ELEMENTWISE_F16, + elementwiseVec4F16: SHADER_MUL_ELEMENTWISE_VEC4_F16, + }, + recordIn, +}; diff --git a/src/shaders/mul/mul-broadcast-f16.wgsl.ts b/src/shaders/mul/mul-broadcast-f16.wgsl.ts new file mode 100644 index 0000000..4d78f1f --- /dev/null +++ b/src/shaders/mul/mul-broadcast-f16.wgsl.ts @@ -0,0 +1,55 @@ +import { TENSOR_META_WGSL } from '../utils/meta'; + +/** + * WGSL compute shader for FP16 binary multiplication with broadcast support. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_MUL_BROADCAST_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +${TENSOR_META_WGSL} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in1_meta: TensorMeta; +@group(0) @binding(5) var in2_meta: TensorMeta; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= out_meta.numel) { + return; + } + + var same = true; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { + same = false; + } + } + if (same) { + output[idx] = input1[idx] * input2[idx]; + return; + } + + var rem = idx; + var l1: u32 = 0u; + var l2: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; + } + output[idx] = input1[l1] * input2[l2]; +} +`; diff --git a/src/shaders/mul/mul-broadcast.wgsl.ts b/src/shaders/mul/mul-broadcast.wgsl.ts new file mode 100644 index 0000000..624522f --- /dev/null +++ b/src/shaders/mul/mul-broadcast.wgsl.ts @@ -0,0 +1,53 @@ +import { TENSOR_META_WGSL } from '../utils/meta'; + +/** + * WGSL compute shader for binary multiplication with broadcast support. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_MUL_BROADCAST = /* wgsl */ ` +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +${TENSOR_META_WGSL} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in1_meta: TensorMeta; +@group(0) @binding(5) var in2_meta: TensorMeta; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= out_meta.numel) { + return; + } + + var same = true; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { + same = false; + } + } + if (same) { + output[idx] = input1[idx] * input2[idx]; + return; + } + + var rem = idx; + var l1: u32 = 0u; + var l2: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; + } + output[idx] = input1[l1] * input2[l2]; +} +`; diff --git a/src/shaders/mul/mul-elementwise-f16.wgsl.ts b/src/shaders/mul/mul-elementwise-f16.wgsl.ts new file mode 100644 index 0000000..23d975b --- /dev/null +++ b/src/shaders/mul/mul-elementwise-f16.wgsl.ts @@ -0,0 +1,31 @@ +/** + * WGSL compute shader for contiguous 1D FP16 elementwise multiplication. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_MUL_ELEMENTWISE_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = input1[idx] * input2[idx]; +} +`; diff --git a/src/shaders/mul/mul-elementwise-vec4-f16.wgsl.ts b/src/shaders/mul/mul-elementwise-vec4-f16.wgsl.ts new file mode 100644 index 0000000..48b9a9a --- /dev/null +++ b/src/shaders/mul/mul-elementwise-vec4-f16.wgsl.ts @@ -0,0 +1,31 @@ +/** + * WGSL compute shader for contiguous 64-bit/128-bit vectorized (vec4) elementwise multiplication. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_MUL_ELEMENTWISE_VEC4_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input1: array>; +@group(0) @binding(1) var input2: array>; +@group(0) @binding(2) var output: array>; + +struct Params { + num_elements: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = input1[idx] * input2[idx]; +} +`; diff --git a/src/shaders/mul/mul-elementwise-vec4.wgsl.ts b/src/shaders/mul/mul-elementwise-vec4.wgsl.ts new file mode 100644 index 0000000..d6f9ec2 --- /dev/null +++ b/src/shaders/mul/mul-elementwise-vec4.wgsl.ts @@ -0,0 +1,29 @@ +/** + * WGSL compute shader for contiguous 128-bit vectorized (vec4) elementwise multiplication. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_MUL_ELEMENTWISE_VEC4 = /* wgsl */ ` +@group(0) @binding(0) var input1: array>; +@group(0) @binding(1) var input2: array>; +@group(0) @binding(2) var output: array>; + +struct Params { + num_elements: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = input1[idx] * input2[idx]; +} +`; diff --git a/src/shaders/mul/mul-elementwise.wgsl.ts b/src/shaders/mul/mul-elementwise.wgsl.ts new file mode 100644 index 0000000..9f913d6 --- /dev/null +++ b/src/shaders/mul/mul-elementwise.wgsl.ts @@ -0,0 +1,29 @@ +/** + * WGSL compute shader for contiguous 1D elementwise multiplication. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_MUL_ELEMENTWISE = /* wgsl */ ` +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = input1[idx] * input2[idx]; +} +`; diff --git a/src/shaders/permute/index.ts b/src/shaders/permute/index.ts new file mode 100644 index 0000000..fc977bf --- /dev/null +++ b/src/shaders/permute/index.ts @@ -0,0 +1,146 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { encodeTensorMeta } from '../utils/meta'; + +import { SHADER_PERMUTE } from './permute.wgsl'; +import { SHADER_PERMUTE_TRANSPOSE2D } from './permute-transpose2d.wgsl'; +import { SHADER_PERMUTE_F16 } from './permute-f16.wgsl'; +import { SHADER_PERMUTE_TRANSPOSE2D_F16 } from './permute-transpose2d-f16.wgsl'; + +const name = 'aten.permute_copy.default'; + +/** + * Positional argument tuple for `aten::permute_copy.default`: + * `[in, dims, out]` + */ +export type PermuteCopyArgs = readonly [ + inTensor: Tensor, + dims: readonly number[], + outTensor: Tensor, +]; + +/** + * Internal dispatch builder for `aten::permute_copy.default`. + * + * @param ctx The execution context. + * @param args Positional argument tuple for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: PermuteCopyArgs): void { + const [inTensor, dims, outTensor] = args; + + if ( + (inTensor.dtype !== 'float32' && inTensor.dtype !== 'float16') || + outTensor.dtype !== inTensor.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + + const device = ctx.device; + const isF16 = inTensor.dtype === 'float16'; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + const rank = inTensor.shape.length; + if (outTensor.shape.length !== rank) { + throw new Error(`${name}: output rank != input rank`); + } + if (dims.length !== rank) { + throw new Error(`${name}: dims count (${dims.length}) != input rank (${rank})`); + } + + const perm = new Uint32Array(8); + const seen = new Set(); + for (let d = 0; d < rank; d++) { + let p = dims[d]; + if (p < 0) { + p += rank; + } + if (p < 0 || p >= rank || seen.has(p)) { + throw new Error(`${name}: invalid permute dims ${dims}`); + } + seen.add(p); + perm[d] = p; + } + for (let d = rank; d < 8; d++) { + perm[d] = d; // identity padding for unused slots + } + for (let d = 0; d < rank; d++) { + if (outTensor.shape[d] !== inTensor.shape[perm[d]]) { + throw new Error(`${name}: output shape does not match permuted input shape`); + } + } + + // Fast path: 2D matrix transpose + if (rank === 2 && perm[0] === 1 && perm[1] === 0) { + const rows = inTensor.shape[0]; + const cols = inTensor.shape[1]; + const paramsBuffer = ctx.uniformBuffer(new Uint32Array([rows, cols])); + + const bundle = createComputeBundle( + device, + isF16 ? SHADER_PERMUTE_TRANSPOSE2D_F16 : SHADER_PERMUTE_TRANSPOSE2D, + [ + { binding: 0, buffer: inTensor.buffer }, + { binding: 1, buffer: outTensor.buffer }, + { binding: 2, buffer: paramsBuffer }, + ] + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX: Math.ceil(cols / 16), + workgroupCountY: Math.ceil(rows / 16), + }); + return; + } + + const wgSize = 64; + const outMetaBuffer = ctx.uniformBuffer(encodeTensorMeta(outTensor)); + const inMetaBuffer = ctx.uniformBuffer(encodeTensorMeta(inTensor)); + const paramsBuffer = ctx.uniformBuffer(perm); + + const bundle = createComputeBundle( + device, + isF16 ? SHADER_PERMUTE_F16 : SHADER_PERMUTE, + [ + { binding: 0, buffer: inTensor.buffer }, + { binding: 1, buffer: outTensor.buffer }, + { binding: 2, buffer: outMetaBuffer }, + { binding: 3, buffer: inMetaBuffer }, + { binding: 4, buffer: paramsBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize } + ); + + const totalWorkgroups = Math.ceil(outTensor.numel / wgSize); + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX: Math.min(totalWorkgroups, 65535), + workgroupCountY: Math.ceil(totalWorkgroups / 65535), + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::permute_copy.default`. + * + * Copies a tensor with its dimensions reordered — a gather where each output + * coordinate maps to the permuted input coordinate. + * Supports float32 and float16 with the `shader-f16` WebGPU feature extension. + */ +export const permuteCopy: Shader = { + name, + code: { + default: SHADER_PERMUTE, + transpose2d: SHADER_PERMUTE_TRANSPOSE2D, + defaultF16: SHADER_PERMUTE_F16, + transpose2dF16: SHADER_PERMUTE_TRANSPOSE2D_F16, + }, + recordIn, +}; diff --git a/src/shaders/permute/permute-f16.wgsl.ts b/src/shaders/permute/permute-f16.wgsl.ts new file mode 100644 index 0000000..4df462b --- /dev/null +++ b/src/shaders/permute/permute-f16.wgsl.ts @@ -0,0 +1,45 @@ +import { TENSOR_META_WGSL } from '../utils/meta'; + +/** + * WGSL compute shader for permuting tensor dimensions (f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_PERMUTE_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +${TENSOR_META_WGSL} +@group(0) @binding(2) var out_meta: TensorMeta; +@group(0) @binding(3) var in_meta: TensorMeta; + +struct Params { + perm: array, 2>, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (out_bufi >= out_meta.numel) { + return; + } + + var rem = out_bufi; + var in_bufi: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + let p = params.perm[d >> 2u][d & 3u]; + in_bufi = in_bufi + coord * in_meta.strides[p >> 2u][p & 3u]; + } + output[out_bufi] = input[in_bufi]; +} +`; diff --git a/src/shaders/permute/permute-transpose2d-f16.wgsl.ts b/src/shaders/permute/permute-transpose2d-f16.wgsl.ts new file mode 100644 index 0000000..6d606ee --- /dev/null +++ b/src/shaders/permute/permute-transpose2d-f16.wgsl.ts @@ -0,0 +1,41 @@ +/** + * WGSL compute shader for 2D matrix transpose using shared-memory tile (16x16 with 17-stride, f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_PERMUTE_TRANSPOSE2D_F16 = /* wgsl */ ` +enable f16; + +struct Transpose2dParams { + rows: u32, + cols: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var params: Transpose2dParams; + +var tile: array, 16>; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let x = wg_id.x * 16u + local_id.x; + let y = wg_id.y * 16u + local_id.y; + + if (x < params.cols && y < params.rows) { + tile[local_id.y][local_id.x] = input[y * params.cols + x]; + } + workgroupBarrier(); + + let out_x = wg_id.y * 16u + local_id.x; + let out_y = wg_id.x * 16u + local_id.y; + + if (out_x < params.rows && out_y < params.cols) { + output[out_y * params.rows + out_x] = tile[local_id.x][local_id.y]; + } +} +`; diff --git a/src/shaders/permute/permute-transpose2d.wgsl.ts b/src/shaders/permute/permute-transpose2d.wgsl.ts new file mode 100644 index 0000000..acd3a4d --- /dev/null +++ b/src/shaders/permute/permute-transpose2d.wgsl.ts @@ -0,0 +1,39 @@ +/** + * WGSL compute shader for 2D matrix transpose using shared-memory tile (16x16 with 17-stride, f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_PERMUTE_TRANSPOSE2D = /* wgsl */ ` +struct Transpose2dParams { + rows: u32, + cols: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var params: Transpose2dParams; + +var tile: array, 16>; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(workgroup_id) wg_id: vec3, + @builtin(local_invocation_id) local_id: vec3 +) { + let x = wg_id.x * 16u + local_id.x; + let y = wg_id.y * 16u + local_id.y; + + if (x < params.cols && y < params.rows) { + tile[local_id.y][local_id.x] = input[y * params.cols + x]; + } + workgroupBarrier(); + + let out_x = wg_id.y * 16u + local_id.x; + let out_y = wg_id.x * 16u + local_id.y; + + if (out_x < params.rows && out_y < params.cols) { + output[out_y * params.rows + out_x] = tile[local_id.x][local_id.y]; + } +} +`; diff --git a/src/shaders/permute/permute.wgsl.ts b/src/shaders/permute/permute.wgsl.ts new file mode 100644 index 0000000..4a21f04 --- /dev/null +++ b/src/shaders/permute/permute.wgsl.ts @@ -0,0 +1,43 @@ +import { TENSOR_META_WGSL } from '../utils/meta'; + +/** + * WGSL compute shader for permuting tensor dimensions (f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_PERMUTE = /* wgsl */ ` +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +${TENSOR_META_WGSL} +@group(0) @binding(2) var out_meta: TensorMeta; +@group(0) @binding(3) var in_meta: TensorMeta; + +struct Params { + perm: array, 2>, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let out_bufi = gid.x + gid.y * (num_workgroups.x * wg_size); + if (out_bufi >= out_meta.numel) { + return; + } + + var rem = out_bufi; + var in_bufi: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + let p = params.perm[d >> 2u][d & 3u]; + in_bufi = in_bufi + coord * in_meta.strides[p >> 2u][p & 3u]; + } + output[out_bufi] = input[in_bufi]; +} +`; diff --git a/src/shaders/prepack/index.ts b/src/shaders/prepack/index.ts new file mode 100644 index 0000000..8dea785 --- /dev/null +++ b/src/shaders/prepack/index.ts @@ -0,0 +1,37 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; + +import { recordFlatCopy } from '../utils/flat-copy'; + +const name = 'et_vk.prepack.default'; + +/** + * Positional argument tuple for `et_vk::prepack.default`: + * `[src, out]` + */ +export type PrepackArgs = readonly [src: Tensor, out: Tensor]; + +/** + * Internal dispatch builder for `et_vk::prepack.default`. + * + * @param ctx The execution context. + * @param args Positional argument tuple for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: PrepackArgs): void { + const [src, out] = args; + recordFlatCopy(ctx, src, out, name); +} + +/** + * WebGPU compute kernel for ExecuTorch `et_vk::prepack.default`. + * + * Materializes a constant into the prepack-output buffer — a single + * write-once flat copy (no layout transformation on the WebGPU path). + * Supports any dtype (float32, float16, etc.) via buffer copy. + */ +export const prepack: Shader = { + name, + code: '/* flat byte copy — no kernel dispatch */', + recordIn, +}; diff --git a/src/shaders/relu/index.ts b/src/shaders/relu/index.ts new file mode 100644 index 0000000..5c08429 --- /dev/null +++ b/src/shaders/relu/index.ts @@ -0,0 +1,94 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { SHADER_RELU } from './relu.wgsl'; +import { SHADER_RELU_VEC4 } from './relu-vec4.wgsl'; +import { SHADER_RELU_F16 } from './relu-f16.wgsl'; +import { SHADER_RELU_VEC4_F16 } from './relu-vec4-f16.wgsl'; + +const name = 'aten.relu.default'; + +/** + * Positional argument tuple for `aten::relu.default`: + * `[inTensor, outTensor]` + */ +export type ReluArgs = readonly [inTensor: Tensor, outTensor: Tensor]; + +/** + * Internal dispatch builder for `aten::relu.default`. + * + * @param ctx The execution context. + * @param args Positional arguments for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, [inTensor, outTensor]: ReluArgs): void { + if ( + (inTensor.dtype !== 'float32' && inTensor.dtype !== 'float16') || + outTensor.dtype !== inTensor.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + if (inTensor.numel !== outTensor.numel) { + throw new Error(`${name}: input numel != output numel`); + } + + const device = ctx.device; + const isF16 = inTensor.dtype === 'float16'; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + const wgSize = 256; + const numel = outTensor.numel; + const useVec4 = numel % 4 === 0; + const effectiveNumel = useVec4 ? numel / 4 : numel; + let code = useVec4 ? SHADER_RELU_VEC4 : SHADER_RELU; + if (isF16) { + code = useVec4 ? SHADER_RELU_VEC4_F16 : SHADER_RELU_F16; + } + + const totalWorkgroups = Math.ceil(effectiveNumel / wgSize); + const workgroupCountX = Math.min(totalWorkgroups, 65535); + const workgroupCountY = Math.ceil(totalWorkgroups / 65535); + + const buffer = new ArrayBuffer(16); + new Uint32Array(buffer)[0] = effectiveNumel; + const uniformBuffer = ctx.uniformBuffer(buffer); + + const bundle = createComputeBundle( + device, + code, + [ + { binding: 0, buffer: inTensor.buffer }, + { binding: 1, buffer: outTensor.buffer }, + { binding: 2, buffer: uniformBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize } + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::relu.default`. + * + * Element-wise rectified linear unit: max(x, 0). Automatically vectorizes + * using 128-bit vec4 or vec4 when tensor size is divisible by 4. + */ +export const relu: Shader = { + name, + code: { + scalar: SHADER_RELU, + vec4: SHADER_RELU_VEC4, + scalarF16: SHADER_RELU_F16, + vec4F16: SHADER_RELU_VEC4_F16, + }, + recordIn, +}; diff --git a/src/shaders/relu/relu-f16.wgsl.ts b/src/shaders/relu/relu-f16.wgsl.ts new file mode 100644 index 0000000..e840ae5 --- /dev/null +++ b/src/shaders/relu/relu-f16.wgsl.ts @@ -0,0 +1,30 @@ +/** + * WGSL compute shader for FP16 ReLU activation (scalar f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_RELU_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = max(input[idx], 0.0h); +} +`; diff --git a/src/shaders/relu/relu-vec4-f16.wgsl.ts b/src/shaders/relu/relu-vec4-f16.wgsl.ts new file mode 100644 index 0000000..4c4ce66 --- /dev/null +++ b/src/shaders/relu/relu-vec4-f16.wgsl.ts @@ -0,0 +1,30 @@ +/** + * WGSL compute shader for FP16 ReLU activation (vec4 64-bit/128-bit vectorized). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_RELU_VEC4_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input: array>; +@group(0) @binding(1) var output: array>; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = max(input[idx], vec4(0.0h)); +} +`; diff --git a/src/shaders/relu/relu-vec4.wgsl.ts b/src/shaders/relu/relu-vec4.wgsl.ts new file mode 100644 index 0000000..9dc34c0 --- /dev/null +++ b/src/shaders/relu/relu-vec4.wgsl.ts @@ -0,0 +1,28 @@ +/** + * WGSL compute shader for ReLU activation (vec4 128-bit vectorized). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_RELU_VEC4 = /* wgsl */ ` +@group(0) @binding(0) var input: array>; +@group(0) @binding(1) var output: array>; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = max(input[idx], vec4(0.0)); +} +`; diff --git a/src/shaders/relu/relu.wgsl.ts b/src/shaders/relu/relu.wgsl.ts new file mode 100644 index 0000000..357378f --- /dev/null +++ b/src/shaders/relu/relu.wgsl.ts @@ -0,0 +1,28 @@ +/** + * WGSL compute shader for ReLU activation (scalar f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_RELU = /* wgsl */ ` +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = max(input[idx], 0.0); +} +`; diff --git a/src/shaders/sigmoid/index.ts b/src/shaders/sigmoid/index.ts new file mode 100644 index 0000000..eda6f2b --- /dev/null +++ b/src/shaders/sigmoid/index.ts @@ -0,0 +1,94 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { SHADER_SIGMOID } from './sigmoid.wgsl'; +import { SHADER_SIGMOID_VEC4 } from './sigmoid-vec4.wgsl'; +import { SHADER_SIGMOID_F16 } from './sigmoid-f16.wgsl'; +import { SHADER_SIGMOID_VEC4_F16 } from './sigmoid-vec4-f16.wgsl'; + +const name = 'aten.sigmoid.default'; + +/** + * Positional argument tuple for `aten::sigmoid.default`: + * `[inTensor, outTensor]` + */ +export type SigmoidArgs = readonly [inTensor: Tensor, outTensor: Tensor]; + +/** + * Internal dispatch builder for `aten::sigmoid.default`. + * + * @param ctx The execution context. + * @param args Positional arguments for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, [inTensor, outTensor]: SigmoidArgs): void { + if ( + (inTensor.dtype !== 'float32' && inTensor.dtype !== 'float16') || + outTensor.dtype !== inTensor.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + if (inTensor.numel !== outTensor.numel) { + throw new Error(`${name}: input numel != output numel`); + } + + const device = ctx.device; + const isF16 = inTensor.dtype === 'float16'; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + const wgSize = 256; + const numel = outTensor.numel; + const useVec4 = numel % 4 === 0; + const effectiveNumel = useVec4 ? numel / 4 : numel; + let code = useVec4 ? SHADER_SIGMOID_VEC4 : SHADER_SIGMOID; + if (isF16) { + code = useVec4 ? SHADER_SIGMOID_VEC4_F16 : SHADER_SIGMOID_F16; + } + + const totalWorkgroups = Math.ceil(effectiveNumel / wgSize); + const workgroupCountX = Math.min(totalWorkgroups, 65535); + const workgroupCountY = Math.ceil(totalWorkgroups / 65535); + + const buffer = new ArrayBuffer(16); + new Uint32Array(buffer)[0] = effectiveNumel; + const uniformBuffer = ctx.uniformBuffer(buffer); + + const bundle = createComputeBundle( + device, + code, + [ + { binding: 0, buffer: inTensor.buffer }, + { binding: 1, buffer: outTensor.buffer }, + { binding: 2, buffer: uniformBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize } + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::sigmoid.default`. + * + * Element-wise sigmoid activation: 1 / (1 + exp(-x)). Automatically + * vectorizes using 128-bit vec4 or vec4 when tensor size is divisible by 4. + */ +export const sigmoid: Shader = { + name, + code: { + scalar: SHADER_SIGMOID, + vec4: SHADER_SIGMOID_VEC4, + scalarF16: SHADER_SIGMOID_F16, + vec4F16: SHADER_SIGMOID_VEC4_F16, + }, + recordIn, +}; diff --git a/src/shaders/sigmoid/sigmoid-f16.wgsl.ts b/src/shaders/sigmoid/sigmoid-f16.wgsl.ts new file mode 100644 index 0000000..cd2913c --- /dev/null +++ b/src/shaders/sigmoid/sigmoid-f16.wgsl.ts @@ -0,0 +1,30 @@ +/** + * WGSL compute shader for FP16 Sigmoid activation (scalar f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SIGMOID_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = 1.0h / (1.0h + exp(-input[idx])); +} +`; diff --git a/src/shaders/sigmoid/sigmoid-vec4-f16.wgsl.ts b/src/shaders/sigmoid/sigmoid-vec4-f16.wgsl.ts new file mode 100644 index 0000000..54a93a7 --- /dev/null +++ b/src/shaders/sigmoid/sigmoid-vec4-f16.wgsl.ts @@ -0,0 +1,30 @@ +/** + * WGSL compute shader for FP16 Sigmoid activation (vec4 64-bit/128-bit vectorized). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SIGMOID_VEC4_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input: array>; +@group(0) @binding(1) var output: array>; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = vec4(1.0h) / (vec4(1.0h) + exp(-input[idx])); +} +`; diff --git a/src/shaders/sigmoid/sigmoid-vec4.wgsl.ts b/src/shaders/sigmoid/sigmoid-vec4.wgsl.ts new file mode 100644 index 0000000..c95eca0 --- /dev/null +++ b/src/shaders/sigmoid/sigmoid-vec4.wgsl.ts @@ -0,0 +1,28 @@ +/** + * WGSL compute shader for Sigmoid activation (vec4 128-bit vectorized). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SIGMOID_VEC4 = /* wgsl */ ` +@group(0) @binding(0) var input: array>; +@group(0) @binding(1) var output: array>; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = vec4(1.0) / (vec4(1.0) + exp(-input[idx])); +} +`; diff --git a/src/shaders/sigmoid/sigmoid.wgsl.ts b/src/shaders/sigmoid/sigmoid.wgsl.ts new file mode 100644 index 0000000..93ddaa5 --- /dev/null +++ b/src/shaders/sigmoid/sigmoid.wgsl.ts @@ -0,0 +1,28 @@ +/** + * WGSL compute shader for Sigmoid activation (scalar f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SIGMOID = /* wgsl */ ` +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = 1.0 / (1.0 + exp(-input[idx])); +} +`; diff --git a/src/shaders/softmax/index.ts b/src/shaders/softmax/index.ts new file mode 100644 index 0000000..dbd9188 --- /dev/null +++ b/src/shaders/softmax/index.ts @@ -0,0 +1,150 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { SHADER_SOFTMAX } from './softmax.wgsl'; +import { SHADER_SOFTMAX_LAST_DIM } from './softmax-last-dim.wgsl'; +import { SHADER_SOFTMAX_LAST_DIM_VEC4 } from './softmax-last-dim-vec4.wgsl'; +import { SHADER_SOFTMAX_F16 } from './softmax-f16.wgsl'; +import { SHADER_SOFTMAX_LAST_DIM_F16 } from './softmax-last-dim-f16.wgsl'; +import { SHADER_SOFTMAX_LAST_DIM_VEC4_F16 } from './softmax-last-dim-vec4-f16.wgsl'; + +const name = 'aten._softmax.default'; + +/** + * Positional argument tuple for `aten::_softmax.default`: + * `[in, dim, half_to_float, out]` + */ +export type SoftmaxArgs = readonly [ + inTensor: Tensor, + dim: number, + halfToFloat: boolean, + outTensor: Tensor, +]; + +/** + * Decomposes dims into [outer, R, inner] for a reduction along `dim`. + * + * @param dims The tensor shape. + * @param dim The reduction dimension (may be negative). + * @param opName Operator name for error messages. + * @returns Tuple of [outer, r, inner] extents. + */ +function decompose( + dims: readonly number[], + dim: number, + opName: string +): [outer: number, r: number, inner: number] { + const ndim = dims.length; + const d = dim < 0 ? dim + ndim : dim; + if (ndim === 0 || d < 0 || d >= ndim) { + throw new Error(`${opName}: dim out of range`); + } + let outer = 1; + let inner = 1; + for (let i = 0; i < d; i++) { + outer *= dims[i]; + } + for (let i = d + 1; i < ndim; i++) { + inner *= dims[i]; + } + return [outer, dims[d], inner]; +} + +/** + * Internal dispatch builder for `aten::_softmax.default`. + * + * @param ctx The execution context. + * @param args Positional argument tuple for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: SoftmaxArgs): void { + const [inTensor, dim, , outTensor] = args; + + if ( + (inTensor.dtype !== 'float32' && inTensor.dtype !== 'float16') || + outTensor.dtype !== inTensor.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + if (inTensor.shape.length === 0) { + throw new Error(`${name}: scalar input unsupported`); + } + if (inTensor.numel !== outTensor.numel) { + throw new Error(`${name}: input/output numel mismatch`); + } + + const device = ctx.device; + const isF16 = inTensor.dtype === 'float16'; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + const [outer, r, inner] = decompose(inTensor.shape, dim, name); + if (outer === 0 || r === 0 || inner === 0) { + throw new Error(`${name}: zero-sized reduction`); + } + + const wgSize = 256; + + let code = isF16 ? SHADER_SOFTMAX_F16 : SHADER_SOFTMAX; + let lines = outer * inner; + let paramsArr = [outer, r, inner, 0]; + + if (inner === 1) { + if (r % 4 === 0) { + code = isF16 ? SHADER_SOFTMAX_LAST_DIM_VEC4_F16 : SHADER_SOFTMAX_LAST_DIM_VEC4; + paramsArr = [outer, r / 4, 1, 0]; + lines = outer; + } else { + code = isF16 ? SHADER_SOFTMAX_LAST_DIM_F16 : SHADER_SOFTMAX_LAST_DIM; + paramsArr = [outer, r, 1, 0]; + lines = outer; + } + } + + const paramsBuffer = ctx.uniformBuffer(new Uint32Array(paramsArr)); + + const bundle = createComputeBundle( + device, + code, + [ + { binding: 0, buffer: inTensor.buffer }, + { binding: 1, buffer: outTensor.buffer }, + { binding: 2, buffer: paramsBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize } + ); + + const totalWorkgroups = Math.ceil(lines / wgSize); + const workgroupCountX = Math.min(totalWorkgroups, 65535); + const workgroupCountY = Math.ceil(totalWorkgroups / 65535); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::_softmax.default`. + * + * Computes softmax along one dimension with a numerically stable online + * (flash-style) max/sum pass per reduction line. + * Supports float32 and float16 with the `shader-f16` WebGPU feature extension. + */ +export const softmax: Shader = { + name, + code: { + default: SHADER_SOFTMAX, + lastDim: SHADER_SOFTMAX_LAST_DIM, + lastDimVec4: SHADER_SOFTMAX_LAST_DIM_VEC4, + defaultF16: SHADER_SOFTMAX_F16, + lastDimF16: SHADER_SOFTMAX_LAST_DIM_F16, + lastDimVec4F16: SHADER_SOFTMAX_LAST_DIM_VEC4_F16, + }, + recordIn, +}; diff --git a/src/shaders/softmax/softmax-f16.wgsl.ts b/src/shaders/softmax/softmax-f16.wgsl.ts new file mode 100644 index 0000000..978f4fa --- /dev/null +++ b/src/shaders/softmax/softmax-f16.wgsl.ts @@ -0,0 +1,50 @@ +/** + * WGSL compute shader for general dimension softmax (f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SOFTMAX_F16 = /* wgsl */ ` +enable f16; + +struct Params { + outer_: u32, + r_: u32, + inner_: u32, + pad_: u32, +}; + +@group(0) @binding(0) var inp: array; +@group(0) @binding(1) var out: array; +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let t = gid.x; + let lines = params.outer_ * params.inner_; + if (t >= lines) { + return; + } + let oo = t / params.inner_; + let ii = t % params.inner_; + let base = oo * params.r_ * params.inner_ + ii; + + var mx: f16 = -65504.0h; + var ssum: f16 = 0.0h; + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let v = inp[base + r * params.inner_]; + if (v > mx) { + ssum = ssum * exp(mx - v); + mx = v; + } + ssum = ssum + exp(v - mx); + } + let inv_ssum = 1.0h / ssum; + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let idx = base + r * params.inner_; + out[idx] = exp(inp[idx] - mx) * inv_ssum; + } +} +`; diff --git a/src/shaders/softmax/softmax-last-dim-f16.wgsl.ts b/src/shaders/softmax/softmax-last-dim-f16.wgsl.ts new file mode 100644 index 0000000..d061e87 --- /dev/null +++ b/src/shaders/softmax/softmax-last-dim-f16.wgsl.ts @@ -0,0 +1,49 @@ +/** + * WGSL compute shader for contiguous last-dimension softmax (inner_ == 1, f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SOFTMAX_LAST_DIM_F16 = /* wgsl */ ` +enable f16; + +struct Params { + outer_: u32, + r_: u32, + inner_: u32, + pad_: u32, +}; + +@group(0) @binding(0) var inp: array; +@group(0) @binding(1) var out: array; +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3 +) { + let row = gid.x + gid.y * (num_workgroups.x * wg_size); + if (row >= params.outer_) { + return; + } + let base = row * params.r_; + var mx: f16 = -65504.0h; + var ssum: f16 = 0.0h; + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let v = inp[base + r]; + if (v > mx) { + ssum = ssum * exp(mx - v); + mx = v; + } + ssum = ssum + exp(v - mx); + } + let inv_ssum = 1.0h / ssum; + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let idx = base + r; + out[idx] = exp(inp[idx] - mx) * inv_ssum; + } +} +`; diff --git a/src/shaders/softmax/softmax-last-dim-vec4-f16.wgsl.ts b/src/shaders/softmax/softmax-last-dim-vec4-f16.wgsl.ts new file mode 100644 index 0000000..d525cfd --- /dev/null +++ b/src/shaders/softmax/softmax-last-dim-vec4-f16.wgsl.ts @@ -0,0 +1,51 @@ +/** + * WGSL compute shader for vectorized contiguous last-dimension softmax (vec4). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SOFTMAX_LAST_DIM_VEC4_F16 = /* wgsl */ ` +enable f16; + +struct Params { + outer_: u32, + r_vecs: u32, + inner_: u32, + pad_: u32, +}; + +@group(0) @binding(0) var inp: array>; +@group(0) @binding(1) var out: array>; +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3 +) { + let row = gid.x + gid.y * (num_workgroups.x * wg_size); + if (row >= params.outer_) { + return; + } + let base = row * params.r_vecs; + var mx: f16 = -65504.0h; + var ssum: f16 = 0.0h; + for (var r: u32 = 0u; r < params.r_vecs; r = r + 1u) { + let v = inp[base + r]; + let row_max = max(max(v.x, v.y), max(v.z, v.w)); + if (row_max > mx) { + ssum = ssum * exp(mx - row_max); + mx = row_max; + } + let e = exp(v - vec4(mx)); + ssum = ssum + e.x + e.y + e.z + e.w; + } + let inv_ssum = 1.0h / ssum; + for (var r: u32 = 0u; r < params.r_vecs; r = r + 1u) { + let idx = base + r; + out[idx] = exp(inp[idx] - vec4(mx)) * inv_ssum; + } +} +`; diff --git a/src/shaders/softmax/softmax-last-dim-vec4.wgsl.ts b/src/shaders/softmax/softmax-last-dim-vec4.wgsl.ts new file mode 100644 index 0000000..0043fac --- /dev/null +++ b/src/shaders/softmax/softmax-last-dim-vec4.wgsl.ts @@ -0,0 +1,49 @@ +/** + * WGSL compute shader for vectorized 128-bit contiguous last-dimension softmax (vec4). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SOFTMAX_LAST_DIM_VEC4 = /* wgsl */ ` +struct Params { + outer_: u32, + r_vecs: u32, + inner_: u32, + pad_: u32, +}; + +@group(0) @binding(0) var inp: array>; +@group(0) @binding(1) var out: array>; +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3 +) { + let row = gid.x + gid.y * (num_workgroups.x * wg_size); + if (row >= params.outer_) { + return; + } + let base = row * params.r_vecs; + var mx: f32 = -3.4e38; + var ssum: f32 = 0.0; + for (var r: u32 = 0u; r < params.r_vecs; r = r + 1u) { + let v = inp[base + r]; + let row_max = max(max(v.x, v.y), max(v.z, v.w)); + if (row_max > mx) { + ssum = ssum * exp(mx - row_max); + mx = row_max; + } + let e = exp(v - vec4(mx)); + ssum = ssum + e.x + e.y + e.z + e.w; + } + let inv_ssum = 1.0 / ssum; + for (var r: u32 = 0u; r < params.r_vecs; r = r + 1u) { + let idx = base + r; + out[idx] = exp(inp[idx] - vec4(mx)) * inv_ssum; + } +} +`; diff --git a/src/shaders/softmax/softmax-last-dim.wgsl.ts b/src/shaders/softmax/softmax-last-dim.wgsl.ts new file mode 100644 index 0000000..1968f80 --- /dev/null +++ b/src/shaders/softmax/softmax-last-dim.wgsl.ts @@ -0,0 +1,48 @@ +/** + * WGSL compute shader for contiguous last-dimension softmax (inner_ == 1, f32). + * Eliminates per-thread integer divisions and modulos. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SOFTMAX_LAST_DIM = /* wgsl */ ` +struct Params { + outer_: u32, + r_: u32, + inner_: u32, + pad_: u32, +}; + +@group(0) @binding(0) var inp: array; +@group(0) @binding(1) var out: array; +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3 +) { + let row = gid.x + gid.y * (num_workgroups.x * wg_size); + if (row >= params.outer_) { + return; + } + let base = row * params.r_; + var mx: f32 = -3.4e38; + var ssum: f32 = 0.0; + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let v = inp[base + r]; + if (v > mx) { + ssum = ssum * exp(mx - v); + mx = v; + } + ssum = ssum + exp(v - mx); + } + let inv_ssum = 1.0 / ssum; + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let idx = base + r; + out[idx] = exp(inp[idx] - mx) * inv_ssum; + } +} +`; diff --git a/src/shaders/softmax/softmax.wgsl.ts b/src/shaders/softmax/softmax.wgsl.ts new file mode 100644 index 0000000..9fa7ca7 --- /dev/null +++ b/src/shaders/softmax/softmax.wgsl.ts @@ -0,0 +1,48 @@ +/** + * WGSL compute shader for general dimension softmax (f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SOFTMAX = /* wgsl */ ` +struct Params { + outer_: u32, + r_: u32, + inner_: u32, + pad_: u32, +}; + +@group(0) @binding(0) var inp: array; +@group(0) @binding(1) var out: array; +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let t = gid.x; + let lines = params.outer_ * params.inner_; + if (t >= lines) { + return; + } + let oo = t / params.inner_; + let ii = t % params.inner_; + let base = oo * params.r_ * params.inner_ + ii; + // Online (flash-style) softmax; -3.4e38 seeds max (Tint overflows -FLT_MAX). + var mx: f32 = -3.4e38; + var ssum: f32 = 0.0; + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let v = inp[base + r * params.inner_]; + if (v > mx) { + ssum = ssum * exp(mx - v); + mx = v; + } + ssum = ssum + exp(v - mx); + } + let inv_ssum = 1.0 / ssum; + for (var r: u32 = 0u; r < params.r_; r = r + 1u) { + let idx = base + r * params.inner_; + out[idx] = exp(inp[idx] - mx) * inv_ssum; + } +} +`; diff --git a/src/shaders/split-with-sizes/index.ts b/src/shaders/split-with-sizes/index.ts new file mode 100644 index 0000000..f02f90c --- /dev/null +++ b/src/shaders/split-with-sizes/index.ts @@ -0,0 +1,143 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { SHADER_SPLIT_WITH_SIZES } from './split-with-sizes.wgsl'; +import { SHADER_SPLIT_WITH_SIZES_BLOCK } from './split-with-sizes-block.wgsl'; +import { SHADER_SPLIT_WITH_SIZES_BLOCK_VEC4 } from './split-with-sizes-block-vec4.wgsl'; +import { SHADER_SPLIT_WITH_SIZES_F16 } from './split-with-sizes-f16.wgsl'; +import { SHADER_SPLIT_WITH_SIZES_BLOCK_F16 } from './split-with-sizes-block-f16.wgsl'; +import { SHADER_SPLIT_WITH_SIZES_BLOCK_VEC4_F16 } from './split-with-sizes-block-vec4-f16.wgsl'; + +const name = 'aten.split_with_sizes_copy.default'; + +/** + * Positional argument tuple for `aten::split_with_sizes_copy.default`: + * `[in, split_sizes, dim, outs]` + */ +export type SplitWithSizesArgs = readonly [ + inTensor: Tensor, + splitSizes: readonly number[], + dim: number, + outs: readonly Tensor[], +]; + +/** + * Internal dispatch builder for `aten::split_with_sizes_copy.default`. + * + * @param ctx The execution context. + * @param args Positional argument tuple for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: SplitWithSizesArgs): void { + const [inTensor, splitSizes, dimArg, outs] = args; + + const ndim = inTensor.shape.length; + let dim = dimArg; + if (dim < 0) { + dim += ndim; + } + if (dim < 0 || dim >= ndim) { + throw new Error(`${name}: dim out of range`); + } + if (outs.length !== splitSizes.length) { + throw new Error(`${name}: outputs (${outs.length}) != sizes (${splitSizes.length})`); + } + + let sizesSum = 0; + for (const s of splitSizes) { + if (s < 0) { + throw new Error(`${name}: negative split size`); + } + sizesSum += s; + } + if (sizesSum !== inTensor.shape[dim]) { + throw new Error(`${name}: sizes must sum to the split dim extent`); + } + + const isF16 = inTensor.dtype === 'float16'; + const device = ctx.device; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + let innerSize = 1; + for (let d = dim + 1; d < ndim; d++) { + innerSize *= inTensor.shape[d]; + } + + const wgSize = 256; + const cIn = inTensor.shape[dim]; + let start = 0; + + for (let i = 0; i < outs.length; i++) { + const out = outs[i]; + if ( + (inTensor.dtype !== 'float32' && inTensor.dtype !== 'float16') || + out.dtype !== inTensor.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + if (out.shape[dim] !== splitSizes[i]) { + throw new Error(`${name}: output ${i} dim ${dim} != split size`); + } + + const cOut = splitSizes[i]; + const isVec4 = innerSize % 4 === 0; + const effectiveInner = isVec4 ? innerSize / 4 : innerSize; + const totalElements = isVec4 ? out.numel / 4 : out.numel; + + const paramsBuffer = ctx.uniformBuffer( + new Uint32Array([cIn, cOut, effectiveInner, start, totalElements, 0, 0, 0]) + ); + start += cOut; + + let code = isVec4 ? SHADER_SPLIT_WITH_SIZES_BLOCK_VEC4 : SHADER_SPLIT_WITH_SIZES_BLOCK; + if (isF16) { + code = isVec4 ? SHADER_SPLIT_WITH_SIZES_BLOCK_VEC4_F16 : SHADER_SPLIT_WITH_SIZES_BLOCK_F16; + } + + const bundle = createComputeBundle( + device, + code, + [ + { binding: 0, buffer: inTensor.buffer }, + { binding: 1, buffer: out.buffer }, + { binding: 2, buffer: paramsBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize } + ); + + const totalWorkgroups = Math.ceil(totalElements / wgSize); + const workgroupCountX = Math.min(totalWorkgroups, 65535); + const workgroupCountY = Math.ceil(totalWorkgroups / 65535); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); + } +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::split_with_sizes_copy.default`. + * + * Splits a tensor into contiguous chunks along one dimension using fast block-gather + * and 128-bit vectorization. + * Supports float32 and float16 with the `shader-f16` WebGPU feature extension. + */ +export const splitWithSizesCopy: Shader = { + name, + code: { + default: SHADER_SPLIT_WITH_SIZES, + block: SHADER_SPLIT_WITH_SIZES_BLOCK, + blockVec4: SHADER_SPLIT_WITH_SIZES_BLOCK_VEC4, + defaultF16: SHADER_SPLIT_WITH_SIZES_F16, + blockF16: SHADER_SPLIT_WITH_SIZES_BLOCK_F16, + blockVec4F16: SHADER_SPLIT_WITH_SIZES_BLOCK_VEC4_F16, + }, + recordIn, +}; diff --git a/src/shaders/split-with-sizes/split-with-sizes-block-f16.wgsl.ts b/src/shaders/split-with-sizes/split-with-sizes-block-f16.wgsl.ts new file mode 100644 index 0000000..ea203c7 --- /dev/null +++ b/src/shaders/split-with-sizes/split-with-sizes-block-f16.wgsl.ts @@ -0,0 +1,44 @@ +/** + * WGSL compute shader for contiguous block-gather tensor slicing (f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SPLIT_WITH_SIZES_BLOCK_F16 = /* wgsl */ ` +enable f16; + +struct BlockParams { + c_in: u32, + c_out: u32, + inner_size: u32, + start: u32, + total_elements: u32, + _p0: u32, + _p1: u32, + _p2: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var params: BlockParams; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3 +) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.total_elements) { + return; + } + let inner_idx = idx % params.inner_size; + let slice_idx = idx / params.inner_size; + let c = slice_idx % params.c_out; + let b = slice_idx / params.c_out; + + let in_idx = (b * params.c_in + params.start + c) * params.inner_size + inner_idx; + output[idx] = input[in_idx]; +} +`; diff --git a/src/shaders/split-with-sizes/split-with-sizes-block-vec4-f16.wgsl.ts b/src/shaders/split-with-sizes/split-with-sizes-block-vec4-f16.wgsl.ts new file mode 100644 index 0000000..d5afde4 --- /dev/null +++ b/src/shaders/split-with-sizes/split-with-sizes-block-vec4-f16.wgsl.ts @@ -0,0 +1,44 @@ +/** + * WGSL compute shader for 128-bit vectorized contiguous block-gather tensor slicing (vec4). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SPLIT_WITH_SIZES_BLOCK_VEC4_F16 = /* wgsl */ ` +enable f16; + +struct BlockParams { + c_in: u32, + c_out: u32, + inner_size: u32, + start: u32, + total_elements: u32, + _p0: u32, + _p1: u32, + _p2: u32, +}; + +@group(0) @binding(0) var input: array>; +@group(0) @binding(1) var output: array>; +@group(0) @binding(2) var params: BlockParams; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3 +) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.total_elements) { + return; + } + let inner_idx = idx % params.inner_size; + let slice_idx = idx / params.inner_size; + let c = slice_idx % params.c_out; + let b = slice_idx / params.c_out; + + let in_idx = (b * params.c_in + params.start + c) * params.inner_size + inner_idx; + output[idx] = input[in_idx]; +} +`; diff --git a/src/shaders/split-with-sizes/split-with-sizes-block-vec4.wgsl.ts b/src/shaders/split-with-sizes/split-with-sizes-block-vec4.wgsl.ts new file mode 100644 index 0000000..e378678 --- /dev/null +++ b/src/shaders/split-with-sizes/split-with-sizes-block-vec4.wgsl.ts @@ -0,0 +1,42 @@ +/** + * WGSL compute shader for 128-bit vectorized contiguous block-gather tensor slicing (vec4). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SPLIT_WITH_SIZES_BLOCK_VEC4 = /* wgsl */ ` +struct BlockParams { + c_in: u32, + c_out: u32, + inner_size: u32, + start: u32, + total_elements: u32, + _p0: u32, + _p1: u32, + _p2: u32, +}; + +@group(0) @binding(0) var input: array>; +@group(0) @binding(1) var output: array>; +@group(0) @binding(2) var params: BlockParams; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3 +) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.total_elements) { + return; + } + let inner_idx = idx % params.inner_size; + let slice_idx = idx / params.inner_size; + let c = slice_idx % params.c_out; + let b = slice_idx / params.c_out; + + let in_idx = (b * params.c_in + params.start + c) * params.inner_size + inner_idx; + output[idx] = input[in_idx]; +} +`; diff --git a/src/shaders/split-with-sizes/split-with-sizes-block.wgsl.ts b/src/shaders/split-with-sizes/split-with-sizes-block.wgsl.ts new file mode 100644 index 0000000..3bfb6f6 --- /dev/null +++ b/src/shaders/split-with-sizes/split-with-sizes-block.wgsl.ts @@ -0,0 +1,42 @@ +/** + * WGSL compute shader for contiguous block-gather tensor slicing (f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SPLIT_WITH_SIZES_BLOCK = /* wgsl */ ` +struct BlockParams { + c_in: u32, + c_out: u32, + inner_size: u32, + start: u32, + total_elements: u32, + _p0: u32, + _p1: u32, + _p2: u32, +}; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; +@group(0) @binding(2) var params: BlockParams; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3 +) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.total_elements) { + return; + } + let inner_idx = idx % params.inner_size; + let slice_idx = idx / params.inner_size; + let c = slice_idx % params.c_out; + let b = slice_idx / params.c_out; + + let in_idx = (b * params.c_in + params.start + c) * params.inner_size + inner_idx; + output[idx] = input[in_idx]; +} +`; diff --git a/src/shaders/split-with-sizes/split-with-sizes-f16.wgsl.ts b/src/shaders/split-with-sizes/split-with-sizes-f16.wgsl.ts new file mode 100644 index 0000000..27ea99e --- /dev/null +++ b/src/shaders/split-with-sizes/split-with-sizes-f16.wgsl.ts @@ -0,0 +1,48 @@ +import { TENSOR_META_WGSL } from '../utils/meta'; + +/** + * WGSL compute shader for general tensor slicing along one dimension (f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SPLIT_WITH_SIZES_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +${TENSOR_META_WGSL} +@group(0) @binding(2) var out_meta: TensorMeta; +@group(0) @binding(3) var in_meta: TensorMeta; + +struct Params { + dim: u32, + start: u32, + step: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let out_bufi = gid.x; + if (out_bufi >= out_meta.numel) { + return; + } + + var rem = out_bufi; + var in_bufi: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + var in_coord = coord; + if (d == params.dim) { + in_coord = params.start + coord * params.step; + } + in_bufi = in_bufi + in_coord * in_meta.strides[d >> 2u][d & 3u]; + } + output[out_bufi] = input[in_bufi]; +} +`; diff --git a/src/shaders/split-with-sizes/split-with-sizes.wgsl.ts b/src/shaders/split-with-sizes/split-with-sizes.wgsl.ts new file mode 100644 index 0000000..7d1bea2 --- /dev/null +++ b/src/shaders/split-with-sizes/split-with-sizes.wgsl.ts @@ -0,0 +1,46 @@ +import { TENSOR_META_WGSL } from '../utils/meta'; + +/** + * WGSL compute shader for general tensor slicing along one dimension (f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SPLIT_WITH_SIZES = /* wgsl */ ` +@group(0) @binding(0) var input: array; +@group(0) @binding(1) var output: array; + +${TENSOR_META_WGSL} +@group(0) @binding(2) var out_meta: TensorMeta; +@group(0) @binding(3) var in_meta: TensorMeta; + +struct Params { + dim: u32, + start: u32, + step: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let out_bufi = gid.x; + if (out_bufi >= out_meta.numel) { + return; + } + + var rem = out_bufi; + var in_bufi: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + var in_coord = coord; + if (d == params.dim) { + in_coord = params.start + coord * params.step; + } + in_bufi = in_bufi + in_coord * in_meta.strides[d >> 2u][d & 3u]; + } + output[out_bufi] = input[in_bufi]; +} +`; diff --git a/src/shaders/squeeze-copy/index.ts b/src/shaders/squeeze-copy/index.ts new file mode 100644 index 0000000..229a65f --- /dev/null +++ b/src/shaders/squeeze-copy/index.ts @@ -0,0 +1,40 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; + +import { recordFlatCopy } from '../utils/flat-copy'; + +const name = 'aten.squeeze_copy.dims'; + +/** + * Positional argument tuple for `aten::squeeze_copy.dims`: + * `[in, dims, out]` + */ +export type SqueezeCopyArgs = readonly [ + inTensor: Tensor, + dims: readonly number[], + outTensor: Tensor, +]; + +/** + * Internal dispatch builder for `aten::squeeze_copy.dims`. + * + * @param ctx The execution context. + * @param args Positional argument tuple for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: SqueezeCopyArgs): void { + const [inTensor, , outTensor] = args; + recordFlatCopy(ctx, inTensor, outTensor, name); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::squeeze_copy.dims`. + * + * Squeeze = numel-preserving flat copy. + * Supports any dtype (float32, float16, etc.) via buffer copy. + */ +export const squeezeCopy: Shader = { + name, + code: '/* flat byte copy — no kernel dispatch */', + recordIn, +}; diff --git a/src/shaders/sub/index.ts b/src/shaders/sub/index.ts new file mode 100644 index 0000000..538d05f --- /dev/null +++ b/src/shaders/sub/index.ts @@ -0,0 +1,146 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { isBroadcastable } from '../utils/broadcasting'; +import { encodeTensorMeta } from '../utils/meta'; + +import { SHADER_SUB_BROADCAST } from './sub-broadcast.wgsl'; +import { SHADER_SUB_BROADCAST_F16 } from './sub-broadcast-f16.wgsl'; +import { SHADER_SUB_ELEMENTWISE } from './sub-elementwise.wgsl'; +import { SHADER_SUB_ELEMENTWISE_F16 } from './sub-elementwise-f16.wgsl'; +import { SHADER_SUB_ELEMENTWISE_VEC4 } from './sub-elementwise-vec4.wgsl'; +import { SHADER_SUB_ELEMENTWISE_VEC4_F16 } from './sub-elementwise-vec4-f16.wgsl'; + +const name = 'aten.sub.Tensor'; + +/** + * Positional argument tuple for `aten::sub.Tensor`: + * Standard ExecuTorch schema: `[in1, in2, alpha, out]` + * Or 3-arg variant: `[in1, in2, out]` + */ +export type SubArgs = + | readonly [in1: Tensor, in2: Tensor, out: Tensor] + | readonly [in1: Tensor, in2: Tensor, alpha: number, out: Tensor]; + +/** + * Internal dispatch builder for `aten::sub.Tensor`. + * + * @param ctx The execution context. + * @param args Positional argument tuple for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: SubArgs): void { + const [in1, in2] = args; + const [alpha, out] = args.length === 4 ? [args[2], args[3]] : [1.0, args[2]]; + + if ( + (in1.dtype !== 'float32' && in1.dtype !== 'float16') || + in2.dtype !== in1.dtype || + out.dtype !== in1.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + if (!isBroadcastable(in1.shape, out.shape) || !isBroadcastable(in2.shape, out.shape)) { + throw new Error(`${name}: Input shapes are not broadcastable to output shape`); + } + + const device = ctx.device; + const isF16 = in1.dtype === 'float16'; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + const sameShape = + in1.shape.length === out.shape.length && + in2.shape.length === out.shape.length && + in1.shape.every((s, i) => s === out.shape[i]) && + in2.shape.every((s, i) => s === out.shape[i]); + + const wgSize = 256; + + if (sameShape) { + const isVec4 = out.numel % 4 === 0; + let code = isVec4 ? SHADER_SUB_ELEMENTWISE_VEC4 : SHADER_SUB_ELEMENTWISE; + if (isF16) { + code = isVec4 ? SHADER_SUB_ELEMENTWISE_VEC4_F16 : SHADER_SUB_ELEMENTWISE_F16; + } + const elements = isVec4 ? out.numel / 4 : out.numel; + const totalWorkgroups = Math.ceil(elements / wgSize); + const workgroupCountX = Math.min(totalWorkgroups, 65535); + const workgroupCountY = Math.ceil(totalWorkgroups / 65535); + const paramsBuffer = ctx.uniformBuffer(new Uint32Array([elements])); + + const bundle = createComputeBundle( + device, + code, + [ + { binding: 0, buffer: in1.buffer }, + { binding: 1, buffer: in2.buffer }, + { binding: 2, buffer: out.buffer }, + { binding: 3, buffer: paramsBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize, alpha } + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); + return; + } + + const outRank = out.shape.length; + const outMetaBuffer = ctx.uniformBuffer(encodeTensorMeta(out)); + const in1MetaBuffer = ctx.uniformBuffer(encodeTensorMeta(in1, outRank)); + const in2MetaBuffer = ctx.uniformBuffer(encodeTensorMeta(in2, outRank)); + + const totalWorkgroups = Math.ceil(out.numel / wgSize); + const workgroupCountX = Math.min(totalWorkgroups, 65535); + const workgroupCountY = Math.ceil(totalWorkgroups / 65535); + + const bundle = createComputeBundle( + device, + isF16 ? SHADER_SUB_BROADCAST_F16 : SHADER_SUB_BROADCAST, + [ + { binding: 0, buffer: in1.buffer }, + { binding: 1, buffer: in2.buffer }, + { binding: 2, buffer: out.buffer }, + { binding: 3, buffer: outMetaBuffer }, + { binding: 4, buffer: in1MetaBuffer }, + { binding: 5, buffer: in2MetaBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize, alpha } + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::sub.Tensor`. + * + * Computes elementwise `output = input1 - alpha * input2` with NumPy/PyTorch-style + * broadcasting (right-aligned, dimensions equal or 1, up to rank 8). Features an + * automatic fast path when input and output shapes match. + */ +export const sub: Shader = { + name, + code: { + default: SHADER_SUB_BROADCAST, + elementwise: SHADER_SUB_ELEMENTWISE, + elementwiseVec4: SHADER_SUB_ELEMENTWISE_VEC4, + defaultF16: SHADER_SUB_BROADCAST_F16, + elementwiseF16: SHADER_SUB_ELEMENTWISE_F16, + elementwiseVec4F16: SHADER_SUB_ELEMENTWISE_VEC4_F16, + }, + recordIn, +}; diff --git a/src/shaders/sub/sub-broadcast-f16.wgsl.ts b/src/shaders/sub/sub-broadcast-f16.wgsl.ts new file mode 100644 index 0000000..17df559 --- /dev/null +++ b/src/shaders/sub/sub-broadcast-f16.wgsl.ts @@ -0,0 +1,60 @@ +import { TENSOR_META_WGSL } from '../utils/meta'; + +/** + * WGSL compute shader for FP16 binary subtraction with broadcast support. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SUB_BROADCAST_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +${TENSOR_META_WGSL} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in1_meta: TensorMeta; +@group(0) @binding(5) var in2_meta: TensorMeta; + +override wg_size: u32 = 64u; +override alpha: f32 = 1.0; + +fn op_f16(a: f16, b: f16) -> f16 { + return a - f16(alpha) * b; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= out_meta.numel) { + return; + } + + var same = true; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { + same = false; + } + } + if (same) { + output[idx] = op_f16(input1[idx], input2[idx]); + return; + } + + var rem = idx; + var l1: u32 = 0u; + var l2: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; + } + output[idx] = op_f16(input1[l1], input2[l2]); +} +`; diff --git a/src/shaders/sub/sub-broadcast.wgsl.ts b/src/shaders/sub/sub-broadcast.wgsl.ts new file mode 100644 index 0000000..cb24b6d --- /dev/null +++ b/src/shaders/sub/sub-broadcast.wgsl.ts @@ -0,0 +1,58 @@ +import { TENSOR_META_WGSL } from '../utils/meta'; + +/** + * WGSL compute shader for binary subtraction with broadcast support. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SUB_BROADCAST = /* wgsl */ ` +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +${TENSOR_META_WGSL} +@group(0) @binding(3) var out_meta: TensorMeta; +@group(0) @binding(4) var in1_meta: TensorMeta; +@group(0) @binding(5) var in2_meta: TensorMeta; + +override wg_size: u32 = 64u; +override alpha: f32 = 1.0; + +fn op(a: f32, b: f32) -> f32 { + return a - alpha * b; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= out_meta.numel) { + return; + } + + var same = true; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || + in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { + same = false; + } + } + if (same) { + output[idx] = op(input1[idx], input2[idx]); + return; + } + + var rem = idx; + var l1: u32 = 0u; + var l2: u32 = 0u; + for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { + let coord = rem / out_meta.strides[d >> 2u][d & 3u]; + rem = rem % out_meta.strides[d >> 2u][d & 3u]; + l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; + l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; + } + output[idx] = op(input1[l1], input2[l2]); +} +`; diff --git a/src/shaders/sub/sub-elementwise-f16.wgsl.ts b/src/shaders/sub/sub-elementwise-f16.wgsl.ts new file mode 100644 index 0000000..5552827 --- /dev/null +++ b/src/shaders/sub/sub-elementwise-f16.wgsl.ts @@ -0,0 +1,32 @@ +/** + * WGSL compute shader for contiguous 1D FP16 elementwise subtraction. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SUB_ELEMENTWISE_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 256u; +override alpha: f32 = 1.0; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = input1[idx] - f16(alpha) * input2[idx]; +} +`; diff --git a/src/shaders/sub/sub-elementwise-vec4-f16.wgsl.ts b/src/shaders/sub/sub-elementwise-vec4-f16.wgsl.ts new file mode 100644 index 0000000..5e989c0 --- /dev/null +++ b/src/shaders/sub/sub-elementwise-vec4-f16.wgsl.ts @@ -0,0 +1,32 @@ +/** + * WGSL compute shader for contiguous 64-bit/128-bit vectorized (vec4) elementwise subtraction. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SUB_ELEMENTWISE_VEC4_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var input1: array>; +@group(0) @binding(1) var input2: array>; +@group(0) @binding(2) var output: array>; + +struct Params { + num_elements: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 256u; +override alpha: f32 = 1.0; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = input1[idx] - f16(alpha) * input2[idx]; +} +`; diff --git a/src/shaders/sub/sub-elementwise-vec4.wgsl.ts b/src/shaders/sub/sub-elementwise-vec4.wgsl.ts new file mode 100644 index 0000000..425a6ef --- /dev/null +++ b/src/shaders/sub/sub-elementwise-vec4.wgsl.ts @@ -0,0 +1,30 @@ +/** + * WGSL compute shader for contiguous 128-bit vectorized (vec4) elementwise subtraction. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SUB_ELEMENTWISE_VEC4 = /* wgsl */ ` +@group(0) @binding(0) var input1: array>; +@group(0) @binding(1) var input2: array>; +@group(0) @binding(2) var output: array>; + +struct Params { + num_elements: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 256u; +override alpha: f32 = 1.0; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = input1[idx] - alpha * input2[idx]; +} +`; diff --git a/src/shaders/sub/sub-elementwise.wgsl.ts b/src/shaders/sub/sub-elementwise.wgsl.ts new file mode 100644 index 0000000..4cee557 --- /dev/null +++ b/src/shaders/sub/sub-elementwise.wgsl.ts @@ -0,0 +1,30 @@ +/** + * WGSL compute shader for contiguous 1D elementwise subtraction. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_SUB_ELEMENTWISE = /* wgsl */ ` +@group(0) @binding(0) var input1: array; +@group(0) @binding(1) var input2: array; +@group(0) @binding(2) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 256u; +override alpha: f32 = 1.0; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let idx = gid.x + gid.y * (num_workgroups.x * wg_size); + if (idx >= params.num_elements) { + return; + } + output[idx] = input1[idx] - alpha * input2[idx]; +} +`; diff --git a/src/shaders/unsqueeze-copy/index.ts b/src/shaders/unsqueeze-copy/index.ts new file mode 100644 index 0000000..4630ac5 --- /dev/null +++ b/src/shaders/unsqueeze-copy/index.ts @@ -0,0 +1,36 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; + +import { recordFlatCopy } from '../utils/flat-copy'; + +const name = 'aten.unsqueeze_copy.default'; + +/** + * Positional argument tuple for `aten::unsqueeze_copy.default`: + * `[in, dim, out]` + */ +export type UnsqueezeCopyArgs = readonly [inTensor: Tensor, dim: number, outTensor: Tensor]; + +/** + * Internal dispatch builder for `aten::unsqueeze_copy.default`. + * + * @param ctx The execution context. + * @param args Positional argument tuple for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: UnsqueezeCopyArgs): void { + const [inTensor, , outTensor] = args; + recordFlatCopy(ctx, inTensor, outTensor, name); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::unsqueeze_copy.default`. + * + * Unsqueeze = numel-preserving flat copy. + * Supports any dtype (float32, float16, etc.) via buffer copy. + */ +export const unsqueezeCopy: Shader = { + name, + code: '/* flat byte copy — no kernel dispatch */', + recordIn, +}; diff --git a/src/shaders/upsample-bilinear2d/index.ts b/src/shaders/upsample-bilinear2d/index.ts new file mode 100644 index 0000000..549c577 --- /dev/null +++ b/src/shaders/upsample-bilinear2d/index.ts @@ -0,0 +1,132 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { SHADER_UPSAMPLE_BILINEAR2D } from './upsample-bilinear2d.wgsl'; +import { SHADER_UPSAMPLE_BILINEAR2D_SPATIAL } from './upsample-bilinear2d-spatial.wgsl'; +import { SHADER_UPSAMPLE_BILINEAR2D_F16 } from './upsample-bilinear2d-f16.wgsl'; +import { SHADER_UPSAMPLE_BILINEAR2D_SPATIAL_F16 } from './upsample-bilinear2d-spatial-f16.wgsl'; + +const name = 'aten.upsample_bilinear2d.vec'; + +/** + * Positional argument tuple for `aten::upsample_bilinear2d.vec`: + * `[in, output_size, align_corners, scales, out]` + */ +export type UpsampleBilinear2dArgs = readonly [ + inTensor: Tensor, + outputSize: readonly number[], + alignCorners: boolean, + scales: readonly number[] | undefined, + outTensor: Tensor, +]; + +/** + * Internal dispatch builder for `aten::upsample_bilinear2d.vec`. + * + * @param ctx The execution context. + * @param args Positional arguments for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: UpsampleBilinear2dArgs): void { + const inTensor = args[0]; + const alignCorners = args[2]; + const outTensor = args[4]; + + if ( + (inTensor.dtype !== 'float32' && inTensor.dtype !== 'float16') || + outTensor.dtype !== inTensor.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + if (inTensor.shape.length !== 4 || outTensor.shape.length !== 4) { + throw new Error(`${name}: expected 4D in/out tensors`); + } + + const device = ctx.device; + const isF16 = inTensor.dtype === 'float16'; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + const [n, c, ih, iw] = inTensor.shape as [number, number, number, number]; + const [outN, outC, oh, ow] = outTensor.shape as [number, number, number, number]; + + if (outN !== n || outC !== c) { + throw new Error(`${name}: N/C mismatch`); + } + if (ih === 0 || iw === 0 || oh === 0 || ow === 0) { + throw new Error(`${name}: zero spatial dim`); + } + + const outNumel = outTensor.numel; + const buffer = new ArrayBuffer(32); + const u32View = new Uint32Array(buffer); + u32View.set([n, c, ih, iw, oh, ow, alignCorners ? 1 : 0, 0]); + + const uniformBuffer = ctx.uniformBuffer(buffer); + + // Use 2D spatial grid when dimensions fit in workgroup limits + if (n * c <= 65535 && Math.ceil(ow / 16) <= 65535 && Math.ceil(oh / 16) <= 65535) { + const bundle = createComputeBundle( + device, + isF16 ? SHADER_UPSAMPLE_BILINEAR2D_SPATIAL_F16 : SHADER_UPSAMPLE_BILINEAR2D_SPATIAL, + [ + { binding: 0, buffer: outTensor.buffer }, + { binding: 1, buffer: inTensor.buffer }, + { binding: 2, buffer: uniformBuffer }, + ] + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX: Math.ceil(ow / 16), + workgroupCountY: Math.ceil(oh / 16), + workgroupCountZ: n * c, + }); + return; + } + + const wgSize = 256; + const totalWorkgroups = Math.ceil(outNumel / wgSize); + const workgroupCountX = Math.min(totalWorkgroups, 65535); + const workgroupCountY = Math.ceil(totalWorkgroups / 65535); + const strideX = workgroupCountX * wgSize; + + const bundle = createComputeBundle( + device, + isF16 ? SHADER_UPSAMPLE_BILINEAR2D_F16 : SHADER_UPSAMPLE_BILINEAR2D, + [ + { binding: 0, buffer: outTensor.buffer }, + { binding: 1, buffer: inTensor.buffer }, + { binding: 2, buffer: uniformBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { stride_x: strideX, wg_size: wgSize } + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::upsample_bilinear2d.vec`. + * + * Implements bilinear 2D upsampling for NCHW float32 and float16 tensors. + * Supports float32 and float16 with the `shader-f16` WebGPU feature extension. + */ +export const upsampleBilinear2d: Shader = { + name, + code: { + default: SHADER_UPSAMPLE_BILINEAR2D, + spatial: SHADER_UPSAMPLE_BILINEAR2D_SPATIAL, + defaultF16: SHADER_UPSAMPLE_BILINEAR2D_F16, + spatialF16: SHADER_UPSAMPLE_BILINEAR2D_SPATIAL_F16, + }, + recordIn, +}; diff --git a/src/shaders/upsample-bilinear2d/upsample-bilinear2d-f16.wgsl.ts b/src/shaders/upsample-bilinear2d/upsample-bilinear2d-f16.wgsl.ts new file mode 100644 index 0000000..2e9fd46 --- /dev/null +++ b/src/shaders/upsample-bilinear2d/upsample-bilinear2d-f16.wgsl.ts @@ -0,0 +1,77 @@ +/** + * WGSL compute shader for bilinear 2D upsampling (f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_UPSAMPLE_BILINEAR2D_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + align_corners: u32, + _p0: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; + +fn src_index(dst: u32, insz: u32, outsz: u32, align: u32) -> f32 { + if (align == 1u) { + if (outsz <= 1u) { + return 0.0; + } + return f32(dst) * f32(insz - 1u) / f32(outsz - 1u); + } + return (f32(dst) + 0.5) * f32(insz) / f32(outsz) - 0.5; +} + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.N * params.C * params.OH * params.OW; + let i = gid.y * stride_x + gid.x; + if (i >= total) { + return; + } + let ow = i % params.OW; + let oh = (i / params.OW) % params.OH; + let c = (i / (params.OW * params.OH)) % params.C; + let n = i / (params.OW * params.OH * params.C); + + let sh = src_index(oh, params.IH, params.OH, params.align_corners); + let sw = src_index(ow, params.IW, params.OW, params.align_corners); + + let h0 = i32(floor(sh)); + let w0 = i32(floor(sw)); + let lh = f16(sh - f32(h0)); + let lw = f16(sw - f32(w0)); + + let ih_max = i32(params.IH) - 1; + let iw_max = i32(params.IW) - 1; + let h0c = u32(clamp(h0, 0, ih_max)); + let h1c = u32(clamp(h0 + 1, 0, ih_max)); + let w0c = u32(clamp(w0, 0, iw_max)); + let w1c = u32(clamp(w0 + 1, 0, iw_max)); + + let base = (n * params.C + c) * params.IH; + let r0 = (base + h0c) * params.IW; + let r1 = (base + h1c) * params.IW; + let v00 = inp[r0 + w0c]; + let v01 = inp[r0 + w1c]; + let v10 = inp[r1 + w0c]; + let v11 = inp[r1 + w1c]; + + let top = v00 + (v01 - v00) * lw; + let bot = v10 + (v11 - v10) * lw; + out[i] = top + (bot - top) * lh; +} +`; diff --git a/src/shaders/upsample-bilinear2d/upsample-bilinear2d-spatial-f16.wgsl.ts b/src/shaders/upsample-bilinear2d/upsample-bilinear2d-spatial-f16.wgsl.ts new file mode 100644 index 0000000..d53c1af --- /dev/null +++ b/src/shaders/upsample-bilinear2d/upsample-bilinear2d-spatial-f16.wgsl.ts @@ -0,0 +1,76 @@ +/** + * WGSL compute shader for bilinear 2D upsampling using 2D spatial grid (f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_UPSAMPLE_BILINEAR2D_SPATIAL_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + align_corners: u32, + _p0: u32, +} +@group(0) @binding(2) var params: Params; + +fn src_index(dst: u32, insz: u32, outsz: u32, align: u32) -> f32 { + if (align == 1u) { + if (outsz <= 1u) { + return 0.0; + } + return f32(dst) * f32(insz - 1u) / f32(outsz - 1u); + } + return (f32(dst) + 0.5) * f32(insz) / f32(outsz) - 0.5; +} + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(workgroup_id) wg_id: vec3 +) { + let ow = gid.x; + let oh = gid.y; + let c_idx = wg_id.z; + + if (ow >= params.OW || oh >= params.OH) { + return; + } + + let sh = src_index(oh, params.IH, params.OH, params.align_corners); + let sw = src_index(ow, params.IW, params.OW, params.align_corners); + + let h0 = i32(floor(sh)); + let w0 = i32(floor(sw)); + let lh = f16(sh - f32(h0)); + let lw = f16(sw - f32(w0)); + + let ih_max = i32(params.IH) - 1; + let iw_max = i32(params.IW) - 1; + let h0c = u32(clamp(h0, 0, ih_max)); + let h1c = u32(clamp(h0 + 1, 0, ih_max)); + let w0c = u32(clamp(w0, 0, iw_max)); + let w1c = u32(clamp(w0 + 1, 0, iw_max)); + + let base = c_idx * params.IH; + let r0 = (base + h0c) * params.IW; + let r1 = (base + h1c) * params.IW; + let v00 = inp[r0 + w0c]; + let v01 = inp[r0 + w1c]; + let v10 = inp[r1 + w0c]; + let v11 = inp[r1 + w1c]; + + let top = v00 + (v01 - v00) * lw; + let bot = v10 + (v11 - v10) * lw; + let out_pos = (c_idx * params.OH + oh) * params.OW + ow; + out[out_pos] = top + (bot - top) * lh; +} +`; diff --git a/src/shaders/upsample-bilinear2d/upsample-bilinear2d-spatial.wgsl.ts b/src/shaders/upsample-bilinear2d/upsample-bilinear2d-spatial.wgsl.ts new file mode 100644 index 0000000..bef1d77 --- /dev/null +++ b/src/shaders/upsample-bilinear2d/upsample-bilinear2d-spatial.wgsl.ts @@ -0,0 +1,74 @@ +/** + * WGSL compute shader for bilinear 2D upsampling using 2D spatial grid (f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_UPSAMPLE_BILINEAR2D_SPATIAL = /* wgsl */ ` +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + align_corners: u32, + _p0: u32, +} +@group(0) @binding(2) var params: Params; + +fn src_index(dst: u32, insz: u32, outsz: u32, align: u32) -> f32 { + if (align == 1u) { + if (outsz <= 1u) { + return 0.0; + } + return f32(dst) * f32(insz - 1u) / f32(outsz - 1u); + } + return (f32(dst) + 0.5) * f32(insz) / f32(outsz) - 0.5; +} + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(workgroup_id) wg_id: vec3 +) { + let ow = gid.x; + let oh = gid.y; + let c_idx = wg_id.z; + + if (ow >= params.OW || oh >= params.OH) { + return; + } + + let sh = src_index(oh, params.IH, params.OH, params.align_corners); + let sw = src_index(ow, params.IW, params.OW, params.align_corners); + + let h0 = i32(floor(sh)); + let w0 = i32(floor(sw)); + let lh = sh - f32(h0); + let lw = sw - f32(w0); + + let ih_max = i32(params.IH) - 1; + let iw_max = i32(params.IW) - 1; + let h0c = u32(clamp(h0, 0, ih_max)); + let h1c = u32(clamp(h0 + 1, 0, ih_max)); + let w0c = u32(clamp(w0, 0, iw_max)); + let w1c = u32(clamp(w0 + 1, 0, iw_max)); + + let base = c_idx * params.IH; + let r0 = (base + h0c) * params.IW; + let r1 = (base + h1c) * params.IW; + let v00 = inp[r0 + w0c]; + let v01 = inp[r0 + w1c]; + let v10 = inp[r1 + w0c]; + let v11 = inp[r1 + w1c]; + + let top = v00 + (v01 - v00) * lw; + let bot = v10 + (v11 - v10) * lw; + let out_pos = (c_idx * params.OH + oh) * params.OW + ow; + out[out_pos] = top + (bot - top) * lh; +} +`; diff --git a/src/shaders/upsample-bilinear2d/upsample-bilinear2d.wgsl.ts b/src/shaders/upsample-bilinear2d/upsample-bilinear2d.wgsl.ts new file mode 100644 index 0000000..0959fbc --- /dev/null +++ b/src/shaders/upsample-bilinear2d/upsample-bilinear2d.wgsl.ts @@ -0,0 +1,75 @@ +/** + * WGSL compute shader for bilinear 2D upsampling (f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_UPSAMPLE_BILINEAR2D = /* wgsl */ ` +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + align_corners: u32, + _p0: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; + +fn src_index(dst: u32, insz: u32, outsz: u32, align: u32) -> f32 { + if (align == 1u) { + if (outsz <= 1u) { + return 0.0; + } + return f32(dst) * f32(insz - 1u) / f32(outsz - 1u); + } + return (f32(dst) + 0.5) * f32(insz) / f32(outsz) - 0.5; +} + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.N * params.C * params.OH * params.OW; + let i = gid.y * stride_x + gid.x; + if (i >= total) { + return; + } + let ow = i % params.OW; + let oh = (i / params.OW) % params.OH; + let c = (i / (params.OW * params.OH)) % params.C; + let n = i / (params.OW * params.OH * params.C); + + let sh = src_index(oh, params.IH, params.OH, params.align_corners); + let sw = src_index(ow, params.IW, params.OW, params.align_corners); + + let h0 = i32(floor(sh)); + let w0 = i32(floor(sw)); + let lh = sh - f32(h0); + let lw = sw - f32(w0); + + let ih_max = i32(params.IH) - 1; + let iw_max = i32(params.IW) - 1; + let h0c = u32(clamp(h0, 0, ih_max)); + let h1c = u32(clamp(h0 + 1, 0, ih_max)); + let w0c = u32(clamp(w0, 0, iw_max)); + let w1c = u32(clamp(w0 + 1, 0, iw_max)); + + let base = (n * params.C + c) * params.IH; + let r0 = (base + h0c) * params.IW; + let r1 = (base + h1c) * params.IW; + let v00 = inp[r0 + w0c]; + let v01 = inp[r0 + w1c]; + let v10 = inp[r1 + w0c]; + let v11 = inp[r1 + w1c]; + + let top = v00 + (v01 - v00) * lw; + let bot = v10 + (v11 - v10) * lw; + out[i] = top + (bot - top) * lh; +} +`; diff --git a/src/shaders/upsample-nearest-2d/index.ts b/src/shaders/upsample-nearest-2d/index.ts new file mode 100644 index 0000000..0c4fa66 --- /dev/null +++ b/src/shaders/upsample-nearest-2d/index.ts @@ -0,0 +1,121 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { SHADER_UPSAMPLE_NEAREST2D } from './upsample-nearest-2d.wgsl'; +import { SHADER_UPSAMPLE_NEAREST2D_SPATIAL } from './upsample-nearest-2d-spatial.wgsl'; +import { SHADER_UPSAMPLE_NEAREST2D_F16 } from './upsample-nearest-2d-f16.wgsl'; +import { SHADER_UPSAMPLE_NEAREST2D_SPATIAL_F16 } from './upsample-nearest-2d-spatial-f16.wgsl'; + +const name = 'aten.upsample_nearest2d.vec'; + +/** + * Positional argument tuple for `aten::upsample_nearest2d.vec`: + * `[in, output_size, scales, out]` + */ +export type UpsampleNearest2dArgs = readonly [ + inTensor: Tensor, + outputSize: undefined, + scales: readonly number[], + outTensor: Tensor, +]; + +/** + * Internal dispatch builder for `aten::upsample_nearest2d.vec`. + * + * @param ctx The execution context. + * @param args Positional argument tuple for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: UpsampleNearest2dArgs): void { + const [inTensor, , , outTensor] = args; + + if ( + (inTensor.dtype !== 'float32' && inTensor.dtype !== 'float16') || + outTensor.dtype !== inTensor.dtype + ) { + throw new Error(`${name}: Tensors must have matching float32 or float16 dtype`); + } + if (inTensor.shape.length !== 4 || outTensor.shape.length !== 4) { + throw new Error(`${name}: input and output must be 4D NCHW`); + } + const [b, c, ih, iw] = inTensor.shape as [number, number, number, number]; + const [ob, oc, oh, ow] = outTensor.shape as [number, number, number, number]; + if (ob !== b || oc !== c) { + throw new Error(`${name}: output batch/channels must match input`); + } + if (oh < ih || ow < iw) { + throw new Error(`${name}: output spatial dims must be >= input spatial dims`); + } + + const device = ctx.device; + const isF16 = inTensor.dtype === 'float16'; + if (isF16 && !device.features.has('shader-f16')) { + throw new Error(`${name}: Device does not support shader-f16 feature`); + } + + const paramsBuffer = ctx.uniformBuffer(new Uint32Array([b, c, ih, iw, oh, ow, 0, 0])); + + // Use 2D spatial grid when dimensions fit in workgroup limits + if (b * c <= 65535 && Math.ceil(ow / 16) <= 65535 && Math.ceil(oh / 16) <= 65535) { + const bundle = createComputeBundle( + device, + isF16 ? SHADER_UPSAMPLE_NEAREST2D_SPATIAL_F16 : SHADER_UPSAMPLE_NEAREST2D_SPATIAL, + [ + { binding: 0, buffer: outTensor.buffer }, + { binding: 1, buffer: inTensor.buffer }, + { binding: 2, buffer: paramsBuffer }, + ] + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX: Math.ceil(ow / 16), + workgroupCountY: Math.ceil(oh / 16), + workgroupCountZ: b * c, + }); + return; + } + + const wgSize = 256; + const total = b * c * oh * ow; + const workgroupCountX = Math.min(Math.ceil(total / wgSize), 65535); + + const bundle = createComputeBundle( + device, + isF16 ? SHADER_UPSAMPLE_NEAREST2D_F16 : SHADER_UPSAMPLE_NEAREST2D, + [ + { binding: 0, buffer: outTensor.buffer }, + { binding: 1, buffer: inTensor.buffer }, + { binding: 2, buffer: paramsBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize, stride_x: workgroupCountX * wgSize } + ); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY: Math.ceil(total / wgSize / 65535), + }); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::upsample_nearest2d.vec`. + * + * Nearest-neighbour 2D upsampling — each output texel copies the input texel + * at `floor(o * I / O)` (ATen convention). + * Supports float32 and float16 with the `shader-f16` WebGPU feature extension. + */ +export const upsampleNearest2d: Shader = { + name, + code: { + default: SHADER_UPSAMPLE_NEAREST2D, + spatial: SHADER_UPSAMPLE_NEAREST2D_SPATIAL, + defaultF16: SHADER_UPSAMPLE_NEAREST2D_F16, + spatialF16: SHADER_UPSAMPLE_NEAREST2D_SPATIAL_F16, + }, + recordIn, +}; diff --git a/src/shaders/upsample-nearest-2d/upsample-nearest-2d-f16.wgsl.ts b/src/shaders/upsample-nearest-2d/upsample-nearest-2d-f16.wgsl.ts new file mode 100644 index 0000000..4098112 --- /dev/null +++ b/src/shaders/upsample-nearest-2d/upsample-nearest-2d-f16.wgsl.ts @@ -0,0 +1,45 @@ +/** + * WGSL compute shader for nearest-neighbour 2D upsampling (f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_UPSAMPLE_NEAREST2D_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.N * params.C * params.OH * params.OW; + let i = gid.y * stride_x + gid.x; + if (i >= total) { + return; + } + let ow = i % params.OW; + let oh = (i / params.OW) % params.OH; + let c = (i / (params.OW * params.OH)) % params.C; + let n = i / (params.OW * params.OH * params.C); + + let ih = (oh * params.IH) / params.OH; + let iw = (ow * params.IW) / params.OW; + let in_idx = ((n * params.C + c) * params.IH + ih) * params.IW + iw; + out[i] = inp[in_idx]; +} +`; diff --git a/src/shaders/upsample-nearest-2d/upsample-nearest-2d-spatial-f16.wgsl.ts b/src/shaders/upsample-nearest-2d/upsample-nearest-2d-spatial-f16.wgsl.ts new file mode 100644 index 0000000..e7f5b5b --- /dev/null +++ b/src/shaders/upsample-nearest-2d/upsample-nearest-2d-spatial-f16.wgsl.ts @@ -0,0 +1,44 @@ +/** + * WGSL compute shader for nearest-neighbour 2D upsampling using 2D spatial grid (f16). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_UPSAMPLE_NEAREST2D_SPATIAL_F16 = /* wgsl */ ` +enable f16; + +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(2) var params: Params; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(workgroup_id) wg_id: vec3 +) { + let ow = gid.x; + let oh = gid.y; + let c_idx = wg_id.z; + + if (ow >= params.OW || oh >= params.OH) { + return; + } + + let ih = (oh * params.IH) / params.OH; + let iw = (ow * params.IW) / params.OW; + let in_idx = (c_idx * params.IH + ih) * params.IW + iw; + let out_idx = (c_idx * params.OH + oh) * params.OW + ow; + out[out_idx] = inp[in_idx]; +} +`; diff --git a/src/shaders/upsample-nearest-2d/upsample-nearest-2d-spatial.wgsl.ts b/src/shaders/upsample-nearest-2d/upsample-nearest-2d-spatial.wgsl.ts new file mode 100644 index 0000000..90038f9 --- /dev/null +++ b/src/shaders/upsample-nearest-2d/upsample-nearest-2d-spatial.wgsl.ts @@ -0,0 +1,42 @@ +/** + * WGSL compute shader for nearest-neighbour 2D upsampling using 2D spatial grid (f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_UPSAMPLE_NEAREST2D_SPATIAL = /* wgsl */ ` +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(2) var params: Params; + +@compute @workgroup_size(16, 16, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(workgroup_id) wg_id: vec3 +) { + let ow = gid.x; + let oh = gid.y; + let c_idx = wg_id.z; + + if (ow >= params.OW || oh >= params.OH) { + return; + } + + let ih = (oh * params.IH) / params.OH; + let iw = (ow * params.IW) / params.OW; + let in_idx = (c_idx * params.IH + ih) * params.IW + iw; + let out_idx = (c_idx * params.OH + oh) * params.OW + ow; + out[out_idx] = inp[in_idx]; +} +`; diff --git a/src/shaders/upsample-nearest-2d/upsample-nearest-2d.wgsl.ts b/src/shaders/upsample-nearest-2d/upsample-nearest-2d.wgsl.ts new file mode 100644 index 0000000..6180001 --- /dev/null +++ b/src/shaders/upsample-nearest-2d/upsample-nearest-2d.wgsl.ts @@ -0,0 +1,43 @@ +/** + * WGSL compute shader for nearest-neighbour 2D upsampling (f32). + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_UPSAMPLE_NEAREST2D = /* wgsl */ ` +@group(0) @binding(0) var out: array; +@group(0) @binding(1) var inp: array; + +struct Params { + N: u32, + C: u32, + IH: u32, + IW: u32, + OH: u32, + OW: u32, + _p0: u32, + _p1: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; +override stride_x: u32 = 4294967295u; + +@compute @workgroup_size(wg_size) +fn main(@builtin(global_invocation_id) gid: vec3) { + let total = params.N * params.C * params.OH * params.OW; + let i = gid.y * stride_x + gid.x; + if (i >= total) { + return; + } + let ow = i % params.OW; + let oh = (i / params.OW) % params.OH; + let c = (i / (params.OW * params.OH)) % params.C; + let n = i / (params.OW * params.OH * params.C); + + let ih = (oh * params.IH) / params.OH; + let iw = (ow * params.IW) / params.OW; + let in_idx = ((n * params.C + c) * params.IH + ih) * params.IW + iw; + out[i] = inp[in_idx]; +} +`; diff --git a/src/shaders/utils/flat-copy-vec4.wgsl.ts b/src/shaders/utils/flat-copy-vec4.wgsl.ts new file mode 100644 index 0000000..009b399 --- /dev/null +++ b/src/shaders/utils/flat-copy-vec4.wgsl.ts @@ -0,0 +1,28 @@ +/** + * WGSL compute shader for flat byte copy using 128-bit vec4 loads and stores. + * Coalesced, high-throughput memory streaming within the active compute pass. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_FLAT_COPY_VEC4 = /* wgsl */ ` +@group(0) @binding(0) var src: array>; +@group(0) @binding(1) var dst: array>; + +struct Params { + count: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3 +) { + let idx = gid.y * 65535u * wg_size + gid.x; + if (idx < params.count) { + dst[idx] = src[idx]; + } +} +`; diff --git a/src/shaders/utils/flat-copy.ts b/src/shaders/utils/flat-copy.ts index c72bc44..e78d3c3 100644 --- a/src/shaders/utils/flat-copy.ts +++ b/src/shaders/utils/flat-copy.ts @@ -1,10 +1,17 @@ import type { Tensor } from '../../tensor'; import type { WgpuExecutionContext } from '../../context'; +import { createComputeBundle } from '../../command'; + +import { SHADER_FLAT_COPY } from './flat-copy.wgsl'; +import { SHADER_FLAT_COPY_VEC4 } from './flat-copy-vec4.wgsl'; /** * Records a flat elementwise copy between two same-numel tensors — the * contiguous-reshape semantics of view/clone-style operators. * + * Uses an in-pass compute shader dispatch (128-bit vec4 memory streaming) + * so that memory moves do not break the active compute pass into multiple passes. + * * Aliased in/out buffers (the same `GPUBuffer`) are already in place: WebGPU * rejects overlapping copyBufferToBuffer, so the copy is skipped entirely. * @@ -32,5 +39,37 @@ export function recordFlatCopy( if (inTensor.buffer === outTensor.buffer) { return; // aliased in/out already in place } - ctx.copy(inTensor, outTensor); + + const totalU32 = Math.ceil(inTensor.byteLength / 4); + const useVec4 = totalU32 % 4 === 0; + const count = useVec4 ? totalU32 / 4 : totalU32; + const code = useVec4 ? SHADER_FLAT_COPY_VEC4 : SHADER_FLAT_COPY; + const wgSize = 256; + + const buffer = new ArrayBuffer(16); + new Uint32Array(buffer)[0] = count; + const uniformBuffer = ctx.uniformBuffer(buffer); + + const bundle = createComputeBundle( + ctx.device, + code, + [ + { binding: 0, buffer: inTensor.buffer }, + { binding: 1, buffer: outTensor.buffer }, + { binding: 2, buffer: uniformBuffer }, + ], + // eslint-disable-next-line camelcase, @typescript-eslint/naming-convention + { wg_size: wgSize } + ); + + const totalWorkgroups = Math.ceil(count / wgSize); + const workgroupCountX = Math.min(totalWorkgroups, 65535); + const workgroupCountY = Math.ceil(totalWorkgroups / 65535); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX, + workgroupCountY, + }); } diff --git a/src/shaders/utils/flat-copy.wgsl.ts b/src/shaders/utils/flat-copy.wgsl.ts new file mode 100644 index 0000000..c6251a2 --- /dev/null +++ b/src/shaders/utils/flat-copy.wgsl.ts @@ -0,0 +1,28 @@ +/** + * WGSL compute shader for flat byte copy using 32-bit u32 loads and stores. + * Coalesced, high-throughput memory streaming within the active compute pass. + * + * Source: wgpu-executorch (https://github.com/barhanc/wgpu-executorch) + * License: MIT (Copyright (c) 2026 Bartosz Hanc ) + */ +export const SHADER_FLAT_COPY = /* wgsl */ ` +@group(0) @binding(0) var src: array; +@group(0) @binding(1) var dst: array; + +struct Params { + count: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256u; + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3 +) { + let idx = gid.y * 65535u * wg_size + gid.x; + if (idx < params.count) { + dst[idx] = src[idx]; + } +} +`; diff --git a/src/shaders/view-copy/index.ts b/src/shaders/view-copy/index.ts new file mode 100644 index 0000000..019fa87 --- /dev/null +++ b/src/shaders/view-copy/index.ts @@ -0,0 +1,36 @@ +import type { Shader } from '../../shader'; +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; + +import { recordFlatCopy } from '../utils/flat-copy'; + +const name = 'aten.view_copy.default'; + +/** + * Positional argument tuple for `aten::view_copy.default`: + * `[in, size, out]` + */ +export type ViewCopyArgs = readonly [inTensor: Tensor, size: readonly number[], outTensor: Tensor]; + +/** + * Internal dispatch builder for `aten::view_copy.default`. + * + * @param ctx The execution context. + * @param args Positional argument tuple for the operator. + */ +function recordIn(ctx: WgpuExecutionContext, args: ViewCopyArgs): void { + const [inTensor, , outTensor] = args; + recordFlatCopy(ctx, inTensor, outTensor, name); +} + +/** + * WebGPU compute kernel for ExecuTorch `aten::view_copy.default`. + * + * Contiguous reshape = flat byte copy. + * Supports any dtype (float32, float16, etc.) via buffer copy. + */ +export const viewCopy: Shader = { + name, + code: '/* flat byte copy — no kernel dispatch */', + recordIn, +}; diff --git a/src/tensor.ts b/src/tensor.ts index 2f0f4c8..5f61d49 100644 --- a/src/tensor.ts +++ b/src/tensor.ts @@ -8,7 +8,14 @@ export type DType = * Supported typed arrays that can be uploaded to a Tensor. */ export type TypedArray = - Float32Array | Int32Array | Int8Array | Uint8Array | Uint16Array | Float64Array | BigInt64Array; + | Float32Array + | Float16Array + | Int32Array + | Int8Array + | Uint8Array + | Uint16Array + | Float64Array + | BigInt64Array; /** * Maximum tensor rank supported by WebGPU kernels (std140 layout limit). diff --git a/tests/benchmark.test.ts b/tests/benchmark.test.ts new file mode 100644 index 0000000..1854c40 --- /dev/null +++ b/tests/benchmark.test.ts @@ -0,0 +1,365 @@ +import { readFileSync } from 'node:fs'; +import { resolve } from 'node:path'; +import { describe, it } from 'vitest'; +import { WgpuExecutionContext, type Program, shaders } from '../src'; +import { setupGPU } from './helpers'; + +const REPO_ROOT = resolve(import.meta.dirname, '..'); + +function makeInput(numel: number): Float32Array { + const arr = new Float32Array(numel); + for (let i = 0; i < numel; i++) { + arr[i] = (i % 251) / 250 - 0.5; + } + return arr; +} + +function makeInputF16(numel: number): Float16Array { + const arr = new Float16Array(numel); + for (let i = 0; i < numel; i++) { + arr[i] = (i % 251) / 250 - 0.5; + } + return arr; +} + +function readArrayBuffer(path: string): ArrayBuffer { + const buf = readFileSync(path); + const copy = new ArrayBuffer(buf.byteLength); + new Uint8Array(copy).set(buf); + return copy; +} + +describe('Benchmarks', () => { + it('benchmark models and hot operators', { timeout: 600_000 }, async () => { + const device = await setupGPU(); + + // 1. DepthART benchmark + { + const ctx = new WgpuExecutionContext(device); + const progPath = resolve(REPO_ROOT, 'example/public/assets/depthart/program.json'); + const weightsPath = resolve(REPO_ROOT, 'example/public/assets/depthart/weights.bin'); + const program = JSON.parse(readFileSync(progPath, 'utf8')) as Program; + const weights = readArrayBuffer(weightsPath); + + const inp = ctx.tensor('float32', [1, 3, 448, 448]).setData(makeInput(3 * 448 * 448)); + const out = ctx.tensor('float32', [1, 1, 448, 448]); + ctx.recordProgram(program, weights, [inp, out]); + + // Warmup + ctx.submit(); + await ctx.sync(); + + const iters = 5; + const start = performance.now(); + for (let i = 0; i < iters; i++) { + ctx.submit(); + } + await ctx.sync(); + const avgMs = (performance.now() - start) / iters; + // eslint-disable-next-line no-console + console.log(`[BENCHMARK] DepthART (1x3x448x448): ${avgMs.toFixed(2)} ms`); + ctx.destroy(); + } + + // 2. YOLO26-N benchmark + { + const ctx = new WgpuExecutionContext(device); + const progPath = resolve(REPO_ROOT, 'example/public/assets/yolo26n/program.json'); + const weightsPath = resolve(REPO_ROOT, 'example/public/assets/yolo26n/weights.bin'); + const program = JSON.parse(readFileSync(progPath, 'utf8')) as Program; + const weights = readArrayBuffer(weightsPath); + + const inp = ctx.tensor('float32', [1, 3, 640, 640]).setData(makeInput(3 * 640 * 640)); + const out = ctx.tensor('float32', [1, 84, 8400]); + ctx.recordProgram(program, weights, [inp, out]); + + // Warmup + ctx.submit(); + await ctx.sync(); + + const iters = 5; + const start = performance.now(); + for (let i = 0; i < iters; i++) { + ctx.submit(); + } + await ctx.sync(); + const avgMs = (performance.now() - start) / iters; + // eslint-disable-next-line no-console + console.log(`[BENCHMARK] YOLO26-N (1x3x640x640): ${avgMs.toFixed(2)} ms`); + ctx.destroy(); + } + + // 3. Conv2D microbenchmark (typical bottleneck layer: 64ch -> 64ch, 56x56, 3x3) + { + const ctx = new WgpuExecutionContext(device); + const inp = ctx.tensor('float32', [1, 64, 56, 56]).setData(makeInput(1 * 64 * 56 * 56)); + const weight = ctx.tensor('float32', [64, 64, 3, 3]).setData(makeInput(64 * 64 * 3 * 3)); + const bias = ctx.tensor('float32', [64]).setData(makeInput(64)); + const out = ctx.tensor('float32', [1, 64, 56, 56]); + + ctx.recordShader(shaders.convolution, [ + inp, + weight, + bias, + [1, 1], // stride + [1, 1], // padding + [1, 1], // dilation + false, // transposed + [0, 0], // output_padding + 1, // groups + out, + ]); + + ctx.submit(); + await ctx.sync(); + + const iters = 20; + const start = performance.now(); + for (let i = 0; i < iters; i++) { + ctx.submit(); + } + await ctx.sync(); + const avgMs = (performance.now() - start) / iters; + // eslint-disable-next-line no-console + console.log(`[BENCHMARK] Conv2D [1,64,56,56]*[64,64,3,3] (FP32): ${avgMs.toFixed(3)} ms`); + ctx.destroy(); + + if (device.features.has('shader-f16')) { + const ctxF16 = new WgpuExecutionContext(device); + const inpF16 = ctxF16 + .tensor('float16', [1, 64, 56, 56]) + .setData(makeInputF16(1 * 64 * 56 * 56)); + const weightF16 = ctxF16 + .tensor('float16', [64, 64, 3, 3]) + .setData(makeInputF16(64 * 64 * 3 * 3)); + const biasF16 = ctxF16.tensor('float16', [64]).setData(makeInputF16(64)); + const outF16 = ctxF16.tensor('float16', [1, 64, 56, 56]); + + ctxF16.recordShader(shaders.convolution, [ + inpF16, + weightF16, + biasF16, + [1, 1], + [1, 1], + [1, 1], + false, + [0, 0], + 1, + outF16, + ]); + + ctxF16.submit(); + await ctxF16.sync(); + + const startF16 = performance.now(); + for (let i = 0; i < iters; i++) { + ctxF16.submit(); + } + await ctxF16.sync(); + const avgMsF16 = (performance.now() - startF16) / iters; + const speedup = avgMs / avgMsF16; + // eslint-disable-next-line no-console + console.log( + `[BENCHMARK] Conv2D [1,64,56,56]*[64,64,3,3] (FP16): ${avgMsF16.toFixed(3)} ms (${speedup.toFixed(2)}x speedup)` + ); + ctxF16.destroy(); + } + } + + // 4. BMM microbenchmark (typical attention shape: [4, 64, 64] * [4, 64, 64]) + { + const ctx = new WgpuExecutionContext(device); + const a = ctx.tensor('float32', [4, 64, 64]).setData(makeInput(4 * 64 * 64)); + const b = ctx.tensor('float32', [4, 64, 64]).setData(makeInput(4 * 64 * 64)); + const out = ctx.tensor('float32', [4, 64, 64]); + + ctx.recordShader(shaders.bmm, [a, b, out]); + + ctx.submit(); + await ctx.sync(); + + const iters = 50; + const start = performance.now(); + for (let i = 0; i < iters; i++) { + ctx.submit(); + } + await ctx.sync(); + const avgMs = (performance.now() - start) / iters; + // eslint-disable-next-line no-console + console.log(`[BENCHMARK] BMM [4,64,64]*[4,64,64] (FP32): ${avgMs.toFixed(3)} ms`); + ctx.destroy(); + + if (device.features.has('shader-f16')) { + const ctxF16 = new WgpuExecutionContext(device); + const aF16 = ctxF16.tensor('float16', [4, 64, 64]).setData(makeInputF16(4 * 64 * 64)); + const bF16 = ctxF16.tensor('float16', [4, 64, 64]).setData(makeInputF16(4 * 64 * 64)); + const outF16 = ctxF16.tensor('float16', [4, 64, 64]); + + ctxF16.recordShader(shaders.bmm, [aF16, bF16, outF16]); + + ctxF16.submit(); + await ctxF16.sync(); + + const startF16 = performance.now(); + for (let i = 0; i < iters; i++) { + ctxF16.submit(); + } + await ctxF16.sync(); + const avgMsF16 = (performance.now() - startF16) / iters; + const speedup = avgMs / avgMsF16; + // eslint-disable-next-line no-console + console.log( + `[BENCHMARK] BMM [4,64,64]*[4,64,64] (FP16): ${avgMsF16.toFixed(3)} ms (${speedup.toFixed(2)}x speedup)` + ); + ctxF16.destroy(); + } + } + + // 5. MM microbenchmark (GEMM [256, 256] * [256, 256]) + { + const ctx = new WgpuExecutionContext(device); + const a = ctx.tensor('float32', [256, 256]).setData(makeInput(256 * 256)); + const b = ctx.tensor('float32', [256, 256]).setData(makeInput(256 * 256)); + const out = ctx.tensor('float32', [256, 256]); + + ctx.recordShader(shaders.mm, [a, b, out]); + + ctx.submit(); + await ctx.sync(); + + const iters = 30; + const start = performance.now(); + for (let i = 0; i < iters; i++) { + ctx.submit(); + } + await ctx.sync(); + const avgMs = (performance.now() - start) / iters; + // eslint-disable-next-line no-console + console.log(`[BENCHMARK] MM [256,256]*[256,256] (FP32): ${avgMs.toFixed(3)} ms`); + ctx.destroy(); + + if (device.features.has('shader-f16')) { + const ctxF16 = new WgpuExecutionContext(device); + const aF16 = ctxF16.tensor('float16', [256, 256]).setData(makeInputF16(256 * 256)); + const bF16 = ctxF16.tensor('float16', [256, 256]).setData(makeInputF16(256 * 256)); + const outF16 = ctxF16.tensor('float16', [256, 256]); + + ctxF16.recordShader(shaders.mm, [aF16, bF16, outF16]); + + ctxF16.submit(); + await ctxF16.sync(); + + const startF16 = performance.now(); + for (let i = 0; i < iters; i++) { + ctxF16.submit(); + } + await ctxF16.sync(); + const avgMsF16 = (performance.now() - startF16) / iters; + const speedup = avgMs / avgMsF16; + // eslint-disable-next-line no-console + console.log( + `[BENCHMARK] MM [256,256]*[256,256] (FP16): ${avgMsF16.toFixed(3)} ms (${speedup.toFixed(2)}x speedup)` + ); + ctxF16.destroy(); + } + } + + // 6. Large MM microbenchmark (GEMM [1024, 1024] * [1024, 1024]) + { + const ctx = new WgpuExecutionContext(device); + const a = ctx.tensor('float32', [1024, 1024]).setData(makeInput(1024 * 1024)); + const b = ctx.tensor('float32', [1024, 1024]).setData(makeInput(1024 * 1024)); + const out = ctx.tensor('float32', [1024, 1024]); + + ctx.recordShader(shaders.mm, [a, b, out]); + + ctx.submit(); + await ctx.sync(); + + const iters = 15; + const start = performance.now(); + for (let i = 0; i < iters; i++) { + ctx.submit(); + } + await ctx.sync(); + const avgMs = (performance.now() - start) / iters; + // eslint-disable-next-line no-console + console.log(`[BENCHMARK] MM [1024,1024]*[1024,1024] (FP32): ${avgMs.toFixed(3)} ms`); + ctx.destroy(); + + if (device.features.has('shader-f16')) { + const ctxF16 = new WgpuExecutionContext(device); + const aF16 = ctxF16.tensor('float16', [1024, 1024]).setData(makeInputF16(1024 * 1024)); + const bF16 = ctxF16.tensor('float16', [1024, 1024]).setData(makeInputF16(1024 * 1024)); + const outF16 = ctxF16.tensor('float16', [1024, 1024]); + + ctxF16.recordShader(shaders.mm, [aF16, bF16, outF16]); + + ctxF16.submit(); + await ctxF16.sync(); + + const startF16 = performance.now(); + for (let i = 0; i < iters; i++) { + ctxF16.submit(); + } + await ctxF16.sync(); + const avgMsF16 = (performance.now() - startF16) / iters; + const speedup = avgMs / avgMsF16; + // eslint-disable-next-line no-console + console.log( + `[BENCHMARK] MM [1024,1024]*[1024,1024] (FP16): ${avgMsF16.toFixed(3)} ms (${speedup.toFixed(2)}x speedup)` + ); + ctxF16.destroy(); + } + } + + // 7. Elementwise microbenchmark (Mul [1, 64, 56, 56]) + { + const ctx = new WgpuExecutionContext(device); + const a = ctx.tensor('float32', [1, 64, 56, 56]).setData(makeInput(64 * 56 * 56)); + const b = ctx.tensor('float32', [1, 64, 56, 56]).setData(makeInput(64 * 56 * 56)); + const out = ctx.tensor('float32', [1, 64, 56, 56]); + + ctx.recordShader(shaders.mul, [a, b, out]); + + ctx.submit(); + await ctx.sync(); + + const iters = 50; + const start = performance.now(); + for (let i = 0; i < iters; i++) { + ctx.submit(); + } + await ctx.sync(); + const avgMs = (performance.now() - start) / iters; + // eslint-disable-next-line no-console + console.log(`[BENCHMARK] Mul [1,64,56,56] (FP32): ${avgMs.toFixed(3)} ms`); + ctx.destroy(); + + if (device.features.has('shader-f16')) { + const ctxF16 = new WgpuExecutionContext(device); + const aF16 = ctxF16.tensor('float16', [1, 64, 56, 56]).setData(makeInputF16(64 * 56 * 56)); + const bF16 = ctxF16.tensor('float16', [1, 64, 56, 56]).setData(makeInputF16(64 * 56 * 56)); + const outF16 = ctxF16.tensor('float16', [1, 64, 56, 56]); + + ctxF16.recordShader(shaders.mul, [aF16, bF16, outF16]); + + ctxF16.submit(); + await ctxF16.sync(); + + const startF16 = performance.now(); + for (let i = 0; i < iters; i++) { + ctxF16.submit(); + } + await ctxF16.sync(); + const avgMsF16 = (performance.now() - startF16) / iters; + const speedup = avgMs / avgMsF16; + // eslint-disable-next-line no-console + console.log( + `[BENCHMARK] Mul [1,64,56,56] (FP16): ${avgMsF16.toFixed(3)} ms (${speedup.toFixed(2)}x speedup)` + ); + ctxF16.destroy(); + } + } + }); +}); diff --git a/tests/helpers/assert.ts b/tests/helpers/assert.ts index 1323739..657340e 100644 --- a/tests/helpers/assert.ts +++ b/tests/helpers/assert.ts @@ -12,7 +12,8 @@ export async function expectTensorClose( expected: number[], tolerance = 1e-5 ): Promise { - const actual = new Float32Array(await tensor.getData()); + const data = await tensor.getData(); + const actual = tensor.dtype === 'float16' ? new Float16Array(data) : new Float32Array(data); if (actual.length !== expected.length) { throw new Error(`Tensor length mismatch: got ${actual.length}, expected ${expected.length}`); @@ -21,6 +22,11 @@ export async function expectTensorClose( let maxDiff = 0; let maxIdx = 0; for (let i = 0; i < actual.length; i++) { + if (!Number.isFinite(actual[i]) || !Number.isFinite(expected[i])) { + throw new Error( + `Non-finite tensor value at index ${i}: got ${actual[i]}, expected ${expected[i]}` + ); + } const diff = Math.abs(actual[i] - expected[i]); if (diff > maxDiff) { maxDiff = diff; diff --git a/tests/helpers/gpu.ts b/tests/helpers/gpu.ts index 6b84c8b..c6f55c0 100644 --- a/tests/helpers/gpu.ts +++ b/tests/helpers/gpu.ts @@ -12,6 +12,10 @@ export async function setupGPU(): Promise { if (!adapter) { throw new Error('No suitable GPU adapter found'); } - cached = await adapter.requestDevice(); + const requiredFeatures: GPUFeatureName[] = []; + if (adapter.features.has('shader-f16')) { + requiredFeatures.push('shader-f16'); + } + cached = await adapter.requestDevice({ requiredFeatures }); return cached; } diff --git a/tests/program/depth-pipeline.test.ts b/tests/program/depth-pipeline.test.ts new file mode 100644 index 0000000..a098f09 --- /dev/null +++ b/tests/program/depth-pipeline.test.ts @@ -0,0 +1,145 @@ +import { readFileSync } from 'node:fs'; +import { resolve } from 'node:path'; +import { expect, it } from 'vitest'; +import { WgpuExecutionContext, type Program } from '../../src'; +import { + clearStats, + depthStatsF16, + drawDepthF16, + preprocessF16, +} from '../../example/src/tasks/depth-estimation/shaders'; +import { setupGPU } from '../helpers'; + +it('renders finite FP16 depth across changing camera frames', async () => { + const device = await setupGPU(); + device.pushErrorScope('validation'); + const ctx = new WgpuExecutionContext(device); + const width = 448; + const height = 448; + const video = device.createTexture({ + size: [width, height], + format: 'rgba8unorm', + usage: GPUTextureUsage.COPY_DST | GPUTextureUsage.TEXTURE_BINDING, + }); + const image = device.createTexture({ + size: [width, height], + format: 'rgba8unorm', + usage: GPUTextureUsage.STORAGE_BINDING | GPUTextureUsage.COPY_SRC, + }); + const inp = ctx.tensor('float16', [1, 3, height, width]); + const out = ctx.tensor('float16', [1, 1, height, width]); + const stats = ctx.storageBuffer(16); + const dir = resolve(import.meta.dirname, '../../example/public/assets/depthart'); + const program = JSON.parse(readFileSync(resolve(dir, 'program-fp16.json'), 'utf8')) as Program; + const weights = Uint8Array.from(readFileSync(resolve(dir, 'weights-fp16.bin'))).buffer; + const reference = new WgpuExecutionContext(device); + const referenceIn = reference.tensor('float32', [1, 3, height, width]); + const referenceOut = reference.tensor('float32', [1, 1, height, width]); + reference.recordProgram( + JSON.parse(readFileSync(resolve(dir, 'program.json'), 'utf8')) as Program, + Uint8Array.from(readFileSync(resolve(dir, 'weights.bin'))).buffer, + [referenceIn, referenceOut] + ); + ctx + .recordShader(preprocessF16, [video, inp]) + .recordProgram(program, weights, [inp, out]) + .recordShader(clearStats, [stats]) + .recordShader(depthStatsF16, [out, stats]) + .recordShader(drawDepthF16, [out, stats, image]); + + const readback = device.createBuffer({ + size: width * height * 4, + usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ, + }); + try { + // A dark circle on white drives the unbalanced decoder above FP16's range; + // uniform frames and the original small reference input do not expose it. + for (const frame of ['circle', 'gradient', 'black', 'white'] as const) { + const pixels = new Uint8Array(width * height * 4); + for (let y = 0; y < height; y++) { + for (let x = 0; x < width; x++) { + const i = (y * width + x) * 4; + pixels[i] = + frame === 'gradient' + ? Math.round((255 * x) / (width - 1)) + : frame === 'white' + ? 255 + : 0; + pixels[i + 1] = frame === 'gradient' ? Math.round((255 * y) / (height - 1)) : pixels[i]; + pixels[i + 2] = frame === 'gradient' ? 128 : pixels[i]; + if (frame === 'circle') { + pixels.fill( + (x - width / 2) ** 2 + (y - height / 2) ** 2 < 120 ** 2 ? 0 : 255, + i, + i + 3 + ); + } + pixels[i + 3] = 255; + } + } + device.queue.writeTexture({ texture: video }, pixels, { bytesPerRow: width * 4 }, [ + width, + height, + ]); + ctx.submit(); + await ctx.sync(); + const input = new Float16Array(await inp.getData()); + expect(input[0]).toBeCloseTo((pixels[0] / 255 - 0.485) / 0.229, 2); + const depth = new Float16Array(await out.getData()); + referenceIn.setData(Float32Array.from(input)); + reference.submit(); + await reference.sync(); + const expected = new Float32Array(await referenceOut.getData()); + let maxError = 0; + let referenceMax = 0; + let depthMin = Infinity; + let depthMax = -Infinity; + for (let i = 0; i < depth.length; i++) { + if (!Number.isFinite(depth[i])) + throw new Error(`${frame}: non-finite depth at ${i}: ${depth[i]}`); + if (!Number.isFinite(expected[i])) + throw new Error(`${frame}: non-finite FP32 reference at ${i}`); + maxError = Math.max(maxError, Math.abs(depth[i] - expected[i])); + referenceMax = Math.max(referenceMax, Math.abs(expected[i])); + depthMin = Math.min(depthMin, depth[i]); + depthMax = Math.max(depthMax, depth[i]); + } + expect(maxError / referenceMax, `${frame}: FP16 versus FP32 depth`).toBeLessThan(0.05); + const encoder = device.createCommandEncoder(); + encoder.copyTextureToBuffer( + { texture: image }, + { buffer: readback, bytesPerRow: width * 4 }, + [width, height] + ); + device.queue.submit([encoder.finish()]); + await readback.mapAsync(GPUMapMode.READ); + const rgba = new Uint8Array(readback.getMappedRange()); + let min = 255; + let max = 0; + for (let i = 0; i < rgba.length; i += 4) { + const gray = Math.round(255 * (1 - (depth[i / 4] - depthMin) / (depthMax - depthMin))); + if ( + Math.abs(rgba[i] - gray) > 1 || + rgba[i + 1] !== rgba[i] || + rgba[i + 2] !== rgba[i] || + rgba[i + 3] !== 255 + ) { + throw new Error( + `${frame}: incorrect rendered pixel ${i / 4}: got ${rgba[i]}, expected ${gray}` + ); + } + min = Math.min(min, rgba[i]); + max = Math.max(max, rgba[i]); + } + readback.unmap(); + expect(max - min, `${frame}: depth contrast`).toBeGreaterThan(200); + } + expect(await device.popErrorScope()).toBeNull(); + } finally { + ctx.destroy(); + reference.destroy(); + video.destroy(); + image.destroy(); + readback.destroy(); + } +}, 600_000); diff --git a/tests/program/depthart.test.ts b/tests/program/depthart.test.ts index 6504df0..bc5d430 100644 --- a/tests/program/depthart.test.ts +++ b/tests/program/depthart.test.ts @@ -24,6 +24,14 @@ function makeInput(): Float32Array { return arr; } +function makeInputF16(): Float16Array { + const arr = new Float16Array(INPUT_NUMEL); + for (let i = 0; i < INPUT_NUMEL; i++) { + arr[i] = (i % 251) / 250 - 0.5; + } + return arr; +} + function readArrayBuffer(path: string): ArrayBuffer { const buf = readFileSync(path); const copy = new ArrayBuffer(buf.byteLength); @@ -67,6 +75,9 @@ describe('depthart program (full ExecuTorch graph)', () => { let beyondTol = 0; const tolerance = 1.0; for (let i = 0; i < actual.length; i++) { + if (!Number.isFinite(actual[i])) { + throw new Error(`Non-finite FP32 depth at index ${i}: ${actual[i]}`); + } const diff = Math.abs(actual[i] - expected[i]); if (diff > maxDiff) { maxDiff = diff; @@ -88,4 +99,46 @@ describe('depthart program (full ExecuTorch graph)', () => { ctx.destroy(); } ); + + it('executes the full operator chain in float16', { timeout: 600_000 }, async () => { + const device = await setupGPU(); + const ctx = new WgpuExecutionContext(device); + + const progPath = resolve(REPO_ROOT, 'example/public/assets/depthart/program-fp16.json'); + const weightsPath = resolve(REPO_ROOT, 'example/public/assets/depthart/weights-fp16.bin'); + const program = JSON.parse(readFileSync(progPath, 'utf8')) as Program; + const weights = readArrayBuffer(weightsPath); + + const inp = ctx.tensor('float16', [...INPUT_SHAPE]).setData(makeInputF16()); + const out = ctx.tensor('float16', [...OUTPUT_SHAPE]); + + ctx.recordProgram(program, weights, [inp, out]); + ctx.submit(); + await ctx.sync(); + + const actual = new Float16Array(await out.getData()); + const expected = new Float32Array(readArrayBuffer(REFERENCE_PATH)); + + if (actual.length !== expected.length) { + throw new Error(`Output length mismatch: got ${actual.length}, expected ${expected.length}`); + } + + // Check against FP32 reference with FP16 tolerance (values reach 1e4) + let maxDiff = 0; + for (let i = 0; i < actual.length; i++) { + if (!Number.isFinite(actual[i])) { + throw new Error(`Non-finite FP16 depth at index ${i}: ${actual[i]}`); + } + const diff = Math.abs(actual[i] - expected[i]); + if (diff > maxDiff) maxDiff = diff; + } + + // High-level correlation check: relative error should be small + const relDiff = maxDiff / 1e4; + if (relDiff > 0.05) { + throw new Error(`FP16 relative difference too high: ${relDiff}`); + } + + ctx.destroy(); + }); }); diff --git a/tests/program/selfie.test.ts b/tests/program/selfie.test.ts new file mode 100644 index 0000000..c19fbd3 --- /dev/null +++ b/tests/program/selfie.test.ts @@ -0,0 +1,38 @@ +import { readFileSync } from 'node:fs'; +import { resolve } from 'node:path'; +import { describe, it } from 'vitest'; + +import { WgpuExecutionContext, type Program } from '../../src'; +import { setupGPU } from '../helpers'; + +const REPO_ROOT = resolve(import.meta.dirname, '../..'); +const PROG_PATH = resolve(REPO_ROOT, 'example/public/assets/selfie-segmentation/program-fp16.json'); +const WEIGHTS_PATH = resolve( + REPO_ROOT, + 'example/public/assets/selfie-segmentation/weights-fp16.bin' +); + +function readArrayBuffer(path: string): ArrayBuffer { + const buf = readFileSync(path); + const copy = new ArrayBuffer(buf.byteLength); + new Uint8Array(copy).set(buf); + return copy; +} + +describe('selfie-segmentation fp16', () => { + it('records and submits selfie-segmentation graph in float16', async () => { + const device = await setupGPU(); + const ctx = new WgpuExecutionContext(device); + + const program = JSON.parse(readFileSync(PROG_PATH, 'utf8')) as Program; + const weights = readArrayBuffer(WEIGHTS_PATH); + + const inp = ctx.tensor('float16', [1, 3, 256, 256]); + const out = ctx.tensor('float16', [1, 1, 256, 256]); + + ctx.recordProgram(program, weights, [inp, out]); + ctx.submit(); + await ctx.sync(); + ctx.destroy(); + }); +}); diff --git a/tests/program/yolo26n.test.ts b/tests/program/yolo26n.test.ts index f2cdaf9..1246eae 100644 --- a/tests/program/yolo26n.test.ts +++ b/tests/program/yolo26n.test.ts @@ -24,6 +24,14 @@ function makeInput(): Float32Array { return arr; } +function makeInputF16(): Float16Array { + const arr = new Float16Array(INPUT_NUMEL); + for (let i = 0; i < INPUT_NUMEL; i++) { + arr[i] = (i % 251) / 250 - 0.5; + } + return arr; +} + function readArrayBuffer(path: string): ArrayBuffer { const buf = readFileSync(path); const copy = new ArrayBuffer(buf.byteLength); @@ -88,4 +96,44 @@ describe('yolo26n program (full ExecuTorch graph)', () => { ctx.destroy(); } ); + + it('executes the full operator chain in float16', { timeout: 600_000 }, async () => { + const device = await setupGPU(); + const ctx = new WgpuExecutionContext(device); + + const progPath = resolve(REPO_ROOT, 'example/public/assets/yolo26n/program-fp16.json'); + const weightsPath = resolve(REPO_ROOT, 'example/public/assets/yolo26n/weights-fp16.bin'); + const program = JSON.parse(readFileSync(progPath, 'utf8')) as Program; + const weights = readArrayBuffer(weightsPath); + + const inp = ctx.tensor('float16', [...INPUT_SHAPE]).setData(makeInputF16()); + const out = ctx.tensor('float16', [...OUTPUT_SHAPE]); + + ctx.recordProgram(program, weights, [inp, out]); + ctx.submit(); + await ctx.sync(); + + const actual = new Float16Array(await out.getData()); + const expected = new Float32Array(readArrayBuffer(REFERENCE_PATH)); + + if (actual.length !== expected.length) { + throw new Error(`Output length mismatch: got ${actual.length}, expected ${expected.length}`); + } + + // Check against FP32 reference with FP16 tolerance + let beyondTol = 0; + const tolerance = 0.5; + for (let i = 0; i < actual.length; i++) { + const diff = Math.abs(actual[i] - expected[i]); + if (diff > tolerance) beyondTol++; + } + + if (beyondTol > actual.length * 0.05) { + throw new Error( + `Output mismatch: ${beyondTol}/${actual.length} elements exceed tolerance ${tolerance}` + ); + } + + ctx.destroy(); + }); }); diff --git a/tests/setup.ts b/tests/setup.ts index a843325..c9756b5 100644 --- a/tests/setup.ts +++ b/tests/setup.ts @@ -1,9 +1,15 @@ import { create, globals } from 'webgpu'; +import { Float16Array as PolyfillFloat16Array } from '@petamoriken/float16'; // Set up Dawn WebGPU globals before any source file is loaded. // This must run before src/tensor.ts which uses GPUBufferUsage at module level. Object.assign(globalThis, globals); +// Polyfill Float16Array for Node runtimes where it is not yet exposed by default (e.g. Node 22). +if (typeof globalThis.Float16Array === 'undefined') { + (globalThis as Record).Float16Array = PolyfillFloat16Array; +} + // Create a global navigator with WebGPU for Node.js. // Use SwiftShader for headless environments (CI) where no GPU is available. // navigator is read-only on globalThis in Node.js, so use defineProperty. diff --git a/tests/shaders/bmm.test.ts b/tests/shaders/bmm.test.ts index d6d96c9..54f2d89 100644 --- a/tests/shaders/bmm.test.ts +++ b/tests/shaders/bmm.test.ts @@ -1,4 +1,4 @@ -import { describe, it } from 'vitest'; +import { describe, it, expect } from 'vitest'; import { WgpuExecutionContext, shaders } from '../../src'; import { setupGPU, expectTensorClose } from '../helpers'; @@ -39,4 +39,63 @@ describe('aten.bmm.default', () => { await expectTensorClose(out, [5, 6, 7, 8, 10, 12, 21, 24], 1e-3); ctx.destroy(); }); + + it('multiplies float16 batched matrices using shader-f16', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const a = ctx.tensor('float16', [2, 2, 2]).setData(new Float16Array([1, 0, 0, 1, 2, 0, 0, 3])); + const b = ctx.tensor('float16', [2, 2, 2]).setData(new Float16Array([5, 6, 7, 8, 5, 6, 7, 8])); + const out = ctx.tensor('float16', [2, 2, 2]); + + ctx.recordShader(shaders.bmm, [a, b, out]); + ctx.submit(); + await ctx.sync(); + + const actual = new Float16Array(await out.getData()); + expect(actual[0]).toBeCloseTo(5, 1); + expect(actual[1]).toBeCloseTo(6, 1); + expect(actual[2]).toBeCloseTo(7, 1); + expect(actual[3]).toBeCloseTo(8, 1); + expect(actual[4]).toBeCloseTo(10, 1); + expect(actual[5]).toBeCloseTo(12, 1); + expect(actual[6]).toBeCloseTo(21, 1); + expect(actual[7]).toBeCloseTo(24, 1); + ctx.destroy(); + }); + + it('multiplies 64x64 batched matrices (tiled BMM path)', async () => { + const device = await setupGPU(); + const ctx = new WgpuExecutionContext(device); + + const [B, N] = [2, 64]; + const aData = new Float32Array(B * N * N); + const bData = new Float32Array(B * N * N); + + // Identity matrices for each batch + for (let b = 0; b < B; b++) { + for (let i = 0; i < N; i++) { + aData[b * (N * N) + i * N + i] = 1.0; + for (let j = 0; j < N; j++) { + bData[b * (N * N) + i * N + j] = (i + j) * 0.1; + } + } + } + + const a = ctx.tensor('float32', [B, N, N]).setData(aData); + const b = ctx.tensor('float32', [B, N, N]).setData(bData); + const out = ctx.tensor('float32', [B, N, N]); + + ctx.recordShader(shaders.bmm, [a, b, out]); + ctx.submit(); + await ctx.sync(); + + // A * I = A + const actual = new Float32Array(await out.getData()); + for (let i = 0; i < actual.length; i++) { + expect(actual[i]).toBeCloseTo(bData[i], 3); + } + ctx.destroy(); + }); }); diff --git a/tests/shaders/conv-with-clamp.test.ts b/tests/shaders/conv-with-clamp.test.ts index ef5df81..2680cb3 100644 --- a/tests/shaders/conv-with-clamp.test.ts +++ b/tests/shaders/conv-with-clamp.test.ts @@ -1,8 +1,94 @@ -import { describe, it } from 'vitest'; +import { describe, it, expect } from 'vitest'; import { WgpuExecutionContext, shaders } from '../../src'; import { setupGPU, expectTensorClose } from '../helpers'; describe('et_vk.conv_with_clamp.default', () => { + for (const dtype of ['float32', 'float16'] as const) { + for (const hasBias of [false, true]) { + it(`computes single-channel pointwise ${dtype} with batches, a vector tail, and bias=${hasBias}`, async () => { + const device = await setupGPU(); + const ctx = new WgpuExecutionContext(device); + const array = dtype === 'float16' ? Float16Array : Float32Array; + const spatial = 260; // 65 vectors: the last workgroup has only one active lane. + const values = new array(2 * 3 * spatial); + for (let i = 0; i < values.length; i++) values[i] = (i % 13) - 6; + const input = ctx.tensor(dtype, [2, 3, 1, spatial]).setData(values); + const weights = ctx.tensor(dtype, [1, 3, 1, 1]).setData(new array([1, -2, 0.5])); + const bias = hasBias ? ctx.tensor(dtype, [1]).setData(new array([3])) : undefined; + const out = ctx.tensor(dtype, [2, 1, 1, spatial]); + const expected = Array.from({ length: 2 * spatial }, (value, i) => { + const base = Math.floor(i / spatial) * 3 * spatial + (i % spatial); + return Math.min( + 10, + Math.max( + -5, + values[base] - + 2 * values[base + spatial] + + 0.5 * values[base + 2 * spatial] + + (hasBias ? 3 : 0) + ) + ); + }); + try { + ctx.recordShader(shaders.convWithClamp, [ + input, + weights, + bias, + [1, 1], + [0, 0], + [1, 1], + false, + [0, 0], + 1, + -5, + 10, + out, + ]); + ctx.submit(); + await ctx.sync(); + await expectTensorClose(out, expected); + } finally { + ctx.destroy(); + } + }); + } + } + + it('widens FP16 pointwise products before accumulation and clamps before narrowing', async () => { + const device = await setupGPU(); + const ctx = new WgpuExecutionContext(device); + const input = ctx + .tensor('float16', [1, 3, 1, 4]) + .setData( + new Float16Array([ + 60000, 60000, -60000, -60000, -60000, -60000, 60000, 60000, 4, 60000, -4, -60000, + ]) + ); + const weights = ctx.tensor('float16', [1, 3, 1, 1]).setData(new Float16Array([2, 2, 2])); + const out = ctx.tensor('float16', [1, 1, 1, 4]); + try { + ctx.recordShader(shaders.convWithClamp, [ + input, + weights, + undefined, + [1, 1], + [0, 0], + [1, 1], + false, + [0, 0], + 1, + -100, + 100, + out, + ]); + ctx.submit(); + await ctx.sync(); + await expectTensorClose(out, [8, 100, -8, -100]); + } finally { + ctx.destroy(); + } + }); + it('computes 2D dense convolution with clamp and bias', async () => { const device = await setupGPU(); const ctx = new WgpuExecutionContext(device); @@ -93,4 +179,159 @@ describe('et_vk.conv_with_clamp.default', () => { await expectTensorClose(out, [0, 6]); ctx.destroy(); }); + + it('computes float16 2D dense convolution with clamp and bias', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const inp = ctx + .tensor('float16', [1, 1, 3, 3]) + .setData(new Float16Array([1, 2, 3, 4, 5, 6, 7, 8, 9])); + const weight = ctx.tensor('float16', [1, 1, 2, 2]).setData(new Float16Array([1, 0, 0, 1])); + const bias = ctx.tensor('float16', [1]).setData(new Float16Array([10])); + const out = ctx.tensor('float16', [1, 1, 2, 2]); + + ctx.recordShader(shaders.convWithClamp, [ + inp, + weight, + bias, + [1, 1], + [0, 0], + [1, 1], + false, + [0, 0], + 1, + 17.0, + 23.0, + out, + ]); + ctx.submit(); + await ctx.sync(); + + await expectTensorClose(out, [17, 18, 22, 23], 1e-2); + ctx.destroy(); + }); + + it('computes float16 64x64 register-tiled dense convolution with clamp', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const B = 1; + const IC = 64; + const IH = 8; + const IW = 8; + const OC = 64; + const KH = 1; + const KW = 1; + + const inData = new Float16Array(B * IC * IH * IW); + for (let i = 0; i < inData.length; i++) inData[i] = ((i % 5) - 2) * 0.5; + const wData = new Float16Array(OC * IC * KH * KW); + for (let i = 0; i < wData.length; i++) wData[i] = ((i % 7) - 3) * 0.25; + + const inp = ctx.tensor('float16', [B, IC, IH, IW]).setData(inData); + const weight = ctx.tensor('float16', [OC, IC, KH, KW]).setData(wData); + const out = ctx.tensor('float16', [B, OC, IH, IW]); + + ctx.recordShader(shaders.convWithClamp, [ + inp, + weight, + undefined, + [1, 1], + [0, 0], + [1, 1], + false, + [0, 0], + 1, + -1.0, + 1.0, + out, + ]); + ctx.submit(); + await ctx.sync(); + + const actual = new Float16Array(await out.getData()); + for (let i = 0; i < actual.length; i++) { + expect(actual[i]).toBeGreaterThanOrEqual(-1.01); + expect(actual[i]).toBeLessThanOrEqual(1.01); + } + ctx.destroy(); + }); + + it('computes float16 depthwise grouped convolution with clamp', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const inp = ctx + .tensor('float16', [1, 2, 2, 2]) + .setData(new Float16Array([-5, -2, 1, 2, 2, 4, 1, 3])); + const weight = ctx + .tensor('float16', [2, 1, 2, 2]) + .setData(new Float16Array([1, 1, 1, 1, 1, 1, 1, 1])); + const out = ctx.tensor('float16', [1, 2, 1, 1]); + + ctx.recordShader(shaders.convWithClamp, [ + inp, + weight, + undefined, + [1, 1], + [0, 0], + [1, 1], + false, + [0, 0], + 2, + 0.0, + 6.0, + out, + ]); + ctx.submit(); + await ctx.sync(); + + await expectTensorClose(out, [0, 6], 1e-2); + ctx.destroy(); + }); + + it('computes 1x1 conv-with-clamp with 64x64 register-tiled path', async () => { + const device = await setupGPU(); + const ctx = new WgpuExecutionContext(device); + + const [B, IC, IH, IW, OC] = [1, 16, 8, 8, 64]; + const inData = new Float32Array(B * IC * IH * IW); + for (let i = 0; i < inData.length; i++) inData[i] = (i % 13) * 0.5 - 2.0; + + const wData = new Float32Array(OC * IC * 1 * 1); + for (let i = 0; i < wData.length; i++) wData[i] = (i % 7) * 0.1; + + const inp = ctx.tensor('float32', [B, IC, IH, IW]).setData(inData); + const weight = ctx.tensor('float32', [OC, IC, 1, 1]).setData(wData); + const out = ctx.tensor('float32', [B, OC, IH, IW]); + + ctx.recordShader(shaders.convWithClamp, [ + inp, + weight, + undefined, + [1, 1], + [0, 0], + [1, 1], + false, + [0, 0], + 1, + 0.0, + 5.0, + out, + ]); + ctx.submit(); + await ctx.sync(); + + // Verify all outputs are clamped to [0.0, 5.0] + const actual = new Float32Array(await out.getData()); + for (let i = 0; i < actual.length; i++) { + expect(actual[i]).toBeGreaterThanOrEqual(0.0); + expect(actual[i]).toBeLessThanOrEqual(5.0001); + } + ctx.destroy(); + }); }); diff --git a/tests/shaders/convolution.test.ts b/tests/shaders/convolution.test.ts index a23fba9..07c1eb7 100644 --- a/tests/shaders/convolution.test.ts +++ b/tests/shaders/convolution.test.ts @@ -1,4 +1,4 @@ -import { describe, it } from 'vitest'; +import { describe, it, expect } from 'vitest'; import { WgpuExecutionContext, shaders } from '../../src'; import { setupGPU, expectTensorClose } from '../helpers'; @@ -223,4 +223,228 @@ describe('aten.convolution.default', () => { ); ctx.destroy(); }); + + it('computes float16 2D dense convolution with bias (f16 GEMM)', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const inp = ctx + .tensor('float16', [1, 1, 3, 3]) + .setData(new Float16Array([1, 2, 3, 4, 5, 6, 7, 8, 9])); + const weight = ctx.tensor('float16', [1, 1, 2, 2]).setData(new Float16Array([1, 0, 0, 1])); + const bias = ctx.tensor('float16', [1]).setData(new Float16Array([10])); + const out = ctx.tensor('float16', [1, 1, 2, 2]); + + ctx.recordShader(shaders.convolution, [ + inp, + weight, + bias, + [1, 1], + [0, 0], + [1, 1], + false, + [0, 0], + 1, + out, + ]); + ctx.submit(); + await ctx.sync(); + + await expectTensorClose(out, [16, 18, 22, 24], 1e-2); + ctx.destroy(); + }); + + it('computes float16 64x64 register-tiled dense convolution', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const B = 1; + const IC = 64; + const IH = 8; + const IW = 8; + const OC = 64; + const KH = 1; + const KW = 1; + + const inData = new Float16Array(B * IC * IH * IW); + for (let i = 0; i < inData.length; i++) inData[i] = ((i % 5) - 2) * 0.5; + const wData = new Float16Array(OC * IC * KH * KW); + for (let i = 0; i < wData.length; i++) wData[i] = ((i % 7) - 3) * 0.25; + + const inp = ctx.tensor('float16', [B, IC, IH, IW]).setData(inData); + const weight = ctx.tensor('float16', [OC, IC, KH, KW]).setData(wData); + const out = ctx.tensor('float16', [B, OC, IH, IW]); + + ctx.recordShader(shaders.convolution, [ + inp, + weight, + undefined, + [1, 1], + [0, 0], + [1, 1], + false, + [0, 0], + 1, + out, + ]); + ctx.submit(); + await ctx.sync(); + + // Verify first few elements manually + const actual = new Float16Array(await out.getData()); + // For oc = 0, spatial (0, 0): sum over ic of in[0, ic, 0, 0] * w[0, ic, 0, 0] + let expected00 = 0; + for (let ic = 0; ic < IC; ic++) { + const inVal = inData[ic * (IH * IW)]; + const wVal = wData[ic]; + expected00 += inVal * wVal; + } + expect(actual[0]).toBeCloseTo(expected00, 1); + ctx.destroy(); + }); + + it('computes float16 depthwise convolution', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const inp = ctx + .tensor('float16', [1, 2, 2, 2]) + .setData(new Float16Array([1, 2, 3, 4, 10, 20, 30, 40])); + const weight = ctx.tensor('float16', [2, 1, 1, 1]).setData(new Float16Array([2, 0.5])); + const out = ctx.tensor('float16', [1, 2, 2, 2]); + + ctx.recordShader(shaders.convolution, [ + inp, + weight, + undefined, + [1, 1], + [0, 0], + [1, 1], + false, + [0, 0], + 2, + out, + ]); + ctx.submit(); + await ctx.sync(); + + await expectTensorClose(out, [2, 4, 6, 8, 5, 10, 15, 20], 1e-2); + ctx.destroy(); + }); + + it('computes 1x1 convolution with Vec4 path', async () => { + const device = await setupGPU(); + const ctx = new WgpuExecutionContext(device); + + const [B, IC, IH, IW, OC] = [1, 4, 2, 2, 4]; + const inData = new Float32Array(B * IC * IH * IW); + for (let i = 0; i < inData.length; i++) inData[i] = i + 1; + + // Identity-like weights for 1x1 + const wData = new Float32Array(OC * IC * 1 * 1); + for (let o = 0; o < OC; o++) { + for (let c = 0; c < IC; c++) { + wData[o * IC + c] = o === c ? 2 : 0; + } + } + + const inp = ctx.tensor('float32', [B, IC, IH, IW]).setData(inData); + const weight = ctx.tensor('float32', [OC, IC, 1, 1]).setData(wData); + const out = ctx.tensor('float32', [B, OC, IH, IW]); + + ctx.recordShader(shaders.convolution, [ + inp, + weight, + undefined, + [1, 1], + [0, 0], + [1, 1], + false, + [0, 0], + 1, + out, + ]); + ctx.submit(); + await ctx.sync(); + + const actual = new Float32Array(await out.getData()); + for (let o = 0; o < OC; o++) { + for (let s = 0; s < IH * IW; s++) { + expect(actual[o * (IH * IW) + s]).toBeCloseTo(inData[o * (IH * IW) + s] * 2, 4); + } + } + ctx.destroy(); + }); + + it('computes 1x1 convolution with 64x64 register-tiled path', async () => { + const device = await setupGPU(); + const ctx = new WgpuExecutionContext(device); + + const [B, IC, IH, IW, OC] = [1, 16, 8, 8, 64]; + const inData = new Float32Array(B * IC * IH * IW); + for (let i = 0; i < inData.length; i++) inData[i] = (i % 13) * 0.1; + + const wData = new Float32Array(OC * IC * 1 * 1); + for (let i = 0; i < wData.length; i++) wData[i] = (i % 7) * 0.1; + + const inp = ctx.tensor('float32', [B, IC, IH, IW]).setData(inData); + const weight = ctx.tensor('float32', [OC, IC, 1, 1]).setData(wData); + const out = ctx.tensor('float32', [B, OC, IH, IW]); + + ctx.recordShader(shaders.convolution, [ + inp, + weight, + undefined, + [1, 1], + [0, 0], + [1, 1], + false, + [0, 0], + 1, + out, + ]); + ctx.submit(); + await ctx.sync(); + + // Verify element (0, 0, 0, 0): dot product of out_ch 0 with in_pos 0 across all IC + let expected0 = 0; + for (let c = 0; c < IC; c++) { + expected0 += inData[c * (IH * IW)] * wData[0 * IC + c]; + } + const actual = new Float32Array(await out.getData()); + expect(actual[0]).toBeCloseTo(expected0, 3); + ctx.destroy(); + }); + + it('computes general grouped convolution (groups=2 fallback)', async () => { + const device = await setupGPU(); + const ctx = new WgpuExecutionContext(device); + + // groups = 2, IC = 4 (2 per group), OC = 4 (2 per group) + const inp = ctx.tensor('float32', [1, 4, 2, 2]).setData(new Float32Array(16).fill(1)); + const weight = ctx.tensor('float32', [4, 2, 1, 1]).setData(new Float32Array(8).fill(2)); + const out = ctx.tensor('float32', [1, 4, 2, 2]); + + ctx.recordShader(shaders.convolution, [ + inp, + weight, + undefined, + [1, 1], + [0, 0], + [1, 1], + false, + [0, 0], + 2, + out, + ]); + ctx.submit(); + await ctx.sync(); + + // Each out channel = sum over 2 input channels of 1 * 2 = 4 + await expectTensorClose(out, new Array(16).fill(4), 1e-3); + ctx.destroy(); + }); }); diff --git a/tests/shaders/fp16.test.ts b/tests/shaders/fp16.test.ts new file mode 100644 index 0000000..1be6e7c --- /dev/null +++ b/tests/shaders/fp16.test.ts @@ -0,0 +1,267 @@ +import { describe, it, expect } from 'vitest'; +import { setupGPU, expectTensorClose } from '../helpers'; +import { WgpuExecutionContext, createComputeBundle, shaders } from '../../src'; + +describe('shader-f16 feature', () => { + it('device has shader-f16 feature and executes WGSL f16 shader', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + + const ctx = new WgpuExecutionContext(device); + const inTensor = ctx.tensor('float16', [4]); + const outTensor = ctx.tensor('float16', [4]); + + const inData = new Float16Array([1.5, -2.0, 3.25, 0.0]); + inTensor.setData(inData); + + const WGSL_FP16 = /* wgsl */ ` + enable f16; + + @group(0) @binding(0) var a: array; + @group(0) @binding(1) var b: array; + + @compute @workgroup_size(4, 1, 1) + fn main(@builtin(global_invocation_id) gid: vec3) { + b[gid.x] = a[gid.x] * 2.0h; + } + `; + + const bundle = createComputeBundle(device, WGSL_FP16, [ + { binding: 0, buffer: inTensor.buffer }, + { binding: 1, buffer: outTensor.buffer }, + ]); + + ctx.addDispatch({ + pipeline: bundle.pipeline, + bindGroup: bundle.bindGroup, + workgroupCountX: 1, + workgroupCountY: 1, + workgroupCountZ: 1, + }); + + ctx.submit(); + await ctx.sync(); + + const raw = await outTensor.getData(); + const result = new Float16Array(raw); + expect(result[0]).toBeCloseTo(3.0); + expect(result[1]).toBeCloseTo(-4.0); + expect(result[2]).toBeCloseTo(6.5); + expect(result[3]).toBeCloseTo(0.0); + + ctx.destroy(); + }); + + it('computes float16 addition with shaders.add (elementwise and broadcast)', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + // 1. Elementwise identical shapes + const a = ctx.tensor('float16', [4]).setData(new Float16Array([1.0, 2.5, 3.0, 4.5])); + const b = ctx.tensor('float16', [4]).setData(new Float16Array([0.5, 1.5, -1.0, 2.0])); + const out1 = ctx.tensor('float16', [4]); + + ctx.recordShader(shaders.add, [a, b, out1]); + ctx.submit(); + await ctx.sync(); + + await expectTensorClose(out1, [1.5, 4.0, 2.0, 6.5], 1e-2); + + // 2. Broadcasted addition with alpha + const x = ctx.tensor('float16', [1, 2, 1]).setData(new Float16Array([10.0, 20.0])); + const y = ctx.tensor('float16', [1, 2, 2]).setData(new Float16Array([1.0, 2.0, 3.0, 4.0])); + const out2 = ctx.tensor('float16', [1, 2, 2]); + + ctx.recordShader(shaders.add, [x, y, 2.0, out2]); + ctx.submit(); + await ctx.sync(); + + // [10 + 2*1, 10 + 2*2, 20 + 2*3, 20 + 2*4] = [12, 14, 26, 28] + await expectTensorClose(out2, [12, 14, 26, 28], 1e-2); + ctx.destroy(); + }); + + it('computes float16 subtraction and multiplication (sub, mul)', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const a = ctx.tensor('float16', [4]).setData(new Float16Array([10.0, 20.0, 30.0, 40.0])); + const b = ctx.tensor('float16', [4]).setData(new Float16Array([1.0, 2.0, 3.0, 4.0])); + const outSub = ctx.tensor('float16', [4]); + const outMul = ctx.tensor('float16', [4]); + + ctx.recordShader(shaders.sub, [a, b, 2.0, outSub]); + ctx.recordShader(shaders.mul, [a, b, outMul]); + ctx.submit(); + await ctx.sync(); + + await expectTensorClose(outSub, [8.0, 16.0, 24.0, 32.0], 1e-2); + await expectTensorClose(outMul, [10.0, 40.0, 90.0, 160.0], 1e-2); + ctx.destroy(); + }); + + it('computes float16 activations (relu, sigmoid, hardswish)', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const inp = ctx.tensor('float16', [4]).setData(new Float16Array([-4.0, -1.0, 0.0, 3.0])); + const outRelu = ctx.tensor('float16', [4]); + const outSigmoid = ctx.tensor('float16', [4]); + const outHardswish = ctx.tensor('float16', [4]); + + ctx.recordShader(shaders.relu, [inp, outRelu]); + ctx.recordShader(shaders.sigmoid, [inp, outSigmoid]); + ctx.recordShader(shaders.hardswish, [inp, outHardswish]); + ctx.submit(); + await ctx.sync(); + + // relu: max(x, 0) -> [0, 0, 0, 3] + await expectTensorClose(outRelu, [0, 0, 0, 3], 1e-2); + + // sigmoid: 1 / (1 + exp(-x)) + const sigm = (x: number) => 1 / (1 + Math.exp(-x)); + await expectTensorClose(outSigmoid, [-4, -1, 0, 3].map(sigm), 1e-2); + + // hardswish: x * clamp(x + 3, 0, 6) / 6 + // -4: -4 * 0 / 6 = 0 + // -1: -1 * 2 / 6 = -0.3333 + // 0: 0 + // 3: 3 * 6 / 6 = 3.0 + await expectTensorClose(outHardswish, [0.0, -0.3333, 0.0, 3.0], 1e-2); + ctx.destroy(); + }); + + it('computes float16 reductions (softmax, meanDim)', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + // Softmax + const smIn = ctx.tensor('float16', [2, 4]).setData(new Float16Array([1, 2, 3, 4, 0, 0, 0, 0])); + const smOut = ctx.tensor('float16', [2, 4]); + + ctx.recordShader(shaders.softmax, [smIn, 1, false, smOut]); + + // MeanDim + const meanIn = ctx + .tensor('float16', [2, 3]) + .setData(new Float16Array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])); + const meanOut = ctx.tensor('float16', [2]); + + ctx.recordShader(shaders.meanDim, [meanIn, [1], false, undefined, meanOut]); + ctx.submit(); + await ctx.sync(); + + const expectedSmRow0 = [1, 2, 3, 4].map((v) => Math.exp(v)); + const sum0 = expectedSmRow0.reduce((a, b) => a + b, 0); + await expectTensorClose( + smOut, + [...expectedSmRow0.map((v) => v / sum0), 0.25, 0.25, 0.25, 0.25], + 1e-2 + ); + + // Row 0 mean = (1+2+3)/3 = 2, Row 1 mean = (4+5+6)/3 = 5 + await expectTensorClose(meanOut, [2.0, 5.0], 1e-2); + ctx.destroy(); + }); + + it('computes float16 maxPool2d, upsampleNearest2d, upsampleBilinear2d', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + // 1. MaxPool2D + const poolIn = ctx + .tensor('float16', [1, 1, 4, 4]) + .setData(new Float16Array([1, 2, 5, 6, 3, 4, 7, 8, 9, 10, 13, 14, 11, 12, 15, 16])); + const poolVals = ctx.tensor('float16', [1, 1, 2, 2]); + const poolIdx = ctx.tensor('int32', [1, 1, 2, 2]); + + ctx.recordShader(shaders.maxPool2d, [ + poolIn, + [2, 2], // kernel + [2, 2], // stride + [0, 0], // pad + [1, 1], // dilation + false, // ceil_mode + [poolVals, poolIdx], + ]); + + // 2. Nearest upsample [1, 1, 2, 2] -> [1, 1, 4, 4] + const nearIn = ctx.tensor('float16', [1, 1, 2, 2]).setData(new Float16Array([10, 20, 30, 40])); + const nearOut = ctx.tensor('float16', [1, 1, 4, 4]); + ctx.recordShader(shaders.upsampleNearest2d, [nearIn, undefined, [2, 2], nearOut]); + + // 3. Bilinear upsample [1, 1, 2, 2] -> [1, 1, 3, 3] with align_corners=true + const biIn = ctx.tensor('float16', [1, 1, 2, 2]).setData(new Float16Array([0, 10, 20, 30])); + const biOut = ctx.tensor('float16', [1, 1, 3, 3]); + ctx.recordShader(shaders.upsampleBilinear2d, [biIn, [3, 3], true, undefined, biOut]); + + ctx.submit(); + await ctx.sync(); + + // Pool results: 2x2 max blocks: [4, 8, 12, 16] + await expectTensorClose(poolVals, [4, 8, 12, 16], 1e-2); + + // Nearest upsample 2x2 -> 4x4 + await expectTensorClose( + nearOut, + [10, 10, 20, 20, 10, 10, 20, 20, 30, 30, 40, 40, 30, 30, 40, 40], + 1e-2 + ); + + // Bilinear upsample 2x2 -> 3x3 align_corners=true + await expectTensorClose(biOut, [0, 5, 10, 10, 15, 20, 20, 25, 30], 1e-2); + + ctx.destroy(); + }); + + it('computes float16 data-movement operators (cat, constantPadNd, split, permute)', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + // 1. Cat + const c1 = ctx.tensor('float16', [1, 2]).setData(new Float16Array([1.0, 2.0])); + const c2 = ctx.tensor('float16', [1, 2]).setData(new Float16Array([3.0, 4.0])); + const catOut = ctx.tensor('float16', [1, 4]); + ctx.recordShader(shaders.cat, [[c1, c2], 1, catOut]); + + // 2. ConstantPadNd + const padIn = ctx.tensor('float16', [1, 1, 2, 2]).setData(new Float16Array([1, 2, 3, 4])); + const padOut = ctx.tensor('float16', [1, 1, 4, 4]); + ctx.recordShader(shaders.constantPadNd, [padIn, [1, 1, 1, 1], 9.0, padOut]); + + // 3. SplitWithSizes + const splitIn = ctx.tensor('float16', [4]).setData(new Float16Array([10, 20, 30, 40])); + const s1 = ctx.tensor('float16', [1]); + const s2 = ctx.tensor('float16', [3]); + ctx.recordShader(shaders.splitWithSizesCopy, [splitIn, [1, 3], 0, [s1, s2]]); + + // 4. Permute (2D transpose) + const permIn = ctx.tensor('float16', [2, 3]).setData(new Float16Array([1, 2, 3, 4, 5, 6])); + const permOut = ctx.tensor('float16', [3, 2]); + ctx.recordShader(shaders.permuteCopy, [permIn, [1, 0], permOut]); + + ctx.submit(); + await ctx.sync(); + + // Cat: [1, 2, 3, 4] + await expectTensorClose(catOut, [1.0, 2.0, 3.0, 4.0], 1e-2); + + // Pad: 2x2 in 4x4 with 1-pad border of 9.0 + await expectTensorClose(padOut, [9, 9, 9, 9, 9, 1, 2, 9, 9, 3, 4, 9, 9, 9, 9, 9], 1e-2); + + // Split: s1=[10], s2=[20, 30, 40] + await expectTensorClose(s1, [10], 1e-2); + await expectTensorClose(s2, [20, 30, 40], 1e-2); + + // Permute 2x3 -> 3x2 transpose + await expectTensorClose(permOut, [1, 4, 2, 5, 3, 6], 1e-2); + + ctx.destroy(); + }); +}); diff --git a/tests/shaders/mm.test.ts b/tests/shaders/mm.test.ts new file mode 100644 index 0000000..54a51b1 --- /dev/null +++ b/tests/shaders/mm.test.ts @@ -0,0 +1,197 @@ +import { describe, it, expect } from 'vitest'; +import { WgpuExecutionContext, shaders } from '../../src'; +import { setupGPU, expectTensorClose } from '../helpers'; + +describe('aten.mm.default', () => { + it('multiplies 2x2 matrices', async () => { + const device = await setupGPU(); + const ctx = new WgpuExecutionContext(device); + + const a = ctx.tensor('float32', [2, 2]).setData(new Float32Array([1, 2, 3, 4])); + const b = ctx.tensor('float32', [2, 2]).setData(new Float32Array([5, 6, 7, 8])); + const out = ctx.tensor('float32', [2, 2]); + + ctx.recordShader(shaders.mm, [a, b, out]); + ctx.submit(); + await ctx.sync(); + + // [1*5+2*7, 1*6+2*8] = [19, 22] + // [3*5+4*7, 3*6+4*8] = [43, 50] + await expectTensorClose(out, [19, 22, 43, 50]); + ctx.destroy(); + }); + + it('multiplies 64x64 matrices (vec4 path)', async () => { + const device = await setupGPU(); + const ctx = new WgpuExecutionContext(device); + + const N = 64; + const aData = new Float32Array(N * N); + const bData = new Float32Array(N * N); + const expected = new Float32Array(N * N); + + for (let i = 0; i < N * N; i++) { + aData[i] = (i % 7) - 3; + bData[i] = (i % 5) - 2; + } + + // CPU reference + for (let r = 0; r < N; r++) { + for (let c = 0; c < N; c++) { + let sum = 0; + for (let k = 0; k < N; k++) { + sum += aData[r * N + k] * bData[k * N + c]; + } + expected[r * N + c] = sum; + } + } + + const a = ctx.tensor('float32', [N, N]).setData(aData); + const b = ctx.tensor('float32', [N, N]).setData(bData); + const out = ctx.tensor('float32', [N, N]); + + ctx.recordShader(shaders.mm, [a, b, out]); + ctx.submit(); + await ctx.sync(); + + const actual = new Float32Array(await out.getData()); + for (let i = 0; i < N * N; i++) { + expect(Math.abs(actual[i] - expected[i])).toBeLessThan(1e-3); + } + ctx.destroy(); + }); + + it('multiplies 128x128 matrices (64x64 multi-workgroup path)', async () => { + const device = await setupGPU(); + const ctx = new WgpuExecutionContext(device); + + const N = 128; + const aData = new Float32Array(N * N); + const bData = new Float32Array(N * N); + const expected = new Float32Array(N * N); + + for (let i = 0; i < N * N; i++) { + aData[i] = (i % 7) - 3; + bData[i] = (i % 5) - 2; + } + + for (let r = 0; r < N; r++) { + for (let c = 0; c < N; c++) { + let sum = 0; + for (let k = 0; k < N; k++) { + sum += aData[r * N + k] * bData[k * N + c]; + } + expected[r * N + c] = sum; + } + } + + const a = ctx.tensor('float32', [N, N]).setData(aData); + const b = ctx.tensor('float32', [N, N]).setData(bData); + const out = ctx.tensor('float32', [N, N]); + + ctx.recordShader(shaders.mm, [a, b, out]); + ctx.submit(); + await ctx.sync(); + + const actual = new Float32Array(await out.getData()); + for (let i = 0; i < N * N; i++) { + expect(Math.abs(actual[i] - expected[i])).toBeLessThan(1e-2); + } + ctx.destroy(); + }); + + it('multiplies float16 2x2 matrices using shader-f16', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const a = ctx.tensor('float16', [2, 2]).setData(new Float16Array([1, 2, 3, 4])); + const b = ctx.tensor('float16', [2, 2]).setData(new Float16Array([5, 6, 7, 8])); + const out = ctx.tensor('float16', [2, 2]); + + ctx.recordShader(shaders.mm, [a, b, out]); + ctx.submit(); + await ctx.sync(); + + const actual = new Float16Array(await out.getData()); + expect(actual[0]).toBeCloseTo(19, 1); + expect(actual[1]).toBeCloseTo(22, 1); + expect(actual[2]).toBeCloseTo(43, 1); + expect(actual[3]).toBeCloseTo(50, 1); + ctx.destroy(); + }); + + it('multiplies float16 64x64 matrices (64x64 register-tiled f16 path)', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const N = 64; + const aData = new Float16Array(N * N); + const bData = new Float16Array(N * N); + const expected = new Float16Array(N * N); + + for (let i = 0; i < N * N; i++) { + aData[i] = (i % 7) - 3; + bData[i] = (i % 5) - 2; + } + + for (let r = 0; r < N; r++) { + for (let c = 0; c < N; c++) { + let sum = 0; + for (let k = 0; k < N; k++) { + sum += aData[r * N + k] * bData[k * N + c]; + } + expected[r * N + c] = sum; + } + } + + const a = ctx.tensor('float16', [N, N]).setData(aData); + const b = ctx.tensor('float16', [N, N]).setData(bData); + const out = ctx.tensor('float16', [N, N]); + + ctx.recordShader(shaders.mm, [a, b, out]); + ctx.submit(); + await ctx.sync(); + + const actual = new Float16Array(await out.getData()); + for (let i = 0; i < N * N; i++) { + expect(Math.abs(actual[i] - expected[i])).toBeLessThan(0.5); + } + ctx.destroy(); + }); + + it('multiplies unaligned odd-shaped matrices (5x7 by 7x3 fallback path)', async () => { + const device = await setupGPU(); + const ctx = new WgpuExecutionContext(device); + + const [M, K, N] = [5, 7, 3]; + const aData = new Float32Array(M * K); + const bData = new Float32Array(K * N); + const expected = new Float32Array(M * N); + + for (let i = 0; i < aData.length; i++) aData[i] = (i % 5) - 2; + for (let i = 0; i < bData.length; i++) bData[i] = (i % 3) + 1; + + for (let r = 0; r < M; r++) { + for (let c = 0; c < N; c++) { + let sum = 0; + for (let k = 0; k < K; k++) { + sum += aData[r * K + k] * bData[k * N + c]; + } + expected[r * N + c] = sum; + } + } + + const a = ctx.tensor('float32', [M, K]).setData(aData); + const b = ctx.tensor('float32', [K, N]).setData(bData); + const out = ctx.tensor('float32', [M, N]); + + ctx.recordShader(shaders.mm, [a, b, out]); + ctx.submit(); + await ctx.sync(); + + await expectTensorClose(out, Array.from(expected), 1e-4); + ctx.destroy(); + }); +}); diff --git a/tests/shaders/mul.test.ts b/tests/shaders/mul.test.ts index 091c50f..4982616 100644 --- a/tests/shaders/mul.test.ts +++ b/tests/shaders/mul.test.ts @@ -43,4 +43,40 @@ describe('aten.mul.Tensor', () => { await expectTensorClose(out, [2, 4, 6, 8, 5, 10, 15, 20]); ctx.destroy(); }); + + it('computes float16 elementwise multiplication', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const a = ctx.tensor('float16', [2, 2]).setData(new Float16Array([1, 2, 3, 4])); + const b = ctx.tensor('float16', [2, 2]).setData(new Float16Array([5, 6, 7, 8])); + const out = ctx.tensor('float16', [2, 2]); + + ctx.recordShader(shaders.mul, [a, b, out]); + ctx.submit(); + await ctx.sync(); + + await expectTensorClose(out, [5, 12, 21, 32], 1e-2); + ctx.destroy(); + }); + + it('computes float16 broadcasted multiplication', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const a = ctx + .tensor('float16', [1, 2, 2, 2]) + .setData(new Float16Array([1, 2, 3, 4, 10, 20, 30, 40])); + const b = ctx.tensor('float16', [1, 2, 1, 1]).setData(new Float16Array([2.0, 0.5])); + const out = ctx.tensor('float16', [1, 2, 2, 2]); + + ctx.recordShader(shaders.mul, [a, b, out]); + ctx.submit(); + await ctx.sync(); + + await expectTensorClose(out, [2, 4, 6, 8, 5, 10, 15, 20], 1e-2); + ctx.destroy(); + }); }); diff --git a/tests/shaders/permute.test.ts b/tests/shaders/permute.test.ts index 29ac276..ffa70d1 100644 --- a/tests/shaders/permute.test.ts +++ b/tests/shaders/permute.test.ts @@ -36,4 +36,36 @@ describe('aten.permute_copy.default', () => { await expectTensorClose(out, [1, 4, 2, 5, 3, 6, 7, 10, 8, 11, 9, 12]); ctx.destroy(); }); + + it('transposes a larger matrix with multi-workgroup shared-memory tiling', async () => { + const device = await setupGPU(); + const ctx = new WgpuExecutionContext(device); + + const R = 64; + const C = 48; + const data = new Float32Array(R * C); + const expected = new Float32Array(C * R); + + for (let r = 0; r < R; r++) { + for (let c = 0; c < C; c++) { + data[r * C + c] = r * 100 + c; + expected[c * R + r] = r * 100 + c; + } + } + + const inp = ctx.tensor('float32', [R, C]).setData(data); + const out = ctx.tensor('float32', [C, R]); + + ctx.recordShader(shaders.permuteCopy, [inp, [1, 0], out]); + ctx.submit(); + await ctx.sync(); + + const actual = new Float32Array(await out.getData()); + for (let i = 0; i < C * R; i++) { + if (actual[i] !== expected[i]) { + throw new Error(`Mismatch at ${i}: got ${actual[i]} expected ${expected[i]}`); + } + } + ctx.destroy(); + }); }); diff --git a/tests/shaders/relu.test.ts b/tests/shaders/relu.test.ts index 236f345..b5bf048 100644 --- a/tests/shaders/relu.test.ts +++ b/tests/shaders/relu.test.ts @@ -20,4 +20,23 @@ describe('aten.relu.default', () => { await expectTensorClose(out, expected); ctx.destroy(); }); + + it('computes element-wise relu correctly for float16', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const values = [-2.0, -0.5, 0.0, 0.5, 2.0]; + const expected = values.map((x) => Math.max(x, 0)); + + const inp = ctx.tensor('float16', [5]).setData(new Float16Array(values)); + const out = ctx.tensor('float16', [5]); + + ctx.recordShader(shaders.relu, [inp, out]); + ctx.submit(); + await ctx.sync(); + + await expectTensorClose(out, expected, 1e-2); + ctx.destroy(); + }); }); diff --git a/tests/shaders/sigmoid.test.ts b/tests/shaders/sigmoid.test.ts index 81e9989..1c496f2 100644 --- a/tests/shaders/sigmoid.test.ts +++ b/tests/shaders/sigmoid.test.ts @@ -20,4 +20,23 @@ describe('aten.sigmoid.default', () => { await expectTensorClose(out, expected); ctx.destroy(); }); + + it('computes element-wise sigmoid correctly for float16', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const values = [-10.0, -2.0, 0.0, 2.0, 10.0]; + const expected = values.map((x) => 1.0 / (1.0 + Math.exp(-x))); + + const inp = ctx.tensor('float16', [5]).setData(new Float16Array(values)); + const out = ctx.tensor('float16', [5]); + + ctx.recordShader(shaders.sigmoid, [inp, out]); + ctx.submit(); + await ctx.sync(); + + await expectTensorClose(out, expected, 1e-2); + ctx.destroy(); + }); }); diff --git a/tests/shaders/softmax.test.ts b/tests/shaders/softmax.test.ts index c2ba84b..f2a413d 100644 --- a/tests/shaders/softmax.test.ts +++ b/tests/shaders/softmax.test.ts @@ -1,4 +1,4 @@ -import { describe, it } from 'vitest'; +import { describe, it, expect } from 'vitest'; import { WgpuExecutionContext, shaders } from '../../src'; import { setupGPU, expectTensorClose } from '../helpers'; @@ -49,4 +49,30 @@ describe('aten._softmax.default', () => { await expectTensorClose(out, [c0[0] / s0, c1[0] / s1, c0[1] / s0, c1[1] / s1]); ctx.destroy(); }); + + it('normalizes along the last dimension with size divisible by 4 (vec4 path)', async () => { + const device = await setupGPU(); + const ctx = new WgpuExecutionContext(device); + + const R = 8; + const inpData = new Float32Array([1, 2, 3, 4, 5, 6, 7, 8, -1, 0, 1, 2, -3, 4, 0, 1]); + const inp = ctx.tensor('float32', [2, R]).setData(inpData); + const out = ctx.tensor('float32', [2, R]); + + ctx.recordShader(shaders.softmax, [inp, -1, false, out]); + ctx.submit(); + await ctx.sync(); + + for (let row = 0; row < 2; row++) { + const slice = Array.from(inpData.slice(row * R, (row + 1) * R)); + const expSlice = slice.map((x) => Math.exp(x)); + const s = expSlice.reduce((a, b) => a + b, 0); + const expected = expSlice.map((x) => x / s); + const actual = new Float32Array(await out.getData()).slice(row * R, (row + 1) * R); + for (let i = 0; i < R; i++) { + expect(Math.abs(actual[i] - expected[i])).toBeLessThan(1e-4); + } + } + ctx.destroy(); + }); }); diff --git a/tests/shaders/sub.test.ts b/tests/shaders/sub.test.ts index 79b601d..0a97443 100644 --- a/tests/shaders/sub.test.ts +++ b/tests/shaders/sub.test.ts @@ -34,4 +34,38 @@ describe('aten.sub.Tensor', () => { await expectTensorClose(out, [8, 18, 28, 6, 16, 26]); ctx.destroy(); }); + + it('computes float16 elementwise subtraction', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const a = ctx.tensor('float16', [2, 3]).setData(new Float16Array([5, 7, 9, 4, 3, 1])); + const b = ctx.tensor('float16', [2, 3]).setData(new Float16Array([1, 2, 3, 4, 5, 6])); + const out = ctx.tensor('float16', [2, 3]); + + ctx.recordShader(shaders.sub, [a, b, out]); + ctx.submit(); + await ctx.sync(); + + await expectTensorClose(out, [4, 5, 6, 0, -2, -5], 1e-2); + ctx.destroy(); + }); + + it('computes float16 alpha-scaled subtraction with broadcast', async () => { + const device = await setupGPU(); + if (!device.features.has('shader-f16')) return; + const ctx = new WgpuExecutionContext(device); + + const a = ctx.tensor('float16', [1, 3]).setData(new Float16Array([10, 20, 30])); + const b = ctx.tensor('float16', [2, 3]).setData(new Float16Array([1, 1, 1, 2, 2, 2])); + const out = ctx.tensor('float16', [2, 3]); + + ctx.recordShader(shaders.sub, [a, b, 2.0, out]); + ctx.submit(); + await ctx.sync(); + + await expectTensorClose(out, [8, 18, 28, 6, 16, 26], 1e-2); + ctx.destroy(); + }); }); diff --git a/utils/benchmark/baseline-main/src/command.ts b/utils/benchmark/baseline-main/src/command.ts new file mode 100644 index 0000000..2bbd943 --- /dev/null +++ b/utils/benchmark/baseline-main/src/command.ts @@ -0,0 +1,144 @@ +/** + * A buffer binding specification for constructing bind groups. + */ +export type BufferBindingSpec = { + /** The shader binding index corresponding to `@binding(n)` in WGSL. */ + readonly binding: number; + /** The underlying WebGPU buffer to bind. */ + readonly buffer: GPUBuffer; + /** Optional byte offset from the start of the buffer (defaults to 0). */ + readonly offset?: number; + /** Optional size in bytes of the buffer slice accessible to the shader. */ + readonly size?: number; +}; + +/** + * A generic WebGPU resource binding specification (e.g. texture view, sampler). + */ +export type ResourceBindingSpec = { + /** The shader binding index corresponding to `@binding(n)` in WGSL. */ + readonly binding: number; + /** The WebGPU binding resource (texture view, sampler, external texture, etc.). */ + readonly resource: GPUBindingResource; +}; + +/** + * Binding specification for constructing bind groups. + */ +export type BindingSpec = BufferBindingSpec | ResourceBindingSpec; + +/** + * Result of creating a compute pipeline and matching bind group. + */ +export type ComputePipelineBundle = { + /** The compiled or cached WebGPU compute pipeline. */ + readonly pipeline: GPUComputePipeline; + /** The bind group pre-configured to match the pipeline's layout. */ + readonly bindGroup: GPUBindGroup; +}; + +/** + * A recorded GPU compute dispatch command. + */ +export type WgpuDispatch = { + /** The compute pipeline to bind for execution. */ + readonly pipeline: GPUComputePipeline; + /** The bind group containing bound buffers and uniforms. */ + readonly bindGroup: GPUBindGroup; + /** Number of workgroups to dispatch in the X dimension. */ + readonly workgroupCountX: number; + /** Number of workgroups to dispatch in the Y dimension. */ + readonly workgroupCountY: number; + /** Optional number of workgroups to dispatch in the Z dimension (defaults to 1). */ + readonly workgroupCountZ?: number; +}; + +/** + * A recorded GPU buffer-to-buffer copy command. + */ +export type WgpuCopy = { + /** Source buffer to copy from. */ + readonly src: GPUBuffer; + /** Destination buffer to copy to. */ + readonly dst: GPUBuffer; + /** Number of bytes to copy. */ + readonly size: number; +}; + +/** + * A recorded GPU command — either a compute dispatch or a buffer copy. + */ +export type WgpuCommand = + | { readonly kind: 'dispatch'; readonly dispatch: WgpuDispatch } + | { readonly kind: 'copy'; readonly copy: WgpuCopy }; + +const cache = new WeakMap>(); + +/** + * Serializes pipeline override constants into a deterministic, + * order-independent cache key string. + * + * @param constants Pipeline override constants. + * @returns Serialized key string. + */ +function serializeConstants(constants?: Record): string { + if (!constants) return ''; + + const entries = Object.entries(constants); + if (entries.length === 0) return ''; + + return entries + .sort(([a], [b]) => a.localeCompare(b)) + .map(([k, v]) => `${k}=${v}`) + .join(';'); +} + +/** + * Builds and caches a compute pipeline along with its bind group from a list of + * bindings. + * + * @param device WebGPU device instance. + * @param code WGSL shader code string. + * @param bindings List of buffer bindings for `@group(0)`. + * @param constants Optional pipeline override constants. + * @param entryPoint Shader entry point function name (default: 'main'). + * @returns The compiled pipeline and its matching bind group. + */ +export function createComputeBundle( + device: GPUDevice, + code: string, + bindings: readonly BindingSpec[], + constants: Record = {}, + entryPoint: string = 'main' +): ComputePipelineBundle { + let deviceMap = cache.get(device); + if (!deviceMap) { + deviceMap = new Map(); + cache.set(device, deviceMap); + } + + const serializedConsts = serializeConstants(constants); + const cacheKey = `${entryPoint}#${serializedConsts}#${code}`; + + let pipeline = deviceMap.get(cacheKey); + if (!pipeline) { + const module = device.createShaderModule({ code }); + const computeStage: GPUProgrammableStage = { module, entryPoint, constants }; + pipeline = device.createComputePipeline({ layout: 'auto', compute: computeStage }); + deviceMap.set(cacheKey, pipeline); + } + + const bindGroupEntries: GPUBindGroupEntry[] = bindings.map((b) => { + if ('buffer' in b) { + return { binding: b.binding, resource: { buffer: b.buffer, offset: b.offset, size: b.size } }; + } + return { binding: b.binding, resource: b.resource }; + }); + + const bindGroup = device.createBindGroup({ + layout: pipeline.getBindGroupLayout(0), + entries: bindGroupEntries, + }); + + return { pipeline, bindGroup }; +} diff --git a/utils/benchmark/baseline-main/src/context.ts b/utils/benchmark/baseline-main/src/context.ts new file mode 100644 index 0000000..26e172f --- /dev/null +++ b/utils/benchmark/baseline-main/src/context.ts @@ -0,0 +1,381 @@ +import type { Shader } from './shader'; +import type { Program } from './program'; +import type { WgpuDispatch, WgpuCommand } from './command'; + +import { Tensor, DTYPE_BYTESIZE, type DType } from './tensor'; +import { resolveArgs } from './program'; +import { shaderRegistry } from './shaders'; + +const UNIFORM_BUFFER_USAGE = GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST; +const STORAGE_BUFFER_USAGE = + GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST; + +/** + * Execution context that manages GPU buffers and the sequence of recorded + * kernel dispatches. + */ +export class WgpuExecutionContext { + readonly #device: GPUDevice; + readonly #commands: WgpuCommand[] = []; + readonly #ownedBuffers = new Set(); + #destroyed = false; + + /** + * Constructs a new WgpuExecutionContext. + * + * @param device The WebGPU device instance. + */ + constructor(device: GPUDevice) { + this.#device = device; + } + + /** The WebGPU device used by this context. */ + get device(): GPUDevice { + return this.#device; + } + + /** + * Transfers ownership of a buffer to the context to ensure it remains + * allocated for the model lifetime and is destroyed on context disposal. + * + * @param buffer The GPUBuffer to take ownership of. + */ + #ownBuffer(buffer: GPUBuffer): void { + this.#ownedBuffers.add(buffer); + } + + /** + * Throws if this context has been destroyed. + * + * @throws {Error} If {@link destroy} has been called. + */ + #assertNotDestroyed(): void { + if (this.#destroyed) { + throw new Error('WgpuExecutionContext is destroyed'); + } + } + + /** + * Destroys all owned buffers and releases resources. + */ + destroy(): void { + if (this.#destroyed) { + return; + } + this.#destroyed = true; + + for (const buffer of this.#ownedBuffers) { + buffer.destroy(); + } + this.#ownedBuffers.clear(); + this.#commands.length = 0; + } + + // ========================================================================= + // Buffer Allocation + // ========================================================================= + + /** + * Creates an empty WebGPU storage buffer whose lifetime is managed by this + * context. + * + * @param size Size in bytes to allocate. + * @returns The newly allocated and owned GPUBuffer with STORAGE | COPY_SRC | + * COPY_DST usage. + */ + storageBuffer(size: number): GPUBuffer { + this.#assertNotDestroyed(); + const alignedSize = Math.max(4, Math.ceil(size / 4) * 4); + const buffer = this.#device.createBuffer({ size: alignedSize, usage: STORAGE_BUFFER_USAGE }); + this.#ownBuffer(buffer); + return buffer; + } + + /** + * Creates and populates a WebGPU uniform buffer whose lifetime is managed by + * this context. + * + * @param data Binary data to copy into the uniform buffer. + * @returns The newly allocated and owned GPUBuffer. + */ + uniformBuffer(data: ArrayBufferView | ArrayBuffer): GPUBuffer { + this.#assertNotDestroyed(); + + const byteLength = data.byteLength; + const alignedSize = Math.max(16, Math.ceil(byteLength / 4) * 4); + const buffer = this.#device.createBuffer({ size: alignedSize, usage: UNIFORM_BUFFER_USAGE }); + + if (data instanceof ArrayBuffer) { + this.#device.queue.writeBuffer(buffer, 0, data, 0, byteLength); + } else { + this.#device.queue.writeBuffer(buffer, 0, data.buffer, data.byteOffset, byteLength); + } + this.#ownBuffer(buffer); + return buffer; + } + + /** + * Creates a Tensor view over a GPU buffer. + * + * Buffer ownership behavior: + * 1. **`buffer` omitted**: The context allocates a fresh `GPUBuffer` via + * {@link storageBuffer}. This buffer is tracked and owned by the context, + * and will be destroyed when {@link destroy} is called on this context. + * 2. **`buffer` provided**: The context wraps the given `GPUBuffer` without + * altering its lifecycle or ownership. If it was already managed (e.g. via + * {@link storageBuffer}), it remains managed; if it is external, its + * lifetime remains with the caller. + * + * @param dtype Element data type of the tensor. + * @param shape Tensor dimensions. + * @param buffer Optional existing GPUBuffer to wrap. If omitted, the context + * allocates and owns a new buffer. + * @returns A new Tensor view. + */ + tensor(dtype: DType, shape: readonly number[], buffer?: GPUBuffer): Tensor { + this.#assertNotDestroyed(); + + if (buffer !== undefined) { + return new Tensor(dtype, shape, this.#device, buffer); + } + + const numel = shape.reduce((a, b) => a * b, 1); + const byteLength = numel * DTYPE_BYTESIZE[dtype]; + const newBuffer = this.storageBuffer(byteLength); + + return new Tensor(dtype, shape, this.#device, newBuffer); + } + + // ========================================================================= + // Dispatch Recording & Execution + // ========================================================================= + + /** + * Records a compute dispatch into the context execution queue. + * + * @param dispatch The dispatch descriptor to record. + */ + addDispatch(dispatch: WgpuDispatch): void { + this.#assertNotDestroyed(); + this.#commands.push({ kind: 'dispatch', dispatch }); + } + + /** + * Records a deferred GPU buffer-to-buffer copy into the context execution + * queue. The copy is encoded and submitted when {@link submit} is called, + * preserving ordering relative to recorded dispatches. + * + * @param src Source tensor to copy from. + * @param dst Destination tensor to copy to. + * @returns This context instance. + */ + copy(src: Tensor, dst: Tensor): this { + this.#assertNotDestroyed(); + if ((src.buffer.usage & GPUBufferUsage.COPY_SRC) === 0) { + throw new Error('copy: source buffer missing COPY_SRC usage'); + } + if ((dst.buffer.usage & GPUBufferUsage.COPY_DST) === 0) { + throw new Error('copy: destination buffer missing COPY_DST usage'); + } + if (dst.byteLength < src.byteLength) { + throw new Error(`copy: destination (${dst.byteLength}B) < source (${src.byteLength}B)`); + } + this.#commands.push({ + kind: 'copy', + copy: { src: src.buffer, dst: dst.buffer, size: src.byteLength }, + }); + return this; + } + + /** + * Records a compute shader dispatch into this context. + * + * @param shader The compute shader to execute. + * @param args Positional arguments expected by the operator. + * @returns This context instance. + */ + recordShader(shader: Shader, args: TArgs): this { + this.#assertNotDestroyed(); + shader.recordIn(this, args); + return this; + } + + /** + * Records a full serialized program — allocates buffers, uploads weights, + * copies inputs, dispatches the operator chain, and copies outputs. + * + * All work is deferred until {@link submit} is called. + * + * @param program The serialized program descriptor (program.json). + * @param weights Raw constant tensor data (weights.bin). + * @param args I/O tensors and scalars: `[in1, ..., out1, ...]` matching + * `input_ids` and `output_ids`. + * @returns This context instance. + */ + recordProgram( + program: Program, + weights: ArrayBuffer, + args: readonly (Tensor | number | boolean | string)[] + ): this { + this.#assertNotDestroyed(); + + const numInputs = program.inputIds.length; + const numOutputs = program.outputIds.length; + if (args.length !== numInputs + numOutputs) { + throw new Error(`program: expected ${numInputs + numOutputs} args, got ${args.length}`); + } + + // Allocate shared pool buffers + const pools = new Map(); + for (const pool of program.memoryPlan.pools) { + pools.set(pool.id, this.storageBuffer(pool.size)); + } + + // Allocate buffers and create Tensor views for all values + const tensors = new Map(); + for (const [i, v] of program.values.entries()) { + if (v.type !== 'tensor') continue; + + let buffer: GPUBuffer; + + if (v.weightsOffset !== undefined && v.weightsLength !== undefined) { + // Constant — dedicated buffer + buffer = this.storageBuffer(v.weightsLength); + } else if (v.memObjId !== undefined) { + // Intermediate - shared pool buffer + if (!pools.has(v.memObjId)) { + throw new Error(`program: no memory pool ${v.memObjId} for value ${i}`); + } + buffer = pools.get(v.memObjId)!; + } else { + // Dedicated buffer - no pool assignment + const numel = v.shape.reduce((a, b) => a * b, 1); + buffer = this.storageBuffer(numel * DTYPE_BYTESIZE[v.dtype]); + } + + tensors.set(i, this.tensor(v.dtype, v.shape, buffer)); + } + + // Upload weights into constant tensors + for (const [i, v] of program.values.entries()) { + if ( + v.type === 'tensor' && // prettier-ignore + v.weightsOffset !== undefined && + v.weightsLength !== undefined + ) { + const t = tensors.get(i); + const w = weights.slice(v.weightsOffset, v.weightsOffset + v.weightsLength); + if (t === undefined) { + throw new Error(`program: no tensor for constant value ${i}`); + } + t.setData(w); + } + } + + // Copy inputs: user args → internal tensors + for (let i = 0; i < numInputs; i++) { + const src = args[i]; + const dst = tensors.get(program.inputIds[i]); + + if (!(src instanceof Tensor)) { + continue; + } + if (dst === undefined) { + throw new Error(`program: no tensor for input value ${program.inputIds[i]}`); + } + this.copy(src, dst); + } + + // Dispatch operator chain + for (const op of program.chain) { + const shader = shaderRegistry[op.name] as Shader | undefined; + if (!shader) { + throw new Error(`program: unknown operator '${op.name}'`); + } + this.recordShader(shader, resolveArgs(op, program.values, tensors)); + } + + // Copy outputs: internal tensors → user args + for (let i = 0; i < numOutputs; i++) { + const src = tensors.get(program.outputIds[i]); + const dst = args[numInputs + i]; + + if (!(dst instanceof Tensor)) { + throw new Error(`program: output ${i} must be a Tensor`); + } + if (src === undefined) { + throw new Error(`program: no tensor for output value ${program.outputIds[i]}`); + } + this.copy(src, dst); + } + + return this; + } + + /** + * Encodes all recorded commands (dispatches and copies) into a single + * command buffer and submits to the GPU queue. + */ + submit(): void { + this.#assertNotDestroyed(); + if (this.#commands.length === 0) { + return; + } + + const encoder = this.#device.createCommandEncoder(); + let pass: GPUComputePassEncoder | undefined; + + for (const cmd of this.#commands) { + switch (cmd.kind) { + case 'copy': { + if (pass) { + pass.end(); + pass = undefined; + } + encoder.copyBufferToBuffer(cmd.copy.src, 0, cmd.copy.dst, 0, cmd.copy.size); + break; + } + + case 'dispatch': { + if (!pass) { + pass = encoder.beginComputePass(); + } + const d = cmd.dispatch; + pass.setPipeline(d.pipeline); + pass.setBindGroup(0, d.bindGroup); + pass.dispatchWorkgroups(d.workgroupCountX, d.workgroupCountY, d.workgroupCountZ ?? 1); + break; + } + + default: { + const exhaustive: never = cmd; + throw new Error(`Unhandled command: ${JSON.stringify(exhaustive)}`); + } + } + } + + if (pass) { + pass.end(); + } + this.#device.queue.submit([encoder.finish()]); + } + + /** + * Waits for all previously submitted GPU work to complete. + * + * @returns A promise that resolves when the GPU queue is idle. + */ + async sync(): Promise { + this.#assertNotDestroyed(); + await this.#device.queue.onSubmittedWorkDone(); + } + + /** + * Clears all recorded commands to prepare for a fresh recording pass while + * keeping owned buffers intact. + */ + reset(): void { + this.#assertNotDestroyed(); + this.#commands.length = 0; + } +} diff --git a/utils/benchmark/baseline-main/src/index.ts b/utils/benchmark/baseline-main/src/index.ts new file mode 100644 index 0000000..b924f9a --- /dev/null +++ b/utils/benchmark/baseline-main/src/index.ts @@ -0,0 +1,7 @@ +export * from './tensor'; +export * from './context'; +export * from './shader'; +export * from './command'; +export * from './program'; + +export * from './shaders'; diff --git a/utils/benchmark/baseline-main/src/program.ts b/utils/benchmark/baseline-main/src/program.ts new file mode 100644 index 0000000..6ddb5ef --- /dev/null +++ b/utils/benchmark/baseline-main/src/program.ts @@ -0,0 +1,224 @@ +import type { Tensor, DType } from './tensor'; + +// =========================================================================== +// Value types +// =========================================================================== + +/** + * A tensor value in a serialized program — declares a GPU buffer with its + * shape, element type, and how the buffer is allocated. + * + * Every tensor value maps to exactly one `GPUBuffer` at runtime. The buffer + * source is determined by the optional fields: + * + * - **Constant** (`weightsOffset` present): data is loaded from `weights.bin` + * into a dedicated buffer during {@link WgpuExecutionContext.recordProgram}. + * These are model weights and stay immutable after upload. + * + * - **Shared pool** (`memObjId` present): the tensor is allocated inside a + * shared memory pool (see {@link ProgramPool}). Multiple tensors in the same + * pool share one `GPUBuffer` because their lifetimes don't overlap — the + * memory planner guarantees only one is live at a time. + * + * - **Dedicated** (neither present): a standalone `GPUBuffer` allocated just + * for this value. Used when the memory planner can't find a pool to share. + */ +export type ProgramTensor = { + readonly type: 'tensor'; + /** Tensor dimensions (e.g., `[1, 3, 520, 520]`). */ + readonly shape: readonly number[]; + /** Element data type. */ + readonly dtype: DType; + /** Shared memory pool index (absent = dedicated buffer). */ + readonly memObjId?: number; + /** Byte offset into weights.bin (absent = not a constant). */ + readonly weightsOffset?: number; + /** Byte length in weights.bin (absent = not a constant). */ + readonly weightsLength?: number; +}; + +/** + * A scalar value passed directly to operators — not backed by a GPU buffer. + * + * Scalars carry per-operator parameters like `alpha` in `aten.add.Tensor`, + * clamp bounds in `conv_with_clamp`, or boolean flags. They are resolved at + * dispatch time and passed to the shader's `recordIn` as plain JS values. + */ +export type ProgramScalar = { + readonly type: 'scalar'; + readonly value: number | boolean | string; +}; + +/** + * An integer list used for operator configuration — not backed by a GPU + * buffer. Common uses: shape hints for `view_copy` (`[1, -1]`), convolution + * padding (`[1, 1]`), stride, dilation. + */ +export type ProgramList = { + readonly type: 'list'; + readonly items: readonly number[]; +}; + +/** + * A list of tensor value indices used by variadic operators — e.g. the input + * tensor list of `aten.cat.default` or the output tensor list of + * `aten.split_with_sizes_copy.default`. + * + * Items are indices into the program's {@link Program.values} table and + * resolve to `Tensor` objects at dispatch time. + */ +export type ProgramTensorList = { + /** Discriminates the tensor-list variant. */ + readonly type: 'tensorlist'; + /** Value indices of the tensors in this list. */ + readonly items: readonly number[]; +}; + +/** + * An empty/unused value slot. Reserved by the serializer for optional + * operator arguments that were omitted (e.g., `memory_format=None` for + * `clone`). Resolves to `undefined` at dispatch time. + */ +export type ProgramNull = { + readonly type: 'null'; +}; + +/** Union of all value types in a serialized program. */ +export type ProgramValue = + ProgramTensor | ProgramScalar | ProgramList | ProgramTensorList | ProgramNull; + +// =========================================================================== +// Program structure +// =========================================================================== + +/** + * A single operator invocation. `name` is the ExecuTorch operator identifier + * (e.g., `"aten.add.Tensor"`, `"et_vk.conv_with_clamp.default"`) and `args` + * are indices into the program's {@link Program.values} table. + * + * The operator is dispatched by looking up `name` in the shader registry and + * calling its `recordIn` with the resolved arguments. + */ +export type ProgramOp = { + /** ExecuTorch operator identifier. */ + readonly name: string; + /** Value indices into `Program.values`. */ + readonly args: readonly number[]; +}; + +/** + * A shared memory pool — a single `GPUBuffer` that multiple tensors use at + * different points in the computation (their lifetimes don't overlap). + * + * Pool size is the maximum byte size of any tensor assigned to it. This is + * computed at export time by the ExecuTorch memory planner and embedded in + * the program JSON. + */ +export type ProgramPool = { + /** Unique pool index (matches {@link ProgramTensor.memObjId}). */ + readonly id: number; + /** Buffer size in bytes. */ + readonly size: number; +}; + +/** + * The serialized program format — the JSON structure produced when converting + * a `.pte` model and consumed by {@link WgpuExecutionContext.recordProgram}. + * + * A program is a flat, ordered list of operator calls (the "chain") operating + * on a table of values (tensors, scalars, lists). Each operator's `args` are + * value indices — the dependency graph is implicit in the value indices, not + * stored as pointers. The order of `chain` defines execution order. + * + * Example: + * ```json + * { + * "version": "1", + * "chain": [{ "name": "aten.add.Tensor", "args": [0, 1, 2, 3] }], + * "values": [ + * { "type": "tensor", "shape": [2, 3], "dtype": "float32", "memObjId": 1 }, + * { "type": "tensor", "shape": [1, 3], "dtype": "float32", "memObjId": 1 }, + * { "type": "scalar", "value": 1.0 }, + * { "type": "tensor", "shape": [2, 3], "dtype": "float32", "memObjId": 0 } + * ], + * "inputIds": [0, 1], + * "outputIds": [3], + * "memoryPlan": { "pools": [{ "id": 0, "size": 24 }, { "id": 1, "size": 36 }] } + * } + * ``` + */ +export type Program = { + /** Schema version (currently `"1"`). */ + readonly version: string; + /** Ordered list of operator invocations — defines execution order. */ + readonly chain: readonly ProgramOp[]; + /** Value table — tensors, scalars, lists, and nulls referenced by `chain`. */ + readonly values: readonly ProgramValue[]; + /** Value indices that are graph inputs (map to `args[0..n]` in `recordProgram`). */ + readonly inputIds: readonly number[]; + /** Value indices that are graph outputs (map to `args[n..]` in `recordProgram`). */ + readonly outputIds: readonly number[]; + /** Memory plan — shared GPU buffer pools for intermediate activations. */ + readonly memoryPlan: { + readonly pools: readonly ProgramPool[]; + }; +}; + +// =========================================================================== +// Argument resolution +// =========================================================================== + +/** + * Resolves an operator's value-index args into concrete values. + * + * For each index in `op.args`, looks up the corresponding entry in `values` + * and returns: + * - `Tensor` for tensor values (from the `tensors` map) + * - `number | boolean | string` for scalar values + * - `number[]` for list values + * - `Tensor[]` for tensor-list values (from the `tensors` map) + * - `undefined` for null values + * + * @param op The operator whose args to resolve. + * @param values The full value table. + * @param tensors Map from value index to its GPU Tensor view. + * @returns Resolved args ready to pass to a shader's `recordIn`. + * @internal + */ +export function resolveArgs( + op: ProgramOp, + values: readonly ProgramValue[], + tensors: Map +): unknown[] { + return op.args.map((idx) => { + const v = values[idx]; + switch (v.type) { + case 'tensor': { + const t = tensors.get(idx); + if (t === undefined) { + throw new Error(`program: no tensor for value ${idx} referenced by '${op.name}'`); + } + return t; + } + case 'scalar': + return v.value; + case 'list': + return [...v.items]; + case 'tensorlist': { + return v.items.map((idx) => { + const item = values[idx]; + if (item?.type !== 'tensor') { + throw new Error(`program: tensorlist item ${idx} is not a tensor (${op.name})`); + } + const t = tensors.get(idx); + if (t === undefined) { + throw new Error(`program: no tensor for value ${idx} referenced by '${op.name}'`); + } + return t; + }); + } + case 'null': + return undefined; + } + }); +} diff --git a/utils/benchmark/baseline-main/src/shader.ts b/utils/benchmark/baseline-main/src/shader.ts new file mode 100644 index 0000000..4622200 --- /dev/null +++ b/utils/benchmark/baseline-main/src/shader.ts @@ -0,0 +1,34 @@ +import type { WgpuExecutionContext } from './context'; + +/** + * Represents a GPU compute shader in the WebGPU execution engine. + * + * @typeParam TArgs The tuple schema of arguments accepted by this shader's + * {@link recordIn} method. + * @typeParam TName Literal string type of the operator name. + */ +export type Shader< + TArgs extends readonly unknown[] = readonly unknown[], + TName extends string = string, +> = { + /** + * Canonical PyTorch / ExecuTorch operator identifier. + */ + readonly name: TName; + + /** + * Raw WebGPU Shading Language source code, either as a single shader string + * or a dictionary of named shader variants. + */ + readonly code: string | Record; + + /** + * Dispatches and records the operator's compute pass into the execution + * context. Validates input/output tensors, allocates uniform buffers, + * compiles or retrieves cached pipelines, and records the dispatch. + * + * @param ctx The execution context to record dispatches into. + * @param args Positional arguments for the operator. + */ + readonly recordIn: (ctx: WgpuExecutionContext, args: TArgs) => void; +}; diff --git a/src/shaders/add.ts b/utils/benchmark/baseline-main/src/shaders/add.ts similarity index 100% rename from src/shaders/add.ts rename to utils/benchmark/baseline-main/src/shaders/add.ts diff --git a/src/shaders/bmm.ts b/utils/benchmark/baseline-main/src/shaders/bmm.ts similarity index 100% rename from src/shaders/bmm.ts rename to utils/benchmark/baseline-main/src/shaders/bmm.ts diff --git a/src/shaders/cat.ts b/utils/benchmark/baseline-main/src/shaders/cat.ts similarity index 100% rename from src/shaders/cat.ts rename to utils/benchmark/baseline-main/src/shaders/cat.ts diff --git a/src/shaders/clone.ts b/utils/benchmark/baseline-main/src/shaders/clone.ts similarity index 100% rename from src/shaders/clone.ts rename to utils/benchmark/baseline-main/src/shaders/clone.ts diff --git a/src/shaders/constant-pad-nd.ts b/utils/benchmark/baseline-main/src/shaders/constant-pad-nd.ts similarity index 100% rename from src/shaders/constant-pad-nd.ts rename to utils/benchmark/baseline-main/src/shaders/constant-pad-nd.ts diff --git a/src/shaders/conv-with-clamp.ts b/utils/benchmark/baseline-main/src/shaders/conv-with-clamp.ts similarity index 100% rename from src/shaders/conv-with-clamp.ts rename to utils/benchmark/baseline-main/src/shaders/conv-with-clamp.ts diff --git a/src/shaders/convolution.ts b/utils/benchmark/baseline-main/src/shaders/convolution.ts similarity index 100% rename from src/shaders/convolution.ts rename to utils/benchmark/baseline-main/src/shaders/convolution.ts diff --git a/src/shaders/hardswish.ts b/utils/benchmark/baseline-main/src/shaders/hardswish.ts similarity index 100% rename from src/shaders/hardswish.ts rename to utils/benchmark/baseline-main/src/shaders/hardswish.ts diff --git a/utils/benchmark/baseline-main/src/shaders/index.ts b/utils/benchmark/baseline-main/src/shaders/index.ts new file mode 100644 index 0000000..8ed11f0 --- /dev/null +++ b/utils/benchmark/baseline-main/src/shaders/index.ts @@ -0,0 +1,60 @@ +// Elementwise arithmetic & activations +import { add } from './add'; +import { hardswish } from './hardswish'; +import { mul } from './mul'; +import { relu } from './relu'; +import { sigmoid } from './sigmoid'; +import { sub } from './sub'; + +// Linear algebra & reductions +import { bmm } from './bmm'; +import { meanDim } from './mean-dim'; +import { mm } from './mm'; +import { softmax } from './softmax'; + +// Convolutions & pooling +import { convWithClamp } from './conv-with-clamp'; +import { convolution } from './convolution'; +import { maxPool2d } from './max-pool-2d'; + +// Resampling +import { upsampleBilinear2d } from './upsample-bilinear2d'; +import { upsampleNearest2d } from './upsample-nearest-2d'; + +// Shape & data movement +import { cat } from './cat'; +import { clone } from './clone'; +import { constantPadNd } from './constant-pad-nd'; +import { permuteCopy } from './permute'; +import { prepack } from './prepack'; +import { splitWithSizesCopy } from './split-with-sizes'; +import { squeezeCopy } from './squeeze-copy'; +import { unsqueezeCopy } from './unsqueeze-copy'; +import { viewCopy } from './view-copy'; + +/** + * Namespace containing all registered WebGPU compute shaders. + */ +// prettier-ignore +export const shaders = { + // Elementwise arithmetic & activations + add, sub, mul, sigmoid, relu, hardswish, + // Linear algebra & reductions + mm, bmm, meanDim, softmax, + // Convolutions & pooling + convolution, convWithClamp, maxPool2d, + // Resampling + upsampleBilinear2d, upsampleNearest2d, + // Shape & data movement + cat, constantPadNd, splitWithSizesCopy, permuteCopy, viewCopy, + clone, squeezeCopy, unsqueezeCopy, prepack, +} as const; + +/** + * Maps ExecuTorch operator names to their shader implementations. + * + * @internal + */ +export const shaderRegistry = Object.fromEntries( + Object.values(shaders).map((shader) => [shader.name, shader]) +); diff --git a/src/shaders/max-pool-2d.ts b/utils/benchmark/baseline-main/src/shaders/max-pool-2d.ts similarity index 100% rename from src/shaders/max-pool-2d.ts rename to utils/benchmark/baseline-main/src/shaders/max-pool-2d.ts diff --git a/src/shaders/mean-dim.ts b/utils/benchmark/baseline-main/src/shaders/mean-dim.ts similarity index 100% rename from src/shaders/mean-dim.ts rename to utils/benchmark/baseline-main/src/shaders/mean-dim.ts diff --git a/src/shaders/mm.ts b/utils/benchmark/baseline-main/src/shaders/mm.ts similarity index 100% rename from src/shaders/mm.ts rename to utils/benchmark/baseline-main/src/shaders/mm.ts diff --git a/src/shaders/mul.ts b/utils/benchmark/baseline-main/src/shaders/mul.ts similarity index 100% rename from src/shaders/mul.ts rename to utils/benchmark/baseline-main/src/shaders/mul.ts diff --git a/src/shaders/permute.ts b/utils/benchmark/baseline-main/src/shaders/permute.ts similarity index 100% rename from src/shaders/permute.ts rename to utils/benchmark/baseline-main/src/shaders/permute.ts diff --git a/src/shaders/prepack.ts b/utils/benchmark/baseline-main/src/shaders/prepack.ts similarity index 100% rename from src/shaders/prepack.ts rename to utils/benchmark/baseline-main/src/shaders/prepack.ts diff --git a/src/shaders/relu.ts b/utils/benchmark/baseline-main/src/shaders/relu.ts similarity index 100% rename from src/shaders/relu.ts rename to utils/benchmark/baseline-main/src/shaders/relu.ts diff --git a/src/shaders/sigmoid.ts b/utils/benchmark/baseline-main/src/shaders/sigmoid.ts similarity index 100% rename from src/shaders/sigmoid.ts rename to utils/benchmark/baseline-main/src/shaders/sigmoid.ts diff --git a/src/shaders/softmax.ts b/utils/benchmark/baseline-main/src/shaders/softmax.ts similarity index 100% rename from src/shaders/softmax.ts rename to utils/benchmark/baseline-main/src/shaders/softmax.ts diff --git a/src/shaders/split-with-sizes.ts b/utils/benchmark/baseline-main/src/shaders/split-with-sizes.ts similarity index 100% rename from src/shaders/split-with-sizes.ts rename to utils/benchmark/baseline-main/src/shaders/split-with-sizes.ts diff --git a/src/shaders/squeeze-copy.ts b/utils/benchmark/baseline-main/src/shaders/squeeze-copy.ts similarity index 100% rename from src/shaders/squeeze-copy.ts rename to utils/benchmark/baseline-main/src/shaders/squeeze-copy.ts diff --git a/src/shaders/sub.ts b/utils/benchmark/baseline-main/src/shaders/sub.ts similarity index 100% rename from src/shaders/sub.ts rename to utils/benchmark/baseline-main/src/shaders/sub.ts diff --git a/src/shaders/unsqueeze-copy.ts b/utils/benchmark/baseline-main/src/shaders/unsqueeze-copy.ts similarity index 100% rename from src/shaders/unsqueeze-copy.ts rename to utils/benchmark/baseline-main/src/shaders/unsqueeze-copy.ts diff --git a/src/shaders/upsample-bilinear2d.ts b/utils/benchmark/baseline-main/src/shaders/upsample-bilinear2d.ts similarity index 100% rename from src/shaders/upsample-bilinear2d.ts rename to utils/benchmark/baseline-main/src/shaders/upsample-bilinear2d.ts diff --git a/src/shaders/upsample-nearest-2d.ts b/utils/benchmark/baseline-main/src/shaders/upsample-nearest-2d.ts similarity index 100% rename from src/shaders/upsample-nearest-2d.ts rename to utils/benchmark/baseline-main/src/shaders/upsample-nearest-2d.ts diff --git a/utils/benchmark/baseline-main/src/shaders/utils/broadcasting.ts b/utils/benchmark/baseline-main/src/shaders/utils/broadcasting.ts new file mode 100644 index 0000000..a3eac20 --- /dev/null +++ b/utils/benchmark/baseline-main/src/shaders/utils/broadcasting.ts @@ -0,0 +1,52 @@ +/** + * Checks whether an input shape is broadcastable to the target output shape. + * Follows NumPy / PyTorch broadcasting semantics: + * - Dimensions are compared from right to left (trailing dimensions first). + * - Two dimensions are compatible if they are equal or if one of them is 1. + * - Input rank cannot exceed the target rank. + * + * @param shape The shape to check. + * @param targetShape The destination broadcast shape. + * @returns True if `shape` can be broadcast to `targetShape`. + */ +export function isBroadcastable(shape: readonly number[], targetShape: readonly number[]): boolean { + if (shape.length > targetShape.length) { + return false; + } + for (let i = 0; i < shape.length; i++) { + const dim = shape[shape.length - 1 - i]; + const targetDim = targetShape[targetShape.length - 1 - i]; + if (dim !== 1 && dim !== targetDim) { + return false; + } + } + return true; +} + +/** + * Computes the broadcasted output shape of two tensor shapes. + * Returns undefined if the shapes are incompatible under NumPy / PyTorch broadcasting rules. + * + * @param a First shape. + * @param b Second shape. + * @returns Broadcasted shape array, or undefined if incompatible. + */ +export function computeBroadcastShape( + a: readonly number[], + b: readonly number[] +): number[] | undefined { + const maxRank = Math.max(a.length, b.length); + const outShape = new Array(maxRank); + + for (let i = 0; i < maxRank; i++) { + const dimA = a[a.length - 1 - i] ?? 1; + const dimB = b[b.length - 1 - i] ?? 1; + + if (dimA !== 1 && dimB !== 1 && dimA !== dimB) { + return undefined; + } + outShape[maxRank - 1 - i] = Math.max(dimA, dimB); + } + + return outShape; +} diff --git a/utils/benchmark/baseline-main/src/shaders/utils/flat-copy.ts b/utils/benchmark/baseline-main/src/shaders/utils/flat-copy.ts new file mode 100644 index 0000000..c72bc44 --- /dev/null +++ b/utils/benchmark/baseline-main/src/shaders/utils/flat-copy.ts @@ -0,0 +1,36 @@ +import type { Tensor } from '../../tensor'; +import type { WgpuExecutionContext } from '../../context'; + +/** + * Records a flat elementwise copy between two same-numel tensors — the + * contiguous-reshape semantics of view/clone-style operators. + * + * Aliased in/out buffers (the same `GPUBuffer`) are already in place: WebGPU + * rejects overlapping copyBufferToBuffer, so the copy is skipped entirely. + * + * Source: ExecuTorch (backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp) + * GitHub: https://github.com/pytorch/executorch/blob/main/backends/webgpu/runtime/ops/view_copy/ViewCopy.cpp + * License: BSD-3-Clause (Copyright (c) Meta Platforms, Inc. and affiliates) + * + * @param ctx The execution context to record the copy into. + * @param inTensor The source tensor. + * @param outTensor The destination tensor (same numel and dtype). + * @param opName Operator name for error messages. + */ +export function recordFlatCopy( + ctx: WgpuExecutionContext, + inTensor: Tensor, + outTensor: Tensor, + opName: string +): void { + if (inTensor.dtype !== outTensor.dtype) { + throw new Error(`${opName}: dtype mismatch (${inTensor.dtype} != ${outTensor.dtype})`); + } + if (inTensor.numel !== outTensor.numel) { + throw new Error(`${opName}: numel mismatch (${inTensor.numel} != ${outTensor.numel})`); + } + if (inTensor.buffer === outTensor.buffer) { + return; // aliased in/out already in place + } + ctx.copy(inTensor, outTensor); +} diff --git a/utils/benchmark/baseline-main/src/shaders/utils/meta.ts b/utils/benchmark/baseline-main/src/shaders/utils/meta.ts new file mode 100644 index 0000000..329af5d --- /dev/null +++ b/utils/benchmark/baseline-main/src/shaders/utils/meta.ts @@ -0,0 +1,82 @@ +import { MAX_NDIM } from '../../tensor'; +import type { Tensor } from '../../tensor'; + +/** + * Total byte size of the std140 TensorMeta uniform buffer (80 bytes). + * Layout: + * - 0..3: ndim (u32) + * - 4..7: numel (u32) + * - 8..15: padding (8 bytes) + * - 16..47: sizes (8 x u32, packaged as 2 x vec4) + * - 48..79: strides (8 x u32, packaged as 2 x vec4) + */ +export const TENSOR_META_BYTE_SIZE = 80; + +/** + * Common WGSL struct definition for std140 TensorMeta uniform buffers. + * Defines ndim, numel, and 8-element sizes/strides packaged as 2 x vec4. + * + * Source: ExecuTorch (backends/webgpu/runtime/ops/TensorMeta.h) + * GitHub: https://github.com/pytorch/executorch/blob/main/backends/webgpu/runtime/ops/TensorMeta.h + * License: BSD-3-Clause (Copyright (c) Meta Platforms, Inc. and affiliates) + */ +export const TENSOR_META_WGSL = /* wgsl */ ` +struct TensorMeta { + ndim: u32, + numel: u32, + sizes: array, 2>, + strides: array, 2>, +} +`; + +/** std140 32-bit word offsets */ +const OFFSET_NDIM = 0; +const OFFSET_NUMEL = 1; +const OFFSET_SIZES = 4; // Byte offset 16 (16 / 4 = 4) +const OFFSET_STRIDES = 12; // Byte offset 48 (48 / 4 = 12) + +/** + * Encodes a Tensor's metadata into an 80-byte std140 ArrayBuffer matching ExecuTorch's + * `TensorMeta` layout (ndim, numel, 8-element sizes, 8-element strides). + * + * If `targetRank` is provided, dimensions are right-aligned to match the broadcast output rank. + * + * @param tensor The Tensor instance to encode. + * @param targetRank Optional target rank to right-align dimensions for broadcasting. + * @returns ArrayBuffer containing the 80-byte std140 uniform data. + */ +export function encodeTensorMeta( + tensor: Tensor, + targetRank: number = tensor.shape.length +): ArrayBuffer { + const shape = tensor.shape; + + if (shape.length > MAX_NDIM || targetRank > MAX_NDIM) { + throw new Error(`Tensor rank exceeds maximum supported rank of ${MAX_NDIM}`); + } + if (shape.length > targetRank) { + throw new Error(`Shape rank (${shape.length}) exceeds targetRank (${targetRank})`); + } + + const buffer = new ArrayBuffer(TENSOR_META_BYTE_SIZE); + const u32 = new Uint32Array(buffer); + + for (let i = 0; i < MAX_NDIM; i++) { + u32[OFFSET_SIZES + i] = 1; + u32[OFFSET_STRIDES + i] = 0; + } + + u32[OFFSET_NDIM] = targetRank; + u32[OFFSET_NUMEL] = tensor.numel; + + let stride = 1; + for (let i = shape.length - 1; i >= 0; i--) { + const slot = targetRank - shape.length + i; + const dim = shape[i]; + u32[OFFSET_SIZES + slot] = dim; + u32[OFFSET_STRIDES + slot] = stride; + stride *= dim; + } + + return buffer; +} diff --git a/src/shaders/view-copy.ts b/utils/benchmark/baseline-main/src/shaders/view-copy.ts similarity index 100% rename from src/shaders/view-copy.ts rename to utils/benchmark/baseline-main/src/shaders/view-copy.ts diff --git a/utils/benchmark/baseline-main/src/tensor.ts b/utils/benchmark/baseline-main/src/tensor.ts new file mode 100644 index 0000000..2f0f4c8 --- /dev/null +++ b/utils/benchmark/baseline-main/src/tensor.ts @@ -0,0 +1,152 @@ +/** + * Supported tensor data types in WebGPU storage buffers. + */ +export type DType = + 'float32' | 'float16' | 'int32' | 'int8' | 'uint8' | 'bool' | 'float64' | 'int64'; + +/** + * Supported typed arrays that can be uploaded to a Tensor. + */ +export type TypedArray = + Float32Array | Int32Array | Int8Array | Uint8Array | Uint16Array | Float64Array | BigInt64Array; + +/** + * Maximum tensor rank supported by WebGPU kernels (std140 layout limit). + */ +export const MAX_NDIM = 8; + +// prettier-ignore +export const DTYPE_BYTESIZE: Record = { + float32: 4, float16: 2, float64: 8, + int32: 4, int64: 8, + int8: 1, uint8: 1, bool: 1, +} as const; + +const STAGING_BUFFER_USAGE = GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST; + +/** + * A non-owning WebGPU-backed Tensor view. + * + * Encapsulates an interpretation of an allocated `GPUBuffer` in device VRAM + * along with its metadata (`dtype`, `shape`, `numel`, `byteLength`). Use + * {@link Tensor.getData} to read raw bytes back into host memory and + * {@link Tensor.setData} to upload host data into it. The underlying + * `GPUBuffer` lifetime is managed externally. + */ +export class Tensor { + readonly #dtype: DType; + readonly #shape: readonly number[]; + readonly #numel: number; + readonly #buffer: GPUBuffer; + readonly #device: GPUDevice; + + constructor(dtype: DType, shape: readonly number[], device: GPUDevice, buffer: GPUBuffer) { + if (shape.length > MAX_NDIM) { + throw new Error(`Tensor rank ${shape.length} exceeds maximum rank of ${MAX_NDIM}`); + } + if (shape.some((dim) => dim <= 0 || !Number.isInteger(dim))) { + throw new Error('Tensor dimensions must be positive integers'); + } + + this.#dtype = dtype; + this.#shape = shape; + this.#numel = shape.reduce((a, b) => a * b, 1); + this.#buffer = buffer; + this.#device = device; + + const requiredBytes = this.#numel * DTYPE_BYTESIZE[dtype]; + if (buffer.size < requiredBytes) { + throw new Error(`GPUBuffer size (${buffer.size}B) < required (${requiredBytes}B)`); + } + if ((buffer.usage & GPUBufferUsage.STORAGE) === 0) { + throw new Error('GPUBuffer must have GPUBufferUsage.STORAGE flag'); + } + } + + /** The element data type of the tensor. */ + get dtype(): DType { + return this.#dtype; + } + + /** Total number of elements across all dimensions. */ + get numel(): number { + return this.#numel; + } + + /** The dimensions of the tensor. */ + get shape(): readonly number[] { + return this.#shape; + } + + /** The underlying WebGPU storage buffer in VRAM. */ + get buffer(): GPUBuffer { + return this.#buffer; + } + + /** The WebGPU device that owns this tensor's buffer. */ + get device(): GPUDevice { + return this.#device; + } + + /** The total byte size of the tensor data. */ + get byteLength(): number { + return this.#numel * DTYPE_BYTESIZE[this.#dtype]; + } + + /** + * Uploads data from a host TypedArray or ArrayBuffer into this tensor's + * storage buffer. A source smaller than the tensor writes a prefix and + * leaves the remaining bytes untouched. Writes are padded to 4-byte + * alignment, so a partial write may zero up to three bytes past the copied + * data. + * + * @param src Source data to copy from. + * @returns This tensor instance. + * @throws {Error} If the source exceeds the tensor's byte length — pass a + * view (e.g. `arr.subarray(0, n)`) to upload part of a larger buffer. + */ + setData(src: TypedArray | ArrayBuffer): this { + const byteLength = this.byteLength; + + if ((this.#buffer.usage & GPUBufferUsage.COPY_DST) === 0) { + throw new Error('setData: buffer missing COPY_DST usage'); + } + if (src.byteLength > byteLength) { + throw new Error(`setData: source ${src.byteLength}B > tensor ${byteLength}B; pass a view`); + } + + const alignedSize = Math.ceil(src.byteLength / 4) * 4; + const bytes = new Uint8Array(alignedSize); + bytes.set( + src instanceof ArrayBuffer + ? new Uint8Array(src, 0, src.byteLength) + : new Uint8Array(src.buffer, src.byteOffset, src.byteLength) + ); + + this.#device.queue.writeBuffer(this.#buffer, 0, bytes.buffer, bytes.byteOffset, alignedSize); + return this; + } + + /** + * Reads raw bytes back from the GPU storage buffer into host memory. + * + * @returns Raw byte buffer containing the tensor data. + */ + async getData(): Promise { + const byteLength = this.byteLength; + const copySize = Math.ceil(byteLength / 4) * 4; + const stagingSize = Math.max(16, copySize); + + const encoder = this.#device.createCommandEncoder(); + const staging = this.#device.createBuffer({ size: stagingSize, usage: STAGING_BUFFER_USAGE }); + + try { + encoder.copyBufferToBuffer(this.#buffer, 0, staging, 0, copySize); + this.#device.queue.submit([encoder.finish()]); + await staging.mapAsync(GPUMapMode.READ); + return staging.getMappedRange(0, stagingSize).slice(0, byteLength); + } finally { + staging.destroy(); + } + } +} diff --git a/utils/benchmark/depth-fp16-results.json b/utils/benchmark/depth-fp16-results.json new file mode 100644 index 0000000..2035456 --- /dev/null +++ b/utils/benchmark/depth-fp16-results.json @@ -0,0 +1,183 @@ +{ + "date": "2026-10-21", + "baselineCommit": "6c3575a828b8721ce303a33b30415b5c0da7bc24", + "method": "Chrome WebGPU, record once; warm up, then time submit() and await sync() per sample; wall-clock medians include submission/completion overhead.", + "kernelWarmup": 5, + "kernelSamples": 25, + "modelWarmup": 3, + "modelSamples": 15, + "references": [ + "https://huggingface.co/kernels/webgpu-kernels/ai.onnx.Conv/blob/v1/build/webgpu/conv-1x1-channel-reduce.wgsl.jinja", + "https://huggingface.co/blog/webgpu-kernels" + ], + "before": { + "gpu": { + "vendor": "apple", + "architecture": "metal-3" + }, + "results": [ + { + "name": "depth output head", + "dtype": "float32", + "medianMs": 2, + "samples": [ + 1.6999999284744263, 1.7000000476837158, 1.8000000715255737, 1.899999976158142, + 1.899999976158142, 1.899999976158142, 1.899999976158142, 2, 2, 2, 2, 2, 2, 2, + 2.100000023841858, 2.100000023841858, 2.1999999284744263, 2.299999952316284, + 2.399999976158142, 2.4000000953674316, 2.5, 2.600000023841858, 2.700000047683716, 3, + 3.0999999046325684 + ] + }, + { + "name": "DepthART", + "dtype": "float32", + "medianMs": 18.399999976158142, + "samples": [ + 18.299999952316284, 18.299999952316284, 18.299999952316284, 18.299999952316284, + 18.300000071525574, 18.399999976158142, 18.399999976158142, 18.399999976158142, + 18.399999976158142, 18.40000009536743, 18.5, 18.5, 18.5, 18.600000023841858, + 18.700000047683716 + ] + }, + { + "name": "depth output head", + "dtype": "float16", + "medianMs": 0.8999999761581421, + "samples": [ + 0.7999999523162842, 0.8000000715255737, 0.8000000715255737, 0.8999999761581421, + 0.8999999761581421, 0.8999999761581421, 0.8999999761581421, 0.8999999761581421, + 0.8999999761581421, 0.8999999761581421, 0.8999999761581421, 0.8999999761581421, + 0.8999999761581421, 0.8999999761581421, 0.8999999761581421, 0.9000000953674316, + 0.9000000953674316, 0.9000000953674316, 1, 1, 1, 1, 1, 1, 1 + ] + }, + { + "name": "DepthART", + "dtype": "float16", + "medianMs": 14, + "samples": [ + 13.799999952316284, 13.899999976158142, 13.900000095367432, 14, 14, 14, 14, 14, + 14.100000023841858, 14.100000023841858, 14.100000023841858, 14.199999928474426, + 14.200000047683716, 14.300000071525574, 14.599999904632568 + ] + } + ] + }, + "afterKernelChange": { + "gpu": { + "vendor": "apple", + "architecture": "metal-3" + }, + "results": [ + { + "name": "depth output head", + "dtype": "float32", + "medianMs": 0.3999999761581421, + "samples": [ + 0.19999992847442627, 0.20000004768371582, 0.20000004768371582, 0.2999999523162842, + 0.2999999523162842, 0.2999999523162842, 0.2999999523162842, 0.2999999523162842, + 0.2999999523162842, 0.2999999523162842, 0.30000007152557373, 0.30000007152557373, + 0.3999999761581421, 0.3999999761581421, 0.3999999761581421, 0.3999999761581421, + 0.3999999761581421, 0.3999999761581421, 0.3999999761581421, 0.40000009536743164, + 0.40000009536743164, 0.40000009536743164, 0.40000009536743164, 0.5, 0.5 + ] + }, + { + "name": "DepthART", + "dtype": "float32", + "medianMs": 17.59999990463257, + "samples": [ + 17.399999976158142, 17.40000009536743, 17.5, 17.5, 17.5, 17.5, 17.5, 17.59999990463257, + 17.59999990463257, 17.600000023841858, 17.699999928474426, 17.700000047683716, + 17.700000047683716, 17.899999976158142, 18.100000023841858 + ] + }, + { + "name": "depth output head", + "dtype": "float16", + "medianMs": 0.2999999523162842, + "samples": [ + 0.19999992847442627, 0.19999992847442627, 0.20000004768371582, 0.20000004768371582, + 0.20000004768371582, 0.20000004768371582, 0.2999999523162842, 0.2999999523162842, + 0.2999999523162842, 0.2999999523162842, 0.2999999523162842, 0.2999999523162842, + 0.2999999523162842, 0.2999999523162842, 0.2999999523162842, 0.2999999523162842, + 0.2999999523162842, 0.30000007152557373, 0.30000007152557373, 0.30000007152557373, + 0.30000007152557373, 0.30000007152557373, 0.30000007152557373, 0.3999999761581421, + 0.40000009536743164 + ] + }, + { + "name": "DepthART", + "dtype": "float16", + "medianMs": 13.399999976158142, + "samples": [ + 13.199999928474426, 13.299999952316284, 13.299999952316284, 13.299999952316284, + 13.300000071525574, 13.300000071525574, 13.399999976158142, 13.399999976158142, + 13.399999976158142, 13.400000095367432, 13.5, 13.5, 13.5, 13.5, 13.600000023841858 + ] + } + ] + }, + "afterWeightRangeFix": { + "gpu": { + "vendor": "apple", + "architecture": "metal-3" + }, + "results": [ + { + "name": "depth output head", + "dtype": "float32", + "medianMs": 0.3999999761581421, + "samples": [ + 0.2999999523162842, 0.3999999761581421, 0.3999999761581421, 0.3999999761581421, + 0.3999999761581421, 0.3999999761581421, 0.3999999761581421, 0.3999999761581421, + 0.3999999761581421, 0.3999999761581421, 0.3999999761581421, 0.3999999761581421, + 0.3999999761581421, 0.40000009536743164, 0.40000009536743164, 0.40000009536743164, 0.5, + 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.6000000238418579 + ] + }, + { + "name": "DepthART", + "dtype": "float32", + "medianMs": 17.600000023841858, + "samples": [ + 17.300000071525574, 17.5, 17.5, 17.600000023841858, 17.600000023841858, + 17.600000023841858, 17.600000023841858, 17.600000023841858, 17.699999928474426, + 17.699999928474426, 17.700000047683716, 17.700000047683716, 17.799999952316284, + 17.899999976158142, 18 + ] + }, + { + "name": "depth output head", + "dtype": "float16", + "medianMs": 0.30000007152557373, + "samples": [ + 0.20000004768371582, 0.2999999523162842, 0.2999999523162842, 0.2999999523162842, + 0.2999999523162842, 0.2999999523162842, 0.2999999523162842, 0.2999999523162842, + 0.2999999523162842, 0.2999999523162842, 0.30000007152557373, 0.30000007152557373, + 0.30000007152557373, 0.30000007152557373, 0.30000007152557373, 0.3999999761581421, + 0.3999999761581421, 0.3999999761581421, 0.3999999761581421, 0.3999999761581421, + 0.40000009536743164, 0.40000009536743164, 0.5, 0.5, 0.5 + ] + }, + { + "name": "DepthART", + "dtype": "float16", + "medianMs": 13.900000095367432, + "samples": [ + 13.299999952316284, 13.399999976158142, 13.399999976158142, 13.400000095367432, 13.5, + 13.700000047683716, 13.899999976158142, 13.900000095367432, 14.199999928474426, + 14.299999952316284, 14.600000023841858, 15.100000023841858, 15.100000023841858, + 15.199999928474426, 15.200000047683716 + ] + } + ] + }, + "correctness": { + "testFiles": 28, + "tests": 118, + "photographicFrame": "FP16 rendered all white before the decoder range fix, full 0..255 grayscale afterward, including repeated FP32/FP16 switches.", + "syntheticRegression": "Black radius-120 circle centered on a 448x448 white frame. Before balancing: 33570 non-finite depth values. After: all finite, within 5% of FP32; every rendered pixel matches CPU min/max normalization to one byte." + }, + "notes": "One Apple Metal GPU. Full-model FP16 timings vary across runs (13.4\u201313.9 ms after versus 14.0 ms before); the 0.9 to 0.3 ms output-head improvement is repeatable. No cross-device speedup claim." +} diff --git a/utils/benchmark/index.html b/utils/benchmark/index.html new file mode 100644 index 0000000..50995a9 --- /dev/null +++ b/utils/benchmark/index.html @@ -0,0 +1,166 @@ + + + + + + wgpu-executorch • WebGPU Hardware Benchmark + + + +
+
+
+

⚡ wgpu-executorch Hardware Benchmark

+

+ Hardware benchmarking on real GPU: Baseline (main branch) vs. Optimized FP32 vs. Native FP16 +

+
+ +
+ +
+
Initializing WebGPU Device...
+
+
+ +
+ + + + + + + + + + + + + + + +
+ + 

+  

+  
+
+
diff --git a/utils/benchmark/main.ts b/utils/benchmark/main.ts
new file mode 100644
index 0000000..ee74d5a
--- /dev/null
+++ b/utils/benchmark/main.ts
@@ -0,0 +1,1026 @@
+import { WgpuExecutionContext, shaders, type Program } from '../../src';
+import {
+  WgpuExecutionContext as BaselineContext,
+  shaders as baselineShaders,
+} from './baseline-main/src';
+
+export interface Stats {
+  median: number;
+  mean: number;
+  std: number;
+  cv: number; // coefficient of variation (std / mean)
+  min: number;
+  max: number;
+  samples: number[];
+}
+
+interface BenchItem {
+  category: string;
+  name: string;
+  shape: string;
+  baselineMs: number | null;
+  baselineStd?: number | null;
+  fp32Ms: number;
+  fp32Std: number;
+  fp16Ms: number | null;
+  fp16Std?: number | null;
+  speedupFp32: string;
+  speedupFp16: string;
+  stats?: {
+    baseline?: Stats | null;
+    fp32: Stats;
+    fp16?: Stats | null;
+  };
+}
+
+function makeF32(numel: number): Float32Array {
+  const arr = new Float32Array(numel);
+  for (let i = 0; i < numel; i++) arr[i] = (i % 251) / 250 - 0.5;
+  return arr;
+}
+
+function makeF16(numel: number): Float16Array {
+  const arr = new Float16Array(numel);
+  for (let i = 0; i < numel; i++) arr[i] = (i % 251) / 250 - 0.5;
+  return arr;
+}
+
+function computeStats(samples: number[]): Stats {
+  const sorted = [...samples].sort((a, b) => a - b);
+  const n = sorted.length;
+  const mid = Math.floor(n / 2);
+  const median = n % 2 === 0 ? (sorted[mid - 1] + sorted[mid]) / 2 : sorted[mid];
+  const mean = samples.reduce((a, b) => a + b, 0) / n;
+  const variance = n > 1 ? samples.reduce((a, b) => a + (b - mean) ** 2, 0) / (n - 1) : 0;
+  const std = Math.sqrt(variance);
+  const cv = mean > 0 ? std / mean : 0;
+  return { median, mean, std, cv, min: sorted[0], max: sorted[n - 1], samples };
+}
+
+async function measureStatsMs(
+  ctx: { submit: () => void; sync: () => Promise },
+  warmup = 3,
+  runs = 9
+): Promise {
+  for (let w = 0; w < warmup; w++) {
+    ctx.submit();
+    await ctx.sync();
+  }
+  const samples: number[] = [];
+  for (let i = 0; i < runs; i++) {
+    const t0 = performance.now();
+    ctx.submit();
+    await ctx.sync();
+    samples.push(performance.now() - t0);
+  }
+  return computeStats(samples);
+}
+
+async function measureBatchedStatsMs(
+  ctx: { submit: () => void; sync: () => Promise },
+  batchSize = 25,
+  warmupBatches = 2,
+  measuredBatches = 5
+): Promise {
+  for (let w = 0; w < warmupBatches; w++) {
+    for (let i = 0; i < batchSize; i++) ctx.submit();
+    await ctx.sync();
+  }
+  const samples: number[] = [];
+  for (let b = 0; b < measuredBatches; b++) {
+    const t0 = performance.now();
+    for (let i = 0; i < batchSize; i++) ctx.submit();
+    await ctx.sync();
+    samples.push((performance.now() - t0) / batchSize);
+  }
+  return computeStats(samples);
+}
+
+async function run() {
+  const statusEl = document.getElementById('status')!;
+  const infoEl = document.getElementById('info')!;
+  const tableEl = document.getElementById('results-table')!;
+  const bodyEl = document.getElementById('results-body')!;
+  const jsonEl = document.getElementById('json-output')!;
+
+  if (!navigator.gpu) {
+    statusEl.innerText = 'WebGPU is not supported in this browser.';
+    return;
+  }
+
+  const adapter = await navigator.gpu.requestAdapter();
+  if (!adapter) {
+    statusEl.innerText = 'Failed to obtain WebGPU adapter.';
+    return;
+  }
+
+  const hasF16 = adapter.features.has('shader-f16');
+  const requiredFeatures: GPUFeatureName[] = [];
+  if (hasF16) requiredFeatures.push('shader-f16');
+
+  const device = await adapter.requestDevice({ requiredFeatures });
+  const info = adapter.info;
+
+  infoEl.innerHTML = `
+    
GPU Hardware: ${info.vendor || 'Unknown'} (${info.architecture || 'Unknown architecture'}) - ${info.device || 'Default device'}

+    
shader-f16 Feature: ${hasF16 ? 'YES (Native 16-bit ALUs & memory bandwidth enabled)' : 'NO'}

+  `;
+
+  const results: BenchItem[] = [];
+  bodyEl.innerHTML = '';
+
+  const addRow = (item: BenchItem) => {
+    results.push(item);
+    const tr = document.createElement('tr');
+    tr.innerHTML = `
+      ${item.category}
+      ${item.name}
+      ${item.shape}
+      
+        ${
+          item.baselineMs !== null
+            ? `
${item.baselineMs.toFixed(3)} ms
${
+                item.baselineStd !== undefined && item.baselineStd !== null
+                  ? `
±${item.baselineStd.toFixed(3)} ms
`
+                  : ''
+              }`
+            : 'N/A'
+        }
+      
+      
+        
${item.fp32Ms.toFixed(3)} ms

+        
±${item.fp32Std.toFixed(3)} ms (${(item.stats?.fp32.cv ? item.stats.fp32.cv * 100 : 0).toFixed(1)}%)

+      
+      
+        ${
+          item.fp16Ms !== null
+            ? `
${item.fp16Ms.toFixed(3)} ms
±${(item.fp16Std ?? 0).toFixed(3)} ms (${(item.stats?.fp16?.cv ? item.stats.fp16.cv * 100 : 0).toFixed(1)}%)
`
+            : 'N/A'
+        }
+      
+      ${item.speedupFp32}
+      ${item.speedupFp16}
+    `;
+    bodyEl.appendChild(tr);
+    tableEl.style.display = 'table';
+  };
+
+  // 0. End-to-End ExecuTorch Models
+  const benchmarkModel = async (
+    modelName: string,
+    displayName: string,
+    inShape: number[],
+    outShape: number[],
+    warmup = 3,
+    runs = 9
+  ) => {
+    statusEl.innerText = `Benchmarking full model: ${displayName} (Baseline - main)...`;
+
+    let baseStats: Stats | null = null;
+    try {
+      const progRes = await fetch(`/assets/${modelName}/program.json`);
+      const program = (await progRes.json()) as Program;
+      const weightsRes = await fetch(`/assets/${modelName}/weights.bin`);
+      const weights = await weightsRes.arrayBuffer();
+
+      const baseCtx = new BaselineContext(device);
+      const inNumel = inShape.reduce((a, b) => a * b, 1);
+      const inp = baseCtx.tensor('float32', inShape).setData(makeF32(inNumel));
+      const out = baseCtx.tensor('float32', outShape);
+
+      baseCtx.recordProgram(program, weights, [inp, out]);
+      baseStats = await measureStatsMs(baseCtx, warmup, runs);
+      baseCtx.destroy();
+    } catch (e) {
+      // eslint-disable-next-line no-console
+      console.warn(`Could not run ${displayName} Baseline:`, e);
+    }
+
+    statusEl.innerText = `Benchmarking full model: ${displayName} (Opt FP32)...`;
+    let fp32Stats: Stats | null = null;
+    try {
+      const progRes = await fetch(`/assets/${modelName}/program.json`);
+      const program = (await progRes.json()) as Program;
+      const weightsRes = await fetch(`/assets/${modelName}/weights.bin`);
+      const weights = await weightsRes.arrayBuffer();
+
+      const ctx32 = new WgpuExecutionContext(device);
+      const inNumel = inShape.reduce((a, b) => a * b, 1);
+      const inp = ctx32.tensor('float32', inShape).setData(makeF32(inNumel));
+      const out = ctx32.tensor('float32', outShape);
+
+      ctx32.recordProgram(program, weights, [inp, out]);
+      fp32Stats = await measureStatsMs(ctx32, warmup, runs);
+      ctx32.destroy();
+    } catch (e) {
+      // eslint-disable-next-line no-console
+      console.warn(`Could not run ${displayName} FP32:`, e);
+      return;
+    }
+
+    let fp16Stats: Stats | null = null;
+    let speedupFp16 = 'N/A';
+    if (hasF16) {
+      statusEl.innerText = `Benchmarking full model: ${displayName} (Native FP16)...`;
+      try {
+        const progRes = await fetch(`/assets/${modelName}/program-fp16.json`);
+        const program = (await progRes.json()) as Program;
+        const weightsRes = await fetch(`/assets/${modelName}/weights-fp16.bin`);
+        const weights = await weightsRes.arrayBuffer();
+
+        const ctx16 = new WgpuExecutionContext(device);
+        const inNumel = inShape.reduce((a, b) => a * b, 1);
+        const inp = ctx16.tensor('float16', inShape).setData(makeF16(inNumel));
+        const out = ctx16.tensor('float16', outShape);
+
+        ctx16.recordProgram(program, weights, [inp, out]);
+        fp16Stats = await measureStatsMs(ctx16, warmup, runs);
+        speedupFp16 = baseStats
+          ? `${(baseStats.median / fp16Stats.median).toFixed(2)}x`
+          : `${(fp32Stats.median / fp16Stats.median).toFixed(2)}x`;
+        ctx16.destroy();
+      } catch (e) {
+        // eslint-disable-next-line no-console
+        console.warn(`Could not run ${displayName} FP16:`, e);
+      }
+    }
+
+    const baselineMs = baseStats?.median ?? null;
+    const fp32Ms = fp32Stats.median;
+    const speedupFp32 = baselineMs ? `${(baselineMs / fp32Ms).toFixed(2)}x` : '1.00x';
+
+    addRow({
+      category: 'End-to-End Model',
+      name: displayName,
+      shape: `[${inShape.join(',')}] → [${outShape.join(',')}]`,
+      baselineMs,
+      baselineStd: baseStats?.std ?? null,
+      fp32Ms,
+      fp32Std: fp32Stats.std,
+      fp16Ms: fp16Stats?.median ?? null,
+      fp16Std: fp16Stats?.std ?? null,
+      speedupFp32,
+      speedupFp16,
+      stats: {
+        baseline: baseStats,
+        fp32: fp32Stats,
+        fp16: fp16Stats,
+      },
+    });
+  };
+
+  await benchmarkModel(
+    'depthart',
+    'DepthART Monocular Depth',
+    [1, 3, 448, 448],
+    [1, 1, 448, 448],
+    4
+  );
+  await benchmarkModel('yolo26n', 'YOLO26-N Object Detection', [1, 3, 640, 640], [1, 84, 8400], 4);
+  await benchmarkModel(
+    'selfie-segmentation',
+    'Selfie Segmentation',
+    [1, 3, 256, 256],
+    [1, 1, 256, 256],
+    5
+  );
+
+  // 1. Matrix Multiplication (MM)
+  const mmConfigs = [
+    [256, 256, 25],
+    [512, 512, 20],
+    [1024, 1024, 15],
+    [2048, 2048, 8],
+  ] as const;
+
+  for (const [M, N, iters] of mmConfigs) {
+    statusEl.innerText = `Benchmarking MM ${M}x${N}...`;
+
+    // Baseline
+    const baseCtx = new BaselineContext(device);
+    const ba = baseCtx.tensor('float32', [M, N]).setData(makeF32(M * N));
+    const bb = baseCtx.tensor('float32', [N, N]).setData(makeF32(N * N));
+    const bout = baseCtx.tensor('float32', [M, N]);
+    baseCtx.recordShader(baselineShaders.mm, [ba, bb, bout]);
+    const baseStats = await measureBatchedStatsMs(baseCtx, iters);
+    baseCtx.destroy();
+
+    // Optimized FP32
+    const ctx32 = new WgpuExecutionContext(device);
+    const a32 = ctx32.tensor('float32', [M, N]).setData(makeF32(M * N));
+    const b32 = ctx32.tensor('float32', [N, N]).setData(makeF32(N * N));
+    const out32 = ctx32.tensor('float32', [M, N]);
+    ctx32.recordShader(shaders.mm, [a32, b32, out32]);
+    const fp32Stats = await measureBatchedStatsMs(ctx32, iters);
+    ctx32.destroy();
+
+    // Native FP16
+    let fp16Stats: Stats | null = null;
+    if (hasF16) {
+      const ctx16 = new WgpuExecutionContext(device);
+      const a16 = ctx16.tensor('float16', [M, N]).setData(makeF16(M * N));
+      const b16 = ctx16.tensor('float16', [N, N]).setData(makeF16(N * N));
+      const out16 = ctx16.tensor('float16', [M, N]);
+      ctx16.recordShader(shaders.mm, [a16, b16, out16]);
+      fp16Stats = await measureBatchedStatsMs(ctx16, iters);
+      ctx16.destroy();
+    }
+
+    const baselineMs = baseStats.median;
+    const fp32Ms = fp32Stats.median;
+    const fp16Ms = fp16Stats?.median ?? null;
+
+    const speedupFp32 = `${(baselineMs / fp32Ms).toFixed(2)}x`;
+    const speedupFp16 = fp16Ms !== null ? `${(baselineMs / fp16Ms).toFixed(2)}x` : 'N/A';
+
+    addRow({
+      category: 'Linear Algebra',
+      name: `MatMul (MM ${M >= 64 ? '64x64' : 'Vec4'})`,
+      shape: `[${M}, ${N}] * [${N}, ${N}]`,
+      baselineMs,
+      baselineStd: baseStats.std,
+      fp32Ms,
+      fp32Std: fp32Stats.std,
+      fp16Ms,
+      fp16Std: fp16Stats?.std ?? null,
+      speedupFp32,
+      speedupFp16,
+      stats: {
+        baseline: baseStats,
+        fp32: fp32Stats,
+        fp16: fp16Stats,
+      },
+    });
+  }
+
+  // 2. Batched Matrix Multiplication (BMM)
+  const bmmConfigs = [
+    [4, 64, 64, 30],
+    [8, 128, 128, 20],
+  ] as const;
+
+  for (const [B, M, N, iters] of bmmConfigs) {
+    statusEl.innerText = `Benchmarking BMM [${B}, ${M}, ${N}]...`;
+
+    // Baseline
+    const baseCtx = new BaselineContext(device);
+    const ba = baseCtx.tensor('float32', [B, M, N]).setData(makeF32(B * M * N));
+    const bb = baseCtx.tensor('float32', [B, N, N]).setData(makeF32(B * N * N));
+    const bout = baseCtx.tensor('float32', [B, M, N]);
+    baseCtx.recordShader(baselineShaders.bmm, [ba, bb, bout]);
+    const baseStats = await measureBatchedStatsMs(baseCtx, iters);
+    baseCtx.destroy();
+
+    // Optimized FP32
+    const ctx32 = new WgpuExecutionContext(device);
+    const a32 = ctx32.tensor('float32', [B, M, N]).setData(makeF32(B * M * N));
+    const b32 = ctx32.tensor('float32', [B, N, N]).setData(makeF32(B * N * N));
+    const out32 = ctx32.tensor('float32', [B, M, N]);
+    ctx32.recordShader(shaders.bmm, [a32, b32, out32]);
+    const fp32Stats = await measureBatchedStatsMs(ctx32, iters);
+    ctx32.destroy();
+
+    // Native FP16
+    let fp16Stats: Stats | null = null;
+    if (hasF16) {
+      const ctx16 = new WgpuExecutionContext(device);
+      const a16 = ctx16.tensor('float16', [B, M, N]).setData(makeF16(B * M * N));
+      const b16 = ctx16.tensor('float16', [B, N, N]).setData(makeF16(B * N * N));
+      const out16 = ctx16.tensor('float16', [B, M, N]);
+      ctx16.recordShader(shaders.bmm, [a16, b16, out16]);
+      fp16Stats = await measureBatchedStatsMs(ctx16, iters);
+      ctx16.destroy();
+    }
+
+    const baselineMs = baseStats.median;
+    const fp32Ms = fp32Stats.median;
+    const fp16Ms = fp16Stats?.median ?? null;
+
+    const speedupFp32 = `${(baselineMs / fp32Ms).toFixed(2)}x`;
+    const speedupFp16 = fp16Ms !== null ? `${(baselineMs / fp16Ms).toFixed(2)}x` : 'N/A';
+
+    addRow({
+      category: 'Linear Algebra',
+      name: 'Batch MatMul (BMM)',
+      shape: `[${B}, ${M}, ${N}] * [${B}, ${N}, ${N}]`,
+      baselineMs,
+      baselineStd: baseStats.std,
+      fp32Ms,
+      fp32Std: fp32Stats.std,
+      fp16Ms,
+      fp16Std: fp16Stats?.std ?? null,
+      speedupFp32,
+      speedupFp16,
+      stats: {
+        baseline: baseStats,
+        fp32: fp32Stats,
+        fp16: fp16Stats,
+      },
+    });
+  }
+
+  // 3. Convolutions (Tile64 3x3, Pointwise 1x1, Depthwise, ConvWithClamp)
+  const convConfigs = [
+    {
+      shader: shaders.convWithClamp,
+      category: 'Convolution',
+      label: 'DepthART output 1x1 (single channel)',
+      inShape: [1, 16, 448, 448],
+      wShape: [1, 16, 1, 1],
+      groups: 1,
+      stride: [1, 1],
+      pad: [0, 0],
+      iters: 25,
+    },
+    {
+      shader: shaders.convolution,
+      category: 'Convolution',
+      label: 'Conv2D 3x3 (Tile64)',
+      inShape: [1, 64, 56, 56],
+      wShape: [64, 64, 3, 3],
+      groups: 1,
+      stride: [1, 1],
+      pad: [1, 1],
+      iters: 20,
+    },
+    {
+      shader: shaders.convolution,
+      category: 'Convolution',
+      label: 'Conv2D 1x1 (Tile64)',
+      inShape: [1, 128, 56, 56],
+      wShape: [128, 128, 1, 1],
+      groups: 1,
+      stride: [1, 1],
+      pad: [0, 0],
+      iters: 20,
+    },
+    {
+      shader: shaders.convolution,
+      category: 'Convolution',
+      label: 'Conv2D 3x3 Depthwise',
+      inShape: [1, 64, 56, 56],
+      wShape: [64, 1, 3, 3],
+      groups: 64,
+      stride: [1, 1],
+      pad: [1, 1],
+      iters: 25,
+    },
+    {
+      shader: shaders.convWithClamp,
+      category: 'Convolution',
+      label: 'ConvWithClamp 3x3 (Tile64)',
+      inShape: [1, 64, 56, 56],
+      wShape: [64, 64, 3, 3],
+      groups: 1,
+      stride: [1, 1],
+      pad: [1, 1],
+      iters: 20,
+    },
+    {
+      shader: shaders.convWithClamp,
+      category: 'Convolution',
+      label: 'ConvWithClamp 1x1 (Tile64)',
+      inShape: [1, 128, 56, 56],
+      wShape: [128, 128, 1, 1],
+      groups: 1,
+      stride: [1, 1],
+      pad: [0, 0],
+      iters: 20,
+    },
+  ];
+
+  for (const c of convConfigs) {
+    statusEl.innerText = `Benchmarking ${c.label}...`;
+
+    const inNumel = c.inShape.reduce((a, b) => a * b, 1);
+    const wNumel = c.wShape.reduce((a, b) => a * b, 1);
+    const outShape = [c.inShape[0], c.wShape[0], c.inShape[2], c.inShape[3]];
+
+    // Baseline
+    const baseCtx = new BaselineContext(device);
+    const bin = baseCtx.tensor('float32', c.inShape).setData(makeF32(inNumel));
+    const bw = baseCtx.tensor('float32', c.wShape).setData(makeF32(wNumel));
+    const bb = baseCtx.tensor('float32', [c.wShape[0]]).setData(makeF32(c.wShape[0]));
+    const bout = baseCtx.tensor('float32', outShape);
+    if (c.shader === shaders.convolution) {
+      baseCtx.recordShader(baselineShaders.convolution, [
+        bin,
+        bw,
+        bb,
+        c.stride,
+        c.pad,
+        [1, 1],
+        false,
+        [0, 0],
+        c.groups,
+        bout,
+      ]);
+    } else {
+      baseCtx.recordShader(baselineShaders.convWithClamp, [
+        bin,
+        bw,
+        bb,
+        c.stride,
+        c.pad,
+        [1, 1],
+        false,
+        [0, 0],
+        c.groups,
+        0.0,
+        6.0,
+        bout,
+      ]);
+    }
+    const baseStats = await measureBatchedStatsMs(baseCtx, c.iters);
+    baseCtx.destroy();
+
+    // FP32
+    const ctx32 = new WgpuExecutionContext(device);
+    const in32 = ctx32.tensor('float32', c.inShape).setData(makeF32(inNumel));
+    const w32 = ctx32.tensor('float32', c.wShape).setData(makeF32(wNumel));
+    const b32 = ctx32.tensor('float32', [c.wShape[0]]).setData(makeF32(c.wShape[0]));
+    const out32 = ctx32.tensor('float32', outShape);
+
+    if (c.shader === shaders.convolution) {
+      ctx32.recordShader(shaders.convolution, [
+        in32,
+        w32,
+        b32,
+        c.stride,
+        c.pad,
+        [1, 1],
+        false,
+        [0, 0],
+        c.groups,
+        out32,
+      ]);
+    } else {
+      ctx32.recordShader(shaders.convWithClamp, [
+        in32,
+        w32,
+        b32,
+        c.stride,
+        c.pad,
+        [1, 1],
+        false,
+        [0, 0],
+        c.groups,
+        0.0,
+        6.0,
+        out32,
+      ]);
+    }
+
+    const fp32Stats = await measureBatchedStatsMs(ctx32, c.iters);
+    ctx32.destroy();
+
+    // FP16
+    let fp16Stats: Stats | null = null;
+    if (hasF16) {
+      const ctx16 = new WgpuExecutionContext(device);
+      const in16 = ctx16.tensor('float16', c.inShape).setData(makeF16(inNumel));
+      const w16 = ctx16.tensor('float16', c.wShape).setData(makeF16(wNumel));
+      const b16 = ctx16.tensor('float16', [c.wShape[0]]).setData(makeF16(c.wShape[0]));
+      const out16 = ctx16.tensor('float16', outShape);
+
+      if (c.shader === shaders.convolution) {
+        ctx16.recordShader(shaders.convolution, [
+          in16,
+          w16,
+          b16,
+          c.stride,
+          c.pad,
+          [1, 1],
+          false,
+          [0, 0],
+          c.groups,
+          out16,
+        ]);
+      } else {
+        ctx16.recordShader(shaders.convWithClamp, [
+          in16,
+          w16,
+          b16,
+          c.stride,
+          c.pad,
+          [1, 1],
+          false,
+          [0, 0],
+          c.groups,
+          0.0,
+          6.0,
+          out16,
+        ]);
+      }
+
+      fp16Stats = await measureBatchedStatsMs(ctx16, c.iters);
+      ctx16.destroy();
+    }
+
+    const baselineMs = baseStats.median;
+    const fp32Ms = fp32Stats.median;
+    const fp16Ms = fp16Stats?.median ?? null;
+
+    const speedupFp32 = `${(baselineMs / fp32Ms).toFixed(2)}x`;
+    const speedupFp16 = fp16Ms !== null ? `${(baselineMs / fp16Ms).toFixed(2)}x` : 'N/A';
+
+    addRow({
+      category: c.category,
+      name: c.label,
+      shape: `[${c.inShape.join(',')}] * [${c.wShape.join(',')}]`,
+      baselineMs,
+      baselineStd: baseStats.std,
+      fp32Ms,
+      fp32Std: fp32Stats.std,
+      fp16Ms,
+      fp16Std: fp16Stats?.std ?? null,
+      speedupFp32,
+      speedupFp16,
+      stats: {
+        baseline: baseStats,
+        fp32: fp32Stats,
+        fp16: fp16Stats,
+      },
+    });
+  }
+
+  // 4. Activations
+  const actConfigs = [
+    { name: 'ReLU (Vec4)', shader: shaders.relu, baseShader: baselineShaders.relu },
+    { name: 'Sigmoid (Vec4)', shader: shaders.sigmoid, baseShader: baselineShaders.sigmoid },
+    { name: 'Hardswish (Vec4)', shader: shaders.hardswish, baseShader: baselineShaders.hardswish },
+  ];
+
+  const actShape = [1, 64, 56, 56];
+  const actNumel = 1 * 64 * 56 * 56;
+  const actIters = 40;
+
+  for (const a of actConfigs) {
+    statusEl.innerText = `Benchmarking ${a.name}...`;
+
+    // Baseline
+    const baseCtx = new BaselineContext(device);
+    const bin = baseCtx.tensor('float32', actShape).setData(makeF32(actNumel));
+    const bout = baseCtx.tensor('float32', actShape);
+    baseCtx.recordShader(a.baseShader, [bin, bout]);
+    const baseStats = await measureBatchedStatsMs(baseCtx, actIters);
+    baseCtx.destroy();
+
+    // FP32
+    const ctx32 = new WgpuExecutionContext(device);
+    const in32 = ctx32.tensor('float32', actShape).setData(makeF32(actNumel));
+    const out32 = ctx32.tensor('float32', actShape);
+    ctx32.recordShader(a.shader, [in32, out32]);
+    const fp32Stats = await measureBatchedStatsMs(ctx32, actIters);
+    ctx32.destroy();
+
+    // FP16
+    let fp16Stats: Stats | null = null;
+    if (hasF16) {
+      const ctx16 = new WgpuExecutionContext(device);
+      const in16 = ctx16.tensor('float16', actShape).setData(makeF16(actNumel));
+      const out16 = ctx16.tensor('float16', actShape);
+      ctx16.recordShader(a.shader, [in16, out16]);
+      fp16Stats = await measureBatchedStatsMs(ctx16, actIters);
+      ctx16.destroy();
+    }
+
+    const baselineMs = baseStats.median;
+    const fp32Ms = fp32Stats.median;
+    const fp16Ms = fp16Stats?.median ?? null;
+
+    const speedupFp32 = `${(baselineMs / fp32Ms).toFixed(2)}x`;
+    const speedupFp16 = fp16Ms !== null ? `${(baselineMs / fp16Ms).toFixed(2)}x` : 'N/A';
+
+    addRow({
+      category: 'Activation',
+      name: a.name,
+      shape: `[${actShape.join(',')}]`,
+      baselineMs,
+      baselineStd: baseStats.std,
+      fp32Ms,
+      fp32Std: fp32Stats.std,
+      fp16Ms,
+      fp16Std: fp16Stats?.std ?? null,
+      speedupFp32,
+      speedupFp16,
+      stats: {
+        baseline: baseStats,
+        fp32: fp32Stats,
+        fp16: fp16Stats,
+      },
+    });
+  }
+
+  // 5. Reductions (Softmax, MeanDim)
+  {
+    statusEl.innerText = 'Benchmarking Softmax (Vec4)...';
+    const smShape = [128, 1024];
+    const smNumel = 128 * 1024;
+    const smIters = 30;
+
+    // Baseline
+    const baseCtx = new BaselineContext(device);
+    const bin = baseCtx.tensor('float32', smShape).setData(makeF32(smNumel));
+    const bout = baseCtx.tensor('float32', smShape);
+    baseCtx.recordShader(baselineShaders.softmax, [bin, 1, false, bout]);
+    const baseStats = await measureBatchedStatsMs(baseCtx, smIters);
+    baseCtx.destroy();
+
+    // FP32
+    const ctx32 = new WgpuExecutionContext(device);
+    const in32 = ctx32.tensor('float32', smShape).setData(makeF32(smNumel));
+    const out32 = ctx32.tensor('float32', smShape);
+    ctx32.recordShader(shaders.softmax, [in32, 1, false, out32]);
+    const fp32Stats = await measureBatchedStatsMs(ctx32, smIters);
+    ctx32.destroy();
+
+    // FP16
+    let fp16Stats: Stats | null = null;
+    if (hasF16) {
+      const ctx16 = new WgpuExecutionContext(device);
+      const in16 = ctx16.tensor('float16', smShape).setData(makeF16(smNumel));
+      const out16 = ctx16.tensor('float16', smShape);
+      ctx16.recordShader(shaders.softmax, [in16, 1, false, out16]);
+      fp16Stats = await measureBatchedStatsMs(ctx16, smIters);
+      ctx16.destroy();
+    }
+
+    const baselineMs = baseStats.median;
+    const fp32Ms = fp32Stats.median;
+    const fp16Ms = fp16Stats?.median ?? null;
+
+    const speedupFp32 = `${(baselineMs / fp32Ms).toFixed(2)}x`;
+    const speedupFp16 = fp16Ms !== null ? `${(baselineMs / fp16Ms).toFixed(2)}x` : 'N/A';
+
+    addRow({
+      category: 'Reduction',
+      name: 'Softmax (Last-Dim Vec4)',
+      shape: `[${smShape.join(',')}]`,
+      baselineMs,
+      baselineStd: baseStats.std,
+      fp32Ms,
+      fp32Std: fp32Stats.std,
+      fp16Ms,
+      fp16Std: fp16Stats?.std ?? null,
+      speedupFp32,
+      speedupFp16,
+      stats: {
+        baseline: baseStats,
+        fp32: fp32Stats,
+        fp16: fp16Stats,
+      },
+    });
+  }
+
+  // 6. Pooling & Resampling
+  {
+    statusEl.innerText = 'Benchmarking MaxPool2D (2D Spatial)...';
+    const pShape = [1, 64, 112, 112];
+    const pNumel = 1 * 64 * 112 * 112;
+    const pIters = 25;
+
+    // Baseline
+    const baseCtx = new BaselineContext(device);
+    const bin = baseCtx.tensor('float32', pShape).setData(makeF32(pNumel));
+    const boutVals = baseCtx.tensor('float32', [1, 64, 56, 56]);
+    const boutIdx = baseCtx.tensor('int32', [1, 64, 56, 56]);
+    baseCtx.recordShader(baselineShaders.maxPool2d, [
+      bin,
+      [2, 2],
+      [2, 2],
+      [0, 0],
+      [1, 1],
+      false,
+      [boutVals, boutIdx],
+    ]);
+    const baseStats = await measureBatchedStatsMs(baseCtx, pIters);
+    baseCtx.destroy();
+
+    // FP32
+    const ctx32 = new WgpuExecutionContext(device);
+    const in32 = ctx32.tensor('float32', pShape).setData(makeF32(pNumel));
+    const outVals32 = ctx32.tensor('float32', [1, 64, 56, 56]);
+    const outIdx32 = ctx32.tensor('int32', [1, 64, 56, 56]);
+    ctx32.recordShader(shaders.maxPool2d, [
+      in32,
+      [2, 2],
+      [2, 2],
+      [0, 0],
+      [1, 1],
+      false,
+      [outVals32, outIdx32],
+    ]);
+    const fp32Stats = await measureBatchedStatsMs(ctx32, pIters);
+    ctx32.destroy();
+
+    // FP16
+    let fp16Stats: Stats | null = null;
+    if (hasF16) {
+      const ctx16 = new WgpuExecutionContext(device);
+      const in16 = ctx16.tensor('float16', pShape).setData(makeF16(pNumel));
+      const outVals16 = ctx16.tensor('float16', [1, 64, 56, 56]);
+      const outIdx16 = ctx16.tensor('int32', [1, 64, 56, 56]);
+      ctx16.recordShader(shaders.maxPool2d, [
+        in16,
+        [2, 2],
+        [2, 2],
+        [0, 0],
+        [1, 1],
+        false,
+        [outVals16, outIdx16],
+      ]);
+      fp16Stats = await measureBatchedStatsMs(ctx16, pIters);
+      ctx16.destroy();
+    }
+
+    const baselineMs = baseStats.median;
+    const fp32Ms = fp32Stats.median;
+    const fp16Ms = fp16Stats?.median ?? null;
+
+    const speedupFp32 = `${(baselineMs / fp32Ms).toFixed(2)}x`;
+    const speedupFp16 = fp16Ms !== null ? `${(baselineMs / fp16Ms).toFixed(2)}x` : 'N/A';
+
+    addRow({
+      category: 'Pooling & Resampling',
+      name: 'MaxPool2D (2D Spatial Grid)',
+      shape: `[${pShape.join(',')}]`,
+      baselineMs,
+      baselineStd: baseStats.std,
+      fp32Ms,
+      fp32Std: fp32Stats.std,
+      fp16Ms,
+      fp16Std: fp16Stats?.std ?? null,
+      speedupFp32,
+      speedupFp16,
+      stats: {
+        baseline: baseStats,
+        fp32: fp32Stats,
+        fp16: fp16Stats,
+      },
+    });
+  }
+
+  // 7. Elementwise Arithmetic (Mul & Add)
+  {
+    statusEl.innerText = 'Benchmarking Elementwise Mul (Vec4)...';
+    const shape = [1, 64, 56, 56];
+    const numel = 1 * 64 * 56 * 56;
+    const iters = 40;
+
+    // Baseline
+    const baseCtx = new BaselineContext(device);
+    const ba = baseCtx.tensor('float32', shape).setData(makeF32(numel));
+    const bb = baseCtx.tensor('float32', shape).setData(makeF32(numel));
+    const bout = baseCtx.tensor('float32', shape);
+    baseCtx.recordShader(baselineShaders.mul, [ba, bb, bout]);
+    const baseStats = await measureBatchedStatsMs(baseCtx, iters);
+    baseCtx.destroy();
+
+    // FP32
+    const ctx32 = new WgpuExecutionContext(device);
+    const a32 = ctx32.tensor('float32', shape).setData(makeF32(numel));
+    const b32 = ctx32.tensor('float32', shape).setData(makeF32(numel));
+    const out32 = ctx32.tensor('float32', shape);
+    ctx32.recordShader(shaders.mul, [a32, b32, out32]);
+    const fp32Stats = await measureBatchedStatsMs(ctx32, iters);
+    ctx32.destroy();
+
+    // FP16
+    let fp16Stats: Stats | null = null;
+    if (hasF16) {
+      const ctx16 = new WgpuExecutionContext(device);
+      const a16 = ctx16.tensor('float16', shape).setData(makeF16(numel));
+      const b16 = ctx16.tensor('float16', shape).setData(makeF16(numel));
+      const out16 = ctx16.tensor('float16', shape);
+      ctx16.recordShader(shaders.mul, [a16, b16, out16]);
+      fp16Stats = await measureBatchedStatsMs(ctx16, iters);
+      ctx16.destroy();
+    }
+
+    const baselineMs = baseStats.median;
+    const fp32Ms = fp32Stats.median;
+    const fp16Ms = fp16Stats?.median ?? null;
+
+    const speedupFp32 = `${(baselineMs / fp32Ms).toFixed(2)}x`;
+    const speedupFp16 = fp16Ms !== null ? `${(baselineMs / fp16Ms).toFixed(2)}x` : 'N/A';
+
+    addRow({
+      category: 'Elementwise Math',
+      name: 'Elementwise Mul (Vec4)',
+      shape: `[${shape.join(',')}]`,
+      baselineMs,
+      baselineStd: baseStats.std,
+      fp32Ms,
+      fp32Std: fp32Stats.std,
+      fp16Ms,
+      fp16Std: fp16Stats?.std ?? null,
+      speedupFp32,
+      speedupFp16,
+      stats: {
+        baseline: baseStats,
+        fp32: fp32Stats,
+        fp16: fp16Stats,
+      },
+    });
+  }
+
+  // 8. Data Movement (Permute 2D Matrix Transpose)
+  {
+    statusEl.innerText = 'Benchmarking Matrix Transpose (Permute)...';
+    const shape = [1024, 1024];
+    const numel = 1024 * 1024;
+    const iters = 25;
+
+    // Baseline
+    const baseCtx = new BaselineContext(device);
+    const ba = baseCtx.tensor('float32', shape).setData(makeF32(numel));
+    const bout = baseCtx.tensor('float32', shape);
+    baseCtx.recordShader(baselineShaders.permuteCopy, [ba, [1, 0], bout]);
+    const baseStats = await measureBatchedStatsMs(baseCtx, iters);
+    baseCtx.destroy();
+
+    // FP32
+    const ctx32 = new WgpuExecutionContext(device);
+    const a32 = ctx32.tensor('float32', shape).setData(makeF32(numel));
+    const out32 = ctx32.tensor('float32', shape);
+    ctx32.recordShader(shaders.permuteCopy, [a32, [1, 0], out32]);
+    const fp32Stats = await measureBatchedStatsMs(ctx32, iters);
+    ctx32.destroy();
+
+    // FP16
+    let fp16Stats: Stats | null = null;
+    if (hasF16) {
+      const ctx16 = new WgpuExecutionContext(device);
+      const a16 = ctx16.tensor('float16', shape).setData(makeF16(numel));
+      const out16 = ctx16.tensor('float16', shape);
+      ctx16.recordShader(shaders.permuteCopy, [a16, [1, 0], out16]);
+      fp16Stats = await measureBatchedStatsMs(ctx16, iters);
+      ctx16.destroy();
+    }
+
+    const baselineMs = baseStats.median;
+    const fp32Ms = fp32Stats.median;
+    const fp16Ms = fp16Stats?.median ?? null;
+
+    const speedupFp32 = `${(baselineMs / fp32Ms).toFixed(2)}x`;
+    const speedupFp16 = fp16Ms !== null ? `${(baselineMs / fp16Ms).toFixed(2)}x` : 'N/A';
+
+    addRow({
+      category: 'Data Movement',
+      name: 'Matrix Transpose (Shared-Memory 16x16)',
+      shape: `[${shape.join(',')}]`,
+      baselineMs,
+      baselineStd: baseStats.std,
+      fp32Ms,
+      fp32Std: fp32Stats.std,
+      fp16Ms,
+      fp16Std: fp16Stats?.std ?? null,
+      speedupFp32,
+      speedupFp16,
+      stats: {
+        baseline: baseStats,
+        fp32: fp32Stats,
+        fp16: fp16Stats,
+      },
+    });
+  }
+
+  statusEl.innerHTML = '✓ Benchmark Complete!';
+  document.title = 'BENCHMARK_DONE';
+  const outJson = JSON.stringify(
+    {
+      gpu: {
+        vendor: info.vendor,
+        architecture: info.architecture,
+        device: info.device,
+        hasF16,
+      },
+      results,
+    },
+    null,
+    2
+  );
+  jsonEl.innerText = outJson;
+  // eslint-disable-next-line no-console
+  console.log('__BENCHMARK_RESULTS__:' + outJson);
+
+  try {
+    await fetch('/api/results', {
+      method: 'POST',
+      // eslint-disable-next-line @typescript-eslint/naming-convention
+      headers: { 'Content-Type': 'application/json' },
+      body: outJson,
+    });
+  } catch {}
+}
+
+const runBtn = document.getElementById('run-btn');
+if (runBtn) {
+  runBtn.addEventListener('click', () => {
+    runBtn.setAttribute('disabled', 'true');
+    run().finally(() => runBtn.removeAttribute('disabled'));
+  });
+}
+
+run();
diff --git a/utils/benchmark/vite.config.ts b/utils/benchmark/vite.config.ts
new file mode 100644
index 0000000..042959f
--- /dev/null
+++ b/utils/benchmark/vite.config.ts
@@ -0,0 +1,94 @@
+/* eslint-disable no-console, @typescript-eslint/naming-convention */
+import { resolve } from 'node:path';
+import fs from 'node:fs';
+import { defineConfig, type Plugin } from 'vite';
+
+function benchmarkReporterPlugin(): Plugin {
+  return {
+    name: 'benchmark-reporter',
+    configureServer(server) {
+      server.middlewares.use('/api/results', (req, res) => {
+        if (req.method === 'POST') {
+          let body = '';
+          req.on('data', (chunk) => {
+            body += chunk;
+          });
+          req.on('end', () => {
+            try {
+              const data = JSON.parse(body);
+              fs.writeFileSync(
+                resolve(import.meta.dirname, '../../benchmark-results.json'),
+                JSON.stringify(data, null, 2)
+              );
+
+              // Print nicely formatted table to console
+              console.log(
+                '\n====================================================================================================================================='
+              );
+              console.log(
+                `GPU HARDWARE: ${data.gpu.vendor.toUpperCase()} (${data.gpu.architecture}) - ${data.gpu.device}`
+              );
+              console.log(
+                `NATIVE FP16:  ${data.gpu.hasF16 ? 'YES (Hardware 16-bit packed ALUs)' : 'NO'}`
+              );
+              console.log(
+                '=====================================================================================================================================\n'
+              );
+
+              console.log(
+                'Category'.padEnd(20) +
+                  'Operator / Model'.padEnd(30) +
+                  'Tensor Shape'.padEnd(26) +
+                  'Baseline (ms)'.padStart(14) +
+                  'Opt FP32 (ms)'.padStart(14) +
+                  'Native FP16'.padStart(14) +
+                  'Opt Speedup'.padStart(13) +
+                  'FP16 Speedup'.padStart(14)
+              );
+              console.log('-'.repeat(145));
+
+              for (const r of data.results) {
+                console.log(
+                  r.category.padEnd(20) +
+                    r.name.padEnd(30) +
+                    r.shape.padEnd(26) +
+                    (r.baselineMs !== undefined && r.baselineMs !== null
+                      ? r.baselineMs.toFixed(3)
+                      : 'N/A'
+                    ).padStart(14) +
+                    r.fp32Ms.toFixed(3).padStart(14) +
+                    (r.fp16Ms ? r.fp16Ms.toFixed(3) : 'N/A').padStart(14) +
+                    (r.speedupFp32 || r.speedup || '1.00x').padStart(13) +
+                    (r.speedupFp16 || 'N/A').padStart(14)
+                );
+              }
+              console.log('-'.repeat(145));
+              console.log('Saved benchmark results to benchmark-results.json\n');
+            } catch (err) {
+              console.error('Failed to parse benchmark results:', err);
+            }
+            res.writeHead(200, { 'Content-Type': 'application/json' });
+            res.end(JSON.stringify({ status: 'ok' }));
+          });
+        } else {
+          res.writeHead(405);
+          res.end();
+        }
+      });
+    },
+  };
+}
+
+export default defineConfig({
+  base: './',
+  publicDir: resolve(import.meta.dirname, '../../example/public'),
+  plugins: [benchmarkReporterPlugin()],
+  server: {
+    port: 5178,
+    host: true,
+  },
+  build: {
+    outDir: resolve(import.meta.dirname, '../../dist-benchmark'),
+    emptyOutDir: true,
+  },
+});
diff --git a/utils/convert_model_to_fp16.mjs b/utils/convert_model_to_fp16.mjs
new file mode 100644
index 0000000..26bd3ad
--- /dev/null
+++ b/utils/convert_model_to_fp16.mjs
@@ -0,0 +1,136 @@
+#!/usr/bin/env node
+import fs from 'node:fs';
+import path from 'node:path';
+import { Float16Array } from '@petamoriken/float16';
+
+const REPO_ROOT = path.resolve(import.meta.dirname, '..');
+const MODELS = ['depthart', 'selfie-segmentation', 'yolo26n'];
+
+for (const model of MODELS) {
+  const dir = path.join(REPO_ROOT, 'example/public/assets', model);
+  const progPath = path.join(dir, 'program.json');
+  const weightsPath = path.join(dir, 'weights.bin');
+
+  if (!fs.existsSync(progPath) || !fs.existsSync(weightsPath)) {
+    console.log(`Skipping ${model}: files not found`);
+    continue;
+  }
+
+  console.log(`Converting ${model} to FP16...`);
+  const program = JSON.parse(fs.readFileSync(progPath, 'utf8'));
+  const oldWeightsBuf = fs.readFileSync(weightsPath);
+
+  const weightScales = new Map();
+  if (model === 'depthart') {
+    // The decoder's penultimate ReLU produces values above 65504 on photographs,
+    // even though the final depth fits FP16. Rebalance the adjacent convolutions:
+    // ReLU(W*x+b)/16, followed by (16*V)*h+c, preserves the FP32 computation and
+    // leaves headroom for FP16 activations and partial sums. Powers of two do not
+    // change the significand when the scaled weights remain representable.
+    const [hidden, output] = program.chain.slice(-2);
+    const hiddenOut = hidden.args.at(-1);
+    if (
+      hidden.name !== 'et_vk.conv_with_clamp.default' ||
+      output.name !== 'et_vk.conv_with_clamp.default' ||
+      output.args[0] !== hiddenOut ||
+      program.values[hidden.args[9]].value !== 0 ||
+      program.values[hidden.args[10]].value < 1e30 ||
+      program.values[output.args[1]].shape.slice(2).some((dim) => dim !== 1) ||
+      program.chain.filter((op) => op.args.slice(0, -1).includes(hiddenOut)).length !== 1
+    ) {
+      throw new Error('DepthART decoder changed: verify FP16 range balancing before conversion');
+    }
+    weightScales.set(hidden.args[1], 1 / 16);
+    weightScales.set(hidden.args[2], 1 / 16);
+    weightScales.set(output.args[1], 16);
+  }
+
+  // We will build a new weights buffer
+  const newWeightsChunks = [];
+  let currentOffset = 0;
+
+  for (let i = 0; i < program.values.length; i++) {
+    const v = program.values[i];
+    if (v.type !== 'tensor') continue;
+
+    if (v.dtype === 'float32') {
+      v.dtype = 'float16';
+    }
+
+    if (v.weightsOffset !== undefined && v.weightsLength !== undefined) {
+      // Align offset to 4 bytes for WebGPU storage buffers
+      const alignedOffset = (currentOffset + 3) & ~3;
+      if (alignedOffset > currentOffset) {
+        newWeightsChunks.push(Buffer.alloc(alignedOffset - currentOffset));
+        currentOffset = alignedOffset;
+      }
+
+      const numel = v.weightsLength / 4;
+      const f32Slice = new Float32Array(
+        oldWeightsBuf.buffer,
+        oldWeightsBuf.byteOffset + v.weightsOffset,
+        numel
+      );
+
+      const scale = weightScales.get(i) ?? 1;
+      const f16Slice = new Float16Array(
+        scale === 1 ? f32Slice : Float32Array.from(f32Slice, (value) => value * scale)
+      );
+      const f16Buf = Buffer.from(f16Slice.buffer, f16Slice.byteOffset, f16Slice.byteLength);
+
+      v.weightsOffset = currentOffset;
+      v.weightsLength = f16Buf.byteLength;
+
+      newWeightsChunks.push(f16Buf);
+      currentOffset += f16Buf.byteLength;
+    }
+  }
+
+  // Recalculate memory pools for intermediate tensors
+  if (program.memoryPlan && program.memoryPlan.pools) {
+    const poolSizes = new Map();
+    for (const pool of program.memoryPlan.pools) {
+      poolSizes.set(pool.id, 0);
+    }
+
+    const DTYPE_BYTES = {
+      float16: 2,
+      float32: 4,
+      int32: 4,
+      int8: 1,
+      uint8: 1,
+      bool: 1,
+      float64: 8,
+      int64: 8,
+    };
+
+    for (const v of program.values) {
+      if (v.type === 'tensor' && v.memObjId !== undefined) {
+        const numel = v.shape.reduce((a, b) => a * b, 1);
+        const bytes = numel * (DTYPE_BYTES[v.dtype] || 2);
+        const alignedBytes = (bytes + 3) & ~3;
+        const cur = poolSizes.get(v.memObjId) || 0;
+        if (alignedBytes > cur) {
+          poolSizes.set(v.memObjId, alignedBytes);
+        }
+      }
+    }
+
+    for (const pool of program.memoryPlan.pools) {
+      pool.size = poolSizes.get(pool.id) || Math.ceil(pool.size / 2);
+    }
+  }
+
+  const newWeightsBuf = Buffer.concat(newWeightsChunks);
+  const outProgPath = path.join(dir, 'program-fp16.json');
+  const outWeightsPath = path.join(dir, 'weights-fp16.bin');
+
+  fs.writeFileSync(outProgPath, JSON.stringify(program, null, 2));
+  fs.writeFileSync(outWeightsPath, newWeightsBuf);
+
+  console.log(
+    `✓ ${model}: weights reduced from ${(oldWeightsBuf.byteLength / 1024 / 1024).toFixed(2)} MB to ${(newWeightsBuf.byteLength / 1024 / 1024).toFixed(2)} MB (50% reduction)`
+  );
+}
+
+console.log('All models converted to FP16 successfully!');