oidn-web 0.4.0 → 0.5.0

package/lib/nativeUNet.js

@@ -1,4 +1,5 @@
 import { Float16Array } from '@petamoriken/float16';
+import { createFinalRgbShader, sharedMemoryBytes as finalRgbSharedMemoryBytes } from './finalRgbShader.js';
 import { optimizeModelGraph, planModelExecution } from './graphOptimizer.js';
 import { OIDNResourceTracker } from './resourceTracker.js';
 const pipelineCachesByDevice = new WeakMap();
@@ -11,13 +12,42 @@ function sharedPipelineCache(device) {
     return cache;
 }
 const WORKGROUP_SIZE = 8;
-const TILED_CONV_WORKGROUP = 8;
-const TILED_CONV_ROWS_PER_THREAD = 4;
-const TILED_CONV_M = TILED_CONV_WORKGROUP * TILED_CONV_ROWS_PER_THREAD;
-const TILED_CONV_N_BLOCKS = TILED_CONV_WORKGROUP;
 const TILED_CONV_K_BLOCKS = 8;
 const SPATIAL_CONV_WORKGROUP = 8;
 const SPATIAL_CONV_PATCH = SPATIAL_CONV_WORKGROUP + 2;
+function gemmTile(options) {
+    const [workgroupX, workgroupY] = options.workgroupSize;
+    return {
+        workgroupX,
+        workgroupY,
+        rowsPerThread: options.rowsPerThread,
+        tileM: workgroupY * options.rowsPerThread,
+        tileNBlocks: workgroupX,
+        // Keep the half partial accumulation grouping independent of tile tuning.
+        tileKBlocks: TILED_CONV_K_BLOCKS
+    };
+}
+// Holes change only workgroup addresses. Cooperative loads still visit the
+// original dense logical tiles and never read the unused padding elements.
+function gemmSharedSizes(options) {
+    const tile = gemmTile(options);
+    const paddedInput = options.sharedLayout === 'padded' || options.sharedLayout === 'padded-input';
+    const paddedWeights = options.sharedLayout === 'padded' || options.sharedLayout === 'padded-weights';
+    return {
+        input: tile.tileM * tile.tileKBlocks + (paddedInput ? tile.workgroupY : 0),
+        weights: tile.tileKBlocks * tile.tileNBlocks * (paddedWeights ? 5 : 4)
+    };
+}
+function gemmSharedInputIndex(index, options) {
+    return options.sharedLayout === 'padded' || options.sharedLayout === 'padded-input'
+        ? `(${index}) + (${index}) / ${options.rowsPerThread * TILED_CONV_K_BLOCKS}u`
+        : index;
+}
+function gemmSharedWeightIndex(index, options) {
+    return options.sharedLayout === 'padded' || options.sharedLayout === 'padded-weights'
+        ? `(${index}) + (${index}) / 4u`
+        : index;
+}
 function roundUp(value, alignment) {
     return Math.ceil(value / alignment) * alignment;
 }
@@ -89,7 +119,7 @@ function createUniformBuffer(device, label, values) {
     data.set(values);
     return createMappedBuffer(device, label, data, GPUBufferUsage.UNIFORM);
 }
