oidn-web 0.4.0 → 0.5.0

package/src/nativeUNet.ts

@@ -1,4 +1,5 @@
 import { Float16Array } from '@petamoriken/float16';
+import { createFinalRgbShader, sharedMemoryBytes as finalRgbSharedMemoryBytes } from './finalRgbShader.js';
 import {
   optimizeModelGraph,
   planModelExecution,
@@ -26,14 +27,42 @@ export type NativeUNetKernel =
   | 'spatial'
   | 'subgroup';
 export type NativeUNetKernelSetting = NativeUNetKernel | 'auto';
+export type NativeUNetGemmWorkgroup = readonly [4 | 8 | 16, 4 | 8];
+
+export interface NativeUNetGemmOptions {
+  /** Internal per-layer tile selection; explicit workgroup sizes remain fixed by default. */
+  tilePolicy?: 'fixed' | 'output-aligned';
+  /** Internal decoder source-branch scheduling experiment. */
+  decoderLoad?: 'per-load' | 'source-first';
+  /** Internal pooling dispatch experiment for the GEMM execution path. */
+  poolLayout?: 'spatial' | 'channels';
+  /** Internal shared-memory padding experiment; does not change storage ABI. */
+  sharedLayout?: 'linear' | 'padded' | 'padded-input' | 'padded-weights';
+  /** Internal loop scheduling experiment; each output keeps the same K order. */
+  accumulationOrder?: 'k-major' | 'row-major';
+  /** Internal final-convolution experiments; retains the original FP16 partial order. */
+  finalLayer?: 'direct' | 'shared-input' | 'shared-input-weights' | 'shared-auto';
+  /** Internal load-layout experiments; packed loads preserve the FP16 storage bits. */
+  loadMode?: 'native' | 'packed-weights' | 'packed-all';
+  /** Incremental addressing (default), or the original analytic path for A/B comparisons. */
+  addressMode?: 'analytic' | 'incremental' | 'base-offset';
+  /** K-major (default) coalesces neighboring output-block loads; output-major preserves the old ABI. */
+  weightLayout?: 'output-major' | 'k-major';
+  /** Register tile height (default 8). The K reduction grouping stays fixed. */
+  rowsPerThread?: 2 | 4 | 8;
+  /** Workgroup width/height (default [8, 8]); width also selects output channel blocks. */
+  workgroupSize?: NativeUNetGemmWorkgroup;
+}
 
 export interface NativeUNetOptions {
   precision?: NativeUNetPrecisionSetting;
   /**
-   * Convolution kernel selection. `auto` uses a model-independent capability
-   * heuristic and falls back to the direct kernel when a tile does not fit.
+   * Convolution kernel selection. `auto` uses implicit GEMM for FP16/FP32
+   * convolutions and the direct kernel for the final output layer.
    */
   kernel?: NativeUNetKernelSetting;
+  /** Optional implicit-GEMM tuning; other convolution kernels ignore it. */
+  gemm?: NativeUNetGemmOptions;
   /** Maximum number of shape-dependent activation plans retained. */
   shapeCacheSize?: number;
 }
@@ -51,6 +80,7 @@ export interface NativeUNetExecutionProfile {
 interface PackedConvBuffers {
   weights: GPUBuffer;
   bias: GPUBuffer;
+  weightLayout: NonNullable<NativeUNetGemmOptions['weightLayout']>;
 }
 
 interface NativePipelineSpec {
@@ -97,15 +127,47 @@ function sharedPipelineCache(device: GPUDevice) {
 }
 
 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: Readonly<Required<NativeUNetGemmOptions>>) {
+  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: Readonly<Required<NativeUNetGemmOptions>>) {
+  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: string, options: Readonly<Required<NativeUNetGemmOptions>>) {
+  return options.sharedLayout === 'padded' || options.sharedLayout === 'padded-input'
+    ? `(${index}) + (${index}) / ${options.rowsPerThread * TILED_CONV_K_BLOCKS}u`
+    : index;
+}
+
+function gemmSharedWeightIndex(index: string, options: Readonly<Required<NativeUNetGemmOptions>>) {
+  return options.sharedLayout === 'padded' || options.sharedLayout === 'padded-weights'
+    ? `(${index}) + (${index}) / 4u`
+    : index;
+}
+
 function roundUp(value: number, alignment: number) {
   return Math.ceil(value / alignment) * alignment;
 }
@@ -213,7 +275,8 @@ function packConvTensors(
   device: GPUDevice,
   id: string,
   tensors: ValidatedConvTensor,
-  precision: NativeUNetPrecision
+  precision: NativeUNetPrecision,
+  weightLayout: NonNullable<NativeUNetGemmOptions['weightLayout']>
 ): PackedConvBuffers {
   const inputBlocks = blocksForChannels(tensors.inputChannels);
   const outputBlocks = blocksForChannels(tensors.outputChannels);
@@ -247,16 +310,18 @@ function packConvTensors(
             const outputChannel = outputBlock * 4 + outputLane;
             for (let inputLane = 0; inputLane < 4; inputLane++) {
               const inputChannel = inputBlock * 4 + inputLane;
-              const packedBlock =
-                ((((outputBlock * tensors.kernelHeight + y) *
+              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 (
@@ -285,13 +350,14 @@ function packConvTensors(
 
   const weights = createMappedBuffer(
     device,
-    `oidn/${id}/weights/${precision}`,
+    `oidn/${id}/weights/${precision}/${weightLayout}`,
     packedWeights,
     GPUBufferUsage.STORAGE
   );
   try {
     return {
       weights,
+      weightLayout,
       bias: createMappedBuffer(
         device,
         `oidn/${id}/bias`,
@@ -382,21 +448,58 @@ ${accumulator}
 `;
 }
 
-function tiledAccumulationCode(precision: NativeUNetPrecision) {
+function gemmPackedLoad(
+  precision: NativeUNetPrecision,
+  gemm: Readonly<Required<NativeUNetGemmOptions>>,
+  weights: boolean
+) {
+  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: string, packed: boolean) {
+  return packed ? `bitcast<vec4<f16>>(${expression})` : expression;
+}
+
+function tiledAccumulationCode(precision: NativeUNetPrecision, gemm: Readonly<Required<NativeUNetGemmOptions>>) {
+  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),
@@ -419,20 +522,20 @@ function tiledAccumulationCode(precision: NativeUNetPrecision) {
         );
       }
     }
-    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),
@@ -690,29 +793,109 @@ fn main(
 `;
 }
 
+/** Precompute spatial predicates once, then advance the original flattened K order. */
+function tiledAddressCacheCode(inputBlocks: number, decoder: boolean, gemm: Readonly<Required<NativeUNetGemmOptions>>, sources?: { blocks: readonly [number, number]; upsampled: 0 | 1 }) {
+  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: number, gemm: Readonly<Required<NativeUNetGemmOptions>>) {
+  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: string, readValue: string, gemm: Readonly<Required<NativeUNetGemmOptions>>) {
+  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: NativeUNetPrecision,
   outputPrecision: NativeUNetPrecision,
   activation: Conv2DNodeSpec['activation'],
   inputBlocks: number,
-  outputBlocks: number
+  outputBlocks: number,
+  gemm: Readonly<Required<NativeUNetGemmOptions>>
 ) {
+  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;
+    workgroupX * workgroupY;
+  const inputTileValues = tileM * tileKBlocks;
   const weightTileValues =
-    TILED_CONV_K_BLOCKS * TILED_CONV_N_BLOCKS * 4;
+    tileKBlocks * tileNBlocks * 4;
   return /* wgsl */ `${shaderPreamble(precision)}
 struct Params {
   inputWidth: u32,
@@ -722,46 +905,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) {
@@ -780,26 +970,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;
@@ -807,18 +1003,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;
@@ -893,7 +1091,8 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 
 function createMaxPoolShader(
   precision: NativeUNetPrecision,
-  outputBlocks: number
+  outputBlocks: number,
+  coalesced: boolean
 ) {
   const valueType = storageVecType(precision);
   return /* wgsl */ `${shaderPreamble(precision)}
@@ -909,11 +1108,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;
   }
@@ -922,15 +1127,15 @@ 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)};
 }
 `;
@@ -938,22 +1143,31 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 
 function createTiledDecoderShader(
   precision: NativeUNetPrecision,
+  outputPrecision: NativeUNetPrecision,
   activation: Conv2DNodeSpec['activation'],
   sourceBlocks: readonly [number, number],
   upsampledSource: 0 | 1,
-  outputBlocks: number
+  outputBlocks: number,
+  gemm: Readonly<Required<NativeUNetGemmOptions>>
 ) {
+  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
+    outputPrecision
   );
   const workgroupThreads =
-    TILED_CONV_WORKGROUP * TILED_CONV_WORKGROUP;
-  const inputTileValues = TILED_CONV_M * TILED_CONV_K_BLOCKS;
+    workgroupX * workgroupY;
+  const inputTileValues = tileM * tileKBlocks;
   const weightTileValues =
-    TILED_CONV_K_BLOCKS * TILED_CONV_N_BLOCKS * 4;
+    tileKBlocks * tileNBlocks * 4;
   const sourceRead = (source: 0 | 1, blockExpression: string) => {
     const sourceX =
       source === upsampledSource
@@ -966,12 +1180,15 @@ function createTiledDecoderShader(
     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)}
@@ -985,47 +1202,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) {
@@ -1048,23 +1279,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;
@@ -1072,18 +1309,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;
@@ -1097,12 +1336,14 @@ fn main(
 
 function createFusedDecoderShader(
   precision: NativeUNetPrecision,
+  outputPrecision: NativeUNetPrecision,
   activation: Conv2DNodeSpec['activation'],
   sourceBlocks: readonly [number, number],
   upsampledSource: 0 | 1,
   outputBlocks: number
 ) {
   const valueType = storageVecType(precision);
+  const outputType = storageVecType(outputPrecision);
   const inputBlocks = sourceBlocks[0] + sourceBlocks[1];
   const sourceCode = (source: 0 | 1, blockOffset: number) => {
     const isUpsampled = source === upsampledSource;
@@ -1124,7 +1365,7 @@ function createFusedDecoderShader(
   };
   const stored = storeExpression(
     activationExpression('acc', activation),
-    precision
+    outputPrecision
   );
   return /* wgsl */ `${shaderPreamble(precision)}
 struct Params {
@@ -1141,7 +1382,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)
@@ -1168,12 +1409,14 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 
 function createSpatialDecoderShader(
   precision: NativeUNetPrecision,
+  outputPrecision: NativeUNetPrecision,
   activation: Conv2DNodeSpec['activation'],
   sourceBlocks: readonly [number, number],
   upsampledSource: 0 | 1,
   outputBlocks: number
 ) {
   const valueType = storageVecType(precision);
+  const outputType = storageVecType(outputPrecision);
   const inputBlocks = sourceBlocks[0] + sourceBlocks[1];
   const patchValues =
     SPATIAL_CONV_PATCH * SPATIAL_CONV_PATCH * inputBlocks;
@@ -1181,7 +1424,7 @@ function createSpatialDecoderShader(
     SPATIAL_CONV_WORKGROUP * SPATIAL_CONV_WORKGROUP;
   const stored = storeExpression(
     activationExpression('acc', activation),
-    precision
+    outputPrecision
   );
   const sourceRead = (source: 0 | 1, blockExpression: string) => {
     const sourceX =
@@ -1211,7 +1454,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}>;
@@ -1361,10 +1604,12 @@ export function resolveNativeUNetPrecision(
 export class NativeUNetExecutor {
   readonly precision: NativeUNetPrecision;
   readonly kernelSetting: NativeUNetKernelSetting;
+  readonly gemm: Readonly<Required<NativeUNetGemmOptions>>;
   readonly maxSpatialInputBlocks: number;
   readonly subgroupsAvailable: boolean;
 
   private _model: UNetModelGraph;
+  private _gemmByOutputBlocks = new Map<number, Readonly<Required<NativeUNetGemmOptions>>>();
   private _packedConvs = new Map<string, PackedConvBuffers>();
   private _pipelineCache: Map<string, GPUComputePipeline>;
   private _pipelinePromises: Map<string, Promise<GPUComputePipeline>>;
@@ -1391,6 +1636,63 @@ export class NativeUNetExecutor {
       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]]) as NativeUNetGemmWorkgroup
+    });
+    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' as GPUFeatureName
     );
@@ -1440,7 +1742,18 @@ 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);
@@ -1513,11 +1826,36 @@ export class NativeUNetExecutor {
       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,
@@ -1527,7 +1865,8 @@ export class NativeUNetExecutor {
               outputPrecision,
               node.activation,
               inputBlocks,
-              outputBlocks
+              outputBlocks,
+              gemm
             )
           : kernel === 'spatial'
             ? createSpatialConvShader(
@@ -1558,11 +1897,12 @@ export class NativeUNetExecutor {
       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') {
@@ -1598,7 +1938,8 @@ 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'
@@ -1607,24 +1948,32 @@ export class NativeUNetExecutor {
       const outputBlocks = blocksForChannels(
         this._model.convChannels.get(node.conv.id)!.outputChannels
       );
+      const gemm = this._gemmForOutput(outputBlocks);
       const key =
-        `decoder-${kernel}/${this.precision}/` +
+        `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,
+              outputPrecision,
               node.conv.activation,
               sourceBlocks,
               upsampledSource,
-              outputBlocks
+              outputBlocks,
+              gemm
             )
           : kernel === 'spatial'
             ? createSpatialDecoderShader(
                 this.precision,
+                outputPrecision,
                 node.conv.activation,
                 sourceBlocks,
                 upsampledSource,
@@ -1632,6 +1981,7 @@ export class NativeUNetExecutor {
               )
             : createFusedDecoderShader(
               this.precision,
+              outputPrecision,
               node.conv.activation,
               sourceBlocks,
               upsampledSource,
@@ -1644,6 +1994,36 @@ export class NativeUNetExecutor {
     );
   }
 
+  // Wider output tiles reuse each input load across twice as many channels.
+  // Require full output blocks and keep the validated base tile as fallback.
+  private _gemmForOutput(outputBlocks: number): Readonly<Required<NativeUNetGemmOptions>> {
+    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]]) as NativeUNetGemmWorkgroup
+    });
+    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;
+  }
+
+  private _coalescedPool() {
+    return this.gemm.poolLayout === 'channels' &&
+      (this.kernelSetting === 'auto' || this.kernelSetting === 'implicit-gemm');
+  }
+
   private _selectConvKernel(
     inputBlocks: number,
     isFinal: boolean
@@ -1661,7 +2041,7 @@ export class NativeUNetExecutor {
     if (this.kernelSetting === 'subgroup') {
       return this.subgroupsAvailable ? 'subgroup' : 'direct';
     }
-    if (this.precision === 'fp32' && !isFinal) return 'implicit-gemm';
+    if (!isFinal) return 'implicit-gemm';
     return 'direct';
   }
 
@@ -1939,6 +2319,13 @@ export class NativeUNetExecutor {
     return execution;
   }
 
+  /** Allocates shape-dependent execution resources before interactive use. */
+  prewarm(shapes: readonly { width: number; height: number }[]) {
+    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')) return false;
@@ -2056,13 +2443,25 @@ export class NativeUNetExecutor {
       );
       pass.setPipeline(execution.nodePipelines[index]);
       pass.setBindGroup(0, execution.nodeBindings[index]);
-      if (execution.nodeKernels[index] === 'implicit-gemm') {
+      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) / TILED_CONV_M
+            (outputShape.width * outputShape.height) / tile.tileM
           ),
           Math.ceil(
-            blocksForChannels(outputShape.channels) / TILED_CONV_N_BLOCKS
+            blocksForChannels(outputShape.channels) / tile.tileNBlocks
           ),
           1
         );
@@ -2135,7 +2534,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: Float32Array,
     width: number,