Perf/kernel optimizations

.gitignore

@@ -1,6 +1,7 @@
 node_modules
 dist
 dist-example
+dist-benchmark
 .vite
 coverage
 *.log

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 <path/to/test.ts>`
+- **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/<operator>/index.ts`).
+  - Submodules must follow the naming pattern `<specializedShaderName>.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/<operator>/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<f32>` (or `vec4<f16>`) 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.

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"
+    }
+  ]
+}

example/index.html

@@ -30,6 +30,22 @@
         <div class="topbar">
           <a class="back" href="#/">← all demos</a>
           <h2 class="task-title"></h2>
+          <div class="precision-control">
+            <span class="precision-label">Precision</span>
+            <label class="precision-switch">
+              <input
+                class="precision-toggle"
+                type="checkbox"
+                role="switch"
+                aria-label="Use FP16 precision"
+                disabled
+              />
+              <span class="precision-track" aria-hidden="true">
+                <span>FP32</span>
+                <span>FP16</span>
+              </span>
+            </label>
+          </div>
         </div>
 
         <div class="stage">