webgpu: fix OOB subgroup-matrix loads for trailing partial tiles

onnxruntime/core/providers/webgpu/math/subgroup_matrix_gemm.cc

@@ -113,6 +113,12 @@ class SubgroupMatrixGemmImpl final : public Gemm::GemmOptImpl {
       return Status::OK();
     }
 
+    // The kernel keeps its operand loads in bounds by shifting a trailing partial
+    // tile back, which is only possible when the tile fits within M and N.
+    if (M < tiling->tile_m || N < tiling->tile_n) {
+      return Status::OK();
+    }
+
     TensorShape output_shape{{static_cast<int64_t>(M), static_cast<int64_t>(N)}};
     auto* output = context.Output(0, output_shape);
     if (output->Shape().Size() == 0) {

onnxruntime/core/providers/webgpu/math/subgroup_matrix_gemm_8x16x16.wgsl.template

@@ -31,17 +31,18 @@
 //
 // Preconditions enforced by the host: K % sg_mat_k == 0; the load row stride is
 // even (Intel f16 subgroupMatrixLoad reads columns as 32-bit pairs), i.e. N even
-// when trans_b == 0 and M even when trans_a == 1. M and N may otherwise be any
-// size; partial tiles are handled by bounds-checked stores at write-out. The
-// matching loads in the last partial M/N tiles do read A rows >= M and B columns
-// >= N, but since K is never partial those over-read values only feed output
-// elements with m >= M or n >= N, which the bounds-checked stores discard - so
-// they can never corrupt a valid result, on any GPU. The out-of-range reads
-// themselves are made safe by WebGPU robust buffer access in the wasm/browser
-// build; the native (Dawn) build enables the disable_robustness device toggle, so
-// there they instead rely on backend-level bounds (e.g. D3D12 returns 0 for
-// out-of-range buffer reads), which holds on the Intel backends this
-// subgroup-matrix path is gated to. Gemm is strictly 2D (no batching).
+// when trans_b == 0 and M even when trans_a == 1; and the tile fits the matrix
+// (kTileM <= M, kTileN <= N - the host falls back otherwise).
+//
+// M and N may otherwise be any size: a trailing partial tile is shifted back to end
+// exactly at M / N so no operand load runs past the end of A or B, and the
+// rows/columns the shifted tile re-covers are skipped at write-out. The shift is
+// required for correctness, not just safety - a subgroupMatrixLoad whose footprint
+// leaves the buffer does not degrade to a per-element zero fill: Tint's robustness
+// transform resets the whole load to offset 0 with the minimum stride, so it
+// silently returns a different, valid tile, which would corrupt the in-range rows
+// of that tile too.
+// Gemm is strictly 2D (no batching).
 
 #param has_c
 #param trans_a
@@ -80,9 +81,13 @@ $MAIN {
     // m_tile * num_n_tile + n_tile.
     let num_n_tile = (uniforms.N + kTileN - 1u) / kTileN;
     let n_tile = workgroup_idx % num_n_tile;   // which kTileN-wide column tile
-    let global_base_n = n_tile * kTileN;
     let m_tile = workgroup_idx / num_n_tile;   // which kTileM-tall row tile
-    let m_base = m_tile * kTileM;
+    let n_tile_start = n_tile * kTileN;
+    let m_tile_start = m_tile * kTileM;
+    // Shift a trailing partial tile back so its operand loads stay inside A and B.
+    // The max() only guards the u32 subtraction; the host guarantees the tile fits.
+    let global_base_n = min(n_tile_start, max(uniforms.N, kTileN) - kTileN);
+    let m_base = min(m_tile_start, max(uniforms.M, kTileM) - kTileM);
     let k_blocks = uniforms.K / kSgMatK;
     let sg_index = local_idx / kSubgroupSize;          // which split-K subgroup (0..kSplitK-1)
     let sg_lane = local_idx % kSubgroupSize;           // lane within the subgroup (0..kSubgroupSize-1)
@@ -539,16 +544,20 @@ $MAIN {
     // Write-out pass: every subgroup writes out whole M-rows from the summed tile
     // in slot 0, striding by kSplitK subgroups (subgroup sg_index handles rows
     // sg_index, sg_index + kSplitK, ...). The subgroup's lanes cooperate on a
-    // row's N elements (strided by the subgroup size). M and N are bounds-checked
-    // so non-multiple dimensions only write valid data. Each element is scaled by
-    // alpha and, when has_c, has beta * C[m, n] added (C broadcast via strides).
+    // row's N elements (strided by the subgroup size). Rows and columns before the
+    // unshifted tile origin belong to the previous tile and are skipped, so a
+    // shifted tile never writes the same output element twice; M and N are
+    // bounds-checked so non-multiple dimensions only write valid data. Each element
+    // is scaled by alpha and, when has_c, has beta * C[m, n] added (C broadcast via
+    // strides).
+    let n_skip = n_tile_start - global_base_n;
     let n_count = min(kTileN, uniforms.N - global_base_n);
     for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) {
         let global_m = m_base + r;
-        if (global_m < uniforms.M) {
+        if (global_m >= m_tile_start && global_m < uniforms.M) {
             let scratch_base = r * kTileN;
             let out_base = global_m * uniforms.N + global_base_n;
-            for (var i: u32 = sg_lane; i < n_count; i = i + kSubgroupSize) {
+            for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) {
                 var val = output_element_t(scratch[scratch_base + i]) * output_element_t(uniforms.alpha);
 #if has_c
                 let global_n = global_base_n + i;

onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul.cc

@@ -154,13 +154,26 @@ class SubgroupMatrixMatMulImpl final : public MatMulOptImpl {
       return Status::OK();
     }
 
+    // N_b is just N rounded up to even - compute it before doing any padding work so
+    // the tile-fit check below can bail out without a wasted pad dispatch.
+    uint32_t N_b = N;
+    if (needs_padded_b) {
+      ORT_RETURN_IF_NOT(N < std::numeric_limits<uint32_t>::max(),
+                        "Cannot pad odd-N B because N+1 exceeds uint32_t range.");
+      N_b = N + 1;
+    }
+
+    // The kernel keeps its operand loads in bounds by shifting a trailing partial
+    // tile back, which is only possible when the tile fits within M and N.
+    if (M < tiling->tile_m || N_b < tiling->tile_n) {
+      return Status::OK();
+    }
+
     // The optimized path will run: now materialize the even-strided B for odd N.
     const Tensor* b_used = b;
-    uint32_t N_b = N;
     if (needs_padded_b) {
-      ORT_RETURN_IF_ERROR(EnsurePaddedB(context, *b, b_shape, N));
+      ORT_RETURN_IF_ERROR(EnsurePaddedB(context, *b, b_shape, N, N_b));
       b_used = padded_b_.get();
-      N_b = padded_b_stride_;
     }
 
     const Tensor* bias = has_bias ? inputs[2] : nullptr;
@@ -211,10 +224,8 @@ class SubgroupMatrixMatMulImpl final : public MatMulOptImpl {
   Status EnsurePaddedB(ComputeContext& context,
                        const Tensor& b,
                        const TensorShape& b_shape,
-                       uint32_t N) const {
-    ORT_RETURN_IF_NOT(N < std::numeric_limits<uint32_t>::max(),
-                      "Cannot pad odd-N B because N+1 exceeds uint32_t range.");
-    const uint32_t n_b = N + 1;
+                       uint32_t N,
+                       uint32_t n_b) const {
     TensorShapeVector padded_dims{b_shape.GetDims().begin(), b_shape.GetDims().end()};
     padded_dims.back() = static_cast<int64_t>(n_b);
     const TensorShape padded_shape{padded_dims};
@@ -241,7 +252,6 @@ class SubgroupMatrixMatMulImpl final : public MatMulOptImpl {
       }
       if (s.IsOK()) {
         padded_b_ = std::move(padded);
-        padded_b_stride_ = n_b;
       }
     });
     // padded_b_ persists the outcome across calls: call_once runs the body only on
@@ -256,7 +266,6 @@ class SubgroupMatrixMatMulImpl final : public MatMulOptImpl {
   // Cached even-strided B for odd N; built once by EnsurePaddedB.
   mutable std::once_flag pad_once_;
   mutable std::unique_ptr<Tensor> padded_b_;
-  mutable uint32_t padded_b_stride_ = 0;
 };
 
 Status GenerateShaderCode8x16x16(ShaderHelper& shader,

onnxruntime/core/providers/webgpu/math/subgroup_matrix_matmul_8x16x16.wgsl.template

@@ -23,11 +23,18 @@
 // distributed round-robin across the subgroups, each accumulates its own partial
 // tile in shared memory, and the partials are summed at write-out.
 //
-// Preconditions enforced by the host: K % sg_mat_k == 0, and B's row stride N_b
-// is even (the load requires 4-byte-aligned row starts). M and N may be any size;
-// partial tiles are handled by bounds-checked stores. Out-of-range A/B loads
-// return zero (WebGPU bounds checking). Output width/stride uses N (not N_b), so
-// any padded B columns in [N, N_b) are read but never written.
+// Preconditions enforced by the host: K % sg_mat_k == 0; B's row stride N_b is
+// even (the load requires 4-byte-aligned row starts); and the tile fits the
+// matrix (kTileM <= M, kTileN <= N - the host falls back otherwise). M and N may
+// otherwise be any size: a trailing partial tile is shifted back to end exactly
+// at M / N_b so no operand load runs past the end of A or B, and the rows/columns
+// the shifted tile re-covers are skipped at write-out. The shift is required for
+// correctness, not just safety - a subgroupMatrixLoad whose footprint leaves the
+// buffer does not degrade to a per-element zero fill: Tint's robustness transform
+// resets the whole load to offset 0 with the minimum stride, so it silently
+// returns a different, valid tile, which would corrupt the in-range rows of that
+// tile too. Output width/stride uses N (not N_b), so a padded B column in
+// [N, N_b) may be read but is never written.
 //
 // Batching: the host dispatches one num_n_tile x num_m_tile tile grid per batch
 // slice; the batch slice is recovered from the flattened workgroup_idx (the
@@ -73,9 +80,16 @@ $MAIN {
     let batch_id = workgroup_idx / tiles_per_slice;     // which batch slice (0 for a shared 2D weight)
     let slice_idx = workgroup_idx % tiles_per_slice;    // tile index within the slice
     let n_tile = slice_idx % uniforms.num_n_tile;       // which kTileN-wide column tile
-    let global_base_n = n_tile * kTileN;
     let m_tile = slice_idx / uniforms.num_n_tile;       // which kTileM-tall row tile
-    let m_base = m_tile * kTileM;
+    let n_tile_start = n_tile * kTileN;
+    let m_tile_start = m_tile * kTileM;
+    // Shift a trailing partial tile back so its operand loads stay inside A and B.
+    // The N shift uses B's padded stride N_b so the shifted origin stays even (the
+    // load needs 4-byte-aligned row starts); the extra column is clipped at
+    // write-out. The max() only guards the u32 subtraction; the host guarantees
+    // the tile fits.
+    let global_base_n = min(n_tile_start, max(uniforms.N_b, kTileN) - kTileN);
+    let m_base = min(m_tile_start, max(uniforms.M, kTileM) - kTileM);
     // Flat-element offsets into A/B/output for this batch slice, derived from
     // M/N/K. For a shared 2D weight batch_id is 0, so B collapses to its base.
     let a_batch_offset = batch_id * uniforms.M * uniforms.K;
@@ -467,15 +481,18 @@ $MAIN {
     // Write-out pass: every subgroup writes out whole M-rows from the summed tile
     // in slot 0, striding by kSplitK subgroups (subgroup sg_index handles rows
     // sg_index, sg_index + kSplitK, ...). The subgroup's lanes cooperate on a
-    // row's N elements (strided by the subgroup size). M and N are bounds-checked
-    // so non-multiple dimensions only write valid data.
+    // row's N elements (strided by the subgroup size). Rows and columns before the
+    // unshifted tile origin belong to the previous tile and are skipped, so a
+    // shifted tile never writes the same output element twice; M and N are
+    // bounds-checked so non-multiple dimensions only write valid data.
+    let n_skip = n_tile_start - global_base_n;
     let n_count = min(kTileN, uniforms.N - global_base_n);
     for (var r: u32 = sg_index; r < kTileM; r = r + kSplitK) {
         let global_m = m_base + r;
-        if (global_m < uniforms.M) {
+        if (global_m >= m_tile_start && global_m < uniforms.M) {
             let scratch_base = r * kTileN;
             let out_base = out_batch_offset + global_m * uniforms.N + global_base_n;
-            for (var i: u32 = sg_lane; i < n_count; i = i + kSubgroupSize) {
+            for (var i: u32 = sg_lane + n_skip; i < n_count; i = i + kSubgroupSize) {
                 write_output(out_base + i, global_base_n + i,
                              output_value_t(scratch[scratch_base + i]));
             }