sync 6fdf6301e2bb

Browse files

Files changed (16) hide show

README.md +184 -0
build/webgpu/bench.json +0 -0
build/webgpu/conv-int-accumulate-spatial.wgsl.jinja +131 -0
build/webgpu/conv-int-im2col-spatial.wgsl.jinja +45 -0
build/webgpu/manifest.json +0 -0
build/webgpu/metadata.json +115 -0
build/webgpu/qlinear-conv-im2col-nhwc.wgsl.jinja +82 -0
build/webgpu/qlinear-conv-nchw-x4.wgsl.jinja +384 -0
build/webgpu/qlinear-conv-nhwc-oc4.wgsl.jinja +394 -0
build/webgpu/qlinear-conv-requantize.wgsl.jinja +302 -0
build/webgpu/quant-dp4a-matmul.wgsl.jinja +573 -0
build/webgpu/quant-exact-matrix.wgsl.jinja +384 -0
build/webgpu/quant-exact-portable.wgsl.jinja +412 -0
build/webgpu/quant-exact-prepare.wgsl.jinja +79 -0
build/webgpu/quant-weight-k-major.wgsl.jinja +42 -0
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

@@ -1,3 +1,187 @@
 ---
 license: apache-2.0
 ---

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.QLinearConv
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Convolution on a quantized input with quantized weights, producing a quantized output, in either tensor layout. `channels_last = 0` reads `X` as `(N, C, D1..Dn)` and writes `Y` as `(N, M, D1..Dn)`; `channels_last = 1` reads `(N, D1..Dn, C)` and writes `(N, D1..Dn, M)`. The weight stays `(M, C/group, k1..kn)` in both layouts. Input and output quantization are per-tensor; weight quantization is per-tensor or per-output-channel. The optional int32 bias is pre-quantized with `scale = x_scale * w_scale` and zero point 0. Spatial ranks 1 to 3 are implemented.
+See the [ONNX Runtime `QLinearConv` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.QLinearConv) for the reference semantics.
+## Inputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `x` | `TX` | — | — | Quantized input tensor, `(N, C, D1 x ... x Dn)` when `channels_last` is 0 and `(N, D1 x ... x Dn, C)` when it is 1. | required |
+| `x_scale` | `TF` | — | — | Per-tensor scale for input `x`. | required |
+| `x_zero_point` | `TX` | — | — | Per-tensor zero point for input `x`. | required |
+| `w` | `TW` | — | — | Quantized weight tensor shaped `(M, C/group, k1 x ... x kn)` in both layouts; `channels_last` never transposes the weight. | required |
+| `w_scale` | `TF` | — | — | Scale for weights `w`; scalar for per-tensor or 1-D of length M for per-output-channel quantization. | required |
+| `w_zero_point` | `TW` | — | — | Zero point for weights `w`; scalar or 1-D of length M matching `w_scale`. | required |
+| `y_scale` | `TF` | — | — | Per-tensor scale for output `y`. | required |
+| `y_zero_point` | `TY` | — | — | Per-tensor zero point for output `y`. | required |
+| `B` | `int32` | `1` | — | Optional 1-D bias of length `M`, pre-quantized with scale `x_scale * w_scale` and zero point 0. | optional |
+## Outputs
+| Name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- |
+| `y` | `TY` | same as `x` | derived | Quantized output tensor; the channel axis follows `channels_last` and the spatial sizes follow the kernel, strides, dilations and padding. | required |
+## Attributes
+Attributes and default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `auto_pad` | `"NOTSET"` | Automatic padding mode. `NOTSET` uses `pads`; `SAME_UPPER` and `SAME_LOWER` choose padding so each output spatial size is `ceil(input / stride)`; `VALID` uses no padding. |
+| `channels_last` | `0` | Tensor layout convention for `x` and `y`: `0` places the channel axis at index 1, `1` places it last. The weight layout is unaffected. Defaults to 0. |
+| `dilations` | — | Optional dilation factors, one positive integer per spatial axis. Omission means all ones. |
+| `group` | `1` | Number of groups that input and output channels are split into; defaults to 1. |
+| `kernel_shape` | — | Optional kernel shape, one positive integer per spatial axis. When present, it must match the spatial dimensions of the weight tensor; omission infers the shape from the weights. |
+| `pads` | — | Optional explicit padding in ONNX order `[begin_axis_0, ..., begin_axis_n, end_axis_0, ..., end_axis_n]`. Omission means all zeros; it cannot be combined with an automatic padding mode. |
+| `strides` | — | Optional stride factors, one positive integer per spatial axis. Omission means all ones. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `TX` | `uint8`, `int8` |
+| `TW` | `uint8`, `int8` |
+| `TY` | `uint8`, `int8` |
+| `TF` | `float32` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `nchw2d_columns_x4` — Channel-first 2-D route without integer dot-product support: one invocation folds four adjacent output columns into a vec4 accumulator for up to four output channels, so each weight load serves sixteen products.
+- `dp4a_pointwise_nchw` — Channel-first 1x1 route: the convolution is exactly one integer GEMM whose A operand is the weight and whose B operand is the activation image, so the shared packed integer dot-product core runs with no materialized column matrix.
+- `matrix_pointwise_nchw` — Exact subgroup-matrix convolution using the same operand preparation. Bounded f32 partial dots preserve byte-integer products; int32 partial accumulation, optional bias and compensated requantization preserve output quantization. Small grids and short contractions use the packed integer kernels.
+- `dp4a_im2col_nchw` — Channel-first route that materializes the column matrix in `(ic, kh, kw)` order and hands it to the shared packed integer dot-product core as the B operand, with the weight as A.
+- `matrix_im2col_nchw` — Exact subgroup-matrix convolution using the same operand preparation. Bounded f32 partial dots preserve byte-integer products; int32 partial accumulation, optional bias and compensated requantization preserve output quantization. Small grids and short contractions use the packed integer kernels.
+- `dp4a_im2col_nchw_bias` — Channel-first route that materializes the column matrix in `(ic, kh, kw)` order and hands it to the shared packed integer dot-product core as the B operand, with the weight as A. A pre-quantized int32 bias is added per output channel before requantization.
+- `matrix_im2col_nchw_bias` — Exact subgroup-matrix convolution using the same operand preparation. Bounded f32 partial dots preserve byte-integer products; int32 partial accumulation, optional bias and compensated requantization preserve output quantization. Small grids and short contractions use the packed integer kernels.
+- `dp4a_pointwise_nhwc` — Channels-last 1x1 route: the activation is already the GEMM's row-major A operand, so only the weight is reordered into K-major `(kh, kw, ic)` columns for the shared packed integer dot-product core, whose output rows land as the channels-last image.
+- `matrix_pointwise_nhwc` — Exact subgroup-matrix convolution using the same operand preparation. Bounded f32 partial dots preserve byte-integer products; int32 partial accumulation, optional bias and compensated requantization preserve output quantization. Small grids and short contractions use the packed integer kernels.
+- `dp4a_pointwise_nhwc_bias` — Channels-last 1x1 route: the activation is already the GEMM's row-major A operand, so only the weight is reordered into K-major `(kh, kw, ic)` columns for the shared packed integer dot-product core, whose output rows land as the channels-last image. A pre-quantized int32 bias is added per output channel before requantization.
+- `matrix_pointwise_nhwc_bias` — Exact subgroup-matrix convolution using the same operand preparation. Bounded f32 partial dots preserve byte-integer products; int32 partial accumulation, optional bias and compensated requantization preserve output quantization. Small grids and short contractions use the packed integer kernels.
+- `dp4a_im2col_nhwc` — Channels-last general route: the activation is already the GEMM's row-major A operand, so only the weight is reordered into K-major `(kh, kw, ic)` columns for the shared packed integer dot-product core, whose output rows land as the channels-last image.
+- `matrix_im2col_nhwc` — Exact subgroup-matrix convolution using the same operand preparation. Bounded f32 partial dots preserve byte-integer products; int32 partial accumulation, optional bias and compensated requantization preserve output quantization. Small grids and short contractions use the packed integer kernels.
+- `dp4a_im2col_nhwc_bias` — Channels-last general route: the activation is already the GEMM's row-major A operand, so only the weight is reordered into K-major `(kh, kw, ic)` columns for the shared packed integer dot-product core, whose output rows land as the channels-last image. A pre-quantized int32 bias is added per output channel before requantization.
+- `matrix_im2col_nhwc_bias` — Exact subgroup-matrix convolution using the same operand preparation. Bounded f32 partial dots preserve byte-integer products; int32 partial accumulation, optional bias and compensated requantization preserve output quantization. Small grids and short contractions use the packed integer kernels.
+- `precast_f16_pointwise_nchw_tensor` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f16_pointwise_nchw_tensor` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f16_pointwise_nchw_row` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f16_pointwise_nchw_row` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f16_im2col_nchw_tensor` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f16_im2col_nchw_tensor` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f16_im2col_nchw_row` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f16_im2col_nchw_row` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f16_im2col_nchw_bias_tensor` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f16_im2col_nchw_bias_tensor` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f16_im2col_nchw_bias_row` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f16_im2col_nchw_bias_row` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f16_pointwise_nhwc_tensor` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f16_pointwise_nhwc_tensor` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f16_pointwise_nhwc_column` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f16_pointwise_nhwc_column` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f16_pointwise_nhwc_bias_tensor` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f16_pointwise_nhwc_bias_tensor` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f16_pointwise_nhwc_bias_column` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f16_pointwise_nhwc_bias_column` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f16_im2col_nhwc_tensor` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f16_im2col_nhwc_tensor` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f16_im2col_nhwc_column` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f16_im2col_nhwc_column` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f16_im2col_nhwc_bias_tensor` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f16_im2col_nhwc_bias_tensor` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f16_im2col_nhwc_bias_column` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f16_im2col_nhwc_bias_column` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f32_pointwise_nchw_tensor` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f32_pointwise_nchw_tensor` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f32_pointwise_nchw_row` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f32_pointwise_nchw_row` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f32_im2col_nchw_tensor` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f32_im2col_nchw_tensor` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f32_im2col_nchw_row` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f32_im2col_nchw_row` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f32_im2col_nchw_bias_tensor` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f32_im2col_nchw_bias_tensor` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f32_im2col_nchw_bias_row` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f32_im2col_nchw_bias_row` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f32_pointwise_nhwc_tensor` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f32_pointwise_nhwc_tensor` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f32_pointwise_nhwc_column` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f32_pointwise_nhwc_column` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f32_pointwise_nhwc_bias_tensor` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f32_pointwise_nhwc_bias_tensor` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f32_pointwise_nhwc_bias_column` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f32_pointwise_nhwc_bias_column` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f32_im2col_nhwc_tensor` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f32_im2col_nhwc_tensor` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f32_im2col_nhwc_column` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f32_im2col_nhwc_column` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f32_im2col_nhwc_bias_tensor` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f32_im2col_nhwc_bias_tensor` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+- `precast_f32_im2col_nhwc_bias_column` — Fuse spatial sampling, weight layout and exact byte centering into padded matrix operands. Both component tiers use bounded `f32` partial dots and `int32` totals before compensated requantization; `f16` exactly represents `[-255, 255]`. Tiles follow device limits, and small workloads use the other routes.
+- `portable_f32_im2col_nhwc_bias_column` — Reuse fused exact byte preparation with a portable vector GEMM. Each workgroup computes a device-bounded tile with bounded exact f32 partial dots, int32 totals, optional per-channel bias and compensated requantization. Exact f16 operand storage is used when supported; f32 requires no optional arithmetic features.
+## Device requirements
+Some implementation variants require `subgroup-matrix`, `shader-f16`, and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
+## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
+- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
+- [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
+- [`conv-int-accumulate-spatial.wgsl.jinja`](build/webgpu/conv-int-accumulate-spatial.wgsl.jinja)
+- [`conv-int-im2col-spatial.wgsl.jinja`](build/webgpu/conv-int-im2col-spatial.wgsl.jinja)
+- [`qlinear-conv-im2col-nhwc.wgsl.jinja`](build/webgpu/qlinear-conv-im2col-nhwc.wgsl.jinja)
+- [`qlinear-conv-nchw-x4.wgsl.jinja`](build/webgpu/qlinear-conv-nchw-x4.wgsl.jinja)
+- [`qlinear-conv-nhwc-oc4.wgsl.jinja`](build/webgpu/qlinear-conv-nhwc-oc4.wgsl.jinja)
+- [`qlinear-conv-requantize.wgsl.jinja`](build/webgpu/qlinear-conv-requantize.wgsl.jinja)
+- [`quant-dp4a-matmul.wgsl.jinja`](build/webgpu/quant-dp4a-matmul.wgsl.jinja)
+- [`quant-exact-matrix.wgsl.jinja`](build/webgpu/quant-exact-matrix.wgsl.jinja)
+- [`quant-exact-portable.wgsl.jinja`](build/webgpu/quant-exact-portable.wgsl.jinja)
+- [`quant-exact-prepare.wgsl.jinja`](build/webgpu/quant-exact-prepare.wgsl.jinja)
+- [`quant-weight-k-major.wgsl.jinja`](build/webgpu/quant-weight-k-major.wgsl.jinja)
+## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
+The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
+Replace each `*Data` placeholder with a typed array containing the corresponding input data.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/com.microsoft.QLinearConv", { version: 1 });
+const { y } = await kernel({
+  x: { data: xData, shape: [1, 1, 1, 1] },
+  x_scale: { data: x_scaleData, shape: [1] },
+  x_zero_point: { data: x_zero_pointData, shape: [1] },
+  w: { data: wData, shape: [2, 1, 1, 1] },
+  w_scale: { data: w_scaleData, shape: [1] },
+  w_zero_point: { data: w_zero_pointData, shape: [1] },
+  y_scale: { data: y_scaleData, shape: [1] },
+  y_zero_point: { data: y_zero_pointData, shape: [1] },
+});
+```

build/webgpu/bench.json ADDED Viewed

The diff for this file is too large to render. See raw diff

build/webgpu/conv-int-accumulate-spatial.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,131 @@

+{% macro quant_read_scalar_zero(fn_name, buffer_name, is_unsigned, present=true) %}
+fn {{ fn_name }}() -> i32 {
+  {% if present %}
+  {% if is_unsigned %}
+  return i32({{ buffer_name }}[0]);
+  {% else %}
+  return {{ buffer_name }}[0];
+  {% endif %}
+  {% else %}
+  return 0;
+  {% endif %}
+}{% endmacro %}
+{% macro quant_read_column_zero(fn_name, buffer_name, is_unsigned, present=true, per_column=false, param_name="col", return_type="i32") %}
+fn {{ fn_name }}({% if per_column %}{{ param_name }}: u32{% endif %}) -> {{ return_type }} {
+  {% if present %}
+  {% if per_column %}
+  {% if is_unsigned %}
+  return i32({{ buffer_name }}[{{ param_name }}]);
+  {% else %}
+  return {{ buffer_name }}[{{ param_name }}];
+  {% endif %}
+  {% else %}
+  {% if is_unsigned %}
+  return i32({{ buffer_name }}[0]);
+  {% else %}
+  return {{ buffer_name }}[0];
+  {% endif %}
+  {% endif %}
+  {% else %}
+  return 0;
+  {% endif %}
+}{% endmacro %}
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) { return; }{% endmacro %}
+{{ env.wgsl.resourceDeclarations }}
+fn read_x(index: u32) -> i32 {
+  {% if xUnsigned %}
+  return i32(x[index]);
+  {% else %}
+  return x[index];
+  {% endif %}
+}
+fn read_w(index: u32) -> i32 {
+  {% if wUnsigned %}
+  return i32(w[index]);
+  {% else %}
+  return w[index];
+  {% endif %}
+}
+{{ quant_read_scalar_zero("read_x_zero", "x_zero_point", xUnsigned, not (xZeroOmitted is defined and xZeroOmitted)) }}
+// A per-output-channel weight zero point applies at every spatial rank.
+{{ quant_read_column_zero("read_w_zero", "w_zero_point", wUnsigned, not (wZeroOmitted is defined and wZeroOmitted), wZeroPerChannel, "oc") }}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "index", guardInline=true) }}
+  {% if channelsLast is defined and channelsLast %}
+  let oc = index % params.outChannels;
+  var t = index / params.outChannels;
+  let ow = t % params.outW;
+  t = t / params.outW;
+  let oh = t % params.outH;
+  t = t / params.outH;
+  {% if spatialRankSpec == 3 %}
+  let od = t % params.outD;
+  t = t / params.outD;
+  {% endif %}
+  let batch = t;
+  {% else %}
+  let ow = index % params.outW;
+  var t = index / params.outW;
+  let oh = t % params.outH;
+  t = t / params.outH;
+  {% if spatialRankSpec == 3 %}
+  let od = t % params.outD;
+  t = t / params.outD;
+  {% endif %}
+  let oc = t % params.outChannels;
+  let batch = t / params.outChannels;
+  {% endif %}
+  let group = oc / params.outChannelsPerGroup;
+  let xzp = read_x_zero();
+  let wzp = read_w_zero({% if wZeroPerChannel %}oc{% endif %});
+  var acc = 0i;
+  for (var ic = 0u; ic < params.weightInChannels; ic = ic + 1u) {
+    let input_channel = group * params.inChannelsPerGroup + ic;
+    {% if spatialRankSpec == 3 %}
+    for (var kd = 0u; kd < params.kernelD; kd = kd + 1u) {
+      let id = i32(od * params.strideD + kd * params.dilationD) - params.padD;
+      if (id < 0 || id >= i32(params.inD)) { continue; }
+    {% endif %}
+{% filter indent(2 if spatialRankSpec == 3 else 0, true) %}
+    for (var kh = 0u; kh < params.kernelH; kh = kh + 1u) {
+      let ih = i32(oh * params.strideH + kh * params.dilationH) - params.padH;
+      if (ih < 0 || ih >= i32(params.inH)) { continue; }
+      for (var kw = 0u; kw < params.kernelW; kw = kw + 1u) {
+        let iw = i32(ow * params.strideW + kw * params.dilationW) - params.padW;
+        if (iw < 0 || iw >= i32(params.inW)) { continue; }
+        {% if spatialRankSpec == 3 %}
+        {% if channelsLast is defined and channelsLast %}
+        let x_index = (((batch * params.inD + u32(id)) * params.inH + u32(ih)) * params.inW + u32(iw)) * params.inChannels + input_channel;
+        {% else %}
+        let x_index = (((batch * params.inChannels + input_channel) * params.inD + u32(id)) * params.inH + u32(ih)) * params.inW + u32(iw);
+        {% endif %}
+        let w_index = (((oc * params.weightInChannels + ic) * params.kernelD + kd) * params.kernelH + kh) * params.kernelW + kw;
+        {% else %}
+        {% if channelsLast is defined and channelsLast %}
+        let x_index = ((batch * params.inH + u32(ih)) * params.inW + u32(iw)) * params.inChannels + input_channel;
+        {% else %}
+        let x_index = ((batch * params.inChannels + input_channel) * params.inH + u32(ih)) * params.inW + u32(iw);
+        {% endif %}
+        let w_index = ((oc * params.weightInChannels + ic) * params.kernelH + kh) * params.kernelW + kw;
+        {% endif %}
+        acc = acc + (read_x(x_index) - xzp) * (read_w(w_index) - wzp);
+      }
+    }
+{% endfilter %}
+    {% if spatialRankSpec == 3 %}
+    }
+    {% endif %}
+  }
+  y[index] = acc;
+}

build/webgpu/conv-int-im2col-spatial.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,45 @@

+/* Widened-int8 im2col for convolution DP4A paths. The materialized K order is
+ * OIHW for 2D input and OIDHW for 3D input. Spatial padding is filled with the
+ * input zero point, so padded taps contribute exactly zero after centering.
+ * Adjacent x invocations cover adjacent output positions for coalesced writes.
+ * Optional extra K rows contain the raw input zero point to align the DP4A
+ * reduction. */
+{{ env.wgsl.resourceDeclarations }}
+{% macro spatial_constants() %}
+const KERNEL_H: u32 = {{ kernelHSpec }}u;
+const KERNEL_W: u32 = {{ kernelWSpec }}u;
+const KERNEL_HW: u32 = KERNEL_H * KERNEL_W;
+const KSIZE: u32 = KERNEL_HW;
+const STRIDE_H: u32 = {{ strideHSpec }}u;
+const STRIDE_W: u32 = {{ strideWSpec }}u;
+const DILATION_H: u32 = {{ dilationHSpec }}u;
+const DILATION_W: u32 = {{ dilationWSpec }}u;
+const PAD_TOP: i32 = {{ padTopSpec }};
+const PAD_LEFT: i32 = {{ padLeftSpec }};
+{% endmacro %}
+{% macro spatial_sample(value="value", input="x", zero="x_zero_point", sampleK="k", sampleScalar="auto") %}
+  let ic = {{ sampleK }} / KSIZE;
+  let kq = {{ sampleK }} - ic * KSIZE;
+  let kh = kq / KERNEL_W;
+  let kw = kq - kh * KERNEL_W;
+  let oh = pos / params.outW;
+  let ow = pos - oh * params.outW;
+  let ih = i32(oh * STRIDE_H + kh * DILATION_H) - PAD_TOP;
+  let iw = i32(ow * STRIDE_W + kw * DILATION_W) - PAD_LEFT;
+  var {{ value }}: {{ bScalar if sampleScalar == "auto" else sampleScalar }} = {{ zero }}[0];
+  if (ih >= 0 && ih < i32(params.inH) && iw >= 0 && iw < i32(params.inW)) {
+    {{ value }} = {{ input }}[{% if sampleChannelsLast is defined and sampleChannelsLast %}((batch * params.inH + u32(ih)) * params.inW + u32(iw)) * params.inChannels + ic{% else %}((batch * params.inChannels + ic) * params.inH + u32(ih)) * params.inW + u32(iw){% endif %}];
+  }
+{% endmacro %}
+{{ spatial_constants() }}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let pos = gid.x;
+  if (pos >= params.outCount) {
+    return;
+  }
+  let k = gid.y;
+  let batch = gid.z;
+{{ spatial_sample() }}
+  cols[(batch * params.kRows + k) * params.outCount + pos] = value;
+}

build/webgpu/manifest.json ADDED Viewed

The diff for this file is too large to render. See raw diff

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,115 @@

