feat: Early return pruning - #2711
Conversation
|
pkg.pr.new packages benchmark commit |
Bundle size comparison (
|
| 🟢 Decreased | ➖ Unchanged | 🔴 Increased (max 0.66%) | ❔ Unknown |
|---|---|---|---|
| 0 | 303 | 21 | 0 |
import * as ... in PR vs import * as ... in target (did bundle size increase?):
| Test | tsdown |
|---|---|
| STATIC_tgpu.ts | 282.66 kB ( |
| tgpu_fn.ts | 282.67 kB ( |
| tgpu_init.ts | 282.67 kB ( |
| tgpu_lazy.ts | 282.67 kB ( |
| tgpu_slot.ts | 282.67 kB ( |
| tgpu_const.ts | 282.67 kB ( |
| tgpu_unroll.ts | 282.67 kB ( |
| tgpu_resolve.ts | 282.67 kB ( |
| tgpu_accessor.ts | 282.67 kB ( |
| tgpu_comptime.ts | 282.67 kB ( |
| tgpu_vertexFn.ts | 282.67 kB ( |
| tgpu_computeFn.ts | 282.67 kB ( |
| tgpu_fragmentFn.ts | 282.67 kB ( |
| tgpu_privateVar.ts | 282.67 kB ( |
| tgpu_vertexLayout.ts | 282.68 kB ( |
| tgpu_workgroupVar.ts | 282.68 kB ( |
| tgpu_initFromDevice.ts | 282.68 kB ( |
| tgpu_bindGroupLayout.ts | 282.68 kB ( |
| tgpu_mutableAccessor.ts | 282.68 kB ( |
| tgpu_resolveWithContext.ts | 282.68 kB ( |
| STATIC_allImports.ts | 308.38 kB ( |
import { ... } in PR vs import * as ... in PR (is the library tree-Shakeable?):
| Test | tsdown |
|---|---|
| tgpu_init.ts | 273.14 kB ( |
| tgpu_initFromDevice.ts | 272.59 kB ( |
| tgpu_resolve.ts | 173.42 kB ( |
| tgpu_resolveWithContext.ts | 173.36 kB ( |
| tgpu_bindGroupLayout.ts | 73.94 kB ( |
| tgpu_mutableAccessor.ts | 68.66 kB ( |
| tgpu_accessor.ts | 68.66 kB ( |
| tgpu_privateVar.ts | 67.35 kB ( |
| tgpu_workgroupVar.ts | 67.35 kB ( |
| tgpu_const.ts | 66.77 kB ( |
| tgpu_lazy.ts | 66.56 kB ( |
| tgpu_fragmentFn.ts | 38.92 kB ( |
| tgpu_fn.ts | 38.87 kB ( |
| tgpu_vertexFn.ts | 38.74 kB ( |
| tgpu_computeFn.ts | 38.44 kB ( |
| tgpu_vertexLayout.ts | 27.57 kB ( |
| tgpu_comptime.ts | 15.18 kB ( |
| tgpu_unroll.ts | 1.75 kB ( |
| tgpu_slot.ts | 1.70 kB ( |
If you wish to run a comparison for other, slower bundlers, run the 'Tree-shake test' from the GitHub Actions menu.
Resolution Time Benchmark---
config:
themeVariables:
xyChart:
plotColorPalette: "#E63946, #3B82F6, #059669"
---
xychart
title "Random Branching (🔴 PR | 🔵 main | 🟢 release)"
x-axis "max depth" [1, 2, 3, 4, 5, 6, 7, 8]
y-axis "time (ms)"
line [0.64, 1.25, 2.83, 4.25, 5.81, 8.59, 13.23, 17.07]
line [0.64, 1.30, 2.88, 4.01, 4.65, 7.35, 14.21, 16.79]
line [0.56, 1.24, 2.58, 3.87, 4.60, 6.92, 12.74, 14.35]
---
config:
themeVariables:
xyChart:
plotColorPalette: "#E63946, #3B82F6, #059669"
---
xychart
title "Linear Recursion (🔴 PR | 🔵 main | 🟢 release)"
x-axis "max depth" [1, 2, 3, 4, 5, 6, 7, 8]
y-axis "time (ms)"
line [0.20, 0.32, 0.44, 0.51, 0.82, 0.77, 0.95, 1.05]
line [0.23, 0.38, 0.52, 0.61, 0.77, 0.84, 1.00, 1.07]
line [0.17, 0.38, 0.39, 0.48, 0.72, 0.80, 0.97, 0.98]
---
config:
themeVariables:
xyChart:
plotColorPalette: "#E63946, #3B82F6, #059669"
---
xychart
title "Full Tree (🔴 PR | 🔵 main | 🟢 release)"
x-axis "max depth" [1, 2, 3, 4, 5, 6, 7, 8]
y-axis "time (ms)"
line [0.63, 1.37, 2.50, 4.68, 8.96, 17.68, 39.97, 76.74]
line [0.73, 1.39, 3.15, 4.84, 7.98, 17.60, 37.36, 76.10]
line [0.68, 1.28, 2.55, 4.04, 7.91, 16.59, 33.18, 71.96]
|
f395a3e to
d0fdeb4
Compare
There was a problem hiding this comment.
Pull request overview
This PR updates the WGSL/GLSL code generators to inline trivially-scoped blocks, track simple control-flow, and prune unreachable statements—resulting in cleaner generated shader code and enabling early-return pruning for compile-time-known branches.
Changes:
- Introduces
ResolvedStatement+ block inlining/pruning logic in the WGSL generator (with control-flow tracking). - Adds
blockDepthto the generation context to support “safe” cosmetic pruning (e.g., redundantreturn;). - Updates a broad set of snapshots/docs to match the new, cleaner generated output (including unroll labeling and block formatting).
Reviewed changes
Copilot reviewed 37 out of 37 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| packages/typegpu/src/tgsl/wgslGenerator.ts | Implements block inlining + control-flow-aware pruning via ResolvedStatement, and updates statement/block generation accordingly. |
| packages/typegpu/src/tgsl/generationHelpers.ts | Extends GenerationCtx with blockDepth used by generator heuristics. |
| packages/typegpu/src/resolutionCtx.ts | Implements blockDepth on the resolution/generation context stack. |
| packages/typegpu-gl/src/glslGenerator.ts | Adapts GLSL generator to new _block(...): ResolvedStatement API and empty-body handling. |
| packages/typegpu-gl/tests/glslGenerator.test.ts | Updates snapshots to match new inlined-block output. |
| packages/typegpu/tests/vector.test.ts | Updates snapshots for simplified return formatting (no redundant braces). |
| packages/typegpu/tests/unroll.test.ts | Updates snapshots for unrolled-loop formatting + new // --- separators. |
| packages/typegpu/tests/tgslFn.test.ts | Updates snapshots for empty function body formatting (fn x() {} style). |
| packages/typegpu/tests/tgsl/wgslGenerator.test.ts | Updates snapshots across many cases to reflect block inlining and pruning behavior. |
| packages/typegpu/tests/tgsl/typeInference.test.ts | Updates snapshots for empty functions ({}) and removed redundant return;. |
| packages/typegpu/tests/tgsl/ternaryOperator.test.ts | Updates snapshot formatting for empty function output. |
| packages/typegpu/tests/tgsl/extensionEnabled.test.ts | Updates snapshots reflecting inlined returns in generated WGSL. |
| packages/typegpu/tests/tgsl/entryFnParamPruning.test.ts | Updates snapshots reflecting inlined single-statement blocks in entry functions. |
| packages/typegpu/tests/struct.test.ts | Updates snapshots removing redundant return; and simplifying empty helper bodies. |
| packages/typegpu/tests/std/environment.test.ts | Updates snapshots reflecting inlined returns. |
| packages/typegpu/tests/slot.test.ts | Updates snapshots reflecting inlined single-statement blocks. |
| packages/typegpu/tests/pipelineInit.test.ts | Updates snapshots for empty compute entry formatting (fn ...() {}) and helper formatting. |
| packages/typegpu/tests/mutabilityTracking.test.ts | Updates snapshots reflecting inlined unrolled blocks and simplified empty loops. |
| packages/typegpu/tests/lazy.test.ts | Updates snapshots for empty overload bodies (fn fill(...) {}). |
| packages/typegpu/tests/computePipeline.test.ts | Updates snapshots for empty pipeline entry formatting and inlined assignments. |
| packages/typegpu/tests/array.test.ts | Updates snapshots removing redundant trailing return;. |
| apps/typegpu-docs/tests/individual-example-tests/*.test.ts | Updates example snapshots to match new unroll labeling, block inlining, and // --- separators. |
| apps/typegpu-docs/src/content/docs/apis/utils.mdx | Updates documentation examples to match the new generated formatting. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
d0fdeb4 to
ce42040
Compare
|
@author — I finished reviewing PR #2711. I found a correctness regression in the early-return pruning logic, so I’m requesting changes. BugIn if (body === '') {
return { code: '', controlFlow: undefined };
}If an inlineable block’s only executable statement is const f = tgpu.fn([d.ptrFn(d.i32)])((out) => {
'use gpu';
{
return;
}
out.$ = 2; // incorrectly preserved
});It should produce an empty/equivalent function; instead it emits Suggested fixPreserve the terminator even when the body is empty: if (body === '') {
return { code: '', controlFlow };
}I verified this one-line change fixes the regression and the existing unit-test suite (167 files / 2190 tests) still passes. Other notes
Please apply the fix, add a regression test, and re-request review.
|
ce42040 to
637d800
Compare
![pullfrog[bot]](https://avatars.githubusercontent.com/in/1768019?s=60&v=4){kind=small}
There was a problem hiding this comment.
I found a correctness regression: an inlineable block whose only statement is return; loses its terminator when the body is pruned to empty, so unreachable code after it is wrongly emitted. Details inline.
I also left a maintainability note about the magic blockDepth thresholds.
Kimi K2 (free via Pullfrog for OSS) ｜ 𝕏
There was a problem hiding this comment.
Latest commit fixes the empty-inline-block regression I called out (preserving controlFlow when the pruned body is empty) and adds regression tests for return/break/continue pruning. The full relevant test suite passes (166 files / 2189 tests). I still think the blockDepth <= 2 / <= 3 thresholds would benefit from a clearer explanation, but that’s non-blocking.
Approving.
Kimi K2 (free via Pullfrog for OSS) ｜ 𝕏
637d800 to
89c5524
Compare
There was a problem hiding this comment.
✅ No new issues found.
Reviewed changes — since the last review, the author squashed the branch to a single commit that adds early-return termination to unrolled loops.
- Stopped generating remaining unrolled iterations once a control-flow terminator is encountered —
wgslGenerator.tsnow propagates the resolved block'scontrolFlowfrom the unrolled loop and breaks the iteration loop. - Added a regression test in
unroll.test.tsthat verifies an earlyreturninside atgpu.unrollstops further iterations and prunes the trailing unreachablereturn.
Targeted tests pass: 2 test files / 122 tests.
Kimi K2 (free via Pullfrog for OSS) ｜ 𝕏
89c5524 to
4f78fbf
Compare
There was a problem hiding this comment.
✅ No new issues found.
Reviewed changes — since the prior review of 89c55242, the branch was force-pushed to a single squash (4f78fbf7) rebased onto current main. I compared the PR's own added lines between the two squashes: the pruning logic in wgslGenerator.ts and the GLSL generator is byte-identical to the previously approved version. The delta is purely rebase adaptation plus snapshot churn.
- Relocated plumbing onto the post-#2808 context shape —
blockDepthmoved intoResolutionCtx(types.ts/resolutionCtx.ts, incl. theItemStateStackImpl.blockDepthgetter, unchanged in behavior) andResolvedStatementmoved toshaderGenerator.ts, re-exported throughtypegpu/~internalfor the GLSL generator. - Adapted the GLSL generator to the new
_block(block, allowInlining, externalMap)signature returningResolvedStatement. - Updated 7 more docs-example snapshots (
caustics,game-of-life,perlin-noise,probability,smoky-triangle,uniformity,vaporrave) to the inlined-block / pruned-return;output now produced on the newer base.
I re-validated the pruning semantics with scratch probe functions (early return; inside if/while/for bodies and nested blocks is correctly preserved so unreachable code stays reachable-only-when-the-return-isn't-hit; trailing returns are stripped only at the function-body or inlineable-block level). Verification: full typegpu suite (175 files / 2320 tests), @typegpu/gl (7 files / 50 tests), and the 7 new example snapshots all pass; lint and typecheck clean.
DeepSeek Flash (free via Pullfrog for OSS) ｜ 𝕏
f7f4b92 to
cc56060
Compare
There was a problem hiding this comment.
✅ No new issues found.
Reviewed changes — the delta since the prior review of 4f78fbf is a single force-pushed re-squash (f7f4b92d, same parent 0a113e7) that refactors the codegen plumbing without changing generated output.
- Extracted
ResolutionCtx.getDedented(code)— replaces the inlinebody.trim().replaceAll('\n ', '\n')used when inlining blocks. SinceINDENT[1]is literally two spaces, output is byte-identical; the extraction is a clarity improvement over a magic string and lives on the internal~internalsurface. - Made
ResolvedStatement.definesInNearestScoperequired and threaded it through_block/_blockStatementand the unroll accumulator — the flag is provably inert at block level (the inline path is only reachable when!definesInNearestScope, and the non-inline branch hardcodesfalse), so the change tightens the type contract without altering behavior.
Targeted codegen tests (wgslGenerator 96, unroll, typeInference — 153 total) pass. Prior feedback on the empty-block controlFlow regression and its tests remains sound.
DeepSeek Flash (free via Pullfrog for OSS) ｜ 𝕏
There was a problem hiding this comment.
ℹ️ No critical issues — one minor doc nit inline.
Reviewed changes — the delta since the prior review of f7f4b92d is a single commit (cc56060b) that renames the ResolvedStatement field controlFlow to endsWithControlFlow to match its clarified JSDoc semantics, plus a small tidy-up of _blockStatement to spread the resolved properties.
- Renamed
controlFlow→endsWithControlFlowacrossshaderGenerator.tsandwgslGenerator.ts(interface,_block,_blockStatement, return/break/continue handlers, unroll accumulator). The rename is mechanically complete — no stray references remain in code, and the GLSL generator only consumes.code, so it's unaffected. - Simplified
_blockStatementto destructurecodeand spread the remaining properties, avoiding the prior triple-field repetition.
The change is behavior-preserving (a pure rename + spread refactor); targeted codegen tests still pass.
DeepSeek Flash (free via Pullfrog for OSS) ｜ 𝕏
74d49fa to
ceb6fe8
Compare
There was a problem hiding this comment.
ℹ️ No behavioral issues found — one intentional-looking branch in the new pruning threshold is actually unreachable (see inline).
Reviewed changes — since the prior review of 74d49fa5, this squashed commit reworks the trailing-return; pruning threshold in _block: it introduces a named functionInitialBlockDepth = 2 constant (replacing the old magic 2/3), tightens the non-inline bound <= 2 → === 2 (equivalent), and changes the inline bound <= 3 → functionInitialBlockDepth - 1.
- Named the pruning depth constant —
functionInitialBlockDepth = 2replaces the magic numbers, and the function-entry invariant message is clarified (still requires depth1on entry). This also addresses the earlier "magic numbers" note. - Tightened the non-inline bound
<= 2→=== 2— equivalent, since non-inline blocks never run below depth 2. - Changed the inline bound
<= 3→<= functionInitialBlockDepth - 1— this is the whole behavioral delta of the PR's latest commit; see the inline note on the dead arm.
I ran the full wgslGenerator (96), unroll, and typeInference suites (153 tests) — all pass, confirming the change is behavior-preserving across the corpus.
DeepSeek Flash (free via Pullfrog for OSS) ｜ 𝕏
2ce6a9b to
38eddbe
Compare
There was a problem hiding this comment.
✅ No new issues found.
Reviewed changes — since the prior review of ceb6fe87, the branch was re-squashed to a single commit (2ce6a9b4) rebased onto current main. The only PR-originated code delta is a simplification of the trailing-return; pruning threshold in _block; everything else in the tree delta (ternary aliasing guard, flush removal, two example thumbnails) is base drift from commits #2834 / #2836 / #2839, not this PR.
- Removed the dead inline-arm from the pruning threshold —
_blocknow uses a plainblockDepth === functionInitialBlockDepthinstead of the ternarywillInline ? blockDepth <= functionInitialBlockDepth - 1 : .... This is exactly the fix suggested in the prior review: the inline arm was unreachable (inline blocks always run at depth ≥ 3), so generated output is unchanged. The earlier dead-arm thread is already[RESOLVED]on GitHub, and the "magic numbers" note was addressed by the namedfunctionInitialBlockDepthconstant in the prior round.
I re-ran the 153 targeted codegen tests (wgslGenerator 96, unroll, typeInference) — all pass, confirming the change is behavior-preserving.
DeepSeek Flash (free via Pullfrog for OSS) ｜ 𝕏
There was a problem hiding this comment.
✅ No new issues found.
Reviewed changes — since my prior review of 2ce6a9b4, the only new delta is a single regression test added to wgslGenerator.test.ts.
- Added
prunes unreachable statementsregression test — asserts that the statementcounter.$\+\+after a trailingreturn;is pruned, so thetgpu.privateVarbacking it is never referenced/declared in the generated WGSL. The inline snapshot (fn main() { const hello = 3; }) would fail if unreachable code after a terminator leaked back into the output, so the test genuinely pins the early-return pruning behavior rather than absorbing it.
I ran the new test — it passes.
DeepSeek Flash (free via Pullfrog for OSS) ｜ 𝕏
aleksanderkatan
left a comment
aleksanderkatan
left a comment
There was a problem hiding this comment.
Great job!
Our resolution starts becoming too smart for our tests. Maybe we should consider having some internal settings for disabling comptime evaluation, branch pruning, return pruning etc.?
| { | ||
| continue; | ||
| } | ||
| continue; |
There was a problem hiding this comment.
This continue looks prunable, will it be prined in the future?
| false, | ||
| )};`; | ||
| return { | ||
| code: stitch`${this.ctx.pre}var ${varName} = ${tryConvertSnippet( |
There was a problem hiding this comment.
Why doesn't this go through emitVarDecl?
| const resolved = | ||
| id.value !== undefined && id.value !== null ? this.ctx.resolveSnippet(id).value : ''; | ||
| return resolved ? `${this.ctx.pre}${resolved};` : ''; | ||
| return { code: resolved ? `${this.ctx.pre}${resolved};` : '', definesInNearestScope: false }; |
There was a problem hiding this comment.
Not sure how we handle raw code snippets, but at this moment, this test breaks:
it('does not break with raw code snippets', () => {
const snip = tgpu['~unstable'].rawCodeSnippet('const a = 1', d.Void);
const fn = () => {
'use gpu';
for (const a of tgpu.unroll([1, 2, 3])) {
snip.$;
}
};
expect(tgpu.resolve([fn])).toMatchInlineSnapshot(`
"fn fn_1() {
// unrolled iteration #0
const a = 1;
// unrolled iteration #1
const a = 1;
// unrolled iteration #2
const a = 1;
// ---
}"
`);
});
38eddbe to
3be11a2
Compare
3be11a2 to
2f06959
Compare
Improvements: