feat: Immediates - #2812
feat: Immediates#2812reczkok wants to merge 2 commits into
Conversation
|
pkg.pr.new packages benchmark commit |
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.89, 1.68, 3.58, 5.27, 6.47, 10.65, 20.33, 20.24]
line [0.84, 1.68, 3.76, 5.58, 6.85, 10.79, 20.66, 22.56]
line [0.88, 1.70, 3.67, 5.85, 6.03, 10.28, 21.53, 21.40]
---
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.28, 0.47, 0.66, 0.77, 0.99, 1.09, 1.24, 1.41]
line [0.32, 0.48, 0.64, 0.70, 0.92, 1.00, 1.18, 1.34]
line [0.27, 0.42, 0.61, 0.71, 0.93, 1.05, 1.23, 1.39]
---
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.73, 2.02, 3.05, 5.98, 10.63, 22.51, 48.64, 99.14]
line [0.83, 1.94, 2.96, 5.85, 10.88, 23.47, 47.52, 97.19]
line [0.83, 1.98, 3.21, 5.79, 10.43, 22.55, 49.04, 96.74]
|
Bundle size comparison (
|
| 🟢 Decreased | ➖ Unchanged | 🔴 Increased (max 1.77%) | ❔ Unknown |
|---|---|---|---|
| 0 | 39 | 286 | 0 |
import * as ... in PR vs import * as ... in target (did bundle size increase?):
| Test | tsdown |
|---|---|
| STATIC_tgpu.ts | 277.66 kB ( |
| tgpu_initFromDevice.ts | 267.04 kB ( |
| tgpu_init.ts | 267.58 kB ( |
| STATIC_allImports.ts | 303.91 kB ( |
| d_ref.ts | 5.03 kB ( |
If you wish to run a comparison for other, slower bundlers, run the 'Tree-shake test' from the GitHub Actions menu.
dbf4cea to
9af8a36
Compare
868be67 to
e991685
Compare
There was a problem hiding this comment.
🟡 Changes recommended
Immediate WGSL emission/validation is incomplete (extension directive handling and size/limit validation), which can cause runtime shader/layout failures.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds experimental “immediates” support to TypeGPU via tgpu['~unstable'].immediateVar, integrating WebGPU’s var<immediate> / pass.setImmediates workflow into resolution, pipelines, pass state, serialization, and docs.
Changes:
- Introduces
immediateVarresource type (schema validation, resolution-time constraints, snapshots/caching, serialization “souls”). - Plumbs immediate usage through resolution and pipeline compilation (pipeline layout
immediateSize, draw/dispatch-timesetImmediateswrites, snapshot dedup behavior). - Adds extensive test coverage and documentation for immediates and fallback patterns.
File summaries
| File | Description |
|---|---|
| packages/typegpu/tests/tgsl/letDeclaration.test.ts | Adds regression test for disallowing handle-space aliasing in const declarations. |
| packages/typegpu/tests/serial.test.ts | Verifies immediate vars and pipeline-held immediate data serialize/restore correctly. |
| packages/typegpu/tests/immediates.test.ts | New test suite covering resolution, validation, runtime writes, overrides, caching, and accessors. |
| packages/typegpu/src/wgslExtensions.ts | Adds constants for WGSL language extensions (e.g. immediate_address_space). |
| packages/typegpu/src/types.ts | Extends ResolutionCtx with registerImmediate to enforce 1 immediate per module. |
| packages/typegpu/src/tgsl/wgslGenerator.ts | Adds immediate address space support; improves error for invalid pointer aliasing; blocks immediate mutation. |
| packages/typegpu/src/tgsl/shaderGenerator.ts | Expands variable scope typing to include immediate. |
| packages/typegpu/src/tgpuUnstable.ts | Exports immediateVar from the unstable barrel. |
| packages/typegpu/src/serial/restore.ts | Adds restore support for immediate-var souls. |
| packages/typegpu/src/resolutionCtx.ts | Tracks and returns usedImmediate; enforces single immediate var per shader module. |
| packages/typegpu/src/indexNamedExports.ts | Exposes MissingImmediatesError, isImmediateVar, and TgpuImmediateVar publicly. |
| packages/typegpu/src/errors.ts | Adds MissingImmediatesError with guidance on providing values. |
| packages/typegpu/src/data/wgslTypes.ts | Adds immediate to AddressSpace. |
| packages/typegpu/src/data/snippet.ts | Adds immediate origin and pointer params mapping. |
| packages/typegpu/src/core/slot/slotTypes.ts | Allows immediates to fulfill accessor/slot inputs. |
| packages/typegpu/src/core/root/rootTypes.ts | Adds enabledWgslLanguageFeatures API for language-feature detection. |
| packages/typegpu/src/core/root/init.ts | Implements enabledWgslLanguageFeatures via navigator.gpu.wgslLanguageFeatures. |
| packages/typegpu/src/core/pipeline/renderPipeline.ts | Supports pipeline-held immediates, soul/priors plumbing, and immediateSize in layout creation. |
| packages/typegpu/src/core/pipeline/priors.ts | Marks immediatesMap as transferable prior state. |
| packages/typegpu/src/core/pipeline/drawState.ts | Adds immediate snapshot state + write/dedup behavior for render/compute passes. |
| packages/typegpu/src/core/pipeline/computePipeline.ts | Supports pipeline-held immediates, soul/priors plumbing, and immediateSize in layout creation. |
| packages/typegpu/src/core/immediate/immediateVar.ts | Implements immediateVar, schema validation, snapshots, caching, and usage validation. |
| packages/typegpu/src/core/commandEncoder/renderPass.ts | Adds typed pass.setImmediates(immediate, value) and resets immediate cache on executeBundles. |
| packages/typegpu/src/core/commandEncoder/computePass.ts | Adds typed pass.setImmediates(immediate, value) for compute passes. |
| packages/typegpu-testing-utility/src/extendedIt.ts | Extends mocks to include setImmediates and wgslLanguageFeatures. |
| packages/typegpu-gl/tests/glslGenerator.test.ts | Adds test ensuring GLSL generation rejects immediates. |
| apps/typegpu-docs/src/content/docs/apis/pipelines.mdx | Documents immediates API, performance notes, constraints, and feature-detection fallback. |
| apps/typegpu-docs/src/content/docs/apis/accessors.mdx | Documents immediates as an accessor fulfillment option + updates comparison table. |
Review details
- Files reviewed: 28/28 changed files
- Comments generated: 3
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
e991685 to
3ed038b
Compare
![pullfrog[bot]](https://avatars.githubusercontent.com/in/1768019?s=60&v=4){kind=small}
There was a problem hiding this comment.
Important
Two things worth addressing before merge: the schema-size constraints are documented but never enforced (cryptic browser OperationError/GPU validation errors for valid-looking schemas like d.f16), and restored pipeline souls can silently substitute the default immediate value when a host serializer doesn't preserve object identity across the soul graph.
Reviewed changes
immediateVarcore — newTgpuImmediateVartype (var<immediate>codegen,$gpuValueOf/$access), schema validation rejecting arrays/atomics/booleans, capture-at-provide-time snapshots with a verbatim byte-copy path for typed arrays, default values, andMissingImmediatesError.- Pass/pipeline state wiring —
pass.setImmediatesandpipeline.with(immediate, value)with a shared-snapshot +(identity, generation)dedup cache inside the typed draw/dispatch state machine (raw-accessed passes andexecuteBundleshandled),immediateSizethreaded intocreatePipelineLayout. - Single-immediate-per-shader rule — enforced in
ResolutionCtx.registerImmediate(stricter than WGSL's per-entry-point limit; deliberate and tested). - Serialization/restore + feature detection —
immediate-varsoul restorer,immediatesMapas a transferable prior,root.enabledWgslLanguageFeatures. - Docs/tests — accessors + pipelines docs, GLSL rejection, ~50 mock and snapshot tests including struct-padding
immediateSize.
⚠️ Restored pipeline souls silently swap pipeline-held immediates for the default
INTERNAL_restoreRenderPipeline/INTERNAL_restoreComputePipeline rebuild usedImmediate (into the core memo) and the immediatesMap keys (new Map(soul.immediates)) from two separate points in the transferred soul graph. When a host serializer re-claims each nested TgpuImmediateVarSoul independently — exactly the model of the repo's own deepRoundTrip harness — the two decode to distinct instances, state.immediates.get(memo.usedImmediate) misses, and writeImmediates falls back to defaultSnapshots: the pipeline-provided value is silently replaced by the default, or the draw throws MissingImmediatesError when there is no default. Bind groups suffer the same split but fail loudly (MissingBindGroupsError); immediates are the first pipeline-held state that degrades silently. The new serial test round-trips the pipeline only by reference (roundTrip), so the split is unpinned.
Technical details
# Restore identity split between `usedImmediate` and `immediatesMap` keys
## Affected sites
- packages/typegpu/src/serial/restore.ts — `'immediate-var': (soul) => immediateVar(...)` recreates a fresh instance per occurrence.
- packages/typegpu/src/core/pipeline/computePipeline.ts:193,197 (`INTERNAL_restoreComputePipeline`): `usedImmediate: soul.usedImmediate` and `immediatesMap: new Map(soul.immediates)` are rebuilt as separate instances.
- packages/typegpu/src/core/pipeline/renderPipeline.ts:397,401 (same pattern for render pipelines).
- packages/typegpu/src/core/immediate/immediateVar.ts:192-196 (`writeImmediates`): `snapshots.get(immediate) ?? defaultSnapshots.get(immediate)` — the map miss silently degrades to the default (or throws `MissingImmediatesError`).
## Required outcome
- After `restoreResource` of a pipeline that holds `.with(immediate, value)`, drawing/dispatching it must write the same bytes it held before the transfer.
- When identity cannot be preserved, degrade loudly (throw) rather than silently substituting the default.
## Suggested approach
- During restore, key the `immediatesMap` with the same instance used for `memo.usedImmediate` (e.g. rebuild the map from `soul.immediates` but re-keyed by the restored `usedImmediate` on a match, or store the map keyed by an identity that survives transfer, like the schema).
- Add a `deepRoundTrip`-style test (serial.test.ts:38 `deepRoundTrip`) for a pipeline holding immediates that asserts the written bytes match the pipeline-held value even when the var appears in several soul fields.
## Open questions for the human
- Does the React Native serialization path (packages/typegpu-react/src/react-native/serialization/register-serializables.ts + transfer-cache) actually preserve object identity across claimed resources, or does it exhibit the split? If it preserves identity, this is latent rather than live, but a non-deduping serializer still corrupts silently.ℹ️ Nitpicks
- In
setImmediateSnapshot, a typed array shorter than the schema skips the fast path and goes throughwriteToArrayBuffer, whosecalculateOffsetstreats the short input as fully consuming its own bytes — it copies fewer bytes with no warning, leaving stale tail bytes in the reused snapshot buffer (longer inputs do get thelogger.warn). Asymmetric; consider warning on short inputs too. (immediateVar.ts:149-165)
DeepSeek Flash (free via Pullfrog for OSS) ｜ 𝕏
| ); | ||
| } | ||
|
|
||
| return sizeOf(immediate.dataType); |
There was a problem hiding this comment.
The schema-size constraints documented in this PR ("size must be a multiple of 4 bytes", "limited by the device's maxImmediateSize") are never enforced. assertValidImmediateSchema accepts d.f16 (sizeOf 2) and d.vec3h (6), so immediateVar(d.f16) constructs happily and only fails later, cryptically: immediateSize not a multiple of 4 / over the limit is a createPipelineLayout validation error, and a non-multiple-of-4 setImmediates data size throws a raw browser OperationError at draw time. validateImmediateUsage already receives the root with device here, so both checks are cheap to add at pipeline creation, and the construction-time % 4 check gives the same friendly early failure the array/atomic/bool checks provide.
There was a problem hiding this comment.
Reviewed changes (incremental since e991685)
- Docs cleanup — removed the
.$name('colorImmediate')/.$name('viewProj')calls from the accessors and fallback examples inaccessors.mdx/pipelines.mdx, matching the rest of the snippets (errors still fall back to<unnamed>, so the dropped names don't lose anything).
No new issues in this delta; the two findings from the full review stand.
DeepSeek Flash (free via Pullfrog for OSS) ｜ 𝕏
aleksanderkatan
left a comment
aleksanderkatan
left a comment
There was a problem hiding this comment.
Looks good! 🏊🔉
I leave some requests, mostly minor ones.
I just wonder, caching immediates intuitively looks to me like it could be vastly simplified, couldn't it be something like this?
const imm = pipeline.usedImmediate;
const value = pass.valueOf(imm) ?? pipeline.valueOf(imm) ?? imm.default
// + some mechanism for cheking if this value was already written,
// if we always copy immediates to new ArrayBuffers,
// then a reference check should sufficeSome parts look like they could have been done in separate PRs, this review took me over 5 hours and I still didn't properly parse everything.
| Note that `pipeline.with(pass)` returns a new pipeline wrapper on every call, so hoist it out of draw loops (`const bound = pipeline.with(pass)`) to let deduplication kick in. For values that change between draws, prefer pass-level state (`pass.setBindGroup`, `pass.setImmediates`) over the allocating `with*` methods. | ||
|
|
There was a problem hiding this comment.
I need to create a new pass in each frame, right?
Did you mean something like "hoist it to the top of draw loops"? Also, I don't understand what deduplication means in this context.
There was a problem hiding this comment.
This will change after the precedence update I'm planning but deduplication refers to .setXYZ commands
| const mainFragment = tgpu.fragmentFn({ out: d.vec4f })(() => tint.$); | ||
| ``` | ||
|
|
||
| The value is provided at execution time, either per-pipeline (`pipeline.with(immediate, value)`) or per-pass (`pass.setImmediates`). |
There was a problem hiding this comment.
I think that "per-pass" usually means that it does not change during a pass, this is not the case, correct?
There was a problem hiding this comment.
I don't think I understand the question
There was a problem hiding this comment.
I just think that the "per-" wording may be misleading, I interpret "X is set per-Y" as "X doesn't change during the scope of Y". For example, color attachment is proivided per pass, but e.g. pipeline and immediates are not provided per-pass, because these can change. I'm not sure if this makes sense
| } | ||
| ``` | ||
|
|
||
| A uniform buffer holds one value per submission, so values that change between draws need one uniform (and bind group) per value on the fallback path. |
There was a problem hiding this comment.
Could you elaborate? I feel like I only know what this sentence means because I already knew that before. Something like this?
| A uniform buffer holds one value per submission, so values that change between draws need one uniform (and bind group) per value on the fallback path. | |
| A uniform buffer's value cannot be changed by a command encoder, so immediate values that change between draws included in a single pass need one uniform (and bind group) per draw. |
| // Implementation | ||
| // -------------- | ||
|
|
||
| function assertValidImmediateSchema(schema: BaseData, rootSchema: BaseData = schema): void { |
There was a problem hiding this comment.
This list would work if we assumed that the schema is already a storable type, but it is not. The following is allowed by the types, and the test fails:
it('rejects other non-host-shareable values', () => {
expect(() =>
tgpu['~unstable'].immediateVar(d.texture2d()),
).toThrowErrorMatchingInlineSnapshot();
expect(() =>
tgpu['~unstable'].immediateVar(d.sampler()),
).toThrowErrorMatchingInlineSnapshot();
expect(() =>
tgpu['~unstable'].immediateVar(d.ptrFn(d.u32)),
).toThrowErrorMatchingInlineSnapshot();
});There already is a StorableData type we could use. Or we could add a new type:
export type ImmediateData = ScalarData | VecData | MatData | WgslStruct;
Also nit: I think it is generally cleaner to case out what is allowed, instead of what is not allowed (small chance of disallowing something that should be allowed > big chance of allowing something that should be disallowed)
|
|
||
| expect(writes).toEqual([[0.25]]); | ||
| }); | ||
|
|
There was a problem hiding this comment.
This is what immediatesMap is for, right? I don't see this tested.
it('allows providing multiple immediates', ({ root, computePassEncoder }) => {
const writes = trackWrites(computePassEncoder.mock.setImmediates);
const imm1 = tgpu['~unstable'].immediateVar(d.f32);
const imm2 = tgpu['~unstable'].immediateVar(d.f32);
const fn = tgpu.computeFn({ workgroupSize: [1] })(() => {
const a = imm1.$;
});
const pipeline = root.createComputePipeline({ compute: fn });
pipeline.with(imm1, 1).with(imm2, 2).dispatchWorkgroups(1);
expect(writes).toEqual([[1]]);
});| `); | ||
| }); | ||
|
|
||
| it('throws when aliasing a value from the handle address space', () => { |
There was a problem hiding this comment.
Why is this test in the letDeclaration.test.ts?
| readonly immediates: ImmediateSnapshotMap = new Map(); | ||
| lastWrittenSnapshot: ImmediateSnapshot | undefined; | ||
| lastWrittenGeneration = 0; |
There was a problem hiding this comment.
These three fields relate to immediates, and their state is correlated (when updating one, we need to remember to update another).
Also, they all appear in both draw states... why is compute state named draw anyway?
I think that this immediate managing functionality should be delegated to another class for clarity
| expect(renderPassEncoder.mock.setImmediates).toBeCalledTimes(2); | ||
| }); | ||
|
|
||
| it('distinguishes immediate variables with identical values', ({ root, renderPassEncoder }) => { |
There was a problem hiding this comment.
Is there a reason to distinguish them or is this test just pinning the current behavior?
| expect(writes).toEqual([[0.5], [0.75]]); | ||
| }); | ||
|
|
||
| it('applies pipeline-level and pass-level immediates in call order', ({ |
There was a problem hiding this comment.
This test's result looks really unintuitive to me.
I'd expect that pass.setImmediates > pipeline.with > default. In this case, the line 409 does absolutely nothing. Is this documented somewhere? What is an intuitive explanation for that?
There was a problem hiding this comment.
I agree but it would require a bigger change for consistency. I have proposed it and would like it to land before this PR, then we can clean this up
If you share what parts you mean I'm happy to split |
I thought about changes like allowing immediates in accessors, implementing glsl side, workletization, the caching mechanism, or maybe even schema validation. |
No description provided.