+{
+  "name": "com.microsoft.QLinearConv",
+  "id": "_com_microsoft_qlinearconv_webgpu_dd39e27",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "deKtgF/RQc9YtJ6rKQbZOEn+82qQqbE7gvVAbfsCrHI=",
+      "conv-int-accumulate-spatial.wgsl.jinja": "ovvyaZB9d/pmohPYt0/x2NLB2wSWMJ0/eoBHcm4i/5E=",
+      "conv-int-im2col-spatial.wgsl.jinja": "23v1AQ7sYZnMKE7b32VbAbjutCCZnsMgXRkXYm4W6/k=",
+      "manifest.json": "TbaAnjGDWtUPENa228Ulfk4p2D8TYPRcs0oyPBIHGoI=",
+      "qlinear-conv-im2col-nhwc.wgsl.jinja": "CfKXjliMWapwWUjjAk3hY51R2rsx1Wcd4LkvdYUWe6I=",
+      "qlinear-conv-nchw-x4.wgsl.jinja": "CrFYcoL6MLszFVzk877uUZu6eMi7h6Wii63GJXWC+x0=",
+      "qlinear-conv-nhwc-oc4.wgsl.jinja": "rTj6WQWmCd0BMzMZJVk/ZssxIDh2HrUEOnuK/owJb3Q=",
+      "qlinear-conv-requantize.wgsl.jinja": "ElB8XipqRr6CjNjGakSOhsdsM/TezfpRQ2zyh+vENPQ=",
+      "quant-dp4a-matmul.wgsl.jinja": "jrfKZZQ/P2yDNDElA9IBQJIxU2p7yJzDjj+dt5D4u1E=",
+      "quant-exact-matrix.wgsl.jinja": "EFugcaUsDIAeUHx17xX66qV/fqfila1p8nlXQCG97Kw=",
+      "quant-exact-portable.wgsl.jinja": "Ym20X8RcLRzXP9DcQSwJBp4z8RRCBc+bwhrofGC2SBc=",
+      "quant-exact-prepare.wgsl.jinja": "hAWISxWqTkqc1z9PRAlaNnKqajQKc6sw6IVWfyYmHns=",
+      "quant-weight-k-major.wgsl.jinja": "B2sKfFD51YrInBxbg8bUKVtU8j/QHUgyoJ7QONaV7iE=",
+      "test.json": "Yi0vCMNVybIz6uyywnEE6TeUR5B22AzjXtgta2kOmK4="
+    }
+  },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.1",
+    "variants": {
+      "spatial_nd2_scalar": ["conv-int-accumulate-spatial.wgsl.jinja", "qlinear-conv-requantize.wgsl.jinja"],
+      "spatial_nd2_scalar_bias": ["conv-int-accumulate-spatial.wgsl.jinja", "qlinear-conv-requantize.wgsl.jinja"],
+      "spatial_nd2_per_channel": ["conv-int-accumulate-spatial.wgsl.jinja", "qlinear-conv-requantize.wgsl.jinja"],
+      "spatial_nd2_per_channel_bias": ["conv-int-accumulate-spatial.wgsl.jinja", "qlinear-conv-requantize.wgsl.jinja"],
+      "spatial_nd3_scalar": ["conv-int-accumulate-spatial.wgsl.jinja", "qlinear-conv-requantize.wgsl.jinja"],
+      "spatial_nd3_scalar_bias": ["conv-int-accumulate-spatial.wgsl.jinja", "qlinear-conv-requantize.wgsl.jinja"],
+      "spatial_nd3_per_channel": ["conv-int-accumulate-spatial.wgsl.jinja", "qlinear-conv-requantize.wgsl.jinja"],
+      "spatial_nd3_per_channel_bias": ["conv-int-accumulate-spatial.wgsl.jinja", "qlinear-conv-requantize.wgsl.jinja"],
+      "nchw2d_columns_x4": ["qlinear-conv-nchw-x4.wgsl.jinja"],
+      "nhwc2d_channels_x4": ["qlinear-conv-nhwc-oc4.wgsl.jinja"],
+      "nhwc2d_depthwise_x4": ["qlinear-conv-nhwc-oc4.wgsl.jinja"],
+      "nhwc2d_channels_x4_bias": ["qlinear-conv-nhwc-oc4.wgsl.jinja"],
+      "nhwc2d_depthwise_x4_bias": ["qlinear-conv-nhwc-oc4.wgsl.jinja"],
+      "dp4a_pointwise_nchw": ["quant-dp4a-matmul.wgsl.jinja"],
+      "matrix_pointwise_nchw": ["quant-exact-matrix.wgsl.jinja"],
+      "dp4a_im2col_nchw": ["conv-int-im2col-spatial.wgsl.jinja", "quant-dp4a-matmul.wgsl.jinja"],
+      "matrix_im2col_nchw": ["conv-int-im2col-spatial.wgsl.jinja", "quant-exact-matrix.wgsl.jinja"],
+      "dp4a_im2col_nchw_bias": ["conv-int-im2col-spatial.wgsl.jinja", "quant-dp4a-matmul.wgsl.jinja"],
+      "matrix_im2col_nchw_bias": ["conv-int-im2col-spatial.wgsl.jinja", "quant-exact-matrix.wgsl.jinja"],
+      "dp4a_pointwise_nhwc": ["quant-dp4a-matmul.wgsl.jinja", "quant-weight-k-major.wgsl.jinja"],
+      "matrix_pointwise_nhwc": ["quant-exact-matrix.wgsl.jinja", "quant-weight-k-major.wgsl.jinja"],
+      "dp4a_pointwise_nhwc_bias": ["quant-dp4a-matmul.wgsl.jinja", "quant-weight-k-major.wgsl.jinja"],
+      "matrix_pointwise_nhwc_bias": ["quant-exact-matrix.wgsl.jinja", "quant-weight-k-major.wgsl.jinja"],
+      "dp4a_im2col_nhwc": ["qlinear-conv-im2col-nhwc.wgsl.jinja", "quant-dp4a-matmul.wgsl.jinja", "quant-weight-k-major.wgsl.jinja"],
+      "matrix_im2col_nhwc": ["qlinear-conv-im2col-nhwc.wgsl.jinja", "quant-exact-matrix.wgsl.jinja", "quant-weight-k-major.wgsl.jinja"],
+      "dp4a_im2col_nhwc_bias": ["qlinear-conv-im2col-nhwc.wgsl.jinja", "quant-dp4a-matmul.wgsl.jinja", "quant-weight-k-major.wgsl.jinja"],
+      "matrix_im2col_nhwc_bias": ["qlinear-conv-im2col-nhwc.wgsl.jinja", "quant-exact-matrix.wgsl.jinja", "quant-weight-k-major.wgsl.jinja"],
+      "precast_f16_pointwise_nchw_tensor": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f16_pointwise_nchw_tensor": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f16_pointwise_nchw_row": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f16_pointwise_nchw_row": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f16_im2col_nchw_tensor": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f16_im2col_nchw_tensor": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f16_im2col_nchw_row": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f16_im2col_nchw_row": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f16_im2col_nchw_bias_tensor": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f16_im2col_nchw_bias_tensor": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f16_im2col_nchw_bias_row": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f16_im2col_nchw_bias_row": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f16_pointwise_nhwc_tensor": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f16_pointwise_nhwc_tensor": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f16_pointwise_nhwc_column": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f16_pointwise_nhwc_column": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f16_pointwise_nhwc_bias_tensor": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f16_pointwise_nhwc_bias_tensor": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f16_pointwise_nhwc_bias_column": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f16_pointwise_nhwc_bias_column": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f16_im2col_nhwc_tensor": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f16_im2col_nhwc_tensor": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f16_im2col_nhwc_column": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f16_im2col_nhwc_column": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f16_im2col_nhwc_bias_tensor": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f16_im2col_nhwc_bias_tensor": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f16_im2col_nhwc_bias_column": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f16_im2col_nhwc_bias_column": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f32_pointwise_nchw_tensor": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f32_pointwise_nchw_tensor": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f32_pointwise_nchw_row": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f32_pointwise_nchw_row": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f32_im2col_nchw_tensor": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f32_im2col_nchw_tensor": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f32_im2col_nchw_row": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f32_im2col_nchw_row": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f32_im2col_nchw_bias_tensor": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f32_im2col_nchw_bias_tensor": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f32_im2col_nchw_bias_row": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f32_im2col_nchw_bias_row": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f32_pointwise_nhwc_tensor": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f32_pointwise_nhwc_tensor": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f32_pointwise_nhwc_column": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f32_pointwise_nhwc_column": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f32_pointwise_nhwc_bias_tensor": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f32_pointwise_nhwc_bias_tensor": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f32_pointwise_nhwc_bias_column": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f32_pointwise_nhwc_bias_column": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f32_im2col_nhwc_tensor": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f32_im2col_nhwc_tensor": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f32_im2col_nhwc_column": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f32_im2col_nhwc_column": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f32_im2col_nhwc_bias_tensor": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f32_im2col_nhwc_bias_tensor": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "precast_f32_im2col_nhwc_bias_column": ["quant-exact-matrix.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"],
+      "portable_f32_im2col_nhwc_bias_column": ["quant-exact-portable.wgsl.jinja", "quant-exact-prepare.wgsl.jinja"]
+    }
+  }
+}

build/webgpu/qlinear-conv-im2col-nhwc.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,82 @@

+// Widened-int8 im2col for a channels-last quantized convolution. The materialized
+// A operand is `[batch, outSpatial, KPadded]` with K enumerating `(kh, kw, ic)`,
+// so the input channel is the fastest axis of both the source activation and the
+// reduction and every gather is a contiguous run. Spatial taps outside the input
+// and the K tail hold the input zero point, so they contribute exactly zero after
+// centering.
+//
+// KPadded is a multiple of four by construction, so the store is one vec4 word per
+// K quadruple, and each invocation writes `WORDS_PER_LANE` of them strided by the
+// lane count, which keeps neighbouring invocations on neighbouring words while
+// giving each lane several stores in flight. When the input channel count is itself
+// a multiple of four a quadruple never crosses a kernel tap, and the gather is the
+// matching aligned vec4 load.
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) { return; }{% endmacro %}
+{{ env.wgsl.resourceDeclarations }}
+const WORDS_PER_LANE: u32 = {{ im2colWords }}u;
+const KERNEL_W: u32 = {{ kernelWSpec }}u;
+const STRIDE_H: u32 = {{ strideHSpec }}u;
+const STRIDE_W: u32 = {{ strideWSpec }}u;
+const DILATION_H: u32 = {{ dilationHSpec }}u;
+const DILATION_W: u32 = {{ dilationWSpec }}u;
+const PAD_TOP: i32 = {{ padTopSpec }};
+const PAD_LEFT: i32 = {{ padLeftSpec }};
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "lane", guardInline=true) }}
+  let zp = x_zero_point[0];
+  for (var i = 0u; i < WORDS_PER_LANE; i = i + 1u) {
+    let word = lane + i * params.count;
+    if (word >= params.totalWords) { break; }
+    let w = word % params.wordsPerRow;
+    let t = word / params.wordsPerRow;
+    let pos = t % params.outCount;
+    let batch = t / params.outCount;
+    let oh = pos / params.outW;
+    let ow = pos - oh * params.outW;
+    let k0 = w * 4u;
+    var value = vec4<{{ xScalar }}>(zp, zp, zp, zp);
+{% if channelsAligned %}
+    // A quadruple of K entries stays inside one kernel tap, so one bounds test and
+    // one aligned vec4 load cover all four input channels.
+    if (k0 < params.kRows) {
+      let ic = k0 % params.inChannels;
+      let khw = k0 / params.inChannels;
+      let kw = khw % KERNEL_W;
+      let kh = khw / KERNEL_W;
+      let ih = i32(oh * STRIDE_H + kh * DILATION_H) - PAD_TOP;
+      let iw = i32(ow * STRIDE_W + kw * DILATION_W) - PAD_LEFT;
+      if (ih >= 0 && ih < i32(params.inH) && iw >= 0 && iw < i32(params.inW)) {
+        value = x[(((batch * params.inH + u32(ih)) * params.inW + u32(iw)) * params.inChannels + ic) / 4u];
+      }
+    }
+{% else %}
+    // An unaligned channel count lets a quadruple straddle two kernel taps and the
+    // K tail, so each lane of the word resolves its own tap.
+{% for j in range(4) %}
+    {
+      let k = k0 + {{ j }}u;
+      if (k < params.kRows) {
+        let ic = k % params.inChannels;
+        let khw = k / params.inChannels;
+        let kw = khw % KERNEL_W;
+        let kh = khw / KERNEL_W;
+        let ih = i32(oh * STRIDE_H + kh * DILATION_H) - PAD_TOP;
+        let iw = i32(ow * STRIDE_W + kw * DILATION_W) - PAD_LEFT;
+        if (ih >= 0 && ih < i32(params.inH) && iw >= 0 && iw < i32(params.inW)) {
+          value[{{ j }}] = x[((batch * params.inH + u32(ih)) * params.inW + u32(iw)) * params.inChannels + ic];
+        }
+      }
+    }
+{% endfor %}
+{% endif %}
+    cols[word] = value;
+  }
+}

build/webgpu/qlinear-conv-nchw-x4.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,384 @@

