You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Add the opt-in loop(index, begin, end[, step]) block utility to repeat kernel blocks, declarations, assignments, and shader-helper calls during C++ shader generation. Kernel loop unrolling is one usage. Generated C++ uses range loops while emitted WGSL contains scoped copies of the body. Ordinary WGSL for loops and loop blocks retain their existing behavior.
Support nested loops, scoped integer indices, parameter and metadata bounds, positive steps, and conditional directives within loop bodies. Evaluate bounds once, use signed counters, and validate int32 ranges and positive steps. Reject shader-dependent bounds, conflicting indices, unbalanced blocks, crossed conditionals, and unsupported break/continue.
Repeated kernel block example (inside $MAIN, with kernel bindings):
#use loop
#param ELEMENTS_PER_INVOCATION
loop(block, 0, ELEMENTS_PER_INVOCATION) {
let index = global_idx * u32(ELEMENTS_PER_INVOCATION) + u32(block);
if (index < uniforms.output_size) {
let value = input[index];
output[index] = max(value, 0.0);
}
}
Nested tile example (inside $MAIN, with kernel-defined bases/strides):
A 2x3 tile emits six scoped copies. Inner bounds may reference outer indices. Parameters affecting expansion must be in the shader cache key.
Document syntax, repeated kernel blocks, nested loops, and limitations. Add regression coverage for helper arguments, conditionals, invalid headers/bounds, scope conflicts, and unchanged ordinary shader loops. Compile and execute generated C++ in both output modes to verify nested indices, empty ranges, arithmetic bounds, and invalid-step errors.
Validation: all 94 generator tests pass with MSVC execution coverage; Ruff format/check and git diff --check pass. No existing kernels migrated; GPU correctness and performance testing are outside this change.
Add the opt-in loop(index, begin, end[, step]) block utility to repeat
kernel blocks, declarations, assignments, and shader-helper calls during
C++ shader generation. Kernel loop unrolling is one usage. Generated C++
uses range loops while emitted WGSL contains scoped copies of the body.
Ordinary WGSL for loops and loop blocks retain their existing behavior.
Support nested loops, scoped integer indices, parameter and metadata
bounds, positive steps, and conditional directives within loop bodies.
Evaluate bounds once, use signed counters, and validate int32 ranges and
positive steps. Reject shader-dependent bounds, conflicting indices,
unbalanced blocks, crossed conditionals, and unsupported break/continue.
Repeated kernel block example (inside $MAIN, with kernel bindings):
#use loop
#param ELEMENTS_PER_INVOCATION
loop(block, 0, ELEMENTS_PER_INVOCATION) {
let index = global_idx * u32(ELEMENTS_PER_INVOCATION) + u32(block);
if (index < uniforms.output_size) {
let value = input[index];
output[index] = max(value, 0.0);
}
}
Nested tile example (inside $MAIN, with kernel-defined bases/strides):
#use loop
#param TILE_ROWS TILE_COLS
loop(row, 0, TILE_ROWS) {
loop(col, 0, TILE_COLS) {
let src = input_base + u32(row) * uniforms.input_stride + u32(col);
let dst = output_base + u32(row) * uniforms.output_stride + u32(col);
output[dst] = input[src];
}
}
A 2x3 tile emits six scoped copies. Inner bounds may reference outer
indices. Parameters affecting expansion must be in the shader cache key.
Document syntax, repeated kernel blocks, nested loops, and limitations.
Add regression coverage for helper arguments, conditionals, invalid
headers/bounds, scope conflicts, and unchanged ordinary shader loops.
Compile and execute generated C++ in both output modes to verify nested
indices, empty ranges, arithmetic bounds, and invalid-step errors.
Validation: all 94 generator tests pass with MSVC execution coverage;
Ruff format/check and git diff --check pass. No existing kernels migrated;
GPU correctness and performance testing are outside this change.
The reason will be displayed to describe this comment to others. Learn more.
🟡 Changes recommended
Loop-index expansion can produce invalid WGSL and uncompilable C++ helper calls, while signed-bound validation incorrectly rejects the minimum int32 value.
Adds an opt-in block-loop utility to WebGPU WGSL templates, allowing generated C++ to expand repeated shader blocks without migrating existing kernels.
Convert integral lvalues and rvalues to strings in generated helper code
without ambiguous overloads. Preserve integer indices in host expressions
and isolate expanded indices in shader text to avoid adjacent minus tokens.
Keep member selectors unchanged when they match a loop index, including
selectors separated from the dot by whitespace or a line break.
Validate signed literals in the bound expression AST so -2147483648 is
accepted while positive 2147483648 remains invalid. Update helper snapshots
and add regressions for selector handling, negative-index expressions,
string-typed helper arguments, and signed-bound limits.
Validation: all 95 template tests pass, including MSVC compilation and
execution in both generator modes. Ruff and git diff --check pass.
Treat loop(...) braces as template syntax rather than automatically
emitting a WGSL block for each iteration. Preserve explicit blocks inside
the body so template authors control shader scope. Retain C++ generation
scopes for host counters and bounds, and omit closing-header indentation
from repeated shader text.
Replace the C++ string loops introduced by microsoft#33395 with loop blocks directly
inside the FlashAttention template main body. Use head_vec_idx and
key_tile_idx for expanded indices. Preserve the NVIDIA shared-memory
head-size selection, rolled fallbacks, FP32 score accumulation, and ordered
QK/PV expressions. The vector-count parameter is derived from head size,
which already participates in cache specialization.
Update repeated-block and nested-loop documentation examples to use
explicit scopes around local declarations. Add compiled regression checks
for flat expansion and explicit per-iteration scopes. Exercise the actual
FlashAttention template generator in both output modes across head sizes
64/128/192/256/512, shared-memory and subgroup paths, expansion selections,
and 16/32 key-tile sizes. Verify inline statements and accumulation order
without helper functions or automatically added shader braces.
Validation: all 96 tests pass with MSVC compilation and execution coverage.
Configured lint, Ruff format/check, and git diff --check pass. Native GPU
correctness and performance validation have not been run in this checkout.
The reason will be displayed to describe this comment to others. Learn more.
The original multiline selector case is fixed, but a template conditional between the dot and selector still causes the selector to be expanded as a loop index:
With FLAG = 1, the previous nonempty source line is #if FLAG, not the preceding shader token .. Both generator modes emit let a = v.\n(0)\n;: the generated C++ compiles, but WGSL validation rejects it with expected identifier, found "(". The FLAG = 0 control emits v.y and validates successfully.
Please preserve member-selector context across template directives and their branches rather than relying only on the previous nonempty source line, and add regression coverage for both branches.
Track shader selector context independently of template directives and restore the entry context for each conditional branch. Merge branch endings at endif and report inconsistent selector context instead of silently expanding a member name as a loop index.
Add coverage for both generator modes, nested and alternative branches, and selectors following endif. Execute generated C++ with both conditional outcomes to verify member names remain intact.
Validation: all 99 WGSL template tests passed, including generated C++ execution; Ruff and git diff --check passed. WGSL validation was not run.
Loop index masks an output receiver used in the body
tools/python/wgsl_template/generator.py:525
With #use loop .rank, loop(output, 0, 1) whose body first uses let n = output.rank; is accepted but emits let n = (0).rank; without registering the output binding. The index pattern masks the helper receiver because the conflict check only knows variables used earlier or in the bounds. Detect receiver/index conflicts before substitution, including first uses inside the body, and add a regression case.
The redeclaration check misses address-space-qualified declarations. Inside loop(k, 0, 2), var<function> k : i32 = 0; is accepted and rewritten to var<function> (0) : i32 = 0;, producing invalid WGSL instead of the intended conflict error. Include the optional var<...> qualifier in this check and cover it with a regression case.
Reject active loop indices that mask enabled shader helper receivers first used inside the body. Recognize address-space-qualified var declarations, including spaced qualifiers, when checking loop index redeclarations.
Add regression cases for receiver conflicts and qualified declarations. Extend generated C++ execution coverage to both offsetToIndices and indicesToOffset with bare integral loop indices in both generator modes, verifying the existing pass_as_string conversion fix.
Validation: all 99 WGSL template tests passed, including generated C++ compilation and execution. Ruff and git diff --check passed.
Redeclaration check misses let/var declarations split across lines
tools/python/wgsl_template/generator.py:599
The redeclaration check only examines one line, so a declaration split as let followed by k = 1; inside loop(k, 0, 2) bypasses it. Generation succeeds and emits let\n(0) = 1; (and the corresponding second copy), which is invalid WGSL, whereas the same declaration on one line is correctly rejected. Track declaration tokens across line boundaries before substituting active indices, and cover split let/var declarations.
Make the shared shader_variable.h string converter recognize integral references by removing the reference before checking is_integral. This fixes generated C++ compilation for bare loop indices passed to GetByOffset and SetByOffset.
Compile and execute the production converter and inline offset wrappers in the regression harness, with only non-inline backends stubbed. Cover both generator modes, mutable and const integer lvalues, integer temporaries, and string_view arguments.
Validation: all 99 WGSL template tests passed, including generated C++ compilation and execution. Ruff and git diff --check passed.
Track declaration tokens across source lines before substituting active loop indices. Reject split let, var, and const declarations, including multiline address-space qualifiers, instead of emitting invalid WGSL with an expanded index as the declared name.
Preserve and merge declaration context across conditional branches. Add regression coverage in both generator modes for split declarations, conditional declarations, and valid declarations whose initializer uses a loop index.
Validation: all 101 WGSL template tests passed, including generated C++ execution. Additional conditional regression cases passed the focused generator suite. Ruff and whitespace checks passed.
Parenthesize bare loop indices in string-typed shader helper arguments so rank-one helpers cannot emit adjacent minus tokens. Preserve integral indices for host expressions and integer-typed helper arguments.
Track square-bracket depth across conditional branches and reject expressions crossing template-loop boundaries. Add regressions for nested loops and conditional branches, plus compiled production rank-one helper coverage in both generator modes.
Validation: all 103 WGSL template tests passed, including generated C++ compilation and execution. Ruff checks, formatting checks, and git diff --check passed.
The reason will be displayed to describe this comment to others. Learn more.
🔵 Needs a closer look
Human assessment is needed because a scope-validation edge case can reject valid declarations after a loop closes and was omitted from the published review.
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
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
Add the opt-in loop(index, begin, end[, step]) block utility to repeat kernel blocks, declarations, assignments, and shader-helper calls during C++ shader generation. Kernel loop unrolling is one usage. Generated C++ uses range loops while emitted WGSL contains scoped copies of the body. Ordinary WGSL for loops and loop blocks retain their existing behavior.
Support nested loops, scoped integer indices, parameter and metadata bounds, positive steps, and conditional directives within loop bodies. Evaluate bounds once, use signed counters, and validate int32 ranges and positive steps. Reject shader-dependent bounds, conflicting indices, unbalanced blocks, crossed conditionals, and unsupported break/continue.
Repeated kernel block example (inside $MAIN, with kernel bindings):
Nested tile example (inside $MAIN, with kernel-defined bases/strides):
A 2x3 tile emits six scoped copies. Inner bounds may reference outer indices. Parameters affecting expansion must be in the shader cache key.
Document syntax, repeated kernel blocks, nested loops, and limitations. Add regression coverage for helper arguments, conditionals, invalid headers/bounds, scope conflicts, and unchanged ordinary shader loops. Compile and execute generated C++ in both output modes to verify nested indices, empty ranges, arithmetic bounds, and invalid-step errors.
Validation: all 94 generator tests pass with MSVC execution coverage; Ruff format/check and git diff --check pass. No existing kernels migrated; GPU correctness and performance testing are outside this change.
Motivation and Context
See above.