-function packConvTensors(device, id, tensors, precision) {
+function packConvTensors(device, id, tensors, precision, weightLayout) {
     const inputBlocks = blocksForChannels(tensors.inputChannels);
     const outputBlocks = blocksForChannels(tensors.outputChannels);
     const packedWeightCount = outputBlocks *
@@ -115,15 +145,18 @@ function packConvTensors(device, id, tensors, precision) {
                         const outputChannel = outputBlock * 4 + outputLane;
                         for (let inputLane = 0; inputLane < 4; inputLane++) {
                             const inputChannel = inputBlock * 4 + inputLane;
-                            const packedBlock = ((((outputBlock * tensors.kernelHeight + y) *
-                                tensors.kernelWidth +
-                                x) *
-                                inputBlocks +
-                                inputBlock) *
-                                16);
+                            const packedBlock = weightLayout === 'k-major'
+                                ? ((((y * tensors.kernelWidth + x) * inputBlocks + inputBlock) *
+                                    outputBlocks + outputBlock) * 16)
+                                : ((((outputBlock * tensors.kernelHeight + y) *
+                                    tensors.kernelWidth +
+                                    x) *
+                                    inputBlocks +
+                                    inputBlock) *
+                                    16);
                             // One vec4 contains the four output lanes for an input lane.
-                            // Direct and tiled kernels share this layout, keeping model
-                            // descriptors independent from the selected kernel.
+                            // Both layouts preserve the vec4's output lanes and the
+                            // convolution's reduction order; only the vec4 address changes.
                             const packedIndex = packedBlock + inputLane * 4 + outputLane;
                             if (outputChannel < tensors.outputChannels &&
                                 inputChannel < tensors.inputChannels) {
@@ -144,10 +177,11 @@ function packConvTensors(device, id, tensors, precision) {
     // in f32. Padded lanes are zero and never escape the final three channels.
     const packedBias = new Float32Array(outputBlocks * 4);
     packedBias.set(tensorFloat32Values(tensors.bias));
-    const weights = createMappedBuffer(device, `oidn/${id}/weights/${precision}`, packedWeights, GPUBufferUsage.STORAGE);
+    const weights = createMappedBuffer(device, `oidn/${id}/weights/${precision}/${weightLayout}`, packedWeights, GPUBufferUsage.STORAGE);
     try {
         return {
             weights,
+            weightLayout,
             bias: createMappedBuffer(device, `oidn/${id}/bias`, packedBias, GPUBufferUsage.STORAGE)
         };
     }
@@ -213,21 +247,52 @@ let weightBase = ${weightBase};
 ${accumulator}
 `;
 }
-function tiledAccumulationCode(precision) {
+function gemmPackedLoad(precision, gemm, weights) {
+    return precision === 'fp16' &&
+        (gemm.loadMode === 'packed-all' ||
+            (weights && gemm.loadMode === 'packed-weights'));
+}
+// vec2<u32> reinterprets exactly the same eight bytes as vec4<f16>.
+function gemmLoadExpression(expression, packed) {
+    return packed ? `bitcast<vec4<f16>>(${expression})` : expression;
+}
+function tiledAccumulationCode(precision, gemm) {
+    const { rowsPerThread, tileNBlocks, tileKBlocks } = gemmTile(gemm);
+    const inputIndex = gemmSharedInputIndex(`tileSpatial * ${tileKBlocks}u + tileK`, gemm);
+    const weightBase = `(tileK * ${tileNBlocks}u + localId.x) * ${gemm.sharedLayout === 'padded' || gemm.sharedLayout === 'padded-weights' ? 5 : 4}u`;
+    if (gemm.accumulationOrder === 'row-major') {
+        const valueType = storageVecType(precision);
+        const target = precision === 'fp16' ? 'partial' : 'acc[row]';
+        return /* wgsl */ `
+    for (var row = 0u; row < ${rowsPerThread}u; row++) {
+      ${precision === 'fp16' ? 'var partial = vec4<f16>(0.0h);' : ''}
+      let tileSpatial = localId.y * ${rowsPerThread}u + row;
+      for (var tileK = 0u; tileK < ${tileKBlocks}u; tileK++) {
+        let weightBase = ${weightBase};
+        let inputValue = inputTile[${inputIndex}];
+        ${target} = fma(weightTile[weightBase], ${valueType}(inputValue.x), ${target});
+        ${target} = fma(weightTile[weightBase + 1u], ${valueType}(inputValue.y), ${target});
+        ${target} = fma(weightTile[weightBase + 2u], ${valueType}(inputValue.z), ${target});
+        ${target} = fma(weightTile[weightBase + 3u], ${valueType}(inputValue.w), ${target});
+      }
+      ${precision === 'fp16' ? 'acc[row] += vec4<f32>(partial);' : ''}
+    }
+`;
+    }
     if (precision === 'fp16') {
         return /* wgsl */ `
-    var partial: array<vec4<f16>, ${TILED_CONV_ROWS_PER_THREAD}>;
-    for (var row = 0u; row < ${TILED_CONV_ROWS_PER_THREAD}u; row++) {
+    var partial: array<vec4<f16>, ${rowsPerThread}>;
+    for (var row = 0u; row < ${rowsPerThread}u; row++) {
       partial[row] = vec4<f16>(0.0h);
     }
-    for (var tileK = 0u; tileK < ${TILED_CONV_K_BLOCKS}u; tileK++) {
+    for (var tileK = 0u; tileK < ${tileKBlocks}u; tileK++) {
       let weightBase =
-        (tileK * ${TILED_CONV_N_BLOCKS}u + localId.x) * 4u;
-      for (var row = 0u; row < ${TILED_CONV_ROWS_PER_THREAD}u; row++) {
+        ${weightBase};
+      for (var row = 0u; row < ${rowsPerThread}u; row++) {
         let tileSpatial =
-          localId.y * ${TILED_CONV_ROWS_PER_THREAD}u + row;
+          localId.y * ${rowsPerThread}u + row;
         let inputValue =
-          inputTile[tileSpatial * ${TILED_CONV_K_BLOCKS}u + tileK];
+          inputTile[${inputIndex}];
         partial[row] = fma(
           weightTile[weightBase],
           vec4<f16>(inputValue.x),
@@ -250,20 +315,20 @@ function tiledAccumulationCode(precision) {
         );
       }
     }
-    for (var row = 0u; row < ${TILED_CONV_ROWS_PER_THREAD}u; row++) {
+    for (var row = 0u; row < ${rowsPerThread}u; row++) {
       acc[row] += vec4<f32>(partial[row]);
     }
 `;
     }
     return /* wgsl */ `
-    for (var tileK = 0u; tileK < ${TILED_CONV_K_BLOCKS}u; tileK++) {
+    for (var tileK = 0u; tileK < ${tileKBlocks}u; tileK++) {
       let weightBase =
-        (tileK * ${TILED_CONV_N_BLOCKS}u + localId.x) * 4u;
-      for (var row = 0u; row < ${TILED_CONV_ROWS_PER_THREAD}u; row++) {
+        ${weightBase};
+      for (var row = 0u; row < ${rowsPerThread}u; row++) {
         let tileSpatial =
-          localId.y * ${TILED_CONV_ROWS_PER_THREAD}u + row;
+          localId.y * ${rowsPerThread}u + row;
         let inputValue =
-          inputTile[tileSpatial * ${TILED_CONV_K_BLOCKS}u + tileK];
+          inputTile[${inputIndex}];
         acc[row] = fma(
           weightTile[weightBase],
           vec4<f32>(inputValue.x),
@@ -474,18 +539,92 @@ fn main(
 }
 `;
 }
+/** Precompute spatial predicates once, then advance the original flattened K order. */
+function tiledAddressCacheCode(inputBlocks, decoder, gemm, sources) {
+    const { workgroupX, workgroupY, tileM, tileKBlocks } = gemmTile(gemm);
+    const width = decoder ? 'params.outputWidth' : 'params.inputWidth';
+    const height = decoder ? 'params.outputHeight' : 'params.inputHeight';
+    const baseOffset = gemm.addressMode === 'base-offset';
+    // Decoder parity travels in the high mask bits, avoiding extra cached arrays.
+    const baseAssignments = sources
+        ? sources.blocks.map((blocks, source) => `cachedBase${source}[load] = (${source === sources.upsampled ? 'y / 2u' : 'y'} * params.source${source}Width + ${source === sources.upsampled ? 'x / 2u' : 'x'}) * ${blocks}u;`).join('\n    ')
+        : `cachedBase0[load] = (y * params.inputWidth + x) * ${inputBlocks}u;`;
+    const loads = tileM * tileKBlocks /
+        (workgroupX * workgroupY);
+    return /* wgsl */ `
+  ${baseOffset ? `var cachedBase0: array<u32, ${loads}>;
+  ${sources ? `var cachedBase1: array<u32, ${loads}>;` : ''}` : `var cachedX: array<u32, ${loads}>;
+  var cachedY: array<u32, ${loads}>;`}
+  var cachedMask: array<u32, ${loads}>;
+  for (var load = 0u; load < ${loads}u; load++) {
+    let loadIndex = localLinear + load * ${workgroupX * workgroupY}u;
+    let spatial = workgroupId.x * ${tileM}u + loadIndex / ${tileKBlocks}u;
+    let x = spatial % params.outputWidth;
+    let y = spatial / params.outputWidth;
+    ${baseOffset ? baseAssignments : `cachedX[load] = x;
+    cachedY[load] = y;`}
+    let columns =
+      select(0u, 0x049u, x > 0u && x - 1u < ${width}) |
+      select(0u, 0x092u, x < ${width}) |
+      select(0u, 0x124u, x + 1u < ${width});
+    let rows =
+      select(0u, 0x007u, y > 0u && y - 1u < ${height}) |
+      select(0u, 0x038u, y < ${height}) |
+      select(0u, 0x1c0u, y + 1u < ${height});
+    cachedMask[load] = select(0u, columns & rows, spatial < spatialCount)${baseOffset && sources ? ' | ((x & 1u) << 9u) | ((y & 1u) << 10u)' : ''};
+  }
+  var channelBlock = (localLinear % ${tileKBlocks}u) % ${inputBlocks}u;
+  var filterPosition = (localLinear % ${tileKBlocks}u) / ${inputBlocks}u;
+`;
+}
+function tiledAddressAdvanceCode(inputBlocks, gemm) {
+    const { tileKBlocks } = gemmTile(gemm);
+    // The remainder is less than inputBlocks, so at most one carry is needed.
+    return /* wgsl */ `
+    channelBlock += ${tileKBlocks % inputBlocks}u;
+    let carry = channelBlock >= ${inputBlocks}u;
+    channelBlock -= select(0u, ${inputBlocks}u, carry);
+    filterPosition += ${Math.floor(tileKBlocks / inputBlocks)}u + select(0u, 1u, carry);
+`;
+}
+function tiledIncrementalInputCode(valueType, readValue, gemm) {
+    const { workgroupX, workgroupY, tileM, tileKBlocks } = gemmTile(gemm);
+    const loads = tileM * tileKBlocks /
+        (workgroupX * workgroupY);
+    return /* wgsl */ `
+    let inputBlock = channelBlock;
+    let kernelY = filterPosition / 3u;
+    let kernelX = filterPosition % 3u;
+    for (var load = 0u; load < ${loads}u; load++) {
+      let loadIndex = localLinear + load * ${workgroupX * workgroupY}u;
+      var value = ${valueType}(0.0);
+      if (filterPosition < 9u && (cachedMask[load] & (1u << filterPosition)) != 0u) {
+        ${gemm.addressMode === 'base-offset' ? '' : `let inputX = i32(cachedX[load]) + i32(kernelX) - 1;
+        let inputY = i32(cachedY[load]) + i32(kernelY) - 1;`}
+        ${readValue}
+      }
+      inputTile[${gemmSharedInputIndex('loadIndex', gemm)}] = value;
+    }
+`;
+}
 /**
- * Implicit-GEMM convolution for FP32. This follows the proven packed WebGPU
- * shape: one 8x8 workgroup computes 32 spatial rows by 8 vec4 output blocks,
- * with four output rows per thread and an eight-vec4 K tile.
+ * Implicit-GEMM convolution for FP16/FP32. This follows the packed WebGPU
+ * shape: configurable spatial rows and vec4 output blocks per workgroup.
+ * The eight-vec4 K tile stays fixed to preserve FP16 rounding across variants.
  */
-function createTiledConvShader(precision, outputPrecision, activation, inputBlocks, outputBlocks) {
+function createTiledConvShader(precision, outputPrecision, activation, inputBlocks, outputBlocks, gemm) {
+    const { workgroupX, workgroupY, rowsPerThread, tileM, tileNBlocks, tileKBlocks } = gemmTile(gemm);
+    const { addressMode, weightLayout } = gemm;
     const inputType = storageVecType(precision);
     const outputType = storageVecType(outputPrecision);
+    const packedInput = gemmPackedLoad(precision, gemm, false);
+    const packedWeights = gemmPackedLoad(precision, gemm, true);
+    const inputStorageType = packedInput ? 'vec2<u32>' : inputType;
+    const weightStorageType = packedWeights ? 'vec2<u32>' : inputType;
     const stored = storeExpression(activationExpression('acc', activation), outputPrecision);
-    const workgroupThreads = TILED_CONV_WORKGROUP * TILED_CONV_WORKGROUP;
-    const inputTileValues = TILED_CONV_M * TILED_CONV_K_BLOCKS;
-    const weightTileValues = TILED_CONV_K_BLOCKS * TILED_CONV_N_BLOCKS * 4;
+    const workgroupThreads = workgroupX * workgroupY;
+    const inputTileValues = tileM * tileKBlocks;
+    const weightTileValues = tileKBlocks * tileNBlocks * 4;
     return /* wgsl */ `${shaderPreamble(precision)}
 struct Params {
   inputWidth: u32,
@@ -495,46 +634,53 @@ struct Params {
   inputBlocks: u32,
   outputBlocks: u32,
 }
-@group(0) @binding(0) var<storage, read> inputData: array<${inputType}>;
-@group(0) @binding(1) var<storage, read> weights: array<${inputType}>;
+@group(0) @binding(0) var<storage, read> inputData: array<${inputStorageType}>;
+@group(0) @binding(1) var<storage, read> weights: array<${weightStorageType}>;
 @group(0) @binding(2) var<storage, read> bias: array<vec4<f32>>;
 @group(0) @binding(3) var<storage, read_write> outputData: array<${outputType}>;
 @group(0) @binding(4) var<uniform> params: Params;
 
-var<workgroup> inputTile: array<${inputType}, ${inputTileValues}>;
-var<workgroup> weightTile: array<${inputType}, ${weightTileValues}>;
+var<workgroup> inputTile: array<${inputType}, ${gemmSharedSizes(gemm).input}>;
+var<workgroup> weightTile: array<${inputType}, ${gemmSharedSizes(gemm).weights}>;
 
-@compute @workgroup_size(${TILED_CONV_WORKGROUP}, ${TILED_CONV_WORKGROUP}, 1)
+@compute @workgroup_size(${workgroupX}, ${workgroupY}, 1)
 fn main(
   @builtin(local_invocation_id) localId: vec3<u32>,
   @builtin(workgroup_id) workgroupId: vec3<u32>
 ) {
   let spatialBase =
-    workgroupId.x * ${TILED_CONV_M}u +
-    localId.y * ${TILED_CONV_ROWS_PER_THREAD}u;
+    workgroupId.x * ${tileM}u +
+    localId.y * ${rowsPerThread}u;
   let outputBlock =
-    workgroupId.y * ${TILED_CONV_N_BLOCKS}u + localId.x;
+    workgroupId.y * ${tileNBlocks}u + localId.x;
   let spatialCount = params.outputWidth * params.outputHeight;
-  var acc: array<vec4<f32>, ${TILED_CONV_ROWS_PER_THREAD}>;
+  var acc: array<vec4<f32>, ${rowsPerThread}>;
   if (outputBlock < ${outputBlocks}u) {
-    for (var row = 0u; row < ${TILED_CONV_ROWS_PER_THREAD}u; row++) {
+    for (var row = 0u; row < ${rowsPerThread}u; row++) {
       acc[row] = bias[outputBlock];
     }
   }
 
   let localLinear =
-    localId.y * ${TILED_CONV_WORKGROUP}u + localId.x;
+    localId.y * ${workgroupX}u + localId.x;
+  ${addressMode !== 'analytic' ? tiledAddressCacheCode(inputBlocks, false, gemm) : ''}
   let totalK = ${inputBlocks * 9}u;
-  for (var kBase = 0u; kBase < totalK; kBase += ${TILED_CONV_K_BLOCKS}u) {
+  for (var kBase = 0u; kBase < totalK; kBase += ${tileKBlocks}u) {
+    ${addressMode !== 'analytic' ? tiledIncrementalInputCode(inputType, `
+        ${addressMode === 'base-offset' ? `let offset = (i32(kernelY) - 1) * i32(params.inputWidth) + i32(kernelX) - 1;
+        let inputIndex = cachedBase0[load] + u32(offset * ${inputBlocks}) + inputBlock;` : `let inputIndex = (u32(inputY) * params.inputWidth + u32(inputX)) *
+          ${inputBlocks}u + inputBlock;`}
+        value = ${gemmLoadExpression('inputData[inputIndex]', packedInput)};
+    `, gemm) : /* wgsl */ `
     for (
       var loadIndex = localLinear;
       loadIndex < ${inputTileValues}u;
       loadIndex += ${workgroupThreads}u
     ) {
-      let tileSpatial = loadIndex / ${TILED_CONV_K_BLOCKS}u;
-      let tileK = loadIndex % ${TILED_CONV_K_BLOCKS}u;
+      let tileSpatial = loadIndex / ${tileKBlocks}u;
+      let tileK = loadIndex % ${tileKBlocks}u;
       let inputSpatialIndex =
-        workgroupId.x * ${TILED_CONV_M}u + tileSpatial;
+        workgroupId.x * ${tileM}u + tileSpatial;
       let kIndex = kBase + tileK;
       var value = ${inputType}(0.0);
       if (inputSpatialIndex < spatialCount && kIndex < totalK) {
@@ -553,26 +699,32 @@ fn main(
           let inputIndex =
             (u32(inputY) * params.inputWidth + u32(inputX)) *
             ${inputBlocks}u + inputBlock;
-          value = inputData[inputIndex];
+          value = ${gemmLoadExpression('inputData[inputIndex]', packedInput)};
         }
       }
-      inputTile[loadIndex] = value;
+      inputTile[${gemmSharedInputIndex('loadIndex', gemm)}] = value;
     }
+    `}
 
     for (
       var loadIndex = localLinear;
       loadIndex < ${weightTileValues}u;
       loadIndex += ${workgroupThreads}u
     ) {
-      let tileK = loadIndex / ${TILED_CONV_N_BLOCKS * 4}u;
-      let outputRemainder = loadIndex % ${TILED_CONV_N_BLOCKS * 4}u;
+      let tileK = loadIndex / ${tileNBlocks * 4}u;
+      let outputRemainder = loadIndex % ${tileNBlocks * 4}u;
       let tileOutputBlock = outputRemainder / 4u;
       let outputLane = outputRemainder % 4u;
       let loadedOutputBlock =
-        workgroupId.y * ${TILED_CONV_N_BLOCKS}u + tileOutputBlock;
+        workgroupId.y * ${tileNBlocks}u + tileOutputBlock;
       let kIndex = kBase + tileK;
       var value = ${inputType}(0.0);
       if (loadedOutputBlock < ${outputBlocks}u && kIndex < totalK) {
+        ${weightLayout === 'k-major' ? /* wgsl */ `
+        let weightIndex = (kIndex * ${outputBlocks}u + loadedOutputBlock) * 4u + outputLane;
+        ` : addressMode !== 'analytic' ? /* wgsl */ `
+        let weightIndex = (loadedOutputBlock * ${inputBlocks * 9}u + kIndex) * 4u + outputLane;
+        ` : /* wgsl */ `
         let inputBlock = kIndex % ${inputBlocks}u;
         let kernelIndex = kIndex / ${inputBlocks}u;
         let kernelY = kernelIndex / 3u;
@@ -580,18 +732,20 @@ fn main(
         let weightIndex =
           ((((loadedOutputBlock * 3u + kernelY) * 3u + kernelX) *
             ${inputBlocks}u + inputBlock) * 4u + outputLane);
-        value = weights[weightIndex];
+        `}
+        value = ${gemmLoadExpression('weights[weightIndex]', packedWeights)};
       }
-      weightTile[loadIndex] = value;
+      weightTile[${gemmSharedWeightIndex('loadIndex', gemm)}] = value;
     }
 
     workgroupBarrier();
-    ${tiledAccumulationCode(precision)}
+    ${tiledAccumulationCode(precision, gemm)}
     workgroupBarrier();
+    ${addressMode !== 'analytic' ? tiledAddressAdvanceCode(inputBlocks, gemm) : ''}
   }
 
   if (outputBlock < ${outputBlocks}u) {
-    for (var row = 0u; row < ${TILED_CONV_ROWS_PER_THREAD}u; row++) {
+    for (var row = 0u; row < ${rowsPerThread}u; row++) {
       let spatialIndex = spatialBase + row;
       if (spatialIndex < spatialCount) {
         let outputIndex = spatialIndex * ${outputBlocks}u + outputBlock;
@@ -653,7 +807,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 }
 `;
 }
-function createMaxPoolShader(precision, outputBlocks) {
+function createMaxPoolShader(precision, outputBlocks, coalesced) {
     const valueType = storageVecType(precision);
     return /* wgsl */ `${shaderPreamble(precision)}
 struct Params {
@@ -668,11 +822,17 @@ struct Params {
 @group(0) @binding(2) var<uniform> params: Params;
 
 @compute @workgroup_size(${WORKGROUP_SIZE}, ${WORKGROUP_SIZE}, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+fn main(
+  @builtin(global_invocation_id) gid: vec3<u32>${coalesced ? ',\n  @builtin(num_workgroups) groupCount: vec3<u32>' : ''}
+) {
+  ${coalesced ? `let linearX = gid.x + gid.z * groupCount.x * ${WORKGROUP_SIZE}u;
+  let outputBlock = linearX % ${outputBlocks}u;
+  let outputX = linearX / ${outputBlocks}u;` : `let outputBlock = gid.z;
+  let outputX = gid.x;`}
   if (
-    gid.x >= params.outputWidth ||
+    outputX >= params.outputWidth ||
     gid.y >= params.outputHeight ||
-    gid.z >= ${outputBlocks}u
+    outputBlock >= ${outputBlocks}u
   ) {
     return;
   }
@@ -681,26 +841,33 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     let inputY = gid.y * 2u + py;
     if (inputY >= params.inputHeight) { continue; }
     for (var px = 0u; px < 2u; px++) {
-      let inputX = gid.x * 2u + px;
+      let inputX = outputX * 2u + px;
       if (inputX >= params.inputWidth) { continue; }
       let inputIndex =
-        (inputY * params.inputWidth + inputX) * ${outputBlocks}u + gid.z;
+        (inputY * params.inputWidth + inputX) * ${outputBlocks}u + outputBlock;
       pooled = max(pooled, vec4<f32>(inputData[inputIndex]));
     }
   }
   let outputIndex =
-    (gid.y * params.outputWidth + gid.x) * ${outputBlocks}u + gid.z;
+    (gid.y * params.outputWidth + outputX) * ${outputBlocks}u + outputBlock;
   outputData[outputIndex] = ${storeExpression('pooled', precision)};
 }
 `;
 }
-function createTiledDecoderShader(precision, activation, sourceBlocks, upsampledSource, outputBlocks) {
+function createTiledDecoderShader(precision, outputPrecision, activation, sourceBlocks, upsampledSource, outputBlocks, gemm) {
+    const { workgroupX, workgroupY, rowsPerThread, tileM, tileNBlocks, tileKBlocks } = gemmTile(gemm);
+    const { addressMode, weightLayout } = gemm;
     const valueType = storageVecType(precision);
+    const outputType = storageVecType(outputPrecision);
+    const packedInput = gemmPackedLoad(precision, gemm, false);
+    const packedWeights = gemmPackedLoad(precision, gemm, true);
+    const inputStorageType = packedInput ? 'vec2<u32>' : valueType;
+    const weightStorageType = packedWeights ? 'vec2<u32>' : valueType;
     const inputBlocks = sourceBlocks[0] + sourceBlocks[1];
-    const stored = storeExpression(activationExpression('acc[row]', activation), precision);
-    const workgroupThreads = TILED_CONV_WORKGROUP * TILED_CONV_WORKGROUP;
-    const inputTileValues = TILED_CONV_M * TILED_CONV_K_BLOCKS;
-    const weightTileValues = TILED_CONV_K_BLOCKS * TILED_CONV_N_BLOCKS * 4;
+    const stored = storeExpression(activationExpression('acc[row]', activation), outputPrecision);
+    const workgroupThreads = workgroupX * workgroupY;
+    const inputTileValues = tileM * tileKBlocks;
+    const weightTileValues = tileKBlocks * tileNBlocks * 4;
     const sourceRead = (source, blockExpression) => {
         const sourceX = source === upsampledSource
             ? 'u32(inputX) / 2u'
@@ -711,12 +878,15 @@ function createTiledDecoderShader(precision, activation, sourceBlocks, upsampled
         return /* wgsl */ `
           {
             let sourceBlock = ${blockExpression};
-            let sourceX = ${sourceX};
+            ${addressMode === 'base-offset' ? `let dx = ${source === upsampledSource ? '(i32((cachedMask[load] >> 9u) & 1u) + i32(kernelX) - 1) >> 1u' : 'i32(kernelX) - 1'};
+            let dy = ${source === upsampledSource ? '(i32((cachedMask[load] >> 10u) & 1u) + i32(kernelY) - 1) >> 1u' : 'i32(kernelY) - 1'};
+            let offset = (dy * i32(params.source${source}Width) + dx) * ${sourceBlocks[source]};
+            let sourceIndex = cachedBase${source}[load] + u32(offset) + sourceBlock;` : `let sourceX = ${sourceX};
             let sourceY = ${sourceY};
             let sourceIndex =
               (sourceY * params.source${source}Width + sourceX) *
-              ${sourceBlocks[source]}u + sourceBlock;
-            value = input${source}[sourceIndex];
+              ${sourceBlocks[source]}u + sourceBlock;`}
+            value = ${gemmLoadExpression(`input${source}[sourceIndex]`, packedInput)};
           }`;
     };
     return /* wgsl */ `${shaderPreamble(precision)}
@@ -730,47 +900,61 @@ struct Params {
   source1Width: u32,
   source1Height: u32,
 }
-@group(0) @binding(0) var<storage, read> input0: array<${valueType}>;
-@group(0) @binding(1) var<storage, read> input1: array<${valueType}>;
-@group(0) @binding(2) var<storage, read> weights: array<${valueType}>;
+@group(0) @binding(0) var<storage, read> input0: array<${inputStorageType}>;
+@group(0) @binding(1) var<storage, read> input1: array<${inputStorageType}>;
+@group(0) @binding(2) var<storage, read> weights: array<${weightStorageType}>;
 @group(0) @binding(3) var<storage, read> bias: array<vec4<f32>>;
-@group(0) @binding(4) var<storage, read_write> outputData: array<${valueType}>;
+@group(0) @binding(4) var<storage, read_write> outputData: array<${outputType}>;
 @group(0) @binding(5) var<uniform> params: Params;
 
-var<workgroup> inputTile: array<${valueType}, ${inputTileValues}>;
-var<workgroup> weightTile: array<${valueType}, ${weightTileValues}>;
+var<workgroup> inputTile: array<${valueType}, ${gemmSharedSizes(gemm).input}>;
+var<workgroup> weightTile: array<${valueType}, ${gemmSharedSizes(gemm).weights}>;
 
-@compute @workgroup_size(${TILED_CONV_WORKGROUP}, ${TILED_CONV_WORKGROUP}, 1)
+@compute @workgroup_size(${workgroupX}, ${workgroupY}, 1)
 fn main(
   @builtin(local_invocation_id) localId: vec3<u32>,
   @builtin(workgroup_id) workgroupId: vec3<u32>
 ) {
   let spatialBase =
-    workgroupId.x * ${TILED_CONV_M}u +
-    localId.y * ${TILED_CONV_ROWS_PER_THREAD}u;
+    workgroupId.x * ${tileM}u +
+    localId.y * ${rowsPerThread}u;
   let outputBlock =
-    workgroupId.y * ${TILED_CONV_N_BLOCKS}u + localId.x;
+    workgroupId.y * ${tileNBlocks}u + localId.x;
   let spatialCount = params.outputWidth * params.outputHeight;
-  var acc: array<vec4<f32>, ${TILED_CONV_ROWS_PER_THREAD}>;
+  var acc: array<vec4<f32>, ${rowsPerThread}>;
   if (outputBlock < ${outputBlocks}u) {
-    for (var row = 0u; row < ${TILED_CONV_ROWS_PER_THREAD}u; row++) {
+    for (var row = 0u; row < ${rowsPerThread}u; row++) {
       acc[row] = bias[outputBlock];
     }
   }
 
   let localLinear =
-    localId.y * ${TILED_CONV_WORKGROUP}u + localId.x;
+    localId.y * ${workgroupX}u + localId.x;
+  ${addressMode !== 'analytic' ? tiledAddressCacheCode(inputBlocks, true, gemm, { blocks: sourceBlocks, upsampled: upsampledSource }) : ''}
   let totalK = ${inputBlocks * 9}u;
-  for (var kBase = 0u; kBase < totalK; kBase += ${TILED_CONV_K_BLOCKS}u) {
+  for (var kBase = 0u; kBase < totalK; kBase += ${tileKBlocks}u) {
+    ${addressMode !== 'analytic' ? (gemm.decoderLoad === 'source-first' ? `
+    if (channelBlock < ${sourceBlocks[0]}u) {
+      ${tiledIncrementalInputCode(valueType, sourceRead(0, 'inputBlock'), gemm)}
+    } else {
+      ${tiledIncrementalInputCode(valueType, sourceRead(1, `inputBlock - ${sourceBlocks[0]}u`), gemm)}
+    }
+    ` : tiledIncrementalInputCode(valueType, `
+        if (inputBlock < ${sourceBlocks[0]}u) {
+          ${sourceRead(0, 'inputBlock')}
+        } else {
+          ${sourceRead(1, `inputBlock - ${sourceBlocks[0]}u`)}
+        }
+    `, gemm)) : /* wgsl */ `
     for (
       var loadIndex = localLinear;
       loadIndex < ${inputTileValues}u;
       loadIndex += ${workgroupThreads}u
     ) {
-      let tileSpatial = loadIndex / ${TILED_CONV_K_BLOCKS}u;
-      let tileK = loadIndex % ${TILED_CONV_K_BLOCKS}u;
+      let tileSpatial = loadIndex / ${tileKBlocks}u;
+      let tileK = loadIndex % ${tileKBlocks}u;
       let outputSpatialIndex =
-        workgroupId.x * ${TILED_CONV_M}u + tileSpatial;
+        workgroupId.x * ${tileM}u + tileSpatial;
       let kIndex = kBase + tileK;
       var value = ${valueType}(0.0);
       if (outputSpatialIndex < spatialCount && kIndex < totalK) {
@@ -793,23 +977,29 @@ fn main(
           }
         }
       }
-      inputTile[loadIndex] = value;
+      inputTile[${gemmSharedInputIndex('loadIndex', gemm)}] = value;
     }
+    `}
 
     for (
       var loadIndex = localLinear;
       loadIndex < ${weightTileValues}u;
       loadIndex += ${workgroupThreads}u
     ) {
-      let tileK = loadIndex / ${TILED_CONV_N_BLOCKS * 4}u;
-      let outputRemainder = loadIndex % ${TILED_CONV_N_BLOCKS * 4}u;
+      let tileK = loadIndex / ${tileNBlocks * 4}u;
+      let outputRemainder = loadIndex % ${tileNBlocks * 4}u;
       let tileOutputBlock = outputRemainder / 4u;
       let outputLane = outputRemainder % 4u;
       let loadedOutputBlock =
-        workgroupId.y * ${TILED_CONV_N_BLOCKS}u + tileOutputBlock;
+        workgroupId.y * ${tileNBlocks}u + tileOutputBlock;
       let kIndex = kBase + tileK;
       var value = ${valueType}(0.0);
       if (loadedOutputBlock < ${outputBlocks}u && kIndex < totalK) {
+        ${weightLayout === 'k-major' ? /* wgsl */ `
+        let weightIndex = (kIndex * ${outputBlocks}u + loadedOutputBlock) * 4u + outputLane;
+        ` : addressMode !== 'analytic' ? /* wgsl */ `
+        let weightIndex = (loadedOutputBlock * ${inputBlocks * 9}u + kIndex) * 4u + outputLane;
+        ` : /* wgsl */ `
         let inputBlock = kIndex % ${inputBlocks}u;
         let kernelIndex = kIndex / ${inputBlocks}u;
         let kernelY = kernelIndex / 3u;
@@ -817,18 +1007,20 @@ fn main(
         let weightIndex =
           ((((loadedOutputBlock * 3u + kernelY) * 3u + kernelX) *
             ${inputBlocks}u + inputBlock) * 4u + outputLane);
-        value = weights[weightIndex];
+        `}
+        value = ${gemmLoadExpression('weights[weightIndex]', packedWeights)};
       }
-      weightTile[loadIndex] = value;
+      weightTile[${gemmSharedWeightIndex('loadIndex', gemm)}] = value;
     }
 
     workgroupBarrier();
-    ${tiledAccumulationCode(precision)}
+    ${tiledAccumulationCode(precision, gemm)}
     workgroupBarrier();
+    ${addressMode !== 'analytic' ? tiledAddressAdvanceCode(inputBlocks, gemm) : ''}
   }
 
   if (outputBlock < ${outputBlocks}u) {
-    for (var row = 0u; row < ${TILED_CONV_ROWS_PER_THREAD}u; row++) {
+    for (var row = 0u; row < ${rowsPerThread}u; row++) {
       let spatialIndex = spatialBase + row;
       if (spatialIndex < spatialCount) {
         let outputIndex = spatialIndex * ${outputBlocks}u + outputBlock;
@@ -839,8 +1031,9 @@ fn main(
 }
 `;
 }
-function createFusedDecoderShader(precision, activation, sourceBlocks, upsampledSource, outputBlocks) {
+function createFusedDecoderShader(precision, outputPrecision, activation, sourceBlocks, upsampledSource, outputBlocks) {
     const valueType = storageVecType(precision);
+    const outputType = storageVecType(outputPrecision);
     const inputBlocks = sourceBlocks[0] + sourceBlocks[1];
     const sourceCode = (source, blockOffset) => {
         const isUpsampled = source === upsampledSource;
@@ -856,7 +1049,7 @@ function createFusedDecoderShader(precision, activation, sourceBlocks, upsampled
       }
 `;
     };
-    const stored = storeExpression(activationExpression('acc', activation), precision);
+    const stored = storeExpression(activationExpression('acc', activation), outputPrecision);
     return /* wgsl */ `${shaderPreamble(precision)}
 struct Params {
   outputWidth: u32,
@@ -872,7 +1065,7 @@ struct Params {
 @group(0) @binding(1) var<storage, read> input1: array<${valueType}>;
 @group(0) @binding(2) var<storage, read> weights: array<${valueType}>;
 @group(0) @binding(3) var<storage, read> bias: array<vec4<f32>>;
-@group(0) @binding(4) var<storage, read_write> outputData: array<${valueType}>;
+@group(0) @binding(4) var<storage, read_write> outputData: array<${outputType}>;
 @group(0) @binding(5) var<uniform> params: Params;
 
 @compute @workgroup_size(${WORKGROUP_SIZE}, ${WORKGROUP_SIZE}, 1)
@@ -896,12 +1089,13 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 }
 `;
 }
-function createSpatialDecoderShader(precision, activation, sourceBlocks, upsampledSource, outputBlocks) {
+function createSpatialDecoderShader(precision, outputPrecision, activation, sourceBlocks, upsampledSource, outputBlocks) {
     const valueType = storageVecType(precision);
+    const outputType = storageVecType(outputPrecision);
     const inputBlocks = sourceBlocks[0] + sourceBlocks[1];
     const patchValues = SPATIAL_CONV_PATCH * SPATIAL_CONV_PATCH * inputBlocks;
     const workgroupThreads = SPATIAL_CONV_WORKGROUP * SPATIAL_CONV_WORKGROUP;
-    const stored = storeExpression(activationExpression('acc', activation), precision);
+    const stored = storeExpression(activationExpression('acc', activation), outputPrecision);
     const sourceRead = (source, blockExpression) => {
         const sourceX = source === upsampledSource ? 'u32(inputX) / 2u' : 'u32(inputX)';
         const sourceY = source === upsampledSource ? 'u32(inputY) / 2u' : 'u32(inputY)';
@@ -927,7 +1121,7 @@ struct Params {
 @group(0) @binding(1) var<storage, read> input1: array<${valueType}>;
 @group(0) @binding(2) var<storage, read> weights: array<${valueType}>;
 @group(0) @binding(3) var<storage, read> bias: array<vec4<f32>>;
-@group(0) @binding(4) var<storage, read_write> outputData: array<${valueType}>;
+@group(0) @binding(4) var<storage, read_write> outputData: array<${outputType}>;
 @group(0) @binding(5) var<uniform> params: Params;
 
 var<workgroup> inputPatch: array<${valueType}, ${patchValues}>;
@@ -1060,9 +1254,11 @@ export class NativeUNetExecutor {
     _device;
     precision;
     kernelSetting;
+    gemm;
     maxSpatialInputBlocks;
     subgroupsAvailable;
     _model;
+    _gemmByOutputBlocks = new Map();
     _packedConvs = new Map();
     _pipelineCache;
     _pipelinePromises;
@@ -1082,6 +1278,60 @@ export class NativeUNetExecutor {
         this._pipelinePromises = pipelineCache.pending;
         this.precision = resolveNativeUNetPrecision(_device, options.precision ?? 'auto');
         this.kernelSetting = options.kernel ?? 'auto';
+        const workgroupSize = options.gemm?.workgroupSize ?? [8, 8];
+        this.gemm = Object.freeze({
+            tilePolicy: options.gemm?.tilePolicy ?? (options.gemm?.workgroupSize ? 'fixed' : 'output-aligned'),
+            decoderLoad: options.gemm?.decoderLoad ?? 'per-load',
+            poolLayout: options.gemm?.poolLayout ?? 'channels',
+            sharedLayout: options.gemm?.sharedLayout ?? (this.precision === 'fp16' ? 'padded-input' : 'padded'),
+            accumulationOrder: options.gemm?.accumulationOrder ?? 'k-major',
+            finalLayer: options.gemm?.finalLayer ?? 'shared-auto',
+            loadMode: options.gemm?.loadMode ?? 'native',
+            addressMode: options.gemm?.addressMode ?? 'incremental',
+            weightLayout: options.gemm?.weightLayout ?? 'k-major',
+            rowsPerThread: options.gemm?.rowsPerThread ?? 8,
+            workgroupSize: Object.freeze([workgroupSize[0], workgroupSize[1]])
+        });
+        if (!['fixed', 'output-aligned'].includes(this.gemm.tilePolicy)) {
+            throw new Error(`Unsupported GEMM tile policy: ${this.gemm.tilePolicy}`);
+        }
+        if (!['per-load', 'source-first'].includes(this.gemm.decoderLoad)) {
+            throw new Error(`Unsupported GEMM decoder load: ${this.gemm.decoderLoad}`);
+        }
+        if (!['spatial', 'channels'].includes(this.gemm.poolLayout)) {
+            throw new Error(`Unsupported GEMM pool layout: ${this.gemm.poolLayout}`);
+        }
+        if (!['linear', 'padded', 'padded-input', 'padded-weights'].includes(this.gemm.sharedLayout)) {
+            throw new Error(`Unsupported GEMM shared layout: ${this.gemm.sharedLayout}`);
+        }
+        if (!['k-major', 'row-major'].includes(this.gemm.accumulationOrder)) {
+            throw new Error(`Unsupported GEMM accumulation order: ${this.gemm.accumulationOrder}`);
+        }
+        if (!['direct', 'shared-input', 'shared-input-weights', 'shared-auto'].includes(this.gemm.finalLayer)) {
+            throw new Error(`Unsupported GEMM final layer: ${this.gemm.finalLayer}`);
+        }
+        if (!['native', 'packed-weights', 'packed-all'].includes(this.gemm.loadMode)) {
+            throw new Error(`Unsupported GEMM load mode: ${this.gemm.loadMode}`);
+        }
+        if (!['analytic', 'incremental', 'base-offset'].includes(this.gemm.addressMode)) {
+            throw new Error(`Unsupported GEMM address mode: ${this.gemm.addressMode}`);
+        }
+        if (!['output-major', 'k-major'].includes(this.gemm.weightLayout)) {
+            throw new Error(`Unsupported GEMM weight layout: ${this.gemm.weightLayout}`);
+        }
+        const tile = gemmTile(this.gemm);
+        const tileStorageBytes = (gemmSharedSizes(this.gemm).input + gemmSharedSizes(this.gemm).weights) *
+            4 * (this.precision === 'fp16' ? 2 : 4);
+        if (![2, 4, 8].includes(tile.rowsPerThread) ||
+            ![4, 8, 16].includes(tile.workgroupX) ||
+            ![4, 8].includes(tile.workgroupY) ||
+            workgroupSize.length !== 2 ||
+            tile.workgroupX > _device.limits.maxComputeWorkgroupSizeX ||
+            tile.workgroupY > _device.limits.maxComputeWorkgroupSizeY ||
+            tile.workgroupX * tile.workgroupY > _device.limits.maxComputeInvocationsPerWorkgroup ||
+            tileStorageBytes > _device.limits.maxComputeWorkgroupStorageSize) {
+            throw new Error('Unsupported GEMM tile configuration for this GPUDevice');
+        }
         this.subgroupsAvailable = _device.features.has('subgroups');
         this.maxSpatialInputBlocks =
             this.precision === 'fp16' &&
@@ -1118,7 +1368,14 @@ export class NativeUNetExecutor {
         }
         try {
             for (const [id, tensors] of model.convTensors) {
-                const packed = packConvTensors(_device, id, tensors, this.precision);
+                // Kernel selection is shape-independent. Pack only the active layout
+                // for this executor, keeping direct/spatial/subgroup and final weights
+                // in their original ABI without duplicating GPU allocations.
+                const kernel = this._selectConvKernel(blocksForChannels(tensors.inputChannels), id === model.spec.output);
+                const weightLayout = kernel === 'implicit-gemm'
+                    ? this.gemm.weightLayout
+                    : 'output-major';
+                const packed = packConvTensors(_device, id, tensors, this.precision, weightLayout);
                 this._resources.track('gpu-buffer', packed.weights);
                 this._resources.track('gpu-buffer', packed.bias);
                 this._packedConvs.set(id, packed);
@@ -1184,15 +1441,38 @@ export class NativeUNetExecutor {
             const outputPrecision = isFinal ? 'fp32' : this.precision;
             const inputBlocks = blocksForChannels(this._model.convChannels.get(node.id).inputChannels);
             const outputBlocks = blocksForChannels(this._model.convChannels.get(node.id).outputChannels);
+            const gemm = this._gemmForOutput(outputBlocks);
             const kernel = this._selectConvKernel(inputBlocks, isFinal);
+            const cacheFinalWeights = gemm.finalLayer === 'shared-input-weights' ||
+                (gemm.finalLayer === 'shared-auto' &&
+                    finalRgbSharedMemoryBytes(this.precision, inputBlocks, true) <=
+                        this._device.limits.maxComputeWorkgroupStorageSize);
+            if (isFinal && this._model.convChannels.get(node.id).outputChannels === 3 &&
+                (this.kernelSetting === 'auto' || this.kernelSetting === 'implicit-gemm') &&
+                gemm.finalLayer !== 'direct' &&
+                this._device.limits.maxComputeWorkgroupSizeX >= 8 &&
+                this._device.limits.maxComputeWorkgroupSizeY >= 8 &&
+                this._device.limits.maxComputeInvocationsPerWorkgroup >= 64 &&
+                finalRgbSharedMemoryBytes(this.precision, inputBlocks, cacheFinalWeights) <=
+                    this._device.limits.maxComputeWorkgroupStorageSize) {
+                return {
+                    key: `conv-final-rgb/${this.precision}/${node.activation}/in${inputBlocks}/weights-${cacheFinalWeights ? 'shared' : 'storage'}`,
+                    kernel: 'direct',
+                    code: createFinalRgbShader(this.precision, node.activation, inputBlocks, cacheFinalWeights)
+                };
+            }
             const key = `conv-${kernel}/${this.precision}/` +
                 `${outputPrecision}/${node.activation}/` +
-                `in${inputBlocks}/out${outputBlocks}`;
+                `in${inputBlocks}/out${outputBlocks}` +
+                (kernel === 'implicit-gemm'
+                    ? `/address-${gemm.addressMode}/weights-${gemm.weightLayout}-v1` +
+                        `/tile-${gemm.workgroupSize.join('x')}-r${gemm.rowsPerThread}/loads-${gemm.loadMode}/shared-${gemm.sharedLayout}/acc-${gemm.accumulationOrder}`
+                    : '');
             return {
                 key,
                 kernel,
                 code: kernel === 'implicit-gemm'
-                    ? createTiledConvShader(this.precision, outputPrecision, node.activation, inputBlocks, outputBlocks)
+                    ? createTiledConvShader(this.precision, outputPrecision, node.activation, inputBlocks, outputBlocks, gemm)
                     : kernel === 'spatial'
                         ? createSpatialConvShader(this.precision, outputPrecision, node.activation, inputBlocks, outputBlocks)
                         : kernel === 'subgroup'
@@ -1202,11 +1482,12 @@ export class NativeUNetExecutor {
         }
         if (node.op === 'maxPool2d') {
             const outputBlocks = blocksForChannels(this._model.channelsByValue.get(node.id));
-            const key = `max-pool/${this.precision}/out${outputBlocks}`;
+            const coalesced = this._coalescedPool();
+            const key = `max-pool/${this.precision}/out${outputBlocks}/${coalesced ? 'channels' : 'spatial'}`;
             return {
                 key,
                 kernel: 'direct',
-                code: createMaxPoolShader(this.precision, outputBlocks)
+                code: createMaxPoolShader(this.precision, outputBlocks, coalesced)
             };
         }
         if (node.op === 'fusedConvReluMaxPool2d') {
@@ -1230,28 +1511,63 @@ export class NativeUNetExecutor {
                 throw new Error(`Native fused decoder ${node.id} has no upsample input`);
             }
             const inputBlocks = sourceBlocks[0] + sourceBlocks[1];
-            const selectedKernel = this._selectConvKernel(inputBlocks, false);
+            const outputPrecision = isFinal ? 'fp32' : this.precision;
+            const selectedKernel = this._selectConvKernel(inputBlocks, isFinal);
             // The subgroup broadcast path currently targets the common standalone
             // convolution layout; fused decoder reads use the direct kernel.
             const kernel = selectedKernel === 'subgroup'
                 ? 'direct'
                 : selectedKernel;
             const outputBlocks = blocksForChannels(this._model.convChannels.get(node.conv.id).outputChannels);
-            const key = `decoder-${kernel}/${this.precision}/` +
+            const gemm = this._gemmForOutput(outputBlocks);
+            const key = `decoder-${kernel}/${this.precision}/${outputPrecision}/` +
                 `${node.conv.activation}/` +
-                `${sourceBlocks.join('+')}/out${outputBlocks}/up${upsampledSource}`;
+                `${sourceBlocks.join('+')}/out${outputBlocks}/up${upsampledSource}` +
+                (kernel === 'implicit-gemm'
+                    ? `/address-${gemm.addressMode}/weights-${gemm.weightLayout}-v1` +
+                        `/tile-${gemm.workgroupSize.join('x')}-r${gemm.rowsPerThread}/loads-${gemm.loadMode}/shared-${gemm.sharedLayout}/acc-${gemm.accumulationOrder}/decoder-${gemm.decoderLoad}`
+                    : '');
             return {
                 key,
                 kernel,
                 code: kernel === 'implicit-gemm'
-                    ? createTiledDecoderShader(this.precision, node.conv.activation, sourceBlocks, upsampledSource, outputBlocks)
+                    ? createTiledDecoderShader(this.precision, outputPrecision, node.conv.activation, sourceBlocks, upsampledSource, outputBlocks, gemm)
                     : kernel === 'spatial'
-                        ? createSpatialDecoderShader(this.precision, node.conv.activation, sourceBlocks, upsampledSource, outputBlocks)
-                        : createFusedDecoderShader(this.precision, node.conv.activation, sourceBlocks, upsampledSource, outputBlocks)
+                        ? createSpatialDecoderShader(this.precision, outputPrecision, node.conv.activation, sourceBlocks, upsampledSource, outputBlocks)
+                        : createFusedDecoderShader(this.precision, outputPrecision, node.conv.activation, sourceBlocks, upsampledSource, outputBlocks)
             };
         }
         throw new Error(`Native OIDN does not implement unfused ${node.op} node ${node.id}`);
     }
+    // Wider output tiles reuse each input load across twice as many channels.
+    // Require full output blocks and keep the validated base tile as fallback.
+    _gemmForOutput(outputBlocks) {
+        if (this.gemm.tilePolicy !== 'output-aligned' ||
+            this.gemm.workgroupSize[0] !== 8 || outputBlocks <= 0 || outputBlocks % 16 !== 0) {
+            return this.gemm;
+        }
+        const cached = this._gemmByOutputBlocks.get(outputBlocks);
+        if (cached)
+            return cached;
+        const candidate = Object.freeze({
+            ...this.gemm,
+            workgroupSize: Object.freeze([16, this.gemm.workgroupSize[1]])
+        });
+        const tile = gemmTile(candidate);
+        const sizes = gemmSharedSizes(candidate);
+        const limits = this._device.limits;
+        const fits = tile.workgroupX <= limits.maxComputeWorkgroupSizeX &&
+            tile.workgroupY <= limits.maxComputeWorkgroupSizeY &&
+            tile.workgroupX * tile.workgroupY <= limits.maxComputeInvocationsPerWorkgroup &&
+            (sizes.input + sizes.weights) * 4 * (this.precision === 'fp16' ? 2 : 4) <= limits.maxComputeWorkgroupStorageSize;
+        const selected = fits ? candidate : this.gemm;
+        this._gemmByOutputBlocks.set(outputBlocks, selected);
+        return selected;
+    }
+    _coalescedPool() {
+        return this.gemm.poolLayout === 'channels' &&
+            (this.kernelSetting === 'auto' || this.kernelSetting === 'implicit-gemm');
+    }
     _selectConvKernel(inputBlocks, isFinal) {
         const spatialFits = this.precision === 'fp16' &&
             inputBlocks <= this.maxSpatialInputBlocks;
@@ -1266,7 +1582,7 @@ export class NativeUNetExecutor {
         if (this.kernelSetting === 'subgroup') {
             return this.subgroupsAvailable ? 'subgroup' : 'direct';
         }
-        if (this.precision === 'fp32' && !isFinal)
+        if (!isFinal)
             return 'implicit-gemm';
         return 'direct';
     }
@@ -1502,6 +1818,12 @@ export class NativeUNetExecutor {
         execution.lastUsed = ++this._clock;
         return execution;
     }
+    /** Allocates shape-dependent execution resources before interactive use. */
+    prewarm(shapes) {
+        for (const shape of shapes) {
+            this._execution(shape.width, shape.height);
+        }
+    }
     /** Captures per-pass GPU timestamps for the next execute call when supported. */
     profileNextExecution() {
         if (!this._device.features.has('timestamp-query'))
@@ -1595,8 +1917,15 @@ export class NativeUNetExecutor {
                 const pass = encoder.beginComputePass(passDescriptor(`oidn/${execution.plan.nodes[index].id}`, index + 1));
                 pass.setPipeline(execution.nodePipelines[index]);
                 pass.setBindGroup(0, execution.nodeBindings[index]);
-                if (execution.nodeKernels[index] === 'implicit-gemm') {
-                    pass.dispatchWorkgroups(Math.ceil((outputShape.width * outputShape.height) / TILED_CONV_M), Math.ceil(blocksForChannels(outputShape.channels) / TILED_CONV_N_BLOCKS), 1);
+                if (node.op === 'maxPool2d' && this._coalescedPool()) {
+                    const groupsX = Math.ceil(outputShape.width * blocksForChannels(outputShape.channels) / WORKGROUP_SIZE);
+                    const maxGroups = this._device.limits.maxComputeWorkgroupsPerDimension;
+                    // Wide channel rows continue in Z instead of exceeding the X limit.
+                    pass.dispatchWorkgroups(Math.min(groupsX, maxGroups), Math.ceil(outputShape.height / WORKGROUP_SIZE), Math.ceil(groupsX / maxGroups));
+                }
+                else if (execution.nodeKernels[index] === 'implicit-gemm') {
+                    const tile = gemmTile(this._gemmForOutput(blocksForChannels(outputShape.channels)));
+                    pass.dispatchWorkgroups(Math.ceil((outputShape.width * outputShape.height) / tile.tileM), Math.ceil(blocksForChannels(outputShape.channels) / tile.tileNBlocks), 1);
                 }
                 else {
                     pass.dispatchWorkgroups(Math.ceil(outputShape.width / WORKGROUP_SIZE), Math.ceil(outputShape.height / WORKGROUP_SIZE), blocksForChannels(outputShape.channels));
@@ -1644,7 +1973,7 @@ export class NativeUNetExecutor {
         }
         return execution.valueBuffers.get(execution.plan.spec.output);
     }
-    /** Compatibility path for ImageData/HDR arrays without TensorFlow.js. */
+    /** Executes interleaved CPU image data through the native GPU runtime. */
     async executeCPU(interleavedInput, width, height) {
         const expectedLength = width * height * this._model.inputChannels;
         if (interleavedInput.length !== expectedLength) {