+{% macro quant_read_scalar_zero(fn_name, buffer_name, is_unsigned, present=true) %}
+fn {{ fn_name }}() -> i32 {
+  {% if is_unsigned %}
+  return i32({{ buffer_name }}[0]);
+  {% else %}
+  return {{ buffer_name }}[0];
+  {% endif %}
+}{% endmacro %}
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% endmacro %}
+{{ env.wgsl.resourceDeclarations }}
+// Rare-boundary fallback: reproduce the reference's binary64 scale division
+// and accumulator product using integer significands; no f64 feature required.
+fn quant_u96_add(a: vec3<u32>, b: vec3<u32>) -> vec3<u32> {
+  let low = a.x + b.x;
+  let middle = a.y + b.y;
+  let with_carry = middle + u32(low < a.x);
+  return vec3<u32>(low, with_carry, a.z + b.z + u32(middle < a.y) + u32(with_carry < middle));
+}
+fn quant_u96_shr(a: vec3<u32>) -> vec3<u32> {
+  return vec3<u32>((a.x >> 1u) | (a.y << 31u), (a.y >> 1u) | (a.z << 31u), a.z >> 1u);
+}
+fn quant_u96_round_shift(value: vec3<u32>, shift: u32) -> vec3<u32> {
+  var q = value;
+  var sticky = false;
+  var guard = false;
+  for (var bit = 0u; bit < shift; bit++) {
+    sticky = sticky || guard;
+    guard = (q.x & 1u) != 0u;
+    q = quant_u96_shr(q);
+  }
+  if (guard && (sticky || (q.x & 1u) != 0u)) {
+    q = quant_u96_add(q, vec3<u32>(1u, 0u, 0u));
+  }
+  return q;
+}
+fn quant_float_significand(bits: u32) -> u32 {
+  return (bits & 0x7fffffu) | select(0x800000u, 0u, (bits & 0x7f800000u) == 0u);
+}
+fn quant_float_exponent(bits: u32) -> i32 {
+  return max(1, i32((bits >> 23u) & 255u)) - 150;
+}
+fn quant_round_binary64(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  let a = quant_float_significand(a_bits);
+  let b = quant_float_significand(b_bits);
+  let denominator = quant_float_significand(y_bits);
+  if (a == 0u || b == 0u) { return 0; }
+  // The exact product of two 24-bit significands occupies at most 48 bits.
+  let p0 = (a & 65535u) * (b & 65535u);
+  let cross = (a >> 16u) * (b & 65535u) + (b >> 16u) * (a & 65535u);
+  let product_low = p0 + (cross << 16u);
+  let product_high = (a >> 16u) * (b >> 16u) + (cross >> 16u) + u32(product_low < p0);
+  var input_bit = select(31 - i32(countLeadingZeros(product_low)), 63 - i32(countLeadingZeros(product_high)), product_high != 0u);
+  var remainder = 0u;
+  var significant_bits = 0u;
+  var leading_exponent = 0;
+  var significand = vec2<u32>(0u);
+  // Long division generates the 53 bits of the binary64 multiplier. The
+  // remainder fits in 25 bits, including its left shift. At most 24 leading
+  // zero bits precede the first significant bit.
+  loop {
+    var incoming = 0u;
+    if (input_bit >= 32) { incoming = (product_high >> u32(input_bit - 32)) & 1u; }
+    else if (input_bit >= 0) { incoming = (product_low >> u32(input_bit)) & 1u; }
+    remainder = (remainder << 1u) | incoming;
+    let digit = u32(remainder >= denominator);
+    if (digit != 0u) { remainder -= denominator; }
+    if (significant_bits != 0u || digit != 0u) {
+      if (significant_bits == 0u) { leading_exponent = input_bit; }
+      significand = vec2<u32>((significand.x << 1u) | digit, (significand.y << 1u) | (significand.x >> 31u));
+      significant_bits++;
+    }
+    input_bit--;
+    if (significant_bits == 53u) { break; }
+  }
+  let twice_remainder = remainder << 1u;
+  if (twice_remainder > denominator || (twice_remainder == denominator && (significand.x & 1u) != 0u)) {
+    let low = significand.x + 1u;
+    significand = vec2<u32>(low, significand.y + u32(low == 0u));
+  }
+  let scale_exponent = quant_float_exponent(a_bits) + quant_float_exponent(b_bits) - quant_float_exponent(y_bits) + leading_exponent - 52;
+  // abs(INT_MIN) is representable as u32. Multiply the rounded multiplier's
+  // significand by the exact accumulator into three little-endian u32 limbs.
+  var magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  var operand = vec3<u32>(significand, 0u);
+  var product = vec3<u32>(0u);
+  loop {
+    if ((magnitude & 1u) != 0u) { product = quant_u96_add(product, operand); }
+    magnitude >>= 1u;
+    if (magnitude == 0u) { break; }
+    operand = vec3<u32>(operand.x << 1u, (operand.y << 1u) | (operand.x >> 31u), (operand.z << 1u) | (operand.y >> 31u));
+  }
+  var product_top = 31 - i32(countLeadingZeros(product.x));
+  if (product.y != 0u) { product_top = 63 - i32(countLeadingZeros(product.y)); }
+  if (product.z != 0u) { product_top = 95 - i32(countLeadingZeros(product.z)); }
+  if (product_top + scale_exponent >= 10) { return select(1024, -1024, acc < 0); }
+  if (product_top + scale_exponent < -2) { return 0; }
+  let product_shift = u32(max(0, product_top - 52));
+  let rounded_product = quant_u96_round_shift(product, product_shift);
+  let result_exponent = scale_exponent + i32(product_shift);
+  var rounded = rounded_product;
+  if (result_exponent < 0) { rounded = quant_u96_round_shift(rounded_product, u32(-result_exponent)); }
+  else { rounded.x <<= u32(result_exponent); }
+  return select(i32(rounded.x), -i32(rounded.x), acc < 0);
+}
+// A power-of-two division and an exactly represented <=48-bit product need
+// no binary64 boundary resolver. This includes ordinary dyadic quantization.
+fn quant_exact_dyadic_product(acc: i32, scale: f32) -> bool {
+  let magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  let significand = quant_float_significand(bitcast<u32>(scale));
+  let integer_bits = 32u - countLeadingZeros(magnitude) - countTrailingZeros(magnitude);
+  let scale_bits = 32u - countLeadingZeros(significand) - countTrailingZeros(significand);
+  return integer_bits + scale_bits <= 48u;
+}
+fn quant_normal_expansion_range(v: f32) -> bool {
+  let exponent = (bitcast<u32>(v) >> 23u) & 255u;
+  // Reserve two binary32 significands plus four carry/error bits above the
+  // minimum normal exponent, so residual arithmetic avoids FTZ underflow.
+  return exponent >= (1u + 2u * 24u + 4u) && exponent < 255u;
+}
+fn round_expansion_ratio(numerator_hi: f32, numerator_lo: f32, denominator: f32, dyadic_scale: f32, dyadic_candidate: bool, acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  let estimate = numerator_hi / denominator;
+  // Anything outside this range saturates after adding an 8-bit zero point;
+  // returning early also keeps floor/i32 conversion away from Inf and overflow.
+  if (estimate <= -1024.0) { return -1024; }
+  if (estimate >= 1024.0) { return 1024; }
+  let lo = floor(estimate);
+  let hi = lo + 1.0;
+  let midpoint = lo + 0.5;
+  let midpoint_hi = fma(midpoint, denominator, 0.0);
+  if ((bitcast<u32>(midpoint_hi) & 0x7f800000u) == 0x7f800000u) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let midpoint_lo = fma(midpoint, denominator, -midpoint_hi);
+  // The high-word subtraction is exact for these similarly sized products
+  // (Sterbenz); adding the residual difference therefore preserves its sign
+  // even when the rounded high words straddle adjacent f32 values.
+  let midpoint_delta = fma(1.0, numerator_hi - midpoint_hi, numerator_lo - midpoint_lo);
+  // Escalate ambiguous expansion signs instead of changing the tie tolerance.
+  // Six FMA/product tails and their additions fit within this conservative
+  // 16*u^2 error envelope (u=2^-24), including binary64 reference rounding.
+  if ((bitcast<u32>(midpoint_delta) & 0x7f800000u) == 0x7f800000u) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  if (abs(midpoint_delta) <= max(abs(numerator_hi), abs(midpoint_hi)) * 0x1p-44) {
+    if (!dyadic_candidate || !quant_exact_dyadic_product(acc, dyadic_scale)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  }
+  if (midpoint_delta < 0.0) {
+    return i32(lo);
+  }
+  if (midpoint_delta > 0.0) {
+    return i32(hi);
+  }
+  let half_lo = floor(lo * 0.5);
+  let lo_is_even = lo - half_lo * 2.0 == 0.0;
+  return i32(select(hi, lo, lo_is_even));
+}
+// Multiplication by a power of two keeps every i32 bit in binary64. Round the
+// exact dyadic value directly, including INT_MIN, half-even ties and saturation.
+fn quant_round_power_two(acc: i32, exponent: i32) -> i32 {
+  let magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  var rounded = 0u;
+  if (exponent >= 0) {
+    if (exponent >= 10) { return select(1024, -1024, acc < 0); }
+    rounded = min(magnitude, 1024u >> u32(exponent)) << u32(exponent);
+  } else {
+    let shift = u32(-exponent);
+    if (shift >= 32u) { return 0; }
+    let whole = magnitude >> shift;
+    let remainder = magnitude & ((1u << shift) - 1u);
+    let half = 1u << (shift - 1u);
+    rounded = min(1024u, whole + u32(remainder > half || (remainder == half && (whole & 1u) != 0u)));
+  }
+  return select(i32(rounded), -i32(rounded), acc < 0);
+}
+fn round_requantized_precise(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  if (((a_bits | b_bits | y_bits) & 0x807fffffu) == 0u && min(a_bits, min(b_bits, y_bits)) != 0u && max(a_bits, max(b_bits, y_bits)) < 0x7f800000u) {
+    let exponent = i32(a_bits >> 23u) + i32(b_bits >> 23u) - i32(y_bits >> 23u) - 127;
+    return quant_round_power_two(acc, exponent);
+  }
+  let acc_f = f32(acc);
+  // Every i32 high part aligned to 256 is exactly representable in f32.
+  // Recover the conversion residual without an out-of-range i32(acc_f)
+  // conversion when INT_MAX rounds up to 2^31.
+  let acc_tail = fma(-1.0, acc_f - f32(acc & -256), f32(acc & 255));
+  // Equal b/y scales cancel exactly. Bypass their product and division so a
+  // rounded intermediate cannot move a near-half value onto the tie.
+  if (bitcast<u32>(b_s) == bitcast<u32>(y_s)) {
+    let direct_hi = fma(acc_f, a_s, 0.0);
+    let direct_lo = fma(acc_tail, a_s, fma(acc_f, a_s, -direct_hi));
+    if (!quant_normal_expansion_range(direct_hi)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+    return round_expansion_ratio(direct_hi, direct_lo, 1.0, a_s, true, acc, a_s, b_s, y_s);
+  }
+  // Keep the scale product as a two-float expansion. Its FMA residual preserves
+  // which side of a rounding midpoint a near-half value occupies.
+  let ab_hi = fma(a_s, b_s, 0.0);
+  let ab_lo = fma(a_s, b_s, -ab_hi);
+  if (!quant_normal_expansion_range(ab_hi) || !quant_normal_expansion_range(y_s) || y_s > 0x1p126) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let product_hi = fma(acc_f, ab_hi, 0.0);
+  if (!quant_normal_expansion_range(product_hi)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let rounded_acc_tail = fma(acc_f, ab_hi, -product_hi) + fma(acc_f, ab_lo, 0.0);
+  let product_tail = fma(acc_tail, ab_hi, rounded_acc_tail) + fma(acc_tail, ab_lo, 0.0);
+  // Compare the expansion directly. Renormalizing it adds a cancelling
+  // subtraction/addition chain that can erase the residual during lowering.
+  let denominator_bits = bitcast<u32>(y_s);
+  let denominator_power_two = (denominator_bits & 0x7fffffu) == 0u;
+  let dyadic_candidate = denominator_power_two && ab_lo == 0.0;
+  return round_expansion_ratio(product_hi, product_tail, y_s, ab_hi, dyadic_candidate, acc, a_s, b_s, y_s);
+}
+// A conservative binary32 error filter accepts only a result safely inside one
+// integer rounding bin. Ambiguous boundaries keep the full expansion/binary64
+// path; this is an exact-result shortcut, not an approximate quantization mode.
+fn round_requantized(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  let ea = (a_bits >> 23u) & 255u;
+  let eb = (b_bits >> 23u) & 255u;
+  let ey = (y_bits >> 23u) & 255u;
+  // Certify the range before arithmetic: overflow need not produce infinity.
+  // Any nonzero i32 converts to a float with biased exponent <=158. These
+  // scale-only bounds cover every accumulator and can be shared across lanes.
+  // Reserve two significands plus carry bits above FTZ, and a factor of two
+  // below overflow for products, division and the expansion midpoint product.
+  if (((a_bits | b_bits | y_bits) & 0x80000000u) != 0u ||
+      min(ea, eb) == 0u || max(ea, eb) >= 255u ||
+      ea + eb < 180u || ea + eb > 347u ||
+      ey < 53u || ey > 242u || ea + eb > ey + 220u) {
+    return quant_round_binary64(acc, a_s, b_s, y_s);
+  }
+  let scale = fma(a_s, b_s, 0.0);
+  let product = fma(f32(acc), scale, 0.0);
+  let estimate = product / y_s;
+  if (estimate <= -1024.0) { return -1024; }
+  if (estimate >= 1024.0) { return 1024; }
+  let nearest = round(estimate);
+  // 32 binary32 unit roundoffs reserve room for the i32 conversion, two
+  // products, quotient, and distance arithmetic, plus binary64 rounding.
+  let margin = max(1.0, abs(estimate)) * 0x1p-19;
+  if (abs(estimate - nearest) < 0.5 - margin) { return i32(nearest); }
+{% if device.adapterInfo.vendor != "apple" %}
+  // A rounded half-integer estimate can hide which side the exact f32-scale
+  // ratio lies on. Resolve that one-ulp neighborhood conservatively outside
+  // the compact fast path; exact power-of-two ratios already have a cheap
+  // and correct path in round_requantized_precise.
+  if (abs(abs(estimate - nearest) - 0.5) <= max(1.0, abs(estimate)) * 0x1p-24) {
+    let dyadic_scales = ((a_bits | b_bits | y_bits) & 0x807fffffu) == 0u &&
+      min(a_bits, min(b_bits, y_bits)) != 0u &&
+      max(a_bits, max(b_bits, y_bits)) < 0x7f800000u;
+    if (!dyadic_scales) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  }
+{% endif %}
+  return round_requantized_precise(acc, a_s, b_s, y_s);
+}
+{{ quant_read_scalar_zero("x_zero", "x_zero_point", xUnsigned) }}
+fn w_zero({% if wZeroPerChannel %}oc: u32{% endif %}) -> i32 {
+{% if wZeroPerChannel %}
+  let i = oc;
+{% else %}
+  let i = 0u;
+{% endif %}
+{% if wUnsigned %}
+  return i32(w_zero_point[i]);
+{% else %}
+  return w_zero_point[i];
+{% endif %}
+}
+fn centered_x(i: u32, zp: i32) -> i32 {
+{% if xUnsigned %}
+  return i32(x[i]) - zp;
+{% else %}
+  return x[i] - zp;
+{% endif %}
+}
+fn centered_w(i: u32, zp: i32) -> i32 {
+{% if wUnsigned %}
+  return i32(w[i]) - zp;
+{% else %}
+  return w[i] - zp;
+{% endif %}
+}
+{{ quant_read_scalar_zero("y_zero", "y_zero_point", yUnsigned) }}
+fn requant(acc: i32{% if wScalePerChannel %}, oc: u32{% endif %}) -> i32 {
+{% if wScalePerChannel %}
+  let ws = w_scale[oc];
+{% else %}
+  let ws = w_scale[0];
+{% endif %}
+  return clamp(round_requantized(acc, x_scale[0], ws, y_scale[0]) + y_zero(), {{ qMin }}, {{ qMax }});
+}
+fn write_y(i: u32, value: i32) {
+{% if yUnsigned %}
+  y[i] = u32(value);
+{% else %}
+  y[i] = value;
+{% endif %}
+}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "flat", "") }}
+  let groups_per_row = (params.outW + 3u) / 4u;
+  let oc_blocks = (params.outChannels + {{ ocTile - 1 }}u) / {{ ocTile }}u;
+  let group_count = params.batch * oc_blocks * params.outH * groups_per_row;
+  if (flat >= group_count) { return; }
+  let ow0 = (flat % groups_per_row) * 4u;
+  let row = flat / groups_per_row;
+  let oh = row % params.outH;
+  let oc0 = ((row / params.outH) % oc_blocks) * {{ ocTile }}u;
+  let n = row / (params.outH * oc_blocks);
+  let group = oc0 / params.outChannelsPerGroup;
+  let in_start = group * params.inChannelsPerGroup;
+  let xzp = x_zero();
+{% for ot in range(ocTile) %}
+  var acc{{ ot }} = vec4<i32>(0);
+{% endfor %}
+  for (var icg = 0u; icg < params.inChannelsPerGroup; icg = icg + 1u) {
+    let ic = in_start + icg;
+    for (var kh = 0u; kh < params.kernelH; kh = kh + 1u) {
+      let ih = i32(oh * params.strideH + kh * params.dilationH) - params.padH;
+      for (var kw = 0u; kw < params.kernelW; kw = kw + 1u) {
+        var x0 = 0i;
+        var x1 = 0i;
+        var x2 = 0i;
+        var x3 = 0i;
+        if (ih >= 0 && ih < i32(params.inH)) {
+          let iw0 = i32(ow0 * params.strideW + kw * params.dilationW) - params.padW;
+          let iw1 = i32((ow0 + 1u) * params.strideW + kw * params.dilationW) - params.padW;
+          let iw2 = i32((ow0 + 2u) * params.strideW + kw * params.dilationW) - params.padW;
+          let iw3 = i32((ow0 + 3u) * params.strideW + kw * params.dilationW) - params.padW;
+          let base = (n * params.inChannels + ic) * params.inH * params.inW + u32(ih) * params.inW;
+          if (ow0 < params.outW && iw0 >= 0 && iw0 < i32(params.inW)) { x0 = centered_x(base + u32(iw0), xzp); }
+          if (ow0 + 1u < params.outW && iw1 >= 0 && iw1 < i32(params.inW)) { x1 = centered_x(base + u32(iw1), xzp); }
+          if (ow0 + 2u < params.outW && iw2 >= 0 && iw2 < i32(params.inW)) { x2 = centered_x(base + u32(iw2), xzp); }
+          if (ow0 + 3u < params.outW && iw3 >= 0 && iw3 < i32(params.inW)) { x3 = centered_x(base + u32(iw3), xzp); }
+        }
+        let xv = vec4<i32>(x0, x1, x2, x3);
+{% for ot in range(ocTile) %}
+        let oc{{ ot }} = oc0 + {{ ot }}u;
+        if (oc{{ ot }} < params.outChannels) {
+          let wi{{ ot }} = (((oc{{ ot }} * params.weightInChannels + icg) * params.kernelH + kh) * params.kernelW + kw);
+          let wc{{ ot }} = centered_w(wi{{ ot }}, w_zero({% if wZeroPerChannel %}oc{{ ot }}{% endif %}));
+          acc{{ ot }} = acc{{ ot }} + xv * vec4<i32>(wc{{ ot }});
+        }
+{% endfor %}
+      }
+    }
+  }
+{% for ot in range(ocTile) %}
+  let out_oc{{ ot }} = oc0 + {{ ot }}u;
+  if (out_oc{{ ot }} < params.outChannels) {
+    let out_base{{ ot }} = ((n * params.outChannels + out_oc{{ ot }}) * params.outH + oh) * params.outW + ow0;
+    write_y(out_base{{ ot }}, requant(acc{{ ot }}.x{% if wScalePerChannel %}, out_oc{{ ot }}{% endif %}));
+    if (ow0 + 1u < params.outW) { write_y(out_base{{ ot }} + 1u, requant(acc{{ ot }}.y{% if wScalePerChannel %}, out_oc{{ ot }}{% endif %})); }
+    if (ow0 + 2u < params.outW) { write_y(out_base{{ ot }} + 2u, requant(acc{{ ot }}.z{% if wScalePerChannel %}, out_oc{{ ot }}{% endif %})); }
+    if (ow0 + 3u < params.outW) { write_y(out_base{{ ot }} + 3u, requant(acc{{ ot }}.w{% if wScalePerChannel %}, out_oc{{ ot }}{% endif %})); }
+  }
+{% endfor %}
+}

build/webgpu/qlinear-conv-nhwc-oc4.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,394 @@

+{% macro quant_read_scalar_zero(fn_name, buffer_name, is_unsigned, present=true) %}
+fn {{ fn_name }}() -> i32 {
+  {% if is_unsigned %}
+  return i32({{ buffer_name }}[0]);
+  {% else %}
+  return {{ buffer_name }}[0];
+  {% endif %}
+}{% endmacro %}
+// Channels-last quantized 2-D convolution without integer dot-product support.
+// One invocation owns four adjacent output channels (a vec4 accumulator whose
+// store is one contiguous run in `(N, H, W, C)`) across `OW_TILE` adjacent output
+// columns, so each weight quadruple is reused by every column in the tile. The
+// grouped arm broadcasts one activation across the quadruple; the depthwise arm
+// (one input and one output channel per group) reads four adjacent activation
+// channels instead, which is the same contiguous run as the store.
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) { return; }{% endmacro %}
+{{ env.wgsl.resourceDeclarations }}
+// Rare-boundary fallback: reproduce the reference's binary64 scale division
+// and accumulator product using integer significands; no f64 feature required.
+fn quant_u96_add(a: vec3<u32>, b: vec3<u32>) -> vec3<u32> {
+  let low = a.x + b.x;
+  let middle = a.y + b.y;
+  let with_carry = middle + u32(low < a.x);
+  return vec3<u32>(low, with_carry, a.z + b.z + u32(middle < a.y) + u32(with_carry < middle));
+}
+fn quant_u96_shr(a: vec3<u32>) -> vec3<u32> {
+  return vec3<u32>((a.x >> 1u) | (a.y << 31u), (a.y >> 1u) | (a.z << 31u), a.z >> 1u);
+}
+fn quant_u96_round_shift(value: vec3<u32>, shift: u32) -> vec3<u32> {
+  var q = value;
+  var sticky = false;
+  var guard = false;
+  for (var bit = 0u; bit < shift; bit++) {
+    sticky = sticky || guard;
+    guard = (q.x & 1u) != 0u;
+    q = quant_u96_shr(q);
+  }
+  if (guard && (sticky || (q.x & 1u) != 0u)) {
+    q = quant_u96_add(q, vec3<u32>(1u, 0u, 0u));
+  }
+  return q;
+}
+fn quant_float_significand(bits: u32) -> u32 {
+  return (bits & 0x7fffffu) | select(0x800000u, 0u, (bits & 0x7f800000u) == 0u);
+}
+fn quant_float_exponent(bits: u32) -> i32 {
+  return max(1, i32((bits >> 23u) & 255u)) - 150;
+}
+fn quant_round_binary64(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  let a = quant_float_significand(a_bits);
+  let b = quant_float_significand(b_bits);
+  let denominator = quant_float_significand(y_bits);
+  if (a == 0u || b == 0u) { return 0; }
+  // The exact product of two 24-bit significands occupies at most 48 bits.
+  let p0 = (a & 65535u) * (b & 65535u);
+  let cross = (a >> 16u) * (b & 65535u) + (b >> 16u) * (a & 65535u);
+  let product_low = p0 + (cross << 16u);
+  let product_high = (a >> 16u) * (b >> 16u) + (cross >> 16u) + u32(product_low < p0);
+  var input_bit = select(31 - i32(countLeadingZeros(product_low)), 63 - i32(countLeadingZeros(product_high)), product_high != 0u);
+  var remainder = 0u;
+  var significant_bits = 0u;
+  var leading_exponent = 0;
+  var significand = vec2<u32>(0u);
+  // Long division generates the 53 bits of the binary64 multiplier. The
+  // remainder fits in 25 bits, including its left shift. At most 24 leading
+  // zero bits precede the first significant bit.
+  loop {
+    var incoming = 0u;
+    if (input_bit >= 32) { incoming = (product_high >> u32(input_bit - 32)) & 1u; }
+    else if (input_bit >= 0) { incoming = (product_low >> u32(input_bit)) & 1u; }
+    remainder = (remainder << 1u) | incoming;
+    let digit = u32(remainder >= denominator);
+    if (digit != 0u) { remainder -= denominator; }
+    if (significant_bits != 0u || digit != 0u) {
+      if (significant_bits == 0u) { leading_exponent = input_bit; }
+      significand = vec2<u32>((significand.x << 1u) | digit, (significand.y << 1u) | (significand.x >> 31u));
+      significant_bits++;
+    }
+    input_bit--;
+    if (significant_bits == 53u) { break; }
+  }
+  let twice_remainder = remainder << 1u;
+  if (twice_remainder > denominator || (twice_remainder == denominator && (significand.x & 1u) != 0u)) {
+    let low = significand.x + 1u;
+    significand = vec2<u32>(low, significand.y + u32(low == 0u));
+  }
+  let scale_exponent = quant_float_exponent(a_bits) + quant_float_exponent(b_bits) - quant_float_exponent(y_bits) + leading_exponent - 52;
+  // abs(INT_MIN) is representable as u32. Multiply the rounded multiplier's
+  // significand by the exact accumulator into three little-endian u32 limbs.
+  var magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  var operand = vec3<u32>(significand, 0u);
+  var product = vec3<u32>(0u);
+  loop {
+    if ((magnitude & 1u) != 0u) { product = quant_u96_add(product, operand); }
+    magnitude >>= 1u;
+    if (magnitude == 0u) { break; }
+    operand = vec3<u32>(operand.x << 1u, (operand.y << 1u) | (operand.x >> 31u), (operand.z << 1u) | (operand.y >> 31u));
+  }
+  var product_top = 31 - i32(countLeadingZeros(product.x));
+  if (product.y != 0u) { product_top = 63 - i32(countLeadingZeros(product.y)); }
+  if (product.z != 0u) { product_top = 95 - i32(countLeadingZeros(product.z)); }
+  if (product_top + scale_exponent >= 10) { return select(1024, -1024, acc < 0); }
+  if (product_top + scale_exponent < -2) { return 0; }
+  let product_shift = u32(max(0, product_top - 52));
+  let rounded_product = quant_u96_round_shift(product, product_shift);
+  let result_exponent = scale_exponent + i32(product_shift);
+  var rounded = rounded_product;
+  if (result_exponent < 0) { rounded = quant_u96_round_shift(rounded_product, u32(-result_exponent)); }
+  else { rounded.x <<= u32(result_exponent); }
+  return select(i32(rounded.x), -i32(rounded.x), acc < 0);
+}
+// A power-of-two division and an exactly represented <=48-bit product need
+// no binary64 boundary resolver. This includes ordinary dyadic quantization.
+fn quant_exact_dyadic_product(acc: i32, scale: f32) -> bool {
+  let magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  let significand = quant_float_significand(bitcast<u32>(scale));
+  let integer_bits = 32u - countLeadingZeros(magnitude) - countTrailingZeros(magnitude);
+  let scale_bits = 32u - countLeadingZeros(significand) - countTrailingZeros(significand);
+  return integer_bits + scale_bits <= 48u;
+}
+fn quant_normal_expansion_range(v: f32) -> bool {
+  let exponent = (bitcast<u32>(v) >> 23u) & 255u;
+  // Reserve two binary32 significands plus four carry/error bits above the
+  // minimum normal exponent, so residual arithmetic avoids FTZ underflow.
+  return exponent >= (1u + 2u * 24u + 4u) && exponent < 255u;
+}
+fn round_expansion_ratio(numerator_hi: f32, numerator_lo: f32, denominator: f32, dyadic_scale: f32, dyadic_candidate: bool, acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  let estimate = numerator_hi / denominator;
+  // Anything outside this range saturates after adding an 8-bit zero point;
+  // returning early also keeps floor/i32 conversion away from Inf and overflow.
+  if (estimate <= -1024.0) { return -1024; }
+  if (estimate >= 1024.0) { return 1024; }
+  let lo = floor(estimate);
+  let hi = lo + 1.0;
+  let midpoint = lo + 0.5;
+  let midpoint_hi = fma(midpoint, denominator, 0.0);
+  if ((bitcast<u32>(midpoint_hi) & 0x7f800000u) == 0x7f800000u) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let midpoint_lo = fma(midpoint, denominator, -midpoint_hi);
+  // The high-word subtraction is exact for these similarly sized products
+  // (Sterbenz); adding the residual difference therefore preserves its sign
+  // even when the rounded high words straddle adjacent f32 values.
+  let midpoint_delta = fma(1.0, numerator_hi - midpoint_hi, numerator_lo - midpoint_lo);
+  // Escalate ambiguous expansion signs instead of changing the tie tolerance.
+  // Six FMA/product tails and their additions fit within this conservative
+  // 16*u^2 error envelope (u=2^-24), including binary64 reference rounding.
+  if ((bitcast<u32>(midpoint_delta) & 0x7f800000u) == 0x7f800000u) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  if (abs(midpoint_delta) <= max(abs(numerator_hi), abs(midpoint_hi)) * 0x1p-44) {
+    if (!dyadic_candidate || !quant_exact_dyadic_product(acc, dyadic_scale)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  }
+  if (midpoint_delta < 0.0) {
+    return i32(lo);
+  }
+  if (midpoint_delta > 0.0) {
+    return i32(hi);
+  }
+  let half_lo = floor(lo * 0.5);
+  let lo_is_even = lo - half_lo * 2.0 == 0.0;
+  return i32(select(hi, lo, lo_is_even));
+}
+// Multiplication by a power of two keeps every i32 bit in binary64. Round the
+// exact dyadic value directly, including INT_MIN, half-even ties and saturation.
+fn quant_round_power_two(acc: i32, exponent: i32) -> i32 {
+  let magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  var rounded = 0u;
+  if (exponent >= 0) {
+    if (exponent >= 10) { return select(1024, -1024, acc < 0); }
+    rounded = min(magnitude, 1024u >> u32(exponent)) << u32(exponent);
+  } else {
+    let shift = u32(-exponent);
+    if (shift >= 32u) { return 0; }
+    let whole = magnitude >> shift;
+    let remainder = magnitude & ((1u << shift) - 1u);
+    let half = 1u << (shift - 1u);
+    rounded = min(1024u, whole + u32(remainder > half || (remainder == half && (whole & 1u) != 0u)));
+  }
+  return select(i32(rounded), -i32(rounded), acc < 0);
+}
+fn round_requantized_precise(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  if (((a_bits | b_bits | y_bits) & 0x807fffffu) == 0u && min(a_bits, min(b_bits, y_bits)) != 0u && max(a_bits, max(b_bits, y_bits)) < 0x7f800000u) {
+    let exponent = i32(a_bits >> 23u) + i32(b_bits >> 23u) - i32(y_bits >> 23u) - 127;
+    return quant_round_power_two(acc, exponent);
+  }
+  let acc_f = f32(acc);
+  // Every i32 high part aligned to 256 is exactly representable in f32.
+  // Recover the conversion residual without an out-of-range i32(acc_f)
+  // conversion when INT_MAX rounds up to 2^31.
+  let acc_tail = fma(-1.0, acc_f - f32(acc & -256), f32(acc & 255));
+  // Equal b/y scales cancel exactly. Bypass their product and division so a
+  // rounded intermediate cannot move a near-half value onto the tie.
+  if (bitcast<u32>(b_s) == bitcast<u32>(y_s)) {
+    let direct_hi = fma(acc_f, a_s, 0.0);
+    let direct_lo = fma(acc_tail, a_s, fma(acc_f, a_s, -direct_hi));
+    if (!quant_normal_expansion_range(direct_hi)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+    return round_expansion_ratio(direct_hi, direct_lo, 1.0, a_s, true, acc, a_s, b_s, y_s);
+  }
+  // Keep the scale product as a two-float expansion. Its FMA residual preserves
+  // which side of a rounding midpoint a near-half value occupies.
+  let ab_hi = fma(a_s, b_s, 0.0);
+  let ab_lo = fma(a_s, b_s, -ab_hi);
+  if (!quant_normal_expansion_range(ab_hi) || !quant_normal_expansion_range(y_s) || y_s > 0x1p126) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let product_hi = fma(acc_f, ab_hi, 0.0);
+  if (!quant_normal_expansion_range(product_hi)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let rounded_acc_tail = fma(acc_f, ab_hi, -product_hi) + fma(acc_f, ab_lo, 0.0);
+  let product_tail = fma(acc_tail, ab_hi, rounded_acc_tail) + fma(acc_tail, ab_lo, 0.0);
+  // Compare the expansion directly. Renormalizing it adds a cancelling
+  // subtraction/addition chain that can erase the residual during lowering.
+  let denominator_bits = bitcast<u32>(y_s);
+  let denominator_power_two = (denominator_bits & 0x7fffffu) == 0u;
+  let dyadic_candidate = denominator_power_two && ab_lo == 0.0;
+  return round_expansion_ratio(product_hi, product_tail, y_s, ab_hi, dyadic_candidate, acc, a_s, b_s, y_s);
+}
+// A conservative binary32 error filter accepts only a result safely inside one
+// integer rounding bin. Ambiguous boundaries keep the full expansion/binary64
+// path; this is an exact-result shortcut, not an approximate quantization mode.
+fn round_requantized(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  let ea = (a_bits >> 23u) & 255u;
+  let eb = (b_bits >> 23u) & 255u;
+  let ey = (y_bits >> 23u) & 255u;
+  // Certify the range before arithmetic: overflow need not produce infinity.
+  // Any nonzero i32 converts to a float with biased exponent <=158. These
+  // scale-only bounds cover every accumulator and can be shared across lanes.
+  // Reserve two significands plus carry bits above FTZ, and a factor of two
+  // below overflow for products, division and the expansion midpoint product.
+  if (((a_bits | b_bits | y_bits) & 0x80000000u) != 0u ||
+      min(ea, eb) == 0u || max(ea, eb) >= 255u ||
+      ea + eb < 180u || ea + eb > 347u ||
+      ey < 53u || ey > 242u || ea + eb > ey + 220u) {
+    return quant_round_binary64(acc, a_s, b_s, y_s);
+  }
+  let scale = fma(a_s, b_s, 0.0);
+  let product = fma(f32(acc), scale, 0.0);
+  let estimate = product / y_s;
+  if (estimate <= -1024.0) { return -1024; }
+  if (estimate >= 1024.0) { return 1024; }
+  let nearest = round(estimate);
+  // 32 binary32 unit roundoffs reserve room for the i32 conversion, two
+  // products, quotient, and distance arithmetic, plus binary64 rounding.
+  let margin = max(1.0, abs(estimate)) * 0x1p-19;
+  if (abs(estimate - nearest) < 0.5 - margin) { return i32(nearest); }
+{% if device.adapterInfo.vendor != "apple" %}
+  // A rounded half-integer estimate can hide which side the exact f32-scale
+  // ratio lies on. Resolve that one-ulp neighborhood conservatively outside
+  // the compact fast path; exact power-of-two ratios already have a cheap
+  // and correct path in round_requantized_precise.
+  if (abs(abs(estimate - nearest) - 0.5) <= max(1.0, abs(estimate)) * 0x1p-24) {
+    let dyadic_scales = ((a_bits | b_bits | y_bits) & 0x807fffffu) == 0u &&
+      min(a_bits, min(b_bits, y_bits)) != 0u &&
+      max(a_bits, max(b_bits, y_bits)) < 0x7f800000u;
+    if (!dyadic_scales) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  }
+{% endif %}
+  return round_requantized_precise(acc, a_s, b_s, y_s);
+}
+const OW_TILE: u32 = {{ owTile }}u;
+fn centered_x(i: u32, zp: i32) -> i32 {
+{% if xUnsigned %}
+  return i32(x[i]) - zp;
+{% else %}
+  return x[i] - zp;
+{% endif %}
+}
+fn read_w(i: u32) -> i32 {
+{% if wUnsigned %}
+  return i32(w[i]);
+{% else %}
+  return w[i];
+{% endif %}
+}
+{{ quant_read_scalar_zero("x_zero", "x_zero_point", xUnsigned) }}
+fn w_zero({% if wZeroPerChannel %}oc: u32{% endif %}) -> i32 {
+{% if wZeroPerChannel %}
+  let i = oc;
+{% else %}
+  let i = 0u;
+{% endif %}
+{% if wUnsigned %}
+  return i32(w_zero_point[i]);
+{% else %}
+  return w_zero_point[i];
+{% endif %}
+}
+{{ quant_read_scalar_zero("y_zero", "y_zero_point", yUnsigned) }}
+fn requant(acc: i32{% if wScalePerChannel %}, oc: u32{% endif %}) -> i32 {
+{% if wScalePerChannel %}
+  let ws = w_scale[oc];
+{% else %}
+  let ws = w_scale[0];
+{% endif %}
+  return clamp(round_requantized(acc, x_scale[0], ws, y_scale[0]) + y_zero(), {{ qMin }}, {{ qMax }});
+}
+fn write_y(i: u32, value: i32) {
+{% if yUnsigned %}
+  y[i] = u32(value);
+{% else %}
+  y[i] = value;
+{% endif %}
+}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "flat", guardInline=true) }}
+  let ocb = flat % params.ocBlocks;
+  var rest = flat / params.ocBlocks;
+  let owg = rest % params.owGroups;
+  rest = rest / params.owGroups;
+  let oh = rest % params.outH;
+  let batch = rest / params.outH;
+  let oc0 = ocb * 4u;
+  let ow0 = owg * OW_TILE;
+  // Four adjacent output channels share one group whenever the group holds a
+  // multiple of four; the depthwise arm instead maps channel `oc0 + j` to input
+  // channel `oc0 + j`, which this same expression produces at unit group widths.
+  let in_start = (oc0 / params.outChannelsPerGroup) * params.inChannelsPerGroup;
+  let w_stride = params.weightInChannels * params.kernelH * params.kernelW;
+  let xzp = x_zero();
+  let wzp = vec4<i32>(w_zero({% if wZeroPerChannel %}oc0{% endif %}), w_zero({% if wZeroPerChannel %}oc0 + 1u{% endif %}), w_zero({% if wZeroPerChannel %}oc0 + 2u{% endif %}), w_zero({% if wZeroPerChannel %}oc0 + 3u{% endif %}));
+{% for ot in range(owTile) %}
+  var acc{{ ot }} = vec4<i32>(0);
+{% endfor %}
+  for (var kh = 0u; kh < params.kernelH; kh = kh + 1u) {
+    let ih = i32(oh * params.strideH + kh * params.dilationH) - params.padH;
+    if (ih < 0 || ih >= i32(params.inH)) { continue; }
+    let row_base = (batch * params.inH + u32(ih)) * params.inW;
+    for (var kw = 0u; kw < params.kernelW; kw = kw + 1u) {
+      for (var icg = 0u; icg < params.inChannelsPerGroup; icg = icg + 1u) {
+        let wb = ((oc0 * params.weightInChannels + icg) * params.kernelH + kh) * params.kernelW + kw;
+        let wv = vec4<i32>(read_w(wb), read_w(wb + w_stride), read_w(wb + 2u * w_stride), read_w(wb + 3u * w_stride)) - wzp;
+{% for ot in range(owTile) %}
+        {
+          let ow = ow0 + {{ ot }}u;
+          let iw = i32(ow * params.strideW + kw * params.dilationW) - params.padW;
+          if (ow < params.outW && iw >= 0 && iw < i32(params.inW)) {
+            let xb = (row_base + u32(iw)) * params.inChannels + in_start + icg;
+{% if depthwise %}
+            let xv = vec4<i32>(centered_x(xb, xzp), centered_x(xb + 1u, xzp), centered_x(xb + 2u, xzp), centered_x(xb + 3u, xzp));
+{% else %}
+            let xv = vec4<i32>(centered_x(xb, xzp));
+{% endif %}
+            acc{{ ot }} = acc{{ ot }} + xv * wv;
+          }
+        }
+{% endfor %}
+      }
+    }
+  }
+{% if hasBias | default(false) %}  let bias4 = vec4<i32>(bias[oc0], bias[oc0 + 1u], bias[oc0 + 2u], bias[oc0 + 3u]);
+{% for ot in range(owTile) %}  acc{{ ot }} = acc{{ ot }} + bias4;
+{% endfor %}{% endif %}{% for ot in range(owTile) %}
+  {
+    let ow = ow0 + {{ ot }}u;
+    if (ow < params.outW) {
+      let ob = ((batch * params.outH + oh) * params.outW + ow) * params.outChannels + oc0;
+      write_y(ob, requant(acc{{ ot }}.x{% if wScalePerChannel %}, oc0{% endif %}));
+      write_y(ob + 1u, requant(acc{{ ot }}.y{% if wScalePerChannel %}, oc0 + 1u{% endif %}));
+      write_y(ob + 2u, requant(acc{{ ot }}.z{% if wScalePerChannel %}, oc0 + 2u{% endif %}));
+      write_y(ob + 3u, requant(acc{{ ot }}.w{% if wScalePerChannel %}, oc0 + 3u{% endif %}));
+    }
+  }
+{% endfor %}
+}

