Skip to content

[webgpu] Add block loop utility to WebGPU WGSL templates - #34097

Open
Jianhui Dai (daijh) wants to merge 9 commits into
microsoft:mainfrom
daijh:webgpu-template-loop
Open

Jianhui Dai (daijh) wants to merge 9 commits into
microsoft:mainfrom
daijh:webgpu-template-loop

Conversation

@daijh

Copy link
Copy Markdown
Contributor

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):

  #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.

Motivation and Context

See above.

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.
Copilot AI balanced review requested due to automatic review settings November 11, 2026 02:21
@azure-pipelines

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

@daijh

Copy link
Copy Markdown
Contributor Author

@daijh

Copy link
Copy Markdown
Contributor Author

Yang Gu (@gyagp)

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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.

4 open issues
What changed in this PR

Adds an opt-in block-loop utility to WebGPU WGSL templates, allowing generated C++ to expand repeated shader blocks without migrating existing kernels.

Changes:

  • Supports nested loops, scoped indices, host-side bounds, positive steps, and conditional directives.
  • Adds validation and regression tests, including generated-C++ execution in both output modes.
  • Documents syntax, examples, limitations, and shader-cache requirements.
File Description
tools/​python/​wgsl_template/​test/​test_generator.py Adds loop validation, generation, and C++ execution tests.
tools/​python/​wgsl_template/​generator.py Implements loop parsing, bounds validation, scoped indices, and expansion.
tools/​python/​wgsl_template/​code_pattern.py Registers the opt-in loop pattern.
onnxruntime/​core/​providers/​webgpu/​wgsl_templates/​README.md Documents loop usage and constraints.

Review effort: Balanced

state.current_bracket_state = 1

elif ptype == "loop-index":
output("expression", str(pattern.replace))

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

Comment thread tools/python/wgsl_template/generator.py Resolved
Comment thread tools/python/wgsl_template/generator.py Resolved
Comment thread tools/python/wgsl_template/generator.py Resolved
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.
@qjia7

Copy link
Copy Markdown
Contributor

Could you use this new feature to replace the code changes in #33395 to test this feature?

@daijh

Copy link
Copy Markdown
Contributor Author

Could you use this new feature to replace the code changes in #33395 to test this feature?

Let me include the #33395 improvement in this PR.

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.
@daijh

Copy link
Copy Markdown
Contributor Author

Could you use this new feature to replace the code changes in #33395 to test this feature?

Done.

@daijh

Copy link
Copy Markdown
Contributor Author

Static kernel checks passed. Since I lack a local NVIDIA GPU, I'll rely on CI for hardware verification.

@daijh

Copy link
Copy Markdown
Contributor Author

/run bots

@qjia7

Copy link
Copy Markdown
Contributor

The request was received. Approve workflow runs. Thanks.

atoms[atom] = f"static_cast<int64_t>({segment.content})"
expression += atom

def render(node: ast.AST) -> str:

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.
Annotate WGSL template loop errors as NoReturn.

Comment thread tools/python/wgsl_template/generator.py Outdated
Comment on lines +337 to +341
for previous in reversed(pass1[: state.current_line]):
if previous.line.strip():
before = previous.line.rstrip()
break
if before.endswith("."):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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:

#use loop
#param FLAG
loop(x, 0, 1) {
  {
    let v = vec2<f32>(1.0);
    let a = v.
#if FLAG
      x
#else
      y
#endif
    ;
  }
}

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done. Preserve WGSL member selectors across template conditionals.

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.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Bare loop indices passed to the offset helpers generate non-compiling C++, while missed receiver and redeclaration conflicts can produce invalid WGSL.

2 open issues
2 previously missed issues

In code that hasn't changed since last review

Medium severity 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.

Medium severity Redeclaration check misses address-space-qualified variables

tools/​python/​wgsl_template/​generator.py:589

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.

3 resolved since last review

Review effort: Balanced

output("expression", str(pattern.replace))
output("code", ")")
return
output("expression", str(pattern.replace))
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.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Needs a closer look

Bare loop indices still fail to compile in the real shader access helpers because their string conversion does not handle integral references.

2 open issues
1 previously missed issue

In code that hasn't changed since last review

Medium severity 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.

Review effort: Balanced

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.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Needs a closer look

Human assessment is needed because negative helper arguments and bracket boundaries crossing template loops can produce invalid WGSL.

2 open issues

Review effort: Balanced

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.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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.

2 open issues

Review effort: Balanced

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants