[WebGPU] Make MatMulNaiveProgram's pipeline cache key cover everything it bakes into WGSL by 4n4ny4 · Pull Request #32048 · microsoft/onnxruntime

Ananya Anand (4n4ny4) · 2026-08-13T00:02:02Z

Description

MatMulNaiveProgram bakes three things into its generated WGSL that its pipeline cache hint did not
declare. Any of them can serve a shader compiled for one configuration to a different one, which
produces wrong results with no error and no warning.

1. The activation kind. Two convolutions with identical shapes and different activations hashed to
the same key, so the first shader compiled was reused for the second.

2. Activation parameters, via a formatting mismatch. Activation::ToString() streamed floats
through an unconfigured std::stringstream ( 6 significant digits) while GetActivationSnippet()
emits them with std::to_string (6 decimal places). Values agreeing to 6 significant digits but
differing as floats therefore produced one key and two different shaders. The maximally separated
colliding pair is 1000015.0 / 1000025.0, 160 float32 ULPs apart, both keyed as 1.00002e+06.
Fixed centrally in ToString() with std::setprecision(std::numeric_limits<float>::max_digits10),
so all six call sites that use it as a hint are corrected at once.

3. is_channels_last. It selects between bias[col] and bias[row + i] in the same program, and
conv.cc varies it, but it was absent from the hint. Added at both MatMulNaiveProgram call sites.

Defects 1 and 3 are the same omission: five of the six programs on main that bake an activation into
WGSL already declared it in their key, and four of those five also declared is_channels_last.
MatMulNaiveProgram was the sole holdout on both counts. Defect 2 is pre-existing and shared
(main's ToString() is byte-identical) which is why the fix is in ToString() rather than at the
call site.

Motivation and Context

Reachable from Conv today: a 1x1 kernel with unit stride and no padding lowers to a matmul, and when N
and K are both under 8 it dispatches MatMulNaiveProgram.

The regression tests put both convolutions in a single graph on purpose. Sequentially created
sessions each get a fresh pipeline cache WebGpuContextFactory refcounts its contexts and destroys
one at refcount zero) so running one activation per session cannot reproduce the bug. Concurrent
sessions are different: contexts are held in a process-global map keyed by context_id, so two live
sessions share one cache. That is what makes defect 3 reachable, and the layout test builds exactly
that configuration.

azure-pipelines · 2026-08-13T00:02:11Z

Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run.

Sushanth Rajasankar (sushraja-msft) · 2026-08-13T00:48:19Z

/azp run

azure-pipelines · 2026-08-13T00:48:30Z

Azure Pipelines: Successfully started running 1 pipeline(s).

Copilot

Pull request overview

Fixes a WebGPU pipeline-cache key collision where MatMulNaiveProgram embeds fused activation logic into generated WGSL but previously did not include the activation in its cache hint, allowing different activations to reuse the wrong compiled shader and silently produce incorrect results. Adds a targeted regression test that forces both activation variants to compile within the same session (shared pipeline cache).

Changes:

Include activation serialization in MatMulNaiveProgram cache hints for both the fused-conv small-matmul path and the MatMul call site.
Add a WebGPU regression test that builds a single graph with two identical 1x1 convs differing only by activation and verifies fusion occurs and results match the unfused baseline.
Document why ASSERT_NO_FATAL_FAILURE wrappers are required in the WebGPU fusion test helper.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/test/optimizer/webgpu_fusion_test_util.h	Adds clarifying comments explaining why fatal-failure propagation wrappers are required in the helper.
onnxruntime/test/optimizer/graph_transform_test.cc	Adds a WebGPU regression test covering pipeline-cache key collisions across different fused activations in one session.
onnxruntime/core/providers/webgpu/nn/conv.cc	Adds activation to the `MatMulNaiveProgram` cache hint for the small-matmul conv lowering path.
onnxruntime/core/providers/webgpu/math/matmul.cc	Adds activation to the `MatMulNaiveProgram` cache hint at the MatMul call site for consistency with shared program usage.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Jiajia Qin (qjia7)

The fix direction looks right. I confirmed that the existing Copilot comment about activation-parameter formatting is valid. Please fix the serialization centrally in Activation::ToString() so all current users benefit, and extend the regression coverage with the same activation kind but different parameters (for example, two LeakyRelu nodes with different alpha values). The test comment saying this is unrelated to activation parameters should be updated as well.

I left one additional inline comment about the remaining layout-dependent shader state.