build/webgpu/qlinear-conv-requantize.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,302 @@

+{% macro quant_read_scalar_zero(fn_name, buffer_name, is_unsigned, present=true) %}
+fn {{ fn_name }}() -> i32 {
+  {% if is_unsigned %}
+  return i32({{ buffer_name }}[0]);
+  {% else %}
+  return {{ buffer_name }}[0];
+  {% endif %}
+}{% endmacro %}
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) { return; }{% endmacro %}
+{{ env.wgsl.resourceDeclarations }}
+// Rare-boundary fallback: reproduce the reference's binary64 scale division
+// and accumulator product using integer significands; no f64 feature required.
+fn quant_u96_add(a: vec3<u32>, b: vec3<u32>) -> vec3<u32> {
+  let low = a.x + b.x;
+  let middle = a.y + b.y;
+  let with_carry = middle + u32(low < a.x);
+  return vec3<u32>(low, with_carry, a.z + b.z + u32(middle < a.y) + u32(with_carry < middle));
+}
+fn quant_u96_shr(a: vec3<u32>) -> vec3<u32> {
+  return vec3<u32>((a.x >> 1u) | (a.y << 31u), (a.y >> 1u) | (a.z << 31u), a.z >> 1u);
+}
+fn quant_u96_round_shift(value: vec3<u32>, shift: u32) -> vec3<u32> {
+  var q = value;
+  var sticky = false;
+  var guard = false;
+  for (var bit = 0u; bit < shift; bit++) {
+    sticky = sticky || guard;
+    guard = (q.x & 1u) != 0u;
+    q = quant_u96_shr(q);
+  }
+  if (guard && (sticky || (q.x & 1u) != 0u)) {
+    q = quant_u96_add(q, vec3<u32>(1u, 0u, 0u));
+  }
+  return q;
+}
+fn quant_float_significand(bits: u32) -> u32 {
+  return (bits & 0x7fffffu) | select(0x800000u, 0u, (bits & 0x7f800000u) == 0u);
+}
+fn quant_float_exponent(bits: u32) -> i32 {
+  return max(1, i32((bits >> 23u) & 255u)) - 150;
+}
+fn quant_round_binary64(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  let a = quant_float_significand(a_bits);
+  let b = quant_float_significand(b_bits);
+  let denominator = quant_float_significand(y_bits);
+  if (a == 0u || b == 0u) { return 0; }
+  // The exact product of two 24-bit significands occupies at most 48 bits.
+  let p0 = (a & 65535u) * (b & 65535u);
+  let cross = (a >> 16u) * (b & 65535u) + (b >> 16u) * (a & 65535u);
+  let product_low = p0 + (cross << 16u);
+  let product_high = (a >> 16u) * (b >> 16u) + (cross >> 16u) + u32(product_low < p0);
+  var input_bit = select(31 - i32(countLeadingZeros(product_low)), 63 - i32(countLeadingZeros(product_high)), product_high != 0u);
+  var remainder = 0u;
+  var significant_bits = 0u;
+  var leading_exponent = 0;
+  var significand = vec2<u32>(0u);
+  // Long division generates the 53 bits of the binary64 multiplier. The
+  // remainder fits in 25 bits, including its left shift. At most 24 leading
+  // zero bits precede the first significant bit.
+  loop {
+    var incoming = 0u;
+    if (input_bit >= 32) { incoming = (product_high >> u32(input_bit - 32)) & 1u; }
+    else if (input_bit >= 0) { incoming = (product_low >> u32(input_bit)) & 1u; }
+    remainder = (remainder << 1u) | incoming;
+    let digit = u32(remainder >= denominator);
+    if (digit != 0u) { remainder -= denominator; }
+    if (significant_bits != 0u || digit != 0u) {
+      if (significant_bits == 0u) { leading_exponent = input_bit; }
+      significand = vec2<u32>((significand.x << 1u) | digit, (significand.y << 1u) | (significand.x >> 31u));
+      significant_bits++;
+    }
+    input_bit--;
+    if (significant_bits == 53u) { break; }
+  }
+  let twice_remainder = remainder << 1u;
+  if (twice_remainder > denominator || (twice_remainder == denominator && (significand.x & 1u) != 0u)) {
+    let low = significand.x + 1u;
+    significand = vec2<u32>(low, significand.y + u32(low == 0u));
+  }
+  let scale_exponent = quant_float_exponent(a_bits) + quant_float_exponent(b_bits) - quant_float_exponent(y_bits) + leading_exponent - 52;
+  // abs(INT_MIN) is representable as u32. Multiply the rounded multiplier's
+  // significand by the exact accumulator into three little-endian u32 limbs.
+  var magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  var operand = vec3<u32>(significand, 0u);
+  var product = vec3<u32>(0u);
+  loop {
+    if ((magnitude & 1u) != 0u) { product = quant_u96_add(product, operand); }
+    magnitude >>= 1u;
+    if (magnitude == 0u) { break; }
+    operand = vec3<u32>(operand.x << 1u, (operand.y << 1u) | (operand.x >> 31u), (operand.z << 1u) | (operand.y >> 31u));
+  }
+  var product_top = 31 - i32(countLeadingZeros(product.x));
+  if (product.y != 0u) { product_top = 63 - i32(countLeadingZeros(product.y)); }
+  if (product.z != 0u) { product_top = 95 - i32(countLeadingZeros(product.z)); }
+  if (product_top + scale_exponent >= 10) { return select(1024, -1024, acc < 0); }
+  if (product_top + scale_exponent < -2) { return 0; }
+  let product_shift = u32(max(0, product_top - 52));
+  let rounded_product = quant_u96_round_shift(product, product_shift);
+  let result_exponent = scale_exponent + i32(product_shift);
+  var rounded = rounded_product;
+  if (result_exponent < 0) { rounded = quant_u96_round_shift(rounded_product, u32(-result_exponent)); }
+  else { rounded.x <<= u32(result_exponent); }
+  return select(i32(rounded.x), -i32(rounded.x), acc < 0);
+}
+// A power-of-two division and an exactly represented <=48-bit product need
+// no binary64 boundary resolver. This includes ordinary dyadic quantization.
+fn quant_exact_dyadic_product(acc: i32, scale: f32) -> bool {
+  let magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  let significand = quant_float_significand(bitcast<u32>(scale));
+  let integer_bits = 32u - countLeadingZeros(magnitude) - countTrailingZeros(magnitude);
+  let scale_bits = 32u - countLeadingZeros(significand) - countTrailingZeros(significand);
+  return integer_bits + scale_bits <= 48u;
+}
+fn quant_normal_expansion_range(v: f32) -> bool {
+  let exponent = (bitcast<u32>(v) >> 23u) & 255u;
+  // Reserve two binary32 significands plus four carry/error bits above the
+  // minimum normal exponent, so residual arithmetic avoids FTZ underflow.
+  return exponent >= (1u + 2u * 24u + 4u) && exponent < 255u;
+}
+fn round_expansion_ratio(numerator_hi: f32, numerator_lo: f32, denominator: f32, dyadic_scale: f32, dyadic_candidate: bool, acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  let estimate = numerator_hi / denominator;
+  // Anything outside this range saturates after adding an 8-bit zero point;
+  // returning early also keeps floor/i32 conversion away from Inf and overflow.
+  if (estimate <= -1024.0) { return -1024; }
+  if (estimate >= 1024.0) { return 1024; }
+  let lo = floor(estimate);
+  let hi = lo + 1.0;
+  let midpoint = lo + 0.5;
+  let midpoint_hi = fma(midpoint, denominator, 0.0);
+  if ((bitcast<u32>(midpoint_hi) & 0x7f800000u) == 0x7f800000u) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let midpoint_lo = fma(midpoint, denominator, -midpoint_hi);
+  // The high-word subtraction is exact for these similarly sized products
+  // (Sterbenz); adding the residual difference therefore preserves its sign
+  // even when the rounded high words straddle adjacent f32 values.
+  let midpoint_delta = fma(1.0, numerator_hi - midpoint_hi, numerator_lo - midpoint_lo);
+  // Escalate ambiguous expansion signs instead of changing the tie tolerance.
+  // Six FMA/product tails and their additions fit within this conservative
+  // 16*u^2 error envelope (u=2^-24), including binary64 reference rounding.
+  if ((bitcast<u32>(midpoint_delta) & 0x7f800000u) == 0x7f800000u) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  if (abs(midpoint_delta) <= max(abs(numerator_hi), abs(midpoint_hi)) * 0x1p-44) {
+    if (!dyadic_candidate || !quant_exact_dyadic_product(acc, dyadic_scale)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  }
+  if (midpoint_delta < 0.0) {
+    return i32(lo);
+  }
+  if (midpoint_delta > 0.0) {
+    return i32(hi);
+  }
+  let half_lo = floor(lo * 0.5);
+  let lo_is_even = lo - half_lo * 2.0 == 0.0;
+  return i32(select(hi, lo, lo_is_even));
+}
+// Multiplication by a power of two keeps every i32 bit in binary64. Round the
+// exact dyadic value directly, including INT_MIN, half-even ties and saturation.
+fn quant_round_power_two(acc: i32, exponent: i32) -> i32 {
+  let magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  var rounded = 0u;
+  if (exponent >= 0) {
+    if (exponent >= 10) { return select(1024, -1024, acc < 0); }
+    rounded = min(magnitude, 1024u >> u32(exponent)) << u32(exponent);
+  } else {
+    let shift = u32(-exponent);
+    if (shift >= 32u) { return 0; }
+    let whole = magnitude >> shift;
+    let remainder = magnitude & ((1u << shift) - 1u);
+    let half = 1u << (shift - 1u);
+    rounded = min(1024u, whole + u32(remainder > half || (remainder == half && (whole & 1u) != 0u)));
+  }
+  return select(i32(rounded), -i32(rounded), acc < 0);
+}
+fn round_requantized_precise(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  if (((a_bits | b_bits | y_bits) & 0x807fffffu) == 0u && min(a_bits, min(b_bits, y_bits)) != 0u && max(a_bits, max(b_bits, y_bits)) < 0x7f800000u) {
+    let exponent = i32(a_bits >> 23u) + i32(b_bits >> 23u) - i32(y_bits >> 23u) - 127;
+    return quant_round_power_two(acc, exponent);
+  }
+  let acc_f = f32(acc);
+  // Every i32 high part aligned to 256 is exactly representable in f32.
+  // Recover the conversion residual without an out-of-range i32(acc_f)
+  // conversion when INT_MAX rounds up to 2^31.
+  let acc_tail = fma(-1.0, acc_f - f32(acc & -256), f32(acc & 255));
+  // Equal b/y scales cancel exactly. Bypass their product and division so a
+  // rounded intermediate cannot move a near-half value onto the tie.
+  if (bitcast<u32>(b_s) == bitcast<u32>(y_s)) {
+    let direct_hi = fma(acc_f, a_s, 0.0);
+    let direct_lo = fma(acc_tail, a_s, fma(acc_f, a_s, -direct_hi));
+    if (!quant_normal_expansion_range(direct_hi)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+    return round_expansion_ratio(direct_hi, direct_lo, 1.0, a_s, true, acc, a_s, b_s, y_s);
+  }
+  // Keep the scale product as a two-float expansion. Its FMA residual preserves
+  // which side of a rounding midpoint a near-half value occupies.
+  let ab_hi = fma(a_s, b_s, 0.0);
+  let ab_lo = fma(a_s, b_s, -ab_hi);
+  if (!quant_normal_expansion_range(ab_hi) || !quant_normal_expansion_range(y_s) || y_s > 0x1p126) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let product_hi = fma(acc_f, ab_hi, 0.0);
+  if (!quant_normal_expansion_range(product_hi)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let rounded_acc_tail = fma(acc_f, ab_hi, -product_hi) + fma(acc_f, ab_lo, 0.0);
+  let product_tail = fma(acc_tail, ab_hi, rounded_acc_tail) + fma(acc_tail, ab_lo, 0.0);
+  // Compare the expansion directly. Renormalizing it adds a cancelling
+  // subtraction/addition chain that can erase the residual during lowering.
+  let denominator_bits = bitcast<u32>(y_s);
+  let denominator_power_two = (denominator_bits & 0x7fffffu) == 0u;
+  let dyadic_candidate = denominator_power_two && ab_lo == 0.0;
+  return round_expansion_ratio(product_hi, product_tail, y_s, ab_hi, dyadic_candidate, acc, a_s, b_s, y_s);
+}
+// A conservative binary32 error filter accepts only a result safely inside one
+// integer rounding bin. Ambiguous boundaries keep the full expansion/binary64
+// path; this is an exact-result shortcut, not an approximate quantization mode.
+fn round_requantized(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  let ea = (a_bits >> 23u) & 255u;
+  let eb = (b_bits >> 23u) & 255u;
+  let ey = (y_bits >> 23u) & 255u;
+  // Certify the range before arithmetic: overflow need not produce infinity.
+  // Any nonzero i32 converts to a float with biased exponent <=158. These
+  // scale-only bounds cover every accumulator and can be shared across lanes.
+  // Reserve two significands plus carry bits above FTZ, and a factor of two
+  // below overflow for products, division and the expansion midpoint product.
+  if (((a_bits | b_bits | y_bits) & 0x80000000u) != 0u ||
+      min(ea, eb) == 0u || max(ea, eb) >= 255u ||
+      ea + eb < 180u || ea + eb > 347u ||
+      ey < 53u || ey > 242u || ea + eb > ey + 220u) {
+    return quant_round_binary64(acc, a_s, b_s, y_s);
+  }
+  let scale = fma(a_s, b_s, 0.0);
+  let product = fma(f32(acc), scale, 0.0);
+  let estimate = product / y_s;
+  if (estimate <= -1024.0) { return -1024; }
+  if (estimate >= 1024.0) { return 1024; }
+  let nearest = round(estimate);
+  // 32 binary32 unit roundoffs reserve room for the i32 conversion, two
+  // products, quotient, and distance arithmetic, plus binary64 rounding.
+  let margin = max(1.0, abs(estimate)) * 0x1p-19;
+  if (abs(estimate - nearest) < 0.5 - margin) { return i32(nearest); }
+{% if device.adapterInfo.vendor != "apple" %}
+  // A rounded half-integer estimate can hide which side the exact f32-scale
+  // ratio lies on. Resolve that one-ulp neighborhood conservatively outside
+  // the compact fast path; exact power-of-two ratios already have a cheap
+  // and correct path in round_requantized_precise.
+  if (abs(abs(estimate - nearest) - 0.5) <= max(1.0, abs(estimate)) * 0x1p-24) {
+    let dyadic_scales = ((a_bits | b_bits | y_bits) & 0x807fffffu) == 0u &&
+      min(a_bits, min(b_bits, y_bits)) != 0u &&
+      max(a_bits, max(b_bits, y_bits)) < 0x7f800000u;
+    if (!dyadic_scales) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  }
+{% endif %}
+  return round_requantized_precise(acc, a_s, b_s, y_s);
+}
+{{ quant_read_scalar_zero("read_y_zero", "y_zero_point", yUnsigned) }}
+fn write_y(index: u32, value: i32) {
+  {% if yUnsigned %}
+  y[index] = u32(value);
+  {% else %}
+  y[index] = value;
+  {% endif %}
+}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, guardInline=true) }}
+{% if hasBias or wScalePerChannel %}
+  let oc = (i / params.outImageSize) % params.outChannels;
+{% endif %}
+  {% if hasBias %}
+  // Optional int32 bias is per output channel, already in (x_scale * w_scale)
+  // units, so it is added to the int32 accumulator before requantization.
+  let raw = accum[i] + B[oc];
+  {% else %}
+  let raw = accum[i];
+  {% endif %}
+{% if wScalePerChannel %}
+  let w_scale_value = w_scale[oc];
+{% else %}
+  let w_scale_value = w_scale[0];
+{% endif %}
+  let q = clamp(round_requantized(raw, x_scale[0], w_scale_value, y_scale[0]) + read_y_zero(), {{ qMin }}, {{ qMax }});
+  write_y(i, q);
+}

