[WebGPU] Webgpu im2col fused activation - #32185
Open
Ananya Anand (4n4ny4) wants to merge 7 commits into
Open
Conversation
…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>
LeakyRelu alpha, Clip min/max and HardSigmoid alpha/beta were baked into the generated WGSL, so every distinct parameter value compiled a separate shader and a separate pipeline. Pass them as uniforms instead: the shader text now depends only on the activation kind, so parameter changes 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. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
The previous regeneration updated the per-template generated headers but
not the aggregate index.h, so generated/nn/im2col_matmul.h referenced
params.param_activation_kind against a struct that did not declare it.
Regenerated with the canonical command documented in test_in_tree_smoke.py:
UPDATE_WGSL_GOLDEN=1 python wgsl_template/test/run_tests.py
Only index.h changed in each variant. index_impl.h, string_table.h and
every file under generated/ were already byte-identical to the generator
output, so the earlier regeneration was correct apart from this one file.
Of the 60 added lines per variant, one is the missing
param_activation_kind member; the other 59 are declarations for the three
math/subgroup_matrix_* templates, which have been on main since 56598d6
without their goldens ever being regenerated. Their generated/math/*.h
bodies are deliberately left out here: they are supplied by microsoft#32115, which
fixes that pre-existing drift on its own. The in-tree golden smoke test
therefore still reports a file-set difference for those three headers
until microsoft#32115 lands, but no longer reports any content mismatch.
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
…factor
The uniforms refactor and the im2col fusion change were landing together, but
they are different kinds of change. Moving activation parameters into uniforms
preserves existing behaviour for kernels that already ran. Enabling fusion on
the im2col path removes a `// TODO: Support fuse` guard and makes a different
kernel run for real fused NHWC convolutions.
That second change cannot be exercised anywhere available. Im2ColMatMulProgram
requires an adapter reporting vendor "intel" and architecture "xe-2lpg",
"xe-2hpg", "xe-3lpg" or "xe-3lpg-xs" -- Lunar Lake, Battlemage or Panther Lake.
No ORT CI agent has one, and no machine available during development does
either, so its three parity tests GTEST_SKIP everywhere. It deserves its own
review rather than riding along with a behaviour-preserving refactor.
Reverted to their state on main:
nn/im2col_matmul.cc
nn/im2col_matmul.h
nn/im2col_matmul.wgsl.template
conv.cc keeps its four activation-uniform hunks but reverts the three that were
only there to carry the Activation into the im2col path: both
CanApplyIm2ColMatMulProgram call sites go back to passing
`activation_.activation_kind_ != ActivationKind::None`, and
ApplyIm2ColMatMulProgram loses its activation argument. Reverting the three
im2col files without these would not compile.
graph_transform_test.cc loses RunWebGpuIm2ColActivationParity, its three tests,
and `#include <fstream>`, which this branch added solely for that helper's
profile reader.
fuse_utils.h loses the comment claiming ActivationKind values are mirrored by
im2col_matmul.wgsl.template, which is no longer true here. Everything else it
gained is still required: WEBGPU_PROGRAM_ACTIVATION_UNIFORM_VARIABLES is used by
six programs, and AppendActivationUniformsData is called at eight sites.
Goldens regenerated with the canonical command:
UPDATE_WGSL_GOLDEN=1 python wgsl_template/test/run_tests.py
Beyond the expected im2col_matmul.h churn, this renumbers __str_N indices in
static-cpp/generated/tensor/{pad,oihw_to_ohwi}.h, because the six activation
string fragments left the shared string table and every later index shifts down
by six. static-cpp-literal is unaffected there since it embeds its strings. The
index_impl.h change is the content-hash comment on the im2col include.
The generated/math/subgroup_matrix_*.h bodies remain deliberately absent; they
are supplied by microsoft#32115. The in-tree golden smoke test therefore still reports a
file-set difference for those three headers, and no content mismatch.
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Im2ColMatMulProgram previously refused any fused activation: CanApplyIm2Col- MatMulProgram returned false whenever the Conv carried one, behind a `// TODO: Support fuse`. Fused NHWC fp16 convolutions therefore fell through to conv2d_mm, grouped_conv or matmul even where im2col was otherwise the better kernel. This adds an activation epilogue to im2col_matmul.wgsl.template covering the same six kinds the other WebGPU Conv kernels support -- Relu, Sigmoid, Clip, HardSigmoid, LeakyRelu, Tanh -- and lifts the guard. The epilogue is only tractable because activation parameters are now uniforms. The template dispatches on a single `activation_kind` int, so one branch per kind is enough; had the parameter values still been baked into the shader text, every distinct alpha or clip bound would have needed its own generated variant. This therefore depends on the uniforms refactor and must land after it. Parameters are read as uniforms.activation_param_0/1, which exist because the program picked up WEBGPU_PROGRAM_ACTIVATION_UNIFORM_VARIABLES and calls AppendActivationUniformsData. Slot usage matches GetActivationUsedUniformCount exactly: none for Relu/Sigmoid/Tanh, one for LeakyRelu, two for Clip and HardSigmoid. Two guards keep the template and the C++ enum from drifting apart. static_asserts in im2col_matmul.cc pin each ActivationKind to the numeric value the template tests, and IsActivationSupported enumerates the kinds explicitly so a newly added enumerator falls to `default: return false` and disables the path rather than silently generating no epilogue. CanApplyIm2ColMatMulProgram now takes the Activation rather than a bool, at both call sites in conv.cc. Keeping ComputeInternal and PrePackInternal in agreement matters: they must reach the same conclusion or prepacked weights are produced for a path that will not read them. Reachability is narrow, and no available hardware can execute it. See the pull request description for the five gates and for what is and is not verified. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Ananya Anand (4n4ny4)
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Im2ColMatMulProgramrefused any fused activation, so fused NHWC fp16 convs fell through to another kernel even where im2col was the better choice. This adds an activation epilogue toim2col_matmul.wgsl.templatefor the same six kinds the other Conv kernels support (Relu, Sigmoid, Clip, HardSigmoid, LeakyRelu, Tanh) and removes the// TODO: Support fuseguard.Depends on #32116 and should land after it. The template reads parameters from
uniforms.activation_param_0/1, which only exist because of the uniforms refactor there. With values still baked into shader text, every distinct alpha or clip bound would have needed its own generated variant, which is why the TODO was there.This path has not been executed anywhere.
IsDeviceSupported()requires vendorintelplus architecturexe-2lpg,xe-2hpg,xe-3lpgorxe-3lpg-xs(Lunar Lake, Battlemage, Panther Lake). No ORT CI agent has one and neither did any machine I had while writing this, so the parity tests skip everywhere and this code has never run. Please weigh it as unexecuted.Reachability is narrow, all five required: qualifying adapter, fp16 only, channels last, group 1, non 1x1 kernel.
Motivation and Context
Every other WebGPU Conv kernel already fuses activations. im2col was the only one that did not.
Four parity tests cover Relu, LeakyRelu, HardSigmoid and Clip. The Clip one is new. Clip is the only two slot activation whose slots mean
{min, max}instead of{alpha, beta}, so it is the one case where mis indexingactivation_param_0/1would still pass the HardSigmoid test. It skips like the others but belongs in the file so it runs as soon as someone with the right hardware builds this.static_asserts pin eachActivationKindto the value the template checks, so reordering the enum breaks the build, andIsActivationSupportedfalls todefault: return falsefor a new enumerator rather than emitting no epilogue.Goldens regenerated with
UPDATE_WGSL_GOLDEN=1 python wgsl_template/test/run_tests.py.generated/math/subgroup_matrix_*.hare still missing because they belong to #32115, so the smoke test reports a file set difference for those three and no content mismatch.