…g it bakes into WGSL MatMulNaiveProgram bakes the activation expression directly into its generated WGSL, but neither of its call sites put the activation into the shader cache hint. Two convolutions with identical shapes and different activations therefore hashed to the same pipeline cache key: whichever compiled first was served to both, and the second silently evaluated the wrong activation. This produces incorrect output with no error and no warning. It is reached from Conv when a 1x1 kernel with unit stride and no padding lowers to a matmul whose N and K are both under 8 -- for example a channel-projection convolution with fewer than 8 channels on each side. Three distinct ways the key could under-distinguish are fixed here. 1. The activation kind. conv.cc is the site that can actually collide. matmul.cc always constructs Activation(), so its hint cannot vary today, but it shares the same program class and is updated to match so a future fused matmul path cannot silently reintroduce the bug. 2. The activation parameters. Activation::ToString streamed its floats through an unconfigured std::stringstream, which formats to 6 significant digits, while GetActivationSnippet emits them with std::to_string, which formats to 6 decimal places. Values that agree to 6 significant digits but differ as floats therefore produced one key and two different shaders: 1000015.0 and 1000025.0 are 160 float32 ULPs apart, yet both format as "1.00002e+06". Setting max_digits10 makes the key round-trip exactly, so it can never under-distinguish regardless of how the shader chooses to spell the value. ToString is shared by all six cache hints that use it, so fixing it there fixes every call site at once. 3. is_channels_last. It selects between bias[col] and bias[row + i] in the same WGSL and varies across conv.cc's calls, but was absent from the hint. The pipeline cache is not per-session -- WebGpuContextFactory keeps contexts in a process-global, reference-counted map -- so two sessions that are alive at the same time with opposite preferred layouts share one cache and can collide on it. Three regression tests, each verified to fail when its own fix alone is reverted: - WebGpuSmallMatMulConvDistinguishesActivationsInPipelineCache puts two convolutions in a single graph so both fused shaders are compiled against one pipeline cache within one session. That detail is what makes the collision observable: running one activation per session does not reproduce it, because each session then needs only one MatMulNaive variant. Relu and Sigmoid disagree on every input. Reverting the conv.cc hint alone fails it with all 384 elements differing, reporting a Relu value where a Sigmoid was expected. - WebGpuSmallMatMulConvDistinguishesActivationParamsInPipelineCache covers the parameter collision end to end with two HardSigmoid activations whose alphas are 1000015 and 1000025. A key collision bounds the alphas' relative difference at ~1e-5, which is below any usable float32 tolerance for an activation like LeakyRelu whose output is alpha*x. HardSigmoid defeats that bound because clamp(alpha * value + beta, 0, 1) pins the output to [0, 1] regardless of alpha's magnitude, turning the difference into a full-scale 0-vs-1 swing. Reverting only the setprecision line fails it with 256 of 384 elements differing by exactly 1. - WebGpuConcurrentLayoutConvsDistinguishedInPipelineCache holds two sessions with opposite preferred layouts alive simultaneously so they share one process-global cache, and shapes them so every other component of the key agrees. Reverting only the is_channels_last hint fails it with 6 of 9 elements differing, each receiving the bias of the wrong channel. - ActivationCacheKeyTest pins the formatter itself. It needs no GPU, so it runs on any CI agent. Each test also asserts the state it depends on -- that both convolutions absorbed their activation, or that the convolution is assigned to the WebGPU EP -- so none can pass trivially if fusion or EP assignment stops firing. Both orderings are exercised in every case so a collision in either direction fails. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Sushanth Rajasankar (sushraja-msft) · 2026-08-14T00:27:19Z

/azp run

azure-pipelines · 2026-08-14T00:27:29Z

Azure Pipelines: Successfully started running 1 pipeline(s).

Sushanth Rajasankar (sushraja-msft)

Other than the test files, your change looks functionally good. Ill take a look at the tests once you address the issue about comments in this PR.

Responds to review feedback on microsoft#32048: comments should state the invariant being enforced and why it matters locally, rather than narrating implementation history or every way a future change could invalidate a test. Deletes one comment that only described a hypothetical regression, and compresses six others that either duplicated the rationale already present at the implementation site or enumerated failure modes the assertion itself makes obvious. Comments only; no code changes. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

Ananya Anand (4n4ny4) · 2026-08-21T18:14:00Z

/azp run

azure-pipelines · 2026-08-21T18:14:07Z

Commenter does not have sufficient privileges for PR 32048 in repo microsoft/onnxruntime

#32116) ### Description Moves fused-activation parameters (LeakyRelu alpha, Clip min/max, HardSigmoid alpha/beta) out of the generated WGSL and into uniforms, for the seven activation kinds that already exist (None, Relu, Sigmoid, Clip, HardSigmoid, LeakyRelu, Tanh). Shader text now depends only on the activation kind, so Activation::ToString() no longer emits parameter values and models that vary a parameter reuse the cached pipeline. Activation uniforms occupy fixed trailing slots in each program's uniform list, appended last so definitions and values stay index-aligned. Tests (9 new): ActivationCacheKeyTest.ParametersDoNotAffectTheKey - the cache-key invariant directly. WebGpuSmallMatMulConvSharesPipelineAcrossParameterValues - two different HardSigmoid alphas share one pipeline and still produce different results. WebGpuConv{Relu,LeakyRelu,HardSigmoid,Clip}FusionMatchesUnfusedResults - execution parity against the unfused graph, since this PR changes runtime behaviour for the parameterized kinds. WebGpuIm2ColConv{Relu,LeakyRelu,HardSigmoid}FusionMatchesUnfusedResults - the same for the im2col path. Golden fixture updates reflect this PR's own template edit. ### Motivation and Context Parameters were baked into the generated WGSL as literals, so every distinct parameter value produced different shader text and therefore a separate shader compile and a separate pipeline. A model with several LeakyRelu slopes paid a compile per slope. Stacked on #32048 --------- Authored-by: Ananya Anand <t-anaanand@microsoft.com>

Jiajia Qin (qjia7) requested a lite review from Copilot August 13, 2026 04:30

Copilot started reviewing on behalf of Jiajia Qin (qjia7) August 13, 2026 04:31 View session

Copilot AI reviewed Aug 13, 2026

View reviewed changes

Comment thread onnxruntime/core/providers/webgpu/nn/conv.cc Outdated

Jiajia Qin (qjia7) reviewed Aug 13, 2026

View reviewed changes

Comment thread onnxruntime/core/providers/webgpu/nn/conv.cc Outdated

Ananya Anand (4n4ny4) force-pushed the webgpu-matmulnaive-activation-cache-key branch from 1ce9e2e to 37d1b05 Compare August 13, 2026 23:56

Ananya Anand (4n4ny4) changed the title ~~[WebGPU] Include the activation in MatMulNaiveProgram's pipeline cache key~~ [WebGPU] Make MatMulNaiveProgram's pipeline cache key cover everything it bakes into WGSL Aug 13, 2026