build/webgpu/quant-dp4a-matmul.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,573 @@

+{% macro quant_read_scalar_zero(fn_name, buffer_name, is_unsigned, present=true) %}
+fn {{ fn_name }}() -> i32 {
+  {% if is_unsigned %}
+  return i32({{ buffer_name }}[0]);
+  {% else %}
+  return {{ buffer_name }}[0];
+  {% endif %}
+}{% endmacro %}
+{% macro quant_read_column_zero(fn_name, buffer_name, is_unsigned, present=true, per_column=false, param_name="col", return_type="i32") %}
+fn {{ fn_name }}({% if per_column %}{{ param_name }}: u32{% endif %}) -> {{ return_type }} {
+  {% if per_column %}
+  {% if is_unsigned %}
+  return i32({{ buffer_name }}[{{ param_name }}]);
+  {% else %}
+  return {{ buffer_name }}[{{ param_name }}];
+  {% endif %}
+  {% else %}
+  {% if is_unsigned %}
+  return i32({{ buffer_name }}[0]);
+  {% else %}
+  return {{ buffer_name }}[0];
+  {% endif %}
+  {% endif %}
+}{% endmacro %}
+{% if not env.device.wgslLanguageFeatures.has("packed_4x8_integer_dot_product") %}
+// Runtime source-compatible fallbacks for the packed 4x8 integer built-ins used
+// by quantized kernels. Native implementations remain unshadowed when the
+// language feature is present. The user functions retain the built-ins'
+// @must_use contract, but cannot be @const; every call is runtime.
+// Inputs are function parameters, so replacing a built-in call with one of
+// these same-name functions preserves WGSL's single evaluation of each
+// argument. Byte extraction is unsigned. Signed lanes use the exact
+// i32(byte ^ 0x80u) - 128i sign extension, avoiding signed-left-shift overflow
+// while preserving every 8-bit two's-complement value.
+@must_use fn dot4I8Packed(e1: u32, e2: u32) -> i32 {
+  let shifts = vec4<u32>(0u, 8u, 16u, 24u);
+  let a_bytes = (vec4<u32>(e1) >> shifts) & vec4<u32>(0xffu);
+  let b_bytes = (vec4<u32>(e2) >> shifts) & vec4<u32>(0xffu);
+  let a = vec4<i32>(a_bytes ^ vec4<u32>(0x80u)) - vec4<i32>(128i);
+  let b = vec4<i32>(b_bytes ^ vec4<u32>(0x80u)) - vec4<i32>(128i);
+  return dot(a, b);
+}
+{% endif %}
+{{ env.wgsl.resourceDeclarations }}
+{% if requant != "none" %}
+// Rare-boundary fallback: reproduce the reference's binary64 scale division
+// and accumulator product using integer significands; no f64 feature required.
+fn quant_u96_add(a: vec3<u32>, b: vec3<u32>) -> vec3<u32> {
+  let low = a.x + b.x;
+  let middle = a.y + b.y;
+  let with_carry = middle + u32(low < a.x);
+  return vec3<u32>(low, with_carry, a.z + b.z + u32(middle < a.y) + u32(with_carry < middle));
+}
+fn quant_u96_shr(a: vec3<u32>) -> vec3<u32> {
+  return vec3<u32>((a.x >> 1u) | (a.y << 31u), (a.y >> 1u) | (a.z << 31u), a.z >> 1u);
+}
+fn quant_u96_round_shift(value: vec3<u32>, shift: u32) -> vec3<u32> {
+  var q = value;
+  var sticky = false;
+  var guard = false;
+  for (var bit = 0u; bit < shift; bit++) {
+    sticky = sticky || guard;
+    guard = (q.x & 1u) != 0u;
+    q = quant_u96_shr(q);
+  }
+  if (guard && (sticky || (q.x & 1u) != 0u)) {
+    q = quant_u96_add(q, vec3<u32>(1u, 0u, 0u));
+  }
+  return q;
+}
+fn quant_float_significand(bits: u32) -> u32 {
+  return (bits & 0x7fffffu) | select(0x800000u, 0u, (bits & 0x7f800000u) == 0u);
+}
+fn quant_float_exponent(bits: u32) -> i32 {
+  return max(1, i32((bits >> 23u) & 255u)) - 150;
+}
+fn quant_round_binary64(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  let a = quant_float_significand(a_bits);
+  let b = quant_float_significand(b_bits);
+  let denominator = quant_float_significand(y_bits);
+  if (a == 0u || b == 0u) { return 0; }
+  // The exact product of two 24-bit significands occupies at most 48 bits.
+  let p0 = (a & 65535u) * (b & 65535u);
+  let cross = (a >> 16u) * (b & 65535u) + (b >> 16u) * (a & 65535u);
+  let product_low = p0 + (cross << 16u);
+  let product_high = (a >> 16u) * (b >> 16u) + (cross >> 16u) + u32(product_low < p0);
+  var input_bit = select(31 - i32(countLeadingZeros(product_low)), 63 - i32(countLeadingZeros(product_high)), product_high != 0u);
+  var remainder = 0u;
+  var significant_bits = 0u;
+  var leading_exponent = 0;
+  var significand = vec2<u32>(0u);
+  // Long division generates the 53 bits of the binary64 multiplier. The
+  // remainder fits in 25 bits, including its left shift. At most 24 leading
+  // zero bits precede the first significant bit.
+  loop {
+    var incoming = 0u;
+    if (input_bit >= 32) { incoming = (product_high >> u32(input_bit - 32)) & 1u; }
+    else if (input_bit >= 0) { incoming = (product_low >> u32(input_bit)) & 1u; }
+    remainder = (remainder << 1u) | incoming;
+    let digit = u32(remainder >= denominator);
+    if (digit != 0u) { remainder -= denominator; }
+    if (significant_bits != 0u || digit != 0u) {
+      if (significant_bits == 0u) { leading_exponent = input_bit; }
+      significand = vec2<u32>((significand.x << 1u) | digit, (significand.y << 1u) | (significand.x >> 31u));
+      significant_bits++;
+    }
+    input_bit--;
+    if (significant_bits == 53u) { break; }
+  }
+  let twice_remainder = remainder << 1u;
+  if (twice_remainder > denominator || (twice_remainder == denominator && (significand.x & 1u) != 0u)) {
+    let low = significand.x + 1u;
+    significand = vec2<u32>(low, significand.y + u32(low == 0u));
+  }
+  let scale_exponent = quant_float_exponent(a_bits) + quant_float_exponent(b_bits) - quant_float_exponent(y_bits) + leading_exponent - 52;
+  // abs(INT_MIN) is representable as u32. Multiply the rounded multiplier's
+  // significand by the exact accumulator into three little-endian u32 limbs.
+  var magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  var operand = vec3<u32>(significand, 0u);
+  var product = vec3<u32>(0u);
+  loop {
+    if ((magnitude & 1u) != 0u) { product = quant_u96_add(product, operand); }
+    magnitude >>= 1u;
+    if (magnitude == 0u) { break; }
+    operand = vec3<u32>(operand.x << 1u, (operand.y << 1u) | (operand.x >> 31u), (operand.z << 1u) | (operand.y >> 31u));
+  }
+  var product_top = 31 - i32(countLeadingZeros(product.x));
+  if (product.y != 0u) { product_top = 63 - i32(countLeadingZeros(product.y)); }
+  if (product.z != 0u) { product_top = 95 - i32(countLeadingZeros(product.z)); }
+  if (product_top + scale_exponent >= 10) { return select(1024, -1024, acc < 0); }
+  if (product_top + scale_exponent < -2) { return 0; }
+  let product_shift = u32(max(0, product_top - 52));
+  let rounded_product = quant_u96_round_shift(product, product_shift);
+  let result_exponent = scale_exponent + i32(product_shift);
+  var rounded = rounded_product;
+  if (result_exponent < 0) { rounded = quant_u96_round_shift(rounded_product, u32(-result_exponent)); }
+  else { rounded.x <<= u32(result_exponent); }
+  return select(i32(rounded.x), -i32(rounded.x), acc < 0);
+}
+// A power-of-two division and an exactly represented <=48-bit product need
+// no binary64 boundary resolver. This includes ordinary dyadic quantization.
+fn quant_exact_dyadic_product(acc: i32, scale: f32) -> bool {
+  let magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  let significand = quant_float_significand(bitcast<u32>(scale));
+  let integer_bits = 32u - countLeadingZeros(magnitude) - countTrailingZeros(magnitude);
+  let scale_bits = 32u - countLeadingZeros(significand) - countTrailingZeros(significand);
+  return integer_bits + scale_bits <= 48u;
+}
+fn quant_normal_expansion_range(v: f32) -> bool {
+  let exponent = (bitcast<u32>(v) >> 23u) & 255u;
+  // Reserve two binary32 significands plus four carry/error bits above the
+  // minimum normal exponent, so residual arithmetic avoids FTZ underflow.
+  return exponent >= (1u + 2u * 24u + 4u) && exponent < 255u;
+}
+fn round_expansion_ratio(numerator_hi: f32, numerator_lo: f32, denominator: f32, dyadic_scale: f32, dyadic_candidate: bool, acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  let estimate = numerator_hi / denominator;
+  // Anything outside this range saturates after adding an 8-bit zero point;
+  // returning early also keeps floor/i32 conversion away from Inf and overflow.
+  if (estimate <= -1024.0) { return -1024; }
+  if (estimate >= 1024.0) { return 1024; }
+  let lo = floor(estimate);
+  let hi = lo + 1.0;
+  let midpoint = lo + 0.5;
+  let midpoint_hi = fma(midpoint, denominator, 0.0);
+  if ((bitcast<u32>(midpoint_hi) & 0x7f800000u) == 0x7f800000u) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let midpoint_lo = fma(midpoint, denominator, -midpoint_hi);
+  // The high-word subtraction is exact for these similarly sized products
+  // (Sterbenz); adding the residual difference therefore preserves its sign
+  // even when the rounded high words straddle adjacent f32 values.
+  let midpoint_delta = fma(1.0, numerator_hi - midpoint_hi, numerator_lo - midpoint_lo);
+  // Escalate ambiguous expansion signs instead of changing the tie tolerance.
+  // Six FMA/product tails and their additions fit within this conservative
+  // 16*u^2 error envelope (u=2^-24), including binary64 reference rounding.
+  if ((bitcast<u32>(midpoint_delta) & 0x7f800000u) == 0x7f800000u) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  if (abs(midpoint_delta) <= max(abs(numerator_hi), abs(midpoint_hi)) * 0x1p-44) {
+    if (!dyadic_candidate || !quant_exact_dyadic_product(acc, dyadic_scale)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  }
+  if (midpoint_delta < 0.0) {
+    return i32(lo);
+  }
+  if (midpoint_delta > 0.0) {
+    return i32(hi);
+  }
+  let half_lo = floor(lo * 0.5);
+  let lo_is_even = lo - half_lo * 2.0 == 0.0;
+  return i32(select(hi, lo, lo_is_even));
+}
+// Multiplication by a power of two keeps every i32 bit in binary64. Round the
+// exact dyadic value directly, including INT_MIN, half-even ties and saturation.
+fn quant_round_power_two(acc: i32, exponent: i32) -> i32 {
+  let magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  var rounded = 0u;
+  if (exponent >= 0) {
+    if (exponent >= 10) { return select(1024, -1024, acc < 0); }
+    rounded = min(magnitude, 1024u >> u32(exponent)) << u32(exponent);
+  } else {
+    let shift = u32(-exponent);
+    if (shift >= 32u) { return 0; }
+    let whole = magnitude >> shift;
+    let remainder = magnitude & ((1u << shift) - 1u);
+    let half = 1u << (shift - 1u);
+    rounded = min(1024u, whole + u32(remainder > half || (remainder == half && (whole & 1u) != 0u)));
+  }
+  return select(i32(rounded), -i32(rounded), acc < 0);
+}
+fn round_requantized_precise(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  if (((a_bits | b_bits | y_bits) & 0x807fffffu) == 0u && min(a_bits, min(b_bits, y_bits)) != 0u && max(a_bits, max(b_bits, y_bits)) < 0x7f800000u) {
+    let exponent = i32(a_bits >> 23u) + i32(b_bits >> 23u) - i32(y_bits >> 23u) - 127;
+    return quant_round_power_two(acc, exponent);
+  }
+  let acc_f = f32(acc);
+  // Every i32 high part aligned to 256 is exactly representable in f32.
+  // Recover the conversion residual without an out-of-range i32(acc_f)
+  // conversion when INT_MAX rounds up to 2^31.
+  let acc_tail = fma(-1.0, acc_f - f32(acc & -256), f32(acc & 255));
+  // Equal b/y scales cancel exactly. Bypass their product and division so a
+  // rounded intermediate cannot move a near-half value onto the tie.
+  if (bitcast<u32>(b_s) == bitcast<u32>(y_s)) {
+    let direct_hi = fma(acc_f, a_s, 0.0);
+    let direct_lo = fma(acc_tail, a_s, fma(acc_f, a_s, -direct_hi));
+    if (!quant_normal_expansion_range(direct_hi)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+    return round_expansion_ratio(direct_hi, direct_lo, 1.0, a_s, true, acc, a_s, b_s, y_s);
+  }
+  // Keep the scale product as a two-float expansion. Its FMA residual preserves
+  // which side of a rounding midpoint a near-half value occupies.
+  let ab_hi = fma(a_s, b_s, 0.0);
+  let ab_lo = fma(a_s, b_s, -ab_hi);
+  if (!quant_normal_expansion_range(ab_hi) || !quant_normal_expansion_range(y_s) || y_s > 0x1p126) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let product_hi = fma(acc_f, ab_hi, 0.0);
+  if (!quant_normal_expansion_range(product_hi)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let rounded_acc_tail = fma(acc_f, ab_hi, -product_hi) + fma(acc_f, ab_lo, 0.0);
+  let product_tail = fma(acc_tail, ab_hi, rounded_acc_tail) + fma(acc_tail, ab_lo, 0.0);
+  // Compare the expansion directly. Renormalizing it adds a cancelling
+  // subtraction/addition chain that can erase the residual during lowering.
+  let denominator_bits = bitcast<u32>(y_s);
+  let denominator_power_two = (denominator_bits & 0x7fffffu) == 0u;
+  let dyadic_candidate = denominator_power_two && ab_lo == 0.0;
+  return round_expansion_ratio(product_hi, product_tail, y_s, ab_hi, dyadic_candidate, acc, a_s, b_s, y_s);
+}
+// A conservative binary32 error filter accepts only a result safely inside one
+// integer rounding bin. Ambiguous boundaries keep the full expansion/binary64
+// path; this is an exact-result shortcut, not an approximate quantization mode.
+fn round_requantized(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  let ea = (a_bits >> 23u) & 255u;
+  let eb = (b_bits >> 23u) & 255u;
+  let ey = (y_bits >> 23u) & 255u;
+  // Certify the range before arithmetic: overflow need not produce infinity.
+  // Any nonzero i32 converts to a float with biased exponent <=158. These
+  // scale-only bounds cover every accumulator and can be shared across lanes.
+  // Reserve two significands plus carry bits above FTZ, and a factor of two
+  // below overflow for products, division and the expansion midpoint product.
+  if (((a_bits | b_bits | y_bits) & 0x80000000u) != 0u ||
+      min(ea, eb) == 0u || max(ea, eb) >= 255u ||
+      ea + eb < 180u || ea + eb > 347u ||
+      ey < 53u || ey > 242u || ea + eb > ey + 220u) {
+    return quant_round_binary64(acc, a_s, b_s, y_s);
+  }
+  let scale = fma(a_s, b_s, 0.0);
+  let product = fma(f32(acc), scale, 0.0);
+  let estimate = product / y_s;
+  if (estimate <= -1024.0) { return -1024; }
+  if (estimate >= 1024.0) { return 1024; }
+  let nearest = round(estimate);
+  // 32 binary32 unit roundoffs reserve room for the i32 conversion, two
+  // products, quotient, and distance arithmetic, plus binary64 rounding.
+  let margin = max(1.0, abs(estimate)) * 0x1p-19;
+  if (abs(estimate - nearest) < 0.5 - margin) { return i32(nearest); }
+{% if device.adapterInfo.vendor != "apple" %}
+  // A rounded half-integer estimate can hide which side the exact f32-scale
+  // ratio lies on. Resolve that one-ulp neighborhood conservatively outside
+  // the compact fast path; exact power-of-two ratios already have a cheap
+  // and correct path in round_requantized_precise.
+  if (abs(abs(estimate - nearest) - 0.5) <= max(1.0, abs(estimate)) * 0x1p-24) {
+    let dyadic_scales = ((a_bits | b_bits | y_bits) & 0x807fffffu) == 0u &&
+      min(a_bits, min(b_bits, y_bits)) != 0u &&
+      max(a_bits, max(b_bits, y_bits)) < 0x7f800000u;
+    if (!dyadic_scales) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  }
+{% endif %}
+  return round_requantized_precise(acc, a_s, b_s, y_s);
+}
+{% endif %}
+const TILE_M: u32 = {{ tileM }}u;
+const TILE_N: u32 = {{ tileN }}u;
+const KT_WORDS: u32 = {{ ktWords }}u;
+const A_TILE_WORDS: u32 = TILE_M * KT_WORDS;
+{% set GEMV_UNROLL = 4 %}
+{% if tileM == 1 %}
+const GEMV_K_UNROLL: u32 = {{ GEMV_UNROLL }}u;
+{% endif %}
+{% set needsBSum = true %}
+var<workgroup> a_tile: array<u32, A_TILE_WORDS>;
+var<workgroup> a_row_sum: array<atomic<i32>, TILE_M>;
+// Pack 4 widened A values into one u32 (byte 0 = lowest k).
+fn pack_a(v: vec4<{{ aScalar }}>) -> u32 {
+{% if aUnsigned %}
+  // u8 -> biased i8: (u - 128) two's-complement low byte == u ^ 0x80.
+  return ((v.x ^ 0x80u) & 0xFFu) | (((v.y ^ 0x80u) & 0xFFu) << 8u) | (((v.z ^ 0x80u) & 0xFFu) << 16u) | (((v.w ^ 0x80u) & 0xFFu) << 24u);
+{% else %}
+  return (u32(v.x) & 0xFFu) | ((u32(v.y) & 0xFFu) << 8u) | ((u32(v.z) & 0xFFu) << 16u) | ((u32(v.w) & 0xFFu) << 24u);
+{% endif %}
+}
+// Sum of the 4 bias-shifted A values of one packed word.
+fn sum4_a(v: vec4<{{ aScalar }}>) -> i32 {
+{% if aUnsigned %}
+  return i32(v.x + v.y + v.z + v.w) - 512;
+{% else %}
+  return v.x + v.y + v.z + v.w;
+{% endif %}
+}
+fn pack_b(b0: {{ bScalar }}, b1: {{ bScalar }}, b2: {{ bScalar }}, b3: {{ bScalar }}) -> u32 {
+{% if bUnsigned %}
+  return ((b0 ^ 0x80u) & 0xFFu) | (((b1 ^ 0x80u) & 0xFFu) << 8u) | (((b2 ^ 0x80u) & 0xFFu) << 16u) | (((b3 ^ 0x80u) & 0xFFu) << 24u);
+{% else %}
+  return (u32(b0) & 0xFFu) | ((u32(b1) & 0xFFu) << 8u) | ((u32(b2) & 0xFFu) << 16u) | ((u32(b3) & 0xFFu) << 24u);
+{% endif %}
+}
+// Centralize the signed packed-int8 dot expression so unrolled callers cannot
+// collide in lowering-generated temporaries.
+fn dot4_packed(a_word: u32, b_word: u32) -> i32 {
+  return dot4I8Packed(a_word, b_word);
+}
+fn sum4_b(b0: {{ bScalar }}, b1: {{ bScalar }}, b2: {{ bScalar }}, b3: {{ bScalar }}) -> i32 {
+{% if bUnsigned %}
+  return i32(b0 + b1 + b2 + b3) - 512;
+{% else %}
+  return b0 + b1 + b2 + b3;
+{% endif %}
+}
+{{ quant_read_column_zero("read_a_zero", "a_zero_point", aUnsigned, per_column=aZeroPerRow, param_name="row") }}
+{{ quant_read_column_zero("read_b_zero", "b_zero_point", bUnsigned, per_column=bZeroPerColumn) }}
+{% if requant != "none" %}
+{{ quant_read_scalar_zero("read_y_zero", "y_zero_point", yUnsigned) }}
+{% endif %}
+@compute @workgroup_size(TILE_N, 1, 1)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  let batch = wg.z;
+  let a_base4 = batch * params.aBatchStride4;
+  let b_base = batch * params.bBatchStride;
+  let y_base = batch * params.yBatchStride;
+  let m_block = wg.y * TILE_M;
+  let n_block = wg.x * TILE_N;
+  let tid = lid.x;
+  let col = n_block + tid;
+  let n_valid = col < params.N;
+{% if aPacked == false %}
+  // Tail variants keep the hot loop on complete words. Packing the one partial
+  // word afterward prevents its bounds checks from infecting every full-word load.
+  let k_words_total = params.K / 4u;
+  let tail_lanes = params.K - k_words_total * 4u;
+{% else %}
+  let k_words_total = params.K / 4u;
+{% endif %}
+  if (tid < TILE_M) {
+    atomicStore(&a_row_sum[tid], 0);
+  }
+  workgroupBarrier();
+{% if tileM == 1 %}
+  var r0: i32 = 0;
+{% for u in range(1, GEMV_UNROLL) %}
+  var r0_{{ u }}: i32 = 0;
+{% endfor %}
+{% else %}
+{% for m in range(tileM) %}
+  var r{{ m }}: i32 = 0;
+{% endfor %}
+{% endif %}
+  var sum_b: i32 = 0;
+{% if tileM == 1 %}
+{% for u in range(1, GEMV_UNROLL) %}
+  var sum_b_{{ u }}: i32 = 0;
+{% endfor %}
+{% endif %}
+  for (var kw0 = 0u; kw0 < k_words_total; kw0 = kw0 + KT_WORDS) {
+    let words = min(KT_WORDS, k_words_total - kw0);
+    // Cooperative vec4 load + pack of the A tile (TILE_M rows x `words` u32).
+    for (var i = tid; i < TILE_M * words; i = i + TILE_N) {
+      let mi = i / words;
+      let ki = i % words;
+      let src_m = m_block + mi;
+      var packed = 0u;
+      if (src_m < params.M) {
+{% if aPacked == false %}
+        let k4 = (kw0 + ki) * 4u;
+        let a_row = a_base4 + src_m * params.K;
+        let v = vec4<{{ aScalar }}>(
+          a[a_row + k4],
+          a[a_row + k4 + 1u],
+          a[a_row + k4 + 2u],
+          a[a_row + k4 + 3u],
+        );
+{% else %}
+        let v = a[a_base4 + src_m * k_words_total + kw0 + ki];
+{% endif %}
+        packed = pack_a(v);
+        atomicAdd(&a_row_sum[mi], sum4_a(v));
+      }
+      a_tile[mi * KT_WORDS + ki] = packed;
+    }
+    workgroupBarrier();
+    if (n_valid) {
+{% if tileM == 1 %}
+      // M=1 has only one output accumulator, so a scalar K loop forms a long
+      // dependent DP4A chain. Interleave several words while preserving the
+      // coalesced mapping of neighboring threads to neighboring B columns.
+      let unrolled_words = words - words % GEMV_K_UNROLL;
+      for (var kk = 0u; kk < unrolled_words; kk += GEMV_K_UNROLL) {
+{% for u in range(GEMV_UNROLL) %}
+        let k4_{{ u }} = (kw0 + kk + {{ u }}u) * 4u;
+        let b_row_{{ u }} = b_base + k4_{{ u }} * params.N + col;
+        let b0_{{ u }} = b[b_row_{{ u }}];
+        let b1_{{ u }} = b[b_row_{{ u }} + params.N];
+        let b2_{{ u }} = b[b_row_{{ u }} + 2u * params.N];
+        let b3_{{ u }} = b[b_row_{{ u }} + 3u * params.N];
+        sum_b{{ "" if u == 0 else "_" ~ u }} += sum4_b(b0_{{ u }}, b1_{{ u }}, b2_{{ u }}, b3_{{ u }});
+        let bp_{{ u }} = pack_b(b0_{{ u }}, b1_{{ u }}, b2_{{ u }}, b3_{{ u }});
+        r0{{ "" if u == 0 else "_" ~ u }} += dot4_packed(a_tile[kk + {{ u }}u], bp_{{ u }});
+{% endfor %}
+      }
+      for (var kk = unrolled_words; kk < words; kk++) {
+        let k4 = (kw0 + kk) * 4u;
+        let b_row = b_base + k4 * params.N + col;
+        let b0 = b[b_row];
+        let b1 = b[b_row + params.N];
+        let b2 = b[b_row + 2u * params.N];
+        let b3 = b[b_row + 3u * params.N];
+        sum_b += sum4_b(b0, b1, b2, b3);
+        let bp = pack_b(b0, b1, b2, b3);
+        r0 += dot4_packed(a_tile[kk], bp);
+      }
+{% else %}
+      for (var kk = 0u; kk < words; kk = kk + 1u) {
+        let k4 = (kw0 + kk) * 4u;
+        let b_row = b_base + k4 * params.N + col;
+        let b0 = b[b_row];
+        let b1 = b[b_row + params.N];
+        let b2 = b[b_row + 2u * params.N];
+        let b3 = b[b_row + 3u * params.N];
+        sum_b = sum_b + sum4_b(b0, b1, b2, b3);
+        let bp = pack_b(b0, b1, b2, b3);
+{% for m in range(tileM) %}
+        r{{ m }} = r{{ m }} + dot4_packed(a_tile[{{ m }}u * KT_WORDS + kk], bp);
+{% endfor %}
+      }
+{% endif %}
+    }
+    workgroupBarrier();
+  }
+{% if aPacked == false %}
+  // Exactly one uniform guarded word handles K % 4. Missing lanes use the raw
+  // bias-domain zero (128 for u8, 0 for i8), so their shifted values, sums, and
+  // dot-product contributions are all zero while params.K remains unpadded for
+  // the zero-point correction below.
+  if (tail_lanes != 0u) {
+    if (tid < TILE_M) {
+      let src_m = m_block + tid;
+      var packed = 0u;
+      if (src_m < params.M) {
+        let k4 = k_words_total * 4u;
+        let a_row = a_base4 + src_m * params.K;
+        let a0 = a[a_row + k4];
+        var a1: {{ aScalar }} = {{ "128u" if aUnsigned else "0" }};
+        var a2: {{ aScalar }} = {{ "128u" if aUnsigned else "0" }};
+        if (tail_lanes > 1u) { a1 = a[a_row + k4 + 1u]; }
+        if (tail_lanes > 2u) { a2 = a[a_row + k4 + 2u]; }
+        let a3: {{ aScalar }} = {{ "128u" if aUnsigned else "0" }};
+        let v = vec4<{{ aScalar }}>(a0, a1, a2, a3);
+        packed = pack_a(v);
+        atomicAdd(&a_row_sum[tid], sum4_a(v));
+      }
+      a_tile[tid * KT_WORDS] = packed;
+    }
+    workgroupBarrier();
+    if (n_valid) {
+      let k4 = k_words_total * 4u;
+      let b_row = b_base + k4 * params.N + col;
+      let b0 = b[b_row];
+      var b1: {{ bScalar }} = {{ "128u" if bUnsigned else "0" }};
+      var b2: {{ bScalar }} = {{ "128u" if bUnsigned else "0" }};
+      if (tail_lanes > 1u) { b1 = b[b_row + params.N]; }
+      if (tail_lanes > 2u) { b2 = b[b_row + 2u * params.N]; }
+      let b3: {{ bScalar }} = {{ "128u" if bUnsigned else "0" }};
+      sum_b = sum_b + sum4_b(b0, b1, b2, b3);
+      let bp = pack_b(b0, b1, b2, b3);
+{% for m in range(tileM) %}
+      r{{ m }} = r{{ m }} + dot4_packed(a_tile[{{ m }}u * KT_WORDS], bp);
+{% endfor %}
+    }
+    workgroupBarrier();
+  }
+{% endif %}
+  if (!n_valid) {
+    return;
+  }
+  // Fold bias + zero-point corrections (exact in modular i32 arithmetic).
+  let cb: i32 = {{ "128" if bUnsigned else "0" }} - read_b_zero({% if bZeroPerColumn %}col{% endif %});
+{% if not aZeroPerRow %}
+  let ca: i32 = {{ "128" if aUnsigned else "0" }} - read_a_zero();
+  let kcc: i32 = i32(params.K) * ca * cb;
+{% endif %}
+{% if needsBSum and tileM == 1 %}
+  let sum_b_total = sum_b{% for u in range(1, GEMV_UNROLL) %} + sum_b_{{ u }}{% endfor %};
+{% endif %}
+{% for m in range(tileM) %}
+  {
+    let mr = m_block + {{ m }}u;
+    if (mr < params.M) {
+{% if aZeroPerRow %}
+      let ca: i32 = {{ "128" if aUnsigned else "0" }} - read_a_zero(mr);
+      let kcc: i32 = i32(params.K) * ca * cb;
+{% endif %}
+      var acc = r{{ m }}{% if tileM == 1 %}{% for u in range(1, GEMV_UNROLL) %} + r0_{{ u }}{% endfor %}{% endif %};
+      acc = acc + cb * atomicLoad(&a_row_sum[{{ m }}u]);
+      acc = acc + ca * {% if tileM == 1 %}sum_b_total{% else %}sum_b{% endif %};
+      acc = acc + kcc;
+{% if hasBias %}
+      acc = acc + bias[{{ "col" if biasPerColumn is defined and biasPerColumn else "mr" }}];
+{% endif %}
+{% if requant == "none" %}
+      y[y_base + mr * params.N + col] = acc;
+{% else %}
+{% set scale_index = "col" if requant == "per_column" else ("mr" if requant == "per_row" else "0u") %}
+      let rounded = round_requantized(acc, f32(a_scale[0]), f32(b_scale[{{ scale_index }}]), f32(y_scale[0]));
+      let q = clamp(rounded + read_y_zero(), {{ qMin }}, {{ qMax }});
+{% if yUnsigned %}
+      y[y_base + mr * params.N + col] = u32(q);
+{% else %}
+      y[y_base + mr * params.N + col] = q;
+{% endif %}
+{% endif %}
+    }
+  }
+{% endfor %}
+}

build/webgpu/quant-exact-matrix.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,384 @@

+{% set precastOperands = precastOperands | default(false) %}
+{% set precastScalar = precastScalar | default("f32") %}
+enable subgroups;
+{% if pinSubgroupSize32 %}
+enable subgroup_size_control;
+{% endif %}
+enable chromium_experimental_subgroup_matrix;
+diagnostic(off, chromium.subgroup_matrix_uniformity);
+{{ env.wgsl.resourceDeclarations }}
+// Rare-boundary fallback: reproduce the reference's binary64 scale division
+// and accumulator product using integer significands; no f64 feature required.
+fn quant_u96_add(a: vec3<u32>, b: vec3<u32>) -> vec3<u32> {
+  let low = a.x + b.x;
+  let middle = a.y + b.y;
+  let with_carry = middle + u32(low < a.x);
+  return vec3<u32>(low, with_carry, a.z + b.z + u32(middle < a.y) + u32(with_carry < middle));
+}
+fn quant_u96_shr(a: vec3<u32>) -> vec3<u32> {
+  return vec3<u32>((a.x >> 1u) | (a.y << 31u), (a.y >> 1u) | (a.z << 31u), a.z >> 1u);
+}
+fn quant_u96_round_shift(value: vec3<u32>, shift: u32) -> vec3<u32> {
+  var q = value;
+  var sticky = false;
+  var guard = false;
+  for (var bit = 0u; bit < shift; bit++) {
+    sticky = sticky || guard;
+    guard = (q.x & 1u) != 0u;
+    q = quant_u96_shr(q);
+  }
+  if (guard && (sticky || (q.x & 1u) != 0u)) {
+    q = quant_u96_add(q, vec3<u32>(1u, 0u, 0u));
+  }
+  return q;
+}
+fn quant_float_significand(bits: u32) -> u32 {
+  return (bits & 0x7fffffu) | select(0x800000u, 0u, (bits & 0x7f800000u) == 0u);
+}
+fn quant_float_exponent(bits: u32) -> i32 {
+  return max(1, i32((bits >> 23u) & 255u)) - 150;
+}
+fn quant_round_binary64(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  let a = quant_float_significand(a_bits);
+  let b = quant_float_significand(b_bits);
+  let denominator = quant_float_significand(y_bits);
+  if (a == 0u || b == 0u) { return 0; }
+  // The exact product of two 24-bit significands occupies at most 48 bits.
+  let p0 = (a & 65535u) * (b & 65535u);
+  let cross = (a >> 16u) * (b & 65535u) + (b >> 16u) * (a & 65535u);
+  let product_low = p0 + (cross << 16u);
+  let product_high = (a >> 16u) * (b >> 16u) + (cross >> 16u) + u32(product_low < p0);
+  var input_bit = select(31 - i32(countLeadingZeros(product_low)), 63 - i32(countLeadingZeros(product_high)), product_high != 0u);
+  var remainder = 0u;
+  var significant_bits = 0u;
+  var leading_exponent = 0;
+  var significand = vec2<u32>(0u);
+  // Long division generates the 53 bits of the binary64 multiplier. The
+  // remainder fits in 25 bits, including its left shift. At most 24 leading
+  // zero bits precede the first significant bit.
+  loop {
+    var incoming = 0u;
+    if (input_bit >= 32) { incoming = (product_high >> u32(input_bit - 32)) & 1u; }
+    else if (input_bit >= 0) { incoming = (product_low >> u32(input_bit)) & 1u; }
+    remainder = (remainder << 1u) | incoming;
+    let digit = u32(remainder >= denominator);
+    if (digit != 0u) { remainder -= denominator; }
+    if (significant_bits != 0u || digit != 0u) {
+      if (significant_bits == 0u) { leading_exponent = input_bit; }
+      significand = vec2<u32>((significand.x << 1u) | digit, (significand.y << 1u) | (significand.x >> 31u));
+      significant_bits++;
+    }
+    input_bit--;
+    if (significant_bits == 53u) { break; }
+  }
+  let twice_remainder = remainder << 1u;
+  if (twice_remainder > denominator || (twice_remainder == denominator && (significand.x & 1u) != 0u)) {
+    let low = significand.x + 1u;
+    significand = vec2<u32>(low, significand.y + u32(low == 0u));
+  }
+  let scale_exponent = quant_float_exponent(a_bits) + quant_float_exponent(b_bits) - quant_float_exponent(y_bits) + leading_exponent - 52;
+  // abs(INT_MIN) is representable as u32. Multiply the rounded multiplier's
+  // significand by the exact accumulator into three little-endian u32 limbs.
+  var magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  var operand = vec3<u32>(significand, 0u);
+  var product = vec3<u32>(0u);
+  loop {
+    if ((magnitude & 1u) != 0u) { product = quant_u96_add(product, operand); }
+    magnitude >>= 1u;
+    if (magnitude == 0u) { break; }
+    operand = vec3<u32>(operand.x << 1u, (operand.y << 1u) | (operand.x >> 31u), (operand.z << 1u) | (operand.y >> 31u));
+  }
+  var product_top = 31 - i32(countLeadingZeros(product.x));
+  if (product.y != 0u) { product_top = 63 - i32(countLeadingZeros(product.y)); }
+  if (product.z != 0u) { product_top = 95 - i32(countLeadingZeros(product.z)); }
+  if (product_top + scale_exponent >= 10) { return select(1024, -1024, acc < 0); }
+  if (product_top + scale_exponent < -2) { return 0; }
+  let product_shift = u32(max(0, product_top - 52));
+  let rounded_product = quant_u96_round_shift(product, product_shift);
+  let result_exponent = scale_exponent + i32(product_shift);
+  var rounded = rounded_product;
+  if (result_exponent < 0) { rounded = quant_u96_round_shift(rounded_product, u32(-result_exponent)); }
+  else { rounded.x <<= u32(result_exponent); }
+  return select(i32(rounded.x), -i32(rounded.x), acc < 0);
+}
+// A power-of-two division and an exactly represented <=48-bit product need
+// no binary64 boundary resolver. This includes ordinary dyadic quantization.
+fn quant_exact_dyadic_product(acc: i32, scale: f32) -> bool {
+  let magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  let significand = quant_float_significand(bitcast<u32>(scale));
+  let integer_bits = 32u - countLeadingZeros(magnitude) - countTrailingZeros(magnitude);
+  let scale_bits = 32u - countLeadingZeros(significand) - countTrailingZeros(significand);
+  return integer_bits + scale_bits <= 48u;
+}
+fn quant_normal_expansion_range(v: f32) -> bool {
+  let exponent = (bitcast<u32>(v) >> 23u) & 255u;
+  // Reserve two binary32 significands plus four carry/error bits above the
+  // minimum normal exponent, so residual arithmetic avoids FTZ underflow.
+  return exponent >= (1u + 2u * 24u + 4u) && exponent < 255u;
+}
+fn round_expansion_ratio(numerator_hi: f32, numerator_lo: f32, denominator: f32, dyadic_scale: f32, dyadic_candidate: bool, acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  let estimate = numerator_hi / denominator;
+  // Anything outside this range saturates after adding an 8-bit zero point;
+  // returning early also keeps floor/i32 conversion away from Inf and overflow.
+  if (estimate <= -1024.0) { return -1024; }
+  if (estimate >= 1024.0) { return 1024; }
+  let lo = floor(estimate);
+  let hi = lo + 1.0;
+  let midpoint = lo + 0.5;
+  let midpoint_hi = fma(midpoint, denominator, 0.0);
+  if ((bitcast<u32>(midpoint_hi) & 0x7f800000u) == 0x7f800000u) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let midpoint_lo = fma(midpoint, denominator, -midpoint_hi);
+  // The high-word subtraction is exact for these similarly sized products
+  // (Sterbenz); adding the residual difference therefore preserves its sign
+  // even when the rounded high words straddle adjacent f32 values.
+  let midpoint_delta = fma(1.0, numerator_hi - midpoint_hi, numerator_lo - midpoint_lo);
+  // Escalate ambiguous expansion signs instead of changing the tie tolerance.
+  // Six FMA/product tails and their additions fit within this conservative
+  // 16*u^2 error envelope (u=2^-24), including binary64 reference rounding.
+  if ((bitcast<u32>(midpoint_delta) & 0x7f800000u) == 0x7f800000u) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  if (abs(midpoint_delta) <= max(abs(numerator_hi), abs(midpoint_hi)) * 0x1p-44) {
+    if (!dyadic_candidate || !quant_exact_dyadic_product(acc, dyadic_scale)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  }
+  if (midpoint_delta < 0.0) {
+    return i32(lo);
+  }
+  if (midpoint_delta > 0.0) {
+    return i32(hi);
+  }
+  let half_lo = floor(lo * 0.5);
+  let lo_is_even = lo - half_lo * 2.0 == 0.0;
+  return i32(select(hi, lo, lo_is_even));
+}
+// Multiplication by a power of two keeps every i32 bit in binary64. Round the
+// exact dyadic value directly, including INT_MIN, half-even ties and saturation.
+fn quant_round_power_two(acc: i32, exponent: i32) -> i32 {
+  let magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  var rounded = 0u;
+  if (exponent >= 0) {
+    if (exponent >= 10) { return select(1024, -1024, acc < 0); }
+    rounded = min(magnitude, 1024u >> u32(exponent)) << u32(exponent);
+  } else {
+    let shift = u32(-exponent);
+    if (shift >= 32u) { return 0; }
+    let whole = magnitude >> shift;
+    let remainder = magnitude & ((1u << shift) - 1u);
+    let half = 1u << (shift - 1u);
+    rounded = min(1024u, whole + u32(remainder > half || (remainder == half && (whole & 1u) != 0u)));
+  }
+  return select(i32(rounded), -i32(rounded), acc < 0);
+}
+fn round_requantized_precise(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  if (((a_bits | b_bits | y_bits) & 0x807fffffu) == 0u && min(a_bits, min(b_bits, y_bits)) != 0u && max(a_bits, max(b_bits, y_bits)) < 0x7f800000u) {
+    let exponent = i32(a_bits >> 23u) + i32(b_bits >> 23u) - i32(y_bits >> 23u) - 127;
+    return quant_round_power_two(acc, exponent);
+  }
+  let acc_f = f32(acc);
+  // Every i32 high part aligned to 256 is exactly representable in f32.
+  // Recover the conversion residual without an out-of-range i32(acc_f)
+  // conversion when INT_MAX rounds up to 2^31.
+  let acc_tail = fma(-1.0, acc_f - f32(acc & -256), f32(acc & 255));
+  // Equal b/y scales cancel exactly. Bypass their product and division so a
+  // rounded intermediate cannot move a near-half value onto the tie.
+  if (bitcast<u32>(b_s) == bitcast<u32>(y_s)) {
+    let direct_hi = fma(acc_f, a_s, 0.0);
+    let direct_lo = fma(acc_tail, a_s, fma(acc_f, a_s, -direct_hi));
+    if (!quant_normal_expansion_range(direct_hi)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+    return round_expansion_ratio(direct_hi, direct_lo, 1.0, a_s, true, acc, a_s, b_s, y_s);
+  }
+  // Keep the scale product as a two-float expansion. Its FMA residual preserves
+  // which side of a rounding midpoint a near-half value occupies.
+  let ab_hi = fma(a_s, b_s, 0.0);
+  let ab_lo = fma(a_s, b_s, -ab_hi);
+  if (!quant_normal_expansion_range(ab_hi) || !quant_normal_expansion_range(y_s) || y_s > 0x1p126) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let product_hi = fma(acc_f, ab_hi, 0.0);
+  if (!quant_normal_expansion_range(product_hi)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let rounded_acc_tail = fma(acc_f, ab_hi, -product_hi) + fma(acc_f, ab_lo, 0.0);
+  let product_tail = fma(acc_tail, ab_hi, rounded_acc_tail) + fma(acc_tail, ab_lo, 0.0);
+  // Compare the expansion directly. Renormalizing it adds a cancelling
+  // subtraction/addition chain that can erase the residual during lowering.
+  let denominator_bits = bitcast<u32>(y_s);
+  let denominator_power_two = (denominator_bits & 0x7fffffu) == 0u;
+  let dyadic_candidate = denominator_power_two && ab_lo == 0.0;
+  return round_expansion_ratio(product_hi, product_tail, y_s, ab_hi, dyadic_candidate, acc, a_s, b_s, y_s);
+}
+// A conservative binary32 error filter accepts only a result safely inside one
+// integer rounding bin. Ambiguous boundaries keep the full expansion/binary64
+// path; this is an exact-result shortcut, not an approximate quantization mode.
+fn round_requantized(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  let ea = (a_bits >> 23u) & 255u;
+  let eb = (b_bits >> 23u) & 255u;
+  let ey = (y_bits >> 23u) & 255u;
+  // Certify the range before arithmetic: overflow need not produce infinity.
+  // Any nonzero i32 converts to a float with biased exponent <=158. These
+  // scale-only bounds cover every accumulator and can be shared across lanes.
+  // Reserve two significands plus carry bits above FTZ, and a factor of two
+  // below overflow for products, division and the expansion midpoint product.
+  if (((a_bits | b_bits | y_bits) & 0x80000000u) != 0u ||
+      min(ea, eb) == 0u || max(ea, eb) >= 255u ||
+      ea + eb < 180u || ea + eb > 347u ||
+      ey < 53u || ey > 242u || ea + eb > ey + 220u) {
+    return quant_round_binary64(acc, a_s, b_s, y_s);
+  }
+  let scale = fma(a_s, b_s, 0.0);
+  let product = fma(f32(acc), scale, 0.0);
+  let estimate = product / y_s;
+  if (estimate <= -1024.0) { return -1024; }
+  if (estimate >= 1024.0) { return 1024; }
+  let nearest = round(estimate);
+  // 32 binary32 unit roundoffs reserve room for the i32 conversion, two
+  // products, quotient, and distance arithmetic, plus binary64 rounding.
+  let margin = max(1.0, abs(estimate)) * 0x1p-19;
+  if (abs(estimate - nearest) < 0.5 - margin) { return i32(nearest); }
+{% if device.adapterInfo.vendor != "apple" %}
+  // A rounded half-integer estimate can hide which side the exact f32-scale
+  // ratio lies on. Resolve that one-ulp neighborhood conservatively outside
+  // the compact fast path; exact power-of-two ratios already have a cheap
+  // and correct path in round_requantized_precise.
+  if (abs(abs(estimate - nearest) - 0.5) <= max(1.0, abs(estimate)) * 0x1p-24) {
+    let dyadic_scales = ((a_bits | b_bits | y_bits) & 0x807fffffu) == 0u &&
+      min(a_bits, min(b_bits, y_bits)) != 0u &&
+      max(a_bits, max(b_bits, y_bits)) < 0x7f800000u;
+    if (!dyadic_scales) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  }
+{% endif %}
+  return round_requantized_precise(acc, a_s, b_s, y_s);
+}
+// Adjusted byte integers fit exactly in f32. Each bounded f32 dot product
+// has absolute magnitude at most 2^24; int32 adds then preserve wraparound.
+const TILE_M: u32 = {{ sgmatTileM }}u;
+const TILE_N: u32 = {{ sgmatTileN }}u;
+const TILE_K: u32 = {{ sgmatTileK }}u;
+{% if precastOperands %}
+const PADDED_K: u32 = {{ preparedK }}u;
+const PADDED_N: u32 = {{ preparedN }}u;
+var<workgroup> matrix_out: array<f32, {{ sgmatSubgroups * 8 * 64 }}>;
+{% else %}
+var<workgroup> tile_A: array<f32, {{ sgmatTileM * sgmatTileK }}>;
+var<workgroup> tile_B: array<f32, {{ sgmatTileK * sgmatTileN }}>;
+{% if not sgmatReuseStaging %}
+var<workgroup> scratch: array<array<array<f32, 64>, 4>, {{ sgmatSubgroups }}>;
+{% endif %}
+{% endif %}
+@compute @workgroup_size({{ sgmatWorkgroup }}){{ " @subgroup_size(32)" if pinSubgroupSize32 else "" }}
+fn main(@builtin(workgroup_id) wid: vec3<u32>,
+{% if not precastOperands or not matrixSubgroupId %}
+        @builtin(local_invocation_index) tid: u32,
+{% endif %}
+        @builtin(subgroup_invocation_id) lane: u32,
+@builtin(subgroup_id) subgroup_index: u32) {
+  let subgroup = subgroup_index;
+  let base_A = (subgroup % {{ sgmatRowGroups }}u) * 16u;
+  let base_B = (subgroup / {{ sgmatRowGroups }}u) * 32u;
+  let row = lane / 4u;
+  let col = (lane % 4u) * 2u;
+  // Bound adjusted byte magnitudes from scalar zero points. A column-wise
+  // B zero point keeps the full byte-difference bound for the whole tile.
+  let max_a = {% if matrixZeroPerRow %}255u{% else %}u32(max(abs({{ 0 if aUnsigned else (0 - 128) }} - i32(a_zero_point[0])), abs({{ 255 if aUnsigned else 127 }} - i32(a_zero_point[0])))){% endif %};
+  let max_b = {% if hasBZero and not bZeroPerColumn %}u32(max(abs({{ 0 if bUnsigned else (0 - 128) }} - i32(b_zero_point[0])), abs({{ 255 if bUnsigned else 127 }} - i32(b_zero_point[0])))){% else %}{{ 255 if hasBZero or bUnsigned else 128 }}u{% endif %};
+  let partial_k = (16777216u / (max_a * max_b) / TILE_K) * TILE_K;
+  var totals: array<i32, 16>;
+  for (var partial = 0u; partial < params.K; partial += partial_k) {
+{% for r in range(2) %}{% for c in range(4) %}
+    var matC{{ r }}{{ c }}: subgroup_matrix_result<f32, 8, 8>;
+{% endfor %}{% endfor %}
+    let partial_end = min(params.K, partial + partial_k);
+    for (var kb = partial; kb < partial_end; kb += TILE_K) {
+{% if not precastOperands %}
+      for (var index = tid; index < TILE_M * TILE_K; index += {{ sgmatWorkgroup }}u) {
+        let m = wid.y * TILE_M + index / TILE_K;
+        let k = kb + index % TILE_K;
+        var value = f32(0.0);
+        if (m < params.M && k < params.K) {
+          value = f32(i32(a[m * params.K + k]) - i32(a_zero_point[{{ "m" if matrixZeroPerRow else "0" }}]));
+        }
+        tile_A[index] = value;
+      }
+      for (var index = tid; index < TILE_K * TILE_N; index += {{ sgmatWorkgroup }}u) {
+        let k = kb + index / TILE_N;
+        let n = wid.x * TILE_N + index % TILE_N;
+        var value = f32(0.0);
+        if (k < params.K && n < params.N) {
+          value = f32(i32(b[{% if matrixBatchB %}wid.z * params.bBatchStride + {% endif %}k * params.N + n]) - i32(b_zero_point[{{ "n" if bZeroPerColumn else "0" }}]));
+        }
+        tile_B[index] = value;
+      }
+      workgroupBarrier();
+{% endif %}
+      for (var step = 0u; step < TILE_K; step += 8u) {
+{% for r in range(2) %}
+        let matA{{ r }} = subgroupMatrixLoad<subgroup_matrix_left<{{ precastScalar if precastOperands else "f32" }}, 8, 8>, row_major>(&{{ "a" if precastOperands else "tile_A" }}, {% if precastOperands %}(wid.y * TILE_M + base_A + {{ r * 8 }}u) * PADDED_K + kb + step{% else %}(base_A + {{ r * 8 }}u) * TILE_K + step{% endif %}, {{ "PADDED_K" if precastOperands else "TILE_K" }});
+{% endfor %}
+{% for c in range(4) %}
+        let matB{{ c }} = subgroupMatrixLoad<subgroup_matrix_right<{{ precastScalar if precastOperands else "f32" }}, 8, 8>, row_major>(&{{ "b" if precastOperands else "tile_B" }}, {% if precastOperands %}{% if matrixBatchB %}wid.z * PADDED_K * PADDED_N + {% endif %}(kb + step) * PADDED_N + wid.x * TILE_N + base_B + {{ c * 8 }}u{% else %}step * TILE_N + base_B + {{ c * 8 }}u{% endif %}, {{ "PADDED_N" if precastOperands else "TILE_N" }});
+{% endfor %}
+{% for r in range(2) %}{% for c in range(4) %}
+        matC{{ r }}{{ c }} = subgroupMatrixMultiplyAccumulate(matA{{ r }}, matB{{ c }}, matC{{ r }}{{ c }});
+{% endfor %}{% endfor %}
+      }
+{% if not precastOperands %}
+      workgroupBarrier();
+{% endif %}
+    }
+{% if precastOperands %}
+    // The prepared operands need no staging barriers. Publish all eight
+    // matrices together before converting the bounded partial to int32.
+{% for r in range(2) %}{% for c in range(4) %}
+    subgroupMatrixStore<row_major>(&matrix_out, subgroup * 512u + {{ r * 256 + c * 64 }}u, matC{{ r }}{{ c }}, 8u);
+{% endfor %}{% endfor %}
+    workgroupBarrier();
+{% for r in range(2) %}{% for c in range(4) %}{% for h in range(2) %}
+    totals[{{ r * 8 + c * 2 + h }}] += i32(matrix_out[subgroup * 512u + {{ r * 256 + c * 64 }}u + row * 8u + col + {{ h }}u]);
+{% endfor %}{% endfor %}{% endfor %}
+    workgroupBarrier();
+{% else %}
+{% if sgmatReuseStaging %}
+    // All staged A reads have completed. Reuse its storage for one result band
+    // at a time: each subgroup needs 4 * 64 floats, exactly its A-tile share.
+{% endif %}
+{% for r in range(2) %}
+{% for c in range(4) %}
+    subgroupMatrixStore<row_major>({% if sgmatReuseStaging %}&tile_A{% else %}&scratch[subgroup][{{ c }}]{% endif %}, {% if sgmatReuseStaging %}subgroup * 256u + {{ c * 64 }}u{% else %}0u{% endif %}, matC{{ r }}{{ c }}, 8u);
+{% endfor %}
+    workgroupBarrier();
+{% for c in range(4) %}{% for h in range(2) %}
+    totals[{{ r * 8 + c * 2 + h }}] += i32({% if sgmatReuseStaging %}tile_A[subgroup * 256u + {{ c * 64 }}u + row * 8u + col + {{ h }}u]{% else %}scratch[subgroup][{{ c }}][row * 8u + col + {{ h }}u]{% endif %});
+{% endfor %}{% endfor %}
+    workgroupBarrier();
+{% endfor %}
+{% endif %}
+  }
+{% for r in range(2) %}{% for c in range(4) %}{% for h in range(2) %}
+  {
+    let m = wid.y * TILE_M + base_A + {{ r * 8 }}u + row;
+    let n = wid.x * TILE_N + base_B + {{ c * 8 + h }}u + col;
+    if (m < params.M && n < params.N) {
+      let acc = totals[{{ r * 8 + c * 2 + h }}]{% if matrixHasBias %} + bias[{{ "m" if matrixBiasPerRow else "n" }}]{% endif %};
+      let rounded = round_requantized(acc, f32(a_scale[0]), f32(b_scale[{{ "m" if matrixScalePerRow else ("n" if bScalePerColumn else "0") }}]), f32(y_scale[0]));
+      let q = clamp(rounded + i32(y_zero_point[0]), {{ qMin }}, {{ qMax }});
+      y[{% if matrixBatchB %}wid.z * params.yBatchStride + {% endif %}m * params.N + n] = {{ "u32(q)" if yUnsigned else "q" }};
+    }
+  }
+{% endfor %}{% endfor %}{% endfor %}
+}

build/webgpu/quant-exact-portable.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,412 @@

+{% set integerStoreRow = matrixScalePerRow or (matrixHasBias and matrixBiasPerRow) %}
+{% set integerStoreColumn = (requantize and bScalePerColumn) or (matrixHasBias and not matrixBiasPerRow) %}
+// Rare-boundary fallback: reproduce the reference's binary64 scale division
+// and accumulator product using integer significands; no f64 feature required.
+fn quant_u96_add(a: vec3<u32>, b: vec3<u32>) -> vec3<u32> {
+  let low = a.x + b.x;
+  let middle = a.y + b.y;
+  let with_carry = middle + u32(low < a.x);
+  return vec3<u32>(low, with_carry, a.z + b.z + u32(middle < a.y) + u32(with_carry < middle));
+}
+fn quant_u96_shr(a: vec3<u32>) -> vec3<u32> {
+  return vec3<u32>((a.x >> 1u) | (a.y << 31u), (a.y >> 1u) | (a.z << 31u), a.z >> 1u);
+}
+fn quant_u96_round_shift(value: vec3<u32>, shift: u32) -> vec3<u32> {
+  var q = value;
+  var sticky = false;
+  var guard = false;
+  for (var bit = 0u; bit < shift; bit++) {
+    sticky = sticky || guard;
+    guard = (q.x & 1u) != 0u;
+    q = quant_u96_shr(q);
+  }
+  if (guard && (sticky || (q.x & 1u) != 0u)) {
+    q = quant_u96_add(q, vec3<u32>(1u, 0u, 0u));
+  }
+  return q;
+}
+fn quant_float_significand(bits: u32) -> u32 {
+  return (bits & 0x7fffffu) | select(0x800000u, 0u, (bits & 0x7f800000u) == 0u);
+}
+fn quant_float_exponent(bits: u32) -> i32 {
+  return max(1, i32((bits >> 23u) & 255u)) - 150;
+}
+fn quant_round_binary64(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  let a = quant_float_significand(a_bits);
+  let b = quant_float_significand(b_bits);
+  let denominator = quant_float_significand(y_bits);
+  if (a == 0u || b == 0u) { return 0; }
+  // The exact product of two 24-bit significands occupies at most 48 bits.
+  let p0 = (a & 65535u) * (b & 65535u);
+  let cross = (a >> 16u) * (b & 65535u) + (b >> 16u) * (a & 65535u);
+  let product_low = p0 + (cross << 16u);
+  let product_high = (a >> 16u) * (b >> 16u) + (cross >> 16u) + u32(product_low < p0);
+  var input_bit = select(31 - i32(countLeadingZeros(product_low)), 63 - i32(countLeadingZeros(product_high)), product_high != 0u);
+  var remainder = 0u;
+  var significant_bits = 0u;
+  var leading_exponent = 0;
+  var significand = vec2<u32>(0u);
+  // Long division generates the 53 bits of the binary64 multiplier. The
+  // remainder fits in 25 bits, including its left shift. At most 24 leading
+  // zero bits precede the first significant bit.
+  loop {
+    var incoming = 0u;
+    if (input_bit >= 32) { incoming = (product_high >> u32(input_bit - 32)) & 1u; }
+    else if (input_bit >= 0) { incoming = (product_low >> u32(input_bit)) & 1u; }
+    remainder = (remainder << 1u) | incoming;
+    let digit = u32(remainder >= denominator);
+    if (digit != 0u) { remainder -= denominator; }
+    if (significant_bits != 0u || digit != 0u) {
+      if (significant_bits == 0u) { leading_exponent = input_bit; }
+      significand = vec2<u32>((significand.x << 1u) | digit, (significand.y << 1u) | (significand.x >> 31u));
+      significant_bits++;
+    }
+    input_bit--;
+    if (significant_bits == 53u) { break; }
+  }
+  let twice_remainder = remainder << 1u;
+  if (twice_remainder > denominator || (twice_remainder == denominator && (significand.x & 1u) != 0u)) {
+    let low = significand.x + 1u;
+    significand = vec2<u32>(low, significand.y + u32(low == 0u));
+  }
+  let scale_exponent = quant_float_exponent(a_bits) + quant_float_exponent(b_bits) - quant_float_exponent(y_bits) + leading_exponent - 52;
+  // abs(INT_MIN) is representable as u32. Multiply the rounded multiplier's
+  // significand by the exact accumulator into three little-endian u32 limbs.
+  var magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  var operand = vec3<u32>(significand, 0u);
+  var product = vec3<u32>(0u);
+  loop {
+    if ((magnitude & 1u) != 0u) { product = quant_u96_add(product, operand); }
+    magnitude >>= 1u;
+    if (magnitude == 0u) { break; }
+    operand = vec3<u32>(operand.x << 1u, (operand.y << 1u) | (operand.x >> 31u), (operand.z << 1u) | (operand.y >> 31u));
+  }
+  var product_top = 31 - i32(countLeadingZeros(product.x));
+  if (product.y != 0u) { product_top = 63 - i32(countLeadingZeros(product.y)); }
+  if (product.z != 0u) { product_top = 95 - i32(countLeadingZeros(product.z)); }
+  if (product_top + scale_exponent >= 10) { return select(1024, -1024, acc < 0); }
+  if (product_top + scale_exponent < -2) { return 0; }
+  let product_shift = u32(max(0, product_top - 52));
+  let rounded_product = quant_u96_round_shift(product, product_shift);
+  let result_exponent = scale_exponent + i32(product_shift);
+  var rounded = rounded_product;
+  if (result_exponent < 0) { rounded = quant_u96_round_shift(rounded_product, u32(-result_exponent)); }
+  else { rounded.x <<= u32(result_exponent); }
+  return select(i32(rounded.x), -i32(rounded.x), acc < 0);
+}
+// A power-of-two division and an exactly represented <=48-bit product need
+// no binary64 boundary resolver. This includes ordinary dyadic quantization.
+fn quant_exact_dyadic_product(acc: i32, scale: f32) -> bool {
+  let magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  let significand = quant_float_significand(bitcast<u32>(scale));
+  let integer_bits = 32u - countLeadingZeros(magnitude) - countTrailingZeros(magnitude);
+  let scale_bits = 32u - countLeadingZeros(significand) - countTrailingZeros(significand);
+  return integer_bits + scale_bits <= 48u;
+}
+fn quant_normal_expansion_range(v: f32) -> bool {
+  let exponent = (bitcast<u32>(v) >> 23u) & 255u;
+  // Reserve two binary32 significands plus four carry/error bits above the
+  // minimum normal exponent, so residual arithmetic avoids FTZ underflow.
+  return exponent >= (1u + 2u * 24u + 4u) && exponent < 255u;
+}
+fn round_expansion_ratio(numerator_hi: f32, numerator_lo: f32, denominator: f32, dyadic_scale: f32, dyadic_candidate: bool, acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  let estimate = numerator_hi / denominator;
+  // Anything outside this range saturates after adding an 8-bit zero point;
+  // returning early also keeps floor/i32 conversion away from Inf and overflow.
+  if (estimate <= -1024.0) { return -1024; }
+  if (estimate >= 1024.0) { return 1024; }
+  let lo = floor(estimate);
+  let hi = lo + 1.0;
+  let midpoint = lo + 0.5;
+  let midpoint_hi = fma(midpoint, denominator, 0.0);
+  if ((bitcast<u32>(midpoint_hi) & 0x7f800000u) == 0x7f800000u) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let midpoint_lo = fma(midpoint, denominator, -midpoint_hi);
+  // The high-word subtraction is exact for these similarly sized products
+  // (Sterbenz); adding the residual difference therefore preserves its sign
+  // even when the rounded high words straddle adjacent f32 values.
+  let midpoint_delta = fma(1.0, numerator_hi - midpoint_hi, numerator_lo - midpoint_lo);
+  // Escalate ambiguous expansion signs instead of changing the tie tolerance.
+  // Six FMA/product tails and their additions fit within this conservative
+  // 16*u^2 error envelope (u=2^-24), including binary64 reference rounding.
+  if ((bitcast<u32>(midpoint_delta) & 0x7f800000u) == 0x7f800000u) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  if (abs(midpoint_delta) <= max(abs(numerator_hi), abs(midpoint_hi)) * 0x1p-44) {
+    if (!dyadic_candidate || !quant_exact_dyadic_product(acc, dyadic_scale)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  }
+  if (midpoint_delta < 0.0) {
+    return i32(lo);
+  }
+  if (midpoint_delta > 0.0) {
+    return i32(hi);
+  }
+  let half_lo = floor(lo * 0.5);
+  let lo_is_even = lo - half_lo * 2.0 == 0.0;
+  return i32(select(hi, lo, lo_is_even));
+}
+// Multiplication by a power of two keeps every i32 bit in binary64. Round the
+// exact dyadic value directly, including INT_MIN, half-even ties and saturation.
+fn quant_round_power_two(acc: i32, exponent: i32) -> i32 {
+  let magnitude = select(u32(acc), 0u - u32(acc), acc < 0);
+  var rounded = 0u;
+  if (exponent >= 0) {
+    if (exponent >= 10) { return select(1024, -1024, acc < 0); }
+    rounded = min(magnitude, 1024u >> u32(exponent)) << u32(exponent);
+  } else {
+    let shift = u32(-exponent);
+    if (shift >= 32u) { return 0; }
+    let whole = magnitude >> shift;
+    let remainder = magnitude & ((1u << shift) - 1u);
+    let half = 1u << (shift - 1u);
+    rounded = min(1024u, whole + u32(remainder > half || (remainder == half && (whole & 1u) != 0u)));
+  }
+  return select(i32(rounded), -i32(rounded), acc < 0);
+}
+fn round_requantized_precise(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  if (((a_bits | b_bits | y_bits) & 0x807fffffu) == 0u && min(a_bits, min(b_bits, y_bits)) != 0u && max(a_bits, max(b_bits, y_bits)) < 0x7f800000u) {
+    let exponent = i32(a_bits >> 23u) + i32(b_bits >> 23u) - i32(y_bits >> 23u) - 127;
+    return quant_round_power_two(acc, exponent);
+  }
+  let acc_f = f32(acc);
+  // Every i32 high part aligned to 256 is exactly representable in f32.
+  // Recover the conversion residual without an out-of-range i32(acc_f)
+  // conversion when INT_MAX rounds up to 2^31.
+  let acc_tail = fma(-1.0, acc_f - f32(acc & -256), f32(acc & 255));
+  // Equal b/y scales cancel exactly. Bypass their product and division so a
+  // rounded intermediate cannot move a near-half value onto the tie.
+  if (bitcast<u32>(b_s) == bitcast<u32>(y_s)) {
+    let direct_hi = fma(acc_f, a_s, 0.0);
+    let direct_lo = fma(acc_tail, a_s, fma(acc_f, a_s, -direct_hi));
+    if (!quant_normal_expansion_range(direct_hi)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+    return round_expansion_ratio(direct_hi, direct_lo, 1.0, a_s, true, acc, a_s, b_s, y_s);
+  }
+  // Keep the scale product as a two-float expansion. Its FMA residual preserves
+  // which side of a rounding midpoint a near-half value occupies.
+  let ab_hi = fma(a_s, b_s, 0.0);
+  let ab_lo = fma(a_s, b_s, -ab_hi);
+  if (!quant_normal_expansion_range(ab_hi) || !quant_normal_expansion_range(y_s) || y_s > 0x1p126) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let product_hi = fma(acc_f, ab_hi, 0.0);
+  if (!quant_normal_expansion_range(product_hi)) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  let rounded_acc_tail = fma(acc_f, ab_hi, -product_hi) + fma(acc_f, ab_lo, 0.0);
+  let product_tail = fma(acc_tail, ab_hi, rounded_acc_tail) + fma(acc_tail, ab_lo, 0.0);
+  // Compare the expansion directly. Renormalizing it adds a cancelling
+  // subtraction/addition chain that can erase the residual during lowering.
+  let denominator_bits = bitcast<u32>(y_s);
+  let denominator_power_two = (denominator_bits & 0x7fffffu) == 0u;
+  let dyadic_candidate = denominator_power_two && ab_lo == 0.0;
+  return round_expansion_ratio(product_hi, product_tail, y_s, ab_hi, dyadic_candidate, acc, a_s, b_s, y_s);
+}
+// A conservative binary32 error filter accepts only a result safely inside one
+// integer rounding bin. Ambiguous boundaries keep the full expansion/binary64
+// path; this is an exact-result shortcut, not an approximate quantization mode.
+fn round_requantized(acc: i32, a_s: f32, b_s: f32, y_s: f32) -> i32 {
+  if (acc == 0) { return 0; }
+  let a_bits = bitcast<u32>(a_s);
+  let b_bits = bitcast<u32>(b_s);
+  let y_bits = bitcast<u32>(y_s);
+  let ea = (a_bits >> 23u) & 255u;
+  let eb = (b_bits >> 23u) & 255u;
+  let ey = (y_bits >> 23u) & 255u;
+  // Certify the range before arithmetic: overflow need not produce infinity.
+  // Any nonzero i32 converts to a float with biased exponent <=158. These
+  // scale-only bounds cover every accumulator and can be shared across lanes.
+  // Reserve two significands plus carry bits above FTZ, and a factor of two
+  // below overflow for products, division and the expansion midpoint product.
+  if (((a_bits | b_bits | y_bits) & 0x80000000u) != 0u ||
+      min(ea, eb) == 0u || max(ea, eb) >= 255u ||
+      ea + eb < 180u || ea + eb > 347u ||
+      ey < 53u || ey > 242u || ea + eb > ey + 220u) {
+    return quant_round_binary64(acc, a_s, b_s, y_s);
+  }
+  let scale = fma(a_s, b_s, 0.0);
+  let product = fma(f32(acc), scale, 0.0);
+  let estimate = product / y_s;
+  if (estimate <= -1024.0) { return -1024; }
+  if (estimate >= 1024.0) { return 1024; }
+  let nearest = round(estimate);
+  // 32 binary32 unit roundoffs reserve room for the i32 conversion, two
+  // products, quotient, and distance arithmetic, plus binary64 rounding.
+  let margin = max(1.0, abs(estimate)) * 0x1p-19;
+  if (abs(estimate - nearest) < 0.5 - margin) { return i32(nearest); }
+{% if device.adapterInfo.vendor != "apple" %}
+  // A rounded half-integer estimate can hide which side the exact f32-scale
+  // ratio lies on. Resolve that one-ulp neighborhood conservatively outside
+  // the compact fast path; exact power-of-two ratios already have a cheap
+  // and correct path in round_requantized_precise.
+  if (abs(abs(estimate - nearest) - 0.5) <= max(1.0, abs(estimate)) * 0x1p-24) {
+    let dyadic_scales = ((a_bits | b_bits | y_bits) & 0x807fffffu) == 0u &&
+      min(a_bits, min(b_bits, y_bits)) != 0u &&
+      max(a_bits, max(b_bits, y_bits)) < 0x7f800000u;
+    if (!dyadic_scales) { return quant_round_binary64(acc, a_s, b_s, y_s); }
+  }
+{% endif %}
+  return round_requantized_precise(acc, a_s, b_s, y_s);
+}
+// Both integer-output and requantizing consumers share the vector GEMM body.
+// Its callback stores scalar outputs so padded operands also cover column tails.
+fn store_integer_result(value: i32, output_index: u32{% if integerStoreColumn %}, column: u32{% endif %}{% if integerStoreRow %}, row: u32{% endif %}) {
+{% if matrixHasBias %}
+  let acc = value + bias[{{ "row" if matrixBiasPerRow else "column" }}];
+{% endif %}
+  let rounded = round_requantized({{ "acc" if matrixHasBias else "value" }}, f32(a_scale[0]), f32(b_scale[{{ "row" if matrixScalePerRow else ("column" if bScalePerColumn else "0") }}]), f32(y_scale[0]));
+  let q = clamp(rounded + i32(y_zero_point[0]), {{ 0 if yUnsigned else (0 - 128) }}, {{ 255 if yUnsigned else 127 }});
+  y[output_index] = {{ "u32(q)" if yUnsigned else "q" }};
+}
+{% set exactIntegerPartialsValue = true %}
+{% set preparedNValue = preparedN | default(0) %}
+{% set integerPartialKValue = integerPartialK | default(0) %}
+{% set requantizeValue = true %}
+{% set bScalePerColumnValue = bScalePerColumn is defined and bScalePerColumn %}
+{% set exactBatchBValue = exactIntegerPartialsValue and (exactBatchB | default(false)) %}
+{% set integerStoreColumnValue = (requantizeValue and bScalePerColumnValue) or (integerStoreColumn | default(false)) %}
+{% set integerStoreRowValue = integerStoreRow | default(false) %}
+{{ env.wgsl.resourceDeclarations }}
+{% set kTileValue = kTile %}
+{% set blockRowsValue = blockRows if blockRows is defined else 128 %}
+{% set blockColsValue = blockCols if blockCols is defined else 64 %}
+{% set accumulatorScalarValue = "f32" %}
+{% set workgroupColsValue = workgroupCols if workgroupCols is defined else 16 %}
+{% set workgroupRowsValue = workgroupRows if workgroupRows is defined else 16 %}
+{% set laneCount = workgroupColsValue * workgroupRowsValue %}
+{% set threadRows = (blockRowsValue / workgroupRowsValue)|int %}
+{% set threadColVectors = (blockColsValue / (workgroupColsValue * 4))|int %}
+{% set threadCols = threadColVectors * 4 %}
+{% set aVectorCount = (blockRowsValue * kTileValue / (4 * laneCount))|int %}
+{% set bVectorCount = (kTileValue * blockColsValue / (4 * laneCount))|int %}
+{% set batchStridesValue = false %}
+// High-intensity matrix multiplication for the large, untransposed layout:
+//   Y[M,N] = alpha * A[M,K] @ B[K,N].
+// B may alternatively be stored physically as [N,K]; that mode retains vec4
+// K loads and transposes only while staging the tile in workgroup memory.
+//
+// Both inputs and the output are four-wide along their contiguous dimension.
+// The output-tile geometry and BK are selected from the shape and device's
+// workgroup-storage limit. f16 operands stay packed in global and workgroup
+// memory, then widen at the FMA so accumulation remains f32.
+const M: u32 = {{ M }}u;
+const K: u32 = {{ K }}u;
+const N: u32 = {{ N }}u;
+const K_V4: u32 = K / 4u;
+const B_N_V4: u32 = {{ preparedNValue }}u / 4u;
+const BK: u32 = {{ kTileValue }}u;
+const BM: u32 = {{ blockRowsValue }}u;
+const BN: u32 = {{ blockColsValue }}u;
+const TM: u32 = {{ threadRows }}u;
+const TN: u32 = {{ threadCols }}u;
+// A is loaded and consumed four K values at a time. Keeping that packing in
+// workgroup memory reduces A-side shared-memory instructions and inner-loop
+// control without changing the four scalar FMA operations or their order.
+var<workgroup> tileA: array<array<vec4<{{ scalar }}>, BK / 4u>, BM>;
+var<workgroup> tileB: array<vec4<{{ scalar }}>, BK * (BN / 4u)>;
+@compute @workgroup_size({{ workgroupColsValue }}, {{ workgroupRowsValue }}, 1)
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+) {
+  let li = lid.y * {{ workgroupColsValue }}u + lid.x;
+  let mBase = wg.y * BM;
+  let nBase = wg.x * BN;
+{% if exactBatchBValue %}
+  let bBatch = wg.z * K * B_N_V4;
+  let yBatch = wg.z * M * N;
+{% endif %}
+{% set aOff = "" %}
+{% set bOff = "bBatch + " if batchStridesValue or exactBatchBValue else "" %}
+{% for i in range(threadRows) %}
+{% for j in range(threadColVectors) %}
+  var acc{{ i }}_{{ j }} = vec4<{{ accumulatorScalarValue }}>(0.0);
+  var total{{ i }}_{{ j }} = vec4<i32>(0);
+{% endfor %}
+{% endfor %}
+  for (var kBase: u32 = 0u; kBase < K; kBase = kBase + BK) {
+{% for av in range(aVectorCount) %}
+    let aVector{{ av }} = li + {{ av * laneCount }}u;
+    let aRow{{ av }} = aVector{{ av }} / (BK / 4u);
+    let aColV4_{{ av }} = aVector{{ av }} % (BK / 4u);
+    let am{{ av }} = mBase + aRow{{ av }};
+    var av4_{{ av }} = vec4<{{ scalar }}>({{ scalar }}(0));
+    if (am{{ av }} < M) {
+      av4_{{ av }} = a[{{ aOff }}am{{ av }} * K_V4 + kBase / 4u + aColV4_{{ av }}];
+    }
+    tileA[aRow{{ av }}][aColV4_{{ av }}] = av4_{{ av }};
+{% endfor %}
+{% for bv in range(bVectorCount) %}
+    let bVector{{ bv }} = li + {{ bv * laneCount }}u;
+    let bRow{{ bv }} = bVector{{ bv }} / (BN / 4u);
+    let bColV4_{{ bv }} = bVector{{ bv }} % (BN / 4u);
+    let bn{{ bv }} = nBase + bColV4_{{ bv }} * 4u;
+    var bv4_{{ bv }} = vec4<{{ scalar }}>({{ scalar }}(0));
+    if (bn{{ bv }} < N) {
+      bv4_{{ bv }} = b[{{ bOff }}(kBase + bRow{{ bv }}) * {{ "B_N_V4" if exactIntegerPartialsValue else "N_V4" }} + bn{{ bv }} / 4u];
+    }
+    tileB[bRow{{ bv }} * (BN / 4u) + bColV4_{{ bv }}] = bv4_{{ bv }};
+{% endfor %}
+    workgroupBarrier();
+    let aRow = lid.y * TM;
+    let bCol4 = lid.x * {{ threadColVectors }}u;
+    for (var kv: u32 = 0u; kv < BK / 4u; kv = kv + 1u) {
+      let kk = kv * 4u;
+{% for j in range(threadColVectors) %}
+      let bv0_{{ j }} = vec4<{{ accumulatorScalarValue }}>(tileB[(kk + 0u) * (BN / 4u) + bCol4 + {{ j }}u]);
+      let bv1_{{ j }} = vec4<{{ accumulatorScalarValue }}>(tileB[(kk + 1u) * (BN / 4u) + bCol4 + {{ j }}u]);
+      let bv2_{{ j }} = vec4<{{ accumulatorScalarValue }}>(tileB[(kk + 2u) * (BN / 4u) + bCol4 + {{ j }}u]);
+      let bv3_{{ j }} = vec4<{{ accumulatorScalarValue }}>(tileB[(kk + 3u) * (BN / 4u) + bCol4 + {{ j }}u]);
+{% endfor %}
+{% for i in range(threadRows) %}
+      let av{{ i }} = tileA[aRow + {{ i }}u][kv];
+{% for j in range(threadColVectors) %}
+      acc{{ i }}_{{ j }} = acc{{ i }}_{{ j }} + vec4<{{ accumulatorScalarValue }}>({{ accumulatorScalarValue }}(av{{ i }}.x)) * bv0_{{ j }};
+      acc{{ i }}_{{ j }} = acc{{ i }}_{{ j }} + vec4<{{ accumulatorScalarValue }}>({{ accumulatorScalarValue }}(av{{ i }}.y)) * bv1_{{ j }};
+      acc{{ i }}_{{ j }} = acc{{ i }}_{{ j }} + vec4<{{ accumulatorScalarValue }}>({{ accumulatorScalarValue }}(av{{ i }}.z)) * bv2_{{ j }};
+      acc{{ i }}_{{ j }} = acc{{ i }}_{{ j }} + vec4<{{ accumulatorScalarValue }}>({{ accumulatorScalarValue }}(av{{ i }}.w)) * bv3_{{ j }};
+{% endfor %}
+{% endfor %}
+    }
+    workgroupBarrier();
+    // Adjusted byte integers and each bounded partial are exact in f32.
+    // Convert before exceeding that bound; int32 additions retain wraparound.
+    if ((kBase + BK) % {{ integerPartialKValue }}u == 0u || kBase + BK == K) {
+{% for i in range(threadRows) %}{% for j in range(threadColVectors) %}
+      total{{ i }}_{{ j }} += vec4<i32>(acc{{ i }}_{{ j }});
+      acc{{ i }}_{{ j }} = vec4<f32>(0.0);
+{% endfor %}{% endfor %}
+    }
+  }
+  let m0 = mBase + lid.y * TM;
+{% for j in range(threadColVectors) %}
+  let n{{ j }} = nBase + lid.x * TN + {{ j * 4 }}u;
+{% endfor %}
+{% for i in range(threadRows) %}
+{% for j in range(threadColVectors) %}
+  if (m0 + {{ i }}u < M && n{{ j }} < N) {
+{% for h in range(4) %}
+    if (n{{ j }} + {{ h }}u < N) {
+      store_integer_result(total{{ i }}_{{ j }}[{{ h }}], {% if exactBatchBValue %}yBatch + {% endif %}(m0 + {{ i }}u) * N + n{{ j }} + {{ h }}u{% if integerStoreColumnValue %}, n{{ j }} + {{ h }}u{% endif %}{% if integerStoreRowValue %}, m0 + {{ i }}u{% endif %});
+    }
+{% endfor %}
+  }
+{% endfor %}
+{% endfor %}
+}

build/webgpu/quant-exact-prepare.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,79 @@

+{{ env.wgsl.resourceDeclarations }}
+{% if prepareSpatialA or prepareSpatialB %}
+{% macro spatial_constants() %}
+const KERNEL_H: u32 = {{ kernelHSpec }}u;
+const KERNEL_W: u32 = {{ kernelWSpec }}u;
+const KERNEL_HW: u32 = KERNEL_H * KERNEL_W;
+const KSIZE: u32 = KERNEL_HW;
+const STRIDE_H: u32 = {{ strideHSpec }}u;
+const STRIDE_W: u32 = {{ strideWSpec }}u;
+const DILATION_H: u32 = {{ dilationHSpec }}u;
+const DILATION_W: u32 = {{ dilationWSpec }}u;
+const PAD_TOP: i32 = {{ padTopSpec }};
+const PAD_LEFT: i32 = {{ padLeftSpec }};
+{% endmacro %}
+{% macro spatial_sample(value="value", input="x", zero="x_zero_point", sampleK="k", sampleScalar="auto") %}
+  let ic = {{ sampleK }} / KSIZE;
+  let kq = {{ sampleK }} - ic * KSIZE;
+  let kh = kq / KERNEL_W;
+  let kw = kq - kh * KERNEL_W;
+  let oh = pos / params.outW;
+  let ow = pos - oh * params.outW;
+  let ih = i32(oh * STRIDE_H + kh * DILATION_H) - PAD_TOP;
+  let iw = i32(ow * STRIDE_W + kw * DILATION_W) - PAD_LEFT;
+  var {{ value }}: {{ bScalar if sampleScalar == "auto" else sampleScalar }} = {{ zero }}[0];
+  if (ih >= 0 && ih < i32(params.inH) && iw >= 0 && iw < i32(params.inW)) {
+    {{ value }} = {{ input }}[{% if sampleChannelsLast is defined and sampleChannelsLast %}((batch * params.inH + u32(ih)) * params.inW + u32(iw)) * params.inChannels + ic{% else %}((batch * params.inChannels + ic) * params.inH + u32(ih)) * params.inW + u32(iw){% endif %}];
+  }
+{% endmacro %}
+{{ spatial_constants() }}
+{% endif %}
+@compute @workgroup_size({{ prepareWorkgroup }})
+fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_index) tid: u32) {
+  let offset = ((wid.y * params.groupsX + wid.x) * {{ prepareWorkgroup }}u + tid) * 4u;
+  if (offset >= params.count) { return; }
+  if (offset < params.aCount) {
+    let m = offset / params.paddedK;
+    let k = offset % params.paddedK;
+{% for h in range(4) %}
+    {
+      var value = 0.0;
+      if (m < params.M && k + {{ h }}u < params.K) {
+{% if prepareSpatialA %}
+        let batch = m / params.outCount;
+        let pos = m % params.outCount;
+        let sourceK = k + {{ h }}u;
+{{ spatial_sample("sample", "a", "a_zero_point", "sourceK", aScalar) }}
+        value = f32(i32(sample) - i32(a_zero_point[0]));
+{% else %}
+        value = f32(i32(a[m * params.K + k + {{ h }}u]) - i32(a_zero_point[{{ "m" if prepareZeroPerRow else "0" }}]));
+{% endif %}
+      }
+      aFloat[offset + {{ h }}u] = {% if precastScalar == "f16" %}f16(value){% else %}value{% endif %};
+    }
+{% endfor %}
+  } else {
+    let bOffset = offset - params.aCount;
+{% if prepareBatchB %}
+    let batch = bOffset / (params.paddedK * params.paddedN);
+    let localOffset = bOffset % (params.paddedK * params.paddedN);
+{% endif %}
+    let k = {{ "localOffset" if prepareBatchB else "bOffset" }} / params.paddedN;
+    let n = {{ "localOffset" if prepareBatchB else "bOffset" }} % params.paddedN;
+{% for h in range(4) %}
+    {
+      var value = 0.0;
+      if (k < params.K && n + {{ h }}u < params.N) {
+{% if prepareSpatialB %}
+        let pos = n + {{ h }}u;
+{{ spatial_sample("sample", "b", "b_zero_point") }}
+        value = f32(i32(sample) - i32(b_zero_point[0]));
+{% else %}
+        value = f32(i32(b[{% if prepareTransposeB %}(n + {{ h }}u) * params.K + k{% else %}{% if prepareBatchB %}batch * params.K * params.N + {% endif %}k * params.N + n + {{ h }}u{% endif %}]) - i32(b_zero_point[{% if bZeroPerColumn %}n + {{ h }}u{% else %}0{% endif %}]));
+{% endif %}
+      }
+      bFloat[bOffset + {{ h }}u] = {% if precastScalar == "f16" %}f16(value){% else %}value{% endif %};
+    }
+{% endfor %}
+  }
+}

build/webgpu/quant-weight-k-major.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,42 @@

+{% macro quant_read_column_zero(fn_name, buffer_name, is_unsigned, present=true, per_column=false, param_name="col", return_type="i32") %}
+fn {{ fn_name }}({% if per_column %}{{ param_name }}: u32{% endif %}) -> {{ return_type }} {
+  {% if per_column %}
+  return {{ buffer_name }}[{{ param_name }}];
+  {% else %}
+  return {{ buffer_name }}[0];
+  {% endif %}
+}{% endmacro %}
+// Reorder a quantized `[M, C/group, kH, kW]` convolution weight into the K-major
+// `[KPadded, M]` operand the shared DP4A GEMM core reads as B. K enumerates
+// `(kh, kw, ic)` so the reduction's fastest axis is the input channel, matching
+// the fastest axis of a channels-last activation. Rows past the source K hold the
+// weight zero point, so padded lanes contribute exactly zero to
+// `(a - a_zero) * (b - b_zero)` while the GEMM keeps its aligned vec4 A path.
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) { return; }{% endmacro %}
+{{ env.wgsl.resourceDeclarations }}
+{{ quant_read_column_zero("read_w_zero", "w_zero_point", false, per_column=wZeroPerChannel, param_name="oc", return_type=wScalar) }}
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d(tunables.WORKGROUP_SIZE, "flat", guardInline=true) }}
+  // Neighbouring invocations cover neighbouring output channels, so both the
+  // K-major store and the requantization that later reads column `oc` stay
+  // contiguous.
+  let oc = flat % params.outChannels;
+  let k = flat / params.outChannels;
+  var value = read_w_zero({% if wZeroPerChannel %}oc{% endif %});
+  if (k < params.kRows) {
+    let ic = k % params.weightInChannels;
+    let khw = k / params.weightInChannels;
+    let kw = khw % params.kernelW;
+    let kh = khw / params.kernelW;
+    value = w[((oc * params.weightInChannels + ic) * params.kernelH + kh) * params.kernelW + kw];
+  }
+  w_k_major[flat] = value;
+}

build/webgpu/test.json ADDED Viewed

The diff for this file is too large to render. See raw diff