feat: Transfer protocol for TypeGPU resources (primarily for worklet support)#2732
Open
reczkok wants to merge 5 commits into
Open
feat: Transfer protocol for TypeGPU resources (primarily for worklet support)#2732reczkok wants to merge 5 commits into
reczkok wants to merge 5 commits into
Conversation
Snapshot protocol for serializing roots, buffers, bindings, textures, bind groups, layouts, pipelines, query sets, samplers, slots, accessors and constants across JS runtimes sharing a GPU device, plus react-native-worklets integration in @typegpu/react (registerSerializables, useFrame on the UI runtime, worklet root context).
|
pkg.pr.new packages benchmark commit |
Bundle size comparison (
|
| 🟢 Decreased | ➖ Unchanged | 🔴 Increased (max 0.28%) | ❔ Unknown |
|---|---|---|---|
| 0 | 300 | 22 | 1 |
import * as ... in PR vs import * as ... in target (did bundle size increase?):
| Test | tsdown |
|---|---|
| d_isInvariantAttrib.ts | 755 B |
import { ... } in PR vs import * as ... in PR (is the library tree-Shakeable?):
| Test | tsdown |
|---|---|
| tgpu_init.ts | 259.65 kB ( |
| tgpu_initFromDevice.ts | 259.13 kB ( |
| tgpu_resolve.ts | 164.62 kB ( |
| tgpu_resolveWithContext.ts | 164.55 kB ( |
| tgpu_bindGroupLayout.ts | 68.97 kB ( |
| tgpu_mutableAccessor.ts | 65.97 kB ( |
| tgpu_accessor.ts | 65.96 kB ( |
| tgpu_privateVar.ts | 65.31 kB ( |
| tgpu_workgroupVar.ts | 65.30 kB ( |
| tgpu_const.ts | 64.55 kB ( |
| tgpu_fn.ts | 38.14 kB ( |
| tgpu_fragmentFn.ts | 38.14 kB ( |
| tgpu_vertexFn.ts | 37.96 kB ( |
| tgpu_computeFn.ts | 37.66 kB ( |
| tgpu_vertexLayout.ts | 26.79 kB ( |
| tgpu_comptime.ts | 14.48 kB ( |
| tgpu_unroll.ts | 1.66 kB ( |
| tgpu_slot.ts | 1.54 kB ( |
| tgpu_lazy.ts | 1.19 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.51, 1.00, 2.10, 3.25, 3.77, 6.08, 11.51, 12.66]
line [0.49, 1.05, 2.67, 3.59, 4.16, 5.81, 11.80, 13.10]
line [0.53, 1.04, 2.33, 3.94, 4.16, 6.27, 12.39, 12.79]
---
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.17, 0.30, 0.46, 0.51, 0.68, 0.68, 0.78, 0.86]
line [0.19, 0.29, 0.37, 0.47, 0.64, 0.65, 0.73, 0.80]
line [0.20, 0.35, 0.41, 0.49, 0.70, 0.68, 0.76, 0.85]
---
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.64, 1.37, 2.13, 3.89, 6.33, 13.76, 33.22, 59.94]
line [0.50, 1.27, 2.04, 3.39, 6.48, 13.41, 29.18, 62.51]
line [0.65, 1.41, 2.68, 3.53, 7.17, 14.81, 30.54, 61.81]
|
iwoplaza
reviewed
Jul 16, 2026
Comment on lines
+233
to
+238
| const view = root | ||
| .createTexture({ size: [2, 2], format: 'rgba8unorm' }) | ||
| .$usage('sampled') | ||
| .createView(); | ||
| expect(snapshotResource(view)).toBeUndefined(); | ||
| expect(isNonTransferableResource(view)).toBe(true); |
Collaborator
There was a problem hiding this comment.
Just curious, why is a view non-transferable? Is it because react-native-webgpu doesn't register serializables for them, or is there a technical limitation?
Contributor
Author
There was a problem hiding this comment.
It is transferable, it's just a lot of schemas to write code for when it's easily derived
Contributor
Author
There was a problem hiding this comment.
We can probably salvage some logic from layouts but I would leave that for a follow up maybe?
Collaborator
There was a problem hiding this comment.
I added an issue for it so we don't forget, we don't need to do it in this PR 👌.
#2740
iwoplaza
changed the title
# Conflicts: # packages/typegpu/src/indexNamedExports.ts # packages/typegpu/src/tgpuBindGroupLayout.ts
Contributor
Author
|
Can't wait what mr. Frog thinks of this one 😆 |
There was a problem hiding this comment.
Pull request overview
This PR introduces a snapshot/restore (“transfer”) protocol for TypeGPU resources to enable safe cross-runtime usage (notably React Native worklets/UI runtime), alongside related runtime type-guard fixes and React Native integration.
Changes:
- Add a resource snapshot/restore registry and schema/layout-entry serialization to transfer many TypeGPU resources by GPUDevice identity.
- Improve WebGPU pipeline encoder type guards to work with JSI HostObjects (avoid
inchecks). - Add React Native worklets integration in
@typegpu/react(optionalreact-native-workletssupport, UI-threaduseFrame, and automatic serializer registration), plus small init-options behavior fix.
Reviewed changes
Copilot reviewed 46 out of 47 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| pnpm-lock.yaml | Locks updated deps (adds react-native-worklets and related transitive updates). |
| packages/typegpu/tests/serializeDataSchema.test.ts | New tests for serializing/deserializing data schemas. |
| packages/typegpu/tests/serial.test.ts | New end-to-end tests for resource snapshot/restore round-trips. |
| packages/typegpu/tests/root.test.ts | Adds test ensuring optionalFeatures aren’t forwarded to requestDevice. |
| packages/typegpu/tests/renderPipeline.test.ts | Adds test coverage for restoring precompiled render pipelines + bind groups. |
| packages/typegpu/tests/internal/typeGuards.test.ts | Adds tests for JSI-safe WebGPU encoder/pass type guards. |
| packages/typegpu/tests/computePipeline.test.ts | Adds test coverage for restoring precompiled compute pipelines + bind groups. |
| packages/typegpu/src/tgpuBindGroupLayout.ts | Bind groups now carry root; constructor updated accordingly. |
| packages/typegpu/src/std/bitcast.ts | Updates getName import path. |
| packages/typegpu/src/serial/types.ts | Defines RestoreContext keyed by GPUDevice -> TgpuRoot. |
| packages/typegpu/src/serial/schema.ts | Adds schema serialization/deserialization for transferable data schemas. |
| packages/typegpu/src/serial/resources.ts | Implements snapshot/restore for buffers, bind groups/layouts, vertex layouts, consts, accessors, render pipeline parts. |
| packages/typegpu/src/serial/registry.ts | Central registry for snapshotters + snapshotResource/restoreResource API. |
| packages/typegpu/src/serial/layoutEntries.ts | Adds bind group layout entry serialization/deserialization. |
| packages/typegpu/src/serial/dataValue.ts | Adds serialization for vec/mat instances via byte snapshots. |
| packages/typegpu/src/resolutionCtx.ts | Updates catch-all bind group construction for resolution. |
| packages/typegpu/src/internal.ts | Exposes new serialization APIs via typegpu/~internal. |
| packages/typegpu/src/indexNamedExports.ts | Exports additional public type guards and types (e.g. UsageLiteral). |
| packages/typegpu/src/data/index.ts | Exports invariant attrib guard/type needed by schema serialization. |
| packages/typegpu/src/core/texture/texture.ts | Adds texture snapshot/restore + ownership semantics for restored raw textures. |
| packages/typegpu/src/core/slot/slot.ts | Adds slot snapshot/restore. |
| packages/typegpu/src/core/sampler/sampler.ts | Adds sampler snapshot/restore (fixed samplers only). |
| packages/typegpu/src/core/root/rootTypes.ts | Adds resourceType: 'root' and guarded-pipeline resourceType typing. |
| packages/typegpu/src/core/root/init.ts | Adds snapshot/restore for roots & guarded compute pipelines; fixes optionalFeatures forwarding; bind group construction passes root. |
| packages/typegpu/src/core/querySet/querySet.ts | Adds query set snapshot/restore; stores root. |
| packages/typegpu/src/core/pipeline/typeGuards.ts | Replaces in checks with typeof … === 'function' for JSI compatibility. |
| packages/typegpu/src/core/pipeline/renderPipeline.ts | Adds snapshot/restore for precompiled render pipelines; factors bind group/vertex buffer collection and timestamp prior restore. |
| packages/typegpu/src/core/pipeline/pipelineUtils.ts | Adds helpers to collect bind group/vertex buffer pairs and restore timestamp priors. |
| packages/typegpu/src/core/pipeline/computePipeline.ts | Adds snapshot/restore for precompiled compute pipelines; supports descriptor-less precompiled core. |
| packages/typegpu/src/core/constant/tgpuConstant.ts | Adds isConst type guard for snapshotter registry. |
| packages/typegpu/src/core/buffer/buffer.ts | Exports UsageLiteral; adds root property; adds INTERNAL_applyBufferUsages. |
| packages/typegpu-react/tsdown.config.ts | Adds build plugin to preserve optional require('react-native-worklets') for Metro. |
| packages/typegpu-react/tests/root-context.test.tsx | Tests passing init options through <Root options=…>. |
| packages/typegpu-react/tests/react-native/use-frame.test.tsx | Tests RN useFrame dispatch behavior with/without worklets. |
| packages/typegpu-react/tests/react-native/register-serializables.test.ts | Tests custom serializer registration and round-trips / failure modes. |
| packages/typegpu-react/src/react-native/worklets-integration.ts | Optional runtime loader for react-native-worklets. |
| packages/typegpu-react/src/react-native/use-frame.ts | RN useFrame that can run worklet callbacks on UI runtime. |
| packages/typegpu-react/src/react-native/use-configure-worklet-context.ts | Worklet-shareable canvas context ref wrapper for UI runtime access. |
| packages/typegpu-react/src/react-native/serialization/transfer-cache.ts | Worklet-side identity cache and root-by-device mapping for restored resources. |
| packages/typegpu-react/src/react-native/serialization/register-serializables.ts | Registers custom serializable for TypeGPU resources in worklets. |
| packages/typegpu-react/src/react-native/index.ts | RN entry installs WebGPU + registers serializables + exports RN-specific hooks. |
| packages/typegpu-react/src/core/use-frame.ts | Extracts reusable startFrameLoop and exports FrameCtx. |
| packages/typegpu-react/src/core/root-context.tsx | Adds <Root options> and disableWorklets plumbing via context. |
| packages/typegpu-react/README.md | Documents RN worklets support. |
| packages/typegpu-react/package.json | Adds conditional exports for RN entry; adds optional react-native-worklets dep. |
| apps/typegpu-docs/src/content/docs/integration/react-native/worklets.mdx | New documentation page for RN worklets integration and transfer rules. |
| apps/typegpu-docs/astro.config.mjs | Adds docs sidebar entry for the new worklets page. |
Files not reviewed (1)
- pnpm-lock.yaml: Generated file
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
![pullfrog[bot]](https://avatars.githubusercontent.com/in/1768019?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
LGTM. Verified tests, typechecks, and lint on the touched packages; manually traced the snapshot/restore paths and the HostObject-compatible type-guard changes. A couple of minor notes attached.
Kimi K2 (free via Pullfrog for OSS) ｜ 𝕏
aleksanderkatan
left a comment
aleksanderkatan
left a comment
Member
There was a problem hiding this comment.
Yay worklets! ⛑️🎨
I didn't have time to properly review the react and react-native part of this PR, I leave some comments regarding the typegpu part.
I guess worklets come at a price. I'm generally a little worried about everything that happens lazily, and about all methods modifying the state like .withs on render pipeline, I'm sure we missed some, and it will be impossible to debug.
Also, I expect many of our future changes to break on worklets. Furthermore, our APIs will be less flexible, every subsequent refactor/API update will be more difficult because it will need to take worklets into account.
But hey, worklets!
Could we also include some internal docs explaining the general idea, intuition behind why we cannot transfer functions, and some weird cases like the catchall and logResources working?
| ## Rules of transfer | ||
|
|
||
| **Definitions are runtime-local.** | ||
| Shader functions (`tgpu.fn`, entry functions), `tgpu.comptime`, and schemas cannot be serialized. |
Member
There was a problem hiding this comment.
Schemas cannot be serialized?
| expect(restoredAccessor.schema).toBe(d.vec3f); | ||
| expect(restoredAccessor.defaultValue).toEqual(d.vec3f(1, 2, 3)); | ||
|
|
||
| const constant = tgpu['~unstable'].const(d.arrayOf(d.f32, 3), [1, 2, 3]); |
Member
There was a problem hiding this comment.
Suggested change
| const constant = tgpu['~unstable'].const(d.arrayOf(d.f32, 3), [1, 2, 3]); | |
| const constant = tgpu.const(d.arrayOf(d.f32, 3), [1, 2, 3]); |
Member
There was a problem hiding this comment.
Free test
it('keeps the initial buffer value', ({ root }) => {
const buf = root.createBuffer(d.u32, 5);
const restoredBuf = roundTrip(buf, root);
const rawBuffer = root.unwrap(restoredBuf);
const writtenBuffer = vi.mocked(rawBuffer.getMappedRange).mock.results[0]?.value;
expect(writtenBuffer).toMatchInlineSnapshot(`
ArrayBuffer [
5,
0,
0,
0,
]
`);
});
Member
There was a problem hiding this comment.
Free tests
it('round-trips buffer accessors', ({ root }) => {
const getter = () => {
'use gpu';
return d.u32(1);
};
const fnAccessor = tgpu.accessor(d.u32, getter);
const restoredFnAccessor = roundTrip(fnAccessor, root);
expect(restoredFnAccessor.resourceType).toBe('accessor');
expect(restoredFnAccessor.schema).toBe(d.u32);
expect(restoredFnAccessor.defaultValue).toBe(getter);
});
it('round-trips function accessors', ({ root }) => {
const getter = () => {
'use gpu';
return d.u32(1);
};
const fnAccessor = tgpu.accessor(d.u32, getter);
const restoredFnAccessor = roundTrip(fnAccessor, root);
expect(restoredFnAccessor.resourceType).toBe('accessor');
expect(restoredFnAccessor.schema).toBe(d.u32);
expect(restoredFnAccessor.defaultValue).toBe(getter);
});
Member
There was a problem hiding this comment.
Free tests
it('does not allow snapshotting destroyed buffers', ({ root }) => {
const buf = root.createBuffer(d.u32, 5);
buf.destroy();
expect(() => snapshotResource(buf)).toThrowErrorMatchingInlineSnapshot(`[Error: This buffer has been destroyed]`);
});
it('does not allow snapshotting destroyed querysets', ({ root }) => {
const querySet = root.createQuerySet('timestamp', 2);
querySet.destroy();
expect(() => roundTrip(querySet, root)).toThrowErrorMatchingInlineSnapshot(
`[Error: This QuerySet has been destroyed.]`,
);
});
it('does not allow snapshotting destroyed textures', ({ root }) => {
const texture = root.createTexture({ size: [2, 2], format: 'rgba8unorm' });
texture.destroy();
expect(() => roundTrip(texture, root)).toThrowErrorMatchingInlineSnapshot(`[Error: This texture has been destroyed]`);
});
Member
There was a problem hiding this comment.
restoredPipeline acts differently - regular pipeline does not throw on drawIndexed, and this one does.
This may be a result of misconfigured mocks or something like this.
I would like you to add a test including index buffers, but I'm not sure how to do that because of #2738
it('...', ({ root }) => {
const indexBuffer = root.createBuffer(d.arrayOf(d.u16, 2)).$usage('index');
const vertexFn = tgpu.vertexFn({ out: { pos: d.builtin.position } })``;
const fragmentFn = tgpu.fragmentFn({ out: { color: d.vec4f } })``;
const pipeline = root
.createRenderPipeline({
vertex: vertexFn,
fragment: fragmentFn,
targets: { color: { format: 'rgba8unorm' } },
})
.withColorAttachment({
color: {
view: {} as unknown as GPUTextureView,
loadOp: 'clear',
storeOp: 'store',
},
})
.withIndexBuffer(indexBuffer);
const restoredPipeline = roundTrip(pipeline, root);
restoredPipeline.drawIndexed(4);
});
Comment on lines
+226
to
+228
| export function INTERNAL_snapshotRoot(root: TgpuRoot): TgpuRootSnapshot { | ||
| return { type: 'root', device: root.device }; | ||
| } |
Member
There was a problem hiding this comment.
ownDevice prop gets lost, I suspect it will always be set to false, right? nameRegistry and logOptions get lost as well
Comment on lines
+464
to
+475
| export interface TgpuRenderPipelineSnapshot { | ||
| readonly type: 'render-pipeline'; | ||
| readonly device: GPUDevice; | ||
| readonly pipeline: GPURenderPipeline; | ||
| readonly fragmentOut: SerializedDataSchema | undefined; | ||
| readonly usedBindGroupLayouts: TgpuBindGroupLayout[]; | ||
| readonly bindGroups: [TgpuBindGroupLayout, TgpuBindGroup | GPUBindGroup][]; | ||
| readonly usedVertexLayouts: TgpuVertexLayout[]; | ||
| readonly vertexBuffers: [TgpuVertexLayout, (TgpuBuffer<BaseData> & VertexFlag) | GPUBuffer][]; | ||
| readonly timestampWrites: TimestampWritesPriors['timestampWrites']; | ||
| readonly performanceCallback: TimestampWritesPriors['performanceCallback']; | ||
| } |
Member
There was a problem hiding this comment.
Do withEncoder/withPass work?
Comment on lines
+12
to
+13
| .replace('__require("react-native-worklets")', 'require("react-native-worklets")') | ||
| .replace(/import \{ __require \} from "[^"]+";\n/, ''); |
Member
There was a problem hiding this comment.
nit
Suggested change
| .replace('__require("react-native-worklets")', 'require("react-native-worklets")') | |
| .replace(/import \{ __require \} from "[^"]+";\n/, ''); | |
| .replaceAll('__require("react-native-worklets")', 'require("react-native-worklets")') | |
| .replaceAll(/import \{ __require \} from "[^"]+";\n/, ''); |
Comment on lines
+52
to
+53
| id = global.__TYPEGPU_REACT_NEXT_TRANSFER_ID__ ?? 0; | ||
| global.__TYPEGPU_REACT_NEXT_TRANSFER_ID__ = id + 1; |
Member
There was a problem hiding this comment.
Will this work when receiving resources from 2 different threads?
Member
|
Tgpu functions are tgpu resources, and they are not snapshottable. However, it('treats functions as non-snapshottable', ({ root }) => {
const fn1 = () => {
'use gpu';
return 0;
};
const fn2 = tgpu.fn(
[],
d.u32,
)(() => {
'use gpu';
return 0;
});
expect(isNonTransferableResource(fn1)).toMatchInlineSnapshot(`false`);
expect(isNonTransferableResource(fn2)).toMatchInlineSnapshot(`false`);
}); |
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.
intotypeof x.foo === 'function'since JSI HostObjects expose methods through property lookup but their has trap reports false (weird but what can you do)Also i'm pretty sure everything in
/serialcan be done in a cleaner fashion (i just didn't come up with it yet - open to suggestions)