feat: Transfer protocol for TypeGPU resources (primarily for worklet support)

apps/typegpu-docs/astro.config.mjs

@@ -243,6 +243,11 @@ export default defineConfig({
               label: 'React Native',
               slug: 'integration/react-native',
             },
+            {
+              label: 'React Native Worklets',
+              slug: 'integration/react-native/worklets',
+              badge: { text: 'experimental', variant: 'caution' },
+            },
             {
               label: 'WESL Interoperability',
               slug: 'integration/wesl-interoperability',

apps/typegpu-docs/src/content/docs/integration/react-native/worklets.mdx

@@ -0,0 +1,117 @@
+---
+title: React Native Worklets
+description: A guide on running TypeGPU render loops on the UI thread with react-native-worklets.
+---
+
+With [react-native-worklets](https://docs.swmansion.com/react-native-worklets/), per-frame GPU work can be scheduled on the UI thread, unaffected by tasks running on the default React Native thread (or RN thread).
+TypeGPU resources created on the RN thread can be captured by worklets directly, they are transferred between runtimes automatically.
+
+## Setup
+
+Follow the [React Native guide](/TypeGPU/integration/react-native/) first, then install `react-native-worklets` and enable its babel plugin with [Bundle Mode](https://docs.swmansion.com/react-native-worklets/docs/bundleMode/):
+
+```diff lang=js title="babel.config.js"
++const workletsPluginOptions = {
++  bundleMode: true,
++  importForwarding: {
++    moduleNames: ['typegpu'],
++    // Directories with your module-scope shader definitions
++    relativePaths: ['my-app/components'],
++  },
++};
+
+module.exports = (api) => {
+  api.cache(true);
+  return {
+    presets: ['babel-preset-expo'],
+    plugins: [
+      'unplugin-typegpu/babel',
++      ['react-native-worklets/plugin', workletsPluginOptions],
+    ],
+  };
+};
+```
+
+No extra imports are needed - `@typegpu/react` detects `react-native-worklets` at runtime and registers the transfer support for TypeGPU resources automatically.
+`useFrame` runs its callback on the UI thread whenever the callback is marked with the `'worklet'` directive; plain callbacks keep running on the RN thread.
+
+After changing the babel config, clear the Metro cache with `npx expo start --clear`.
+
+## Example
+
+Create resources on the RN thread, then use them freely inside a `useFrame` worklet:
+
+```tsx
+import { useMemo } from 'react';
+import { Canvas } from 'react-native-webgpu';
+import tgpu, { common, d } from 'typegpu';
+import { useConfigureContext, useFrame, useRoot, useUniform } from '@typegpu/react';
+
+export function Pulse() {
+  const root = useRoot();
+  const color = useUniform(d.vec3f, { initial: d.vec3f(0.114, 0.447, 0.941) });
+
+  const pipeline = useMemo(
+    () =>
+      root.createRenderPipeline({
+        vertex: common.fullScreenTriangle,
+        fragment: () => {
+          'use gpu';
+          return d.vec4f(color.$, 1);
+        },
+      }),
+    [root, color],
+  );
+
+  const { ref, ctxRef } = useConfigureContext({ alphaMode: 'premultiplied' });
+
+  // Runs each frame on the UI thread
+  useFrame(({ elapsedSeconds }) => {
+    'worklet';
+    const ctx = ctxRef.current;
+    if (!ctx) return;
+
+    color.write(d.vec3f(0.5 + Math.sin(elapsedSeconds) * 0.5, 0.447, 0.941));
+    pipeline.withColorAttachment({ view: ctx }).draw(3);
+    ctx.present?.();
+  });
+
+  return <Canvas ref={ref} style={{ aspectRatio: 1 }} transparent />;
+}
+```
+
+The `color` uniform and `pipeline` captured by the worklet are transferred to the UI runtime on first use.
+Both runtimes share the same underlying GPU objects, and transferring the same resource again yields the same object back.
+
+This works for buffers (including `createUniform`/`createMutable`/`createReadonly`), textures, samplers, bind groups and their layouts, vertex layouts, query sets, pipelines, roots, slots, accessors, consts, and vector/matrix instances.
+
+## Opting out
+
+To keep everything on the RN thread even with `react-native-worklets` installed, pass `disableWorklets` to the `Root` provider:
+
+```tsx
+<Root disableWorklets>
+  <Pulse />
+</Root>
+```
+
+`useFrame` then runs its callbacks on the RN thread, and `useConfigureContext` returns a plain object of the same shape, so the code above keeps working unchanged.
+
+## Rules of transfer
+
+**Definitions are runtime-local.**
+Shader functions (`tgpu.fn`, entry functions), `tgpu.comptime`, and schemas cannot be serialized.
+Make sure to create them on the runtime they will be used on, or keep them at module scope in files covered by `importForwarding`; worklets then re-import them natively on the UI runtime instead of transferring them.
+
+**Functions crossing runtimes must be worklets.**
+Any plain function reachable from a transferred resource (e.g. a `withPerformanceCallback` callback) has to be marked with the `'worklet'` directive.
+
+:::note
+Definitions created dynamically (inside components or worklets) cannot cross runtimes and throw.
+If you create pipelines on the UI thread, memoize them - resolving a pipeline every frame is wasted work.
+:::
+
+:::caution
+Render pipelines created on the UI runtime must specify an explicit target format (e.g. `targets: { format: 'bgra8unorm' }`), the default relies on
+`navigator.gpu.getPreferredCanvasFormat()`, and `navigator` is not available on worklet runtimes by default.
+:::

packages/typegpu-react/README.md

@@ -46,6 +46,12 @@ const App = (props: Props) => {
 };
 ```
 
+# React Native
+
+When `react-native-worklets` is installed, per-frame GPU work can run on the UI thread - `useFrame` picks it up automatically for worklet callbacks.
+TypeGPU resources captured by worklets are transferred between runtimes automatically.
+See the [React Native Worklets guide](https://typegpu.com/integration/react-native/worklets).
+
 ## TypeGPU is created by Software Mansion
 
 [![swm](https://logo.swmansion.com/logo?color=white&variant=desktop&width=150&tag=typegpu-github 'Software Mansion')](https://swmansion.com)

packages/typegpu-react/package.json

@@ -17,7 +17,11 @@
   "sideEffects": false,
   "types": "./src/browser/index.ts",
   "exports": {
-    ".": "./src/browser/index.ts",
+    ".": {
+      "react-native": "./src/react-native/index.ts",
+      "browser": "./src/browser/index.ts",
+      "default": "./src/browser/index.ts"
+    },
     "./package.json": "./package.json"
   },
   "publishConfig": {
@@ -53,6 +57,7 @@
     "react": "catalog:",
     "react-dom": "catalog:",
     "react-native": "0.84.1",
+    "react-native-worklets": "0.10.2",
     "tsdown": "catalog:build",
     "typegpu": "workspace:*",
     "typegpu-testing-utility": "workspace:^",
@@ -63,6 +68,7 @@
     "react": "^19.0.0",
     "react-native": "*",
     "react-native-webgpu": "*",
+    "react-native-worklets": "*",
     "typegpu": "workspace:^"
   },
   "peerDependenciesMeta": {
@@ -71,6 +77,9 @@
     },
     "react-native-webgpu": {
       "optional": true
+    },
+    "react-native-worklets": {
+      "optional": true
     }
   }
 }

packages/typegpu-react/src/core/root-context.tsx

@@ -9,7 +9,7 @@ import React, {
   useRef,
   useState,
 } from 'react';
-import { tgpu, type TgpuRoot } from 'typegpu';
+import { tgpu, type InitOptions, type TgpuRoot } from 'typegpu';
 import { useDeferredCleanup } from './helper-hooks.ts';
 import { useBailOnServer } from './use-bail-on-server.ts';
 
@@ -83,6 +83,11 @@ interface RootContext {
 class OwnRootContext implements RootContext {
   #result: RootContextResult | undefined;
   #destroyed: boolean = false;
+  readonly #options: InitOptions | undefined;
+
+  constructor(options?: InitOptions) {
+    this.#options = options;
+  }
 
   initOrGetRoot(): RootContextResult {
     if (this.#destroyed) {
@@ -91,7 +96,7 @@ class OwnRootContext implements RootContext {
     }
 
     if (!this.#result) {
-      const promise = tgpu.init().then(
+      const promise = tgpu.init(this.#options).then(
         (root) => {
           if (this.#destroyed) {
             root.destroy();
@@ -156,14 +161,30 @@ const globalRootContextValue = new OwnRootContext();
 
 const rootContext = createContext<RootContext | null>(null);
 
+const workletsDisabledContext = createContext(false);
+
+/** @internal Reads the `disableWorklets` flag from the nearest <Root> provider */
+export function useWorkletsDisabled(): boolean {
+  return useContext(workletsDisabledContext);
+}
+
 export interface RootProps {
+  /** Options used when this provider creates its own root, ignored when `root` is provided */
+  options?: InitOptions | undefined;
   /**
    * An existing root to provide. If undefined (default), a new root will be initialized for
    * this provider's children.
    *
    * @default undefined
    */
   root?: TgpuRoot | undefined;
+  /**
+   * (React Native only) When true, `useFrame` runs on the RN thread even if
+   * `react-native-worklets` is installed. Ignored on the web
+   *
+   * @default false
+   */
+  disableWorklets?: boolean | undefined;
   children?: ReactNode | undefined;
 }
 
@@ -183,8 +204,8 @@ function WarnSuspense() {
   return null;
 }
 
-export const Root = ({ children, root }: RootProps) => {
-  const [ownCtx] = useState(() => new OwnRootContext());
+export const Root = ({ children, options, root, disableWorklets = false }: RootProps) => {
+  const [ownCtx] = useState(() => new OwnRootContext(options));
   const existingRootCtx = useMemo(() => {
     if (root) {
       return new ExistingRootContext(root);
@@ -198,7 +219,9 @@ export const Root = ({ children, root }: RootProps) => {
 
   return (
     <rootContext.Provider value={existingRootCtx ?? ownCtx}>
-      <Suspense fallback={<WarnSuspense />}>{children}</Suspense>
+      <workletsDisabledContext.Provider value={disableWorklets}>
+        <Suspense fallback={<WarnSuspense />}>{children}</Suspense>
+      </workletsDisabledContext.Provider>
     </rootContext.Provider>
   );
 };

packages/typegpu-react/src/core/use-frame.ts

@@ -1,6 +1,6 @@
 import { useEffect, useRef } from 'react';
 
-interface FrameCtx {
+export interface FrameCtx {
   /**
    * Time elapsed since the last frame
    */
@@ -11,39 +11,41 @@ interface FrameCtx {
   readonly elapsedSeconds: number;
 }
 
+export function startFrameLoop(cb: (ctx: FrameCtx) => void): () => void {
+  let frameId: number | undefined;
+  let startTime: number | undefined;
+  let lastTime: number | undefined;
+
+  const loop = () => {
+    frameId = requestAnimationFrame(loop);
+
+    const now = performance.now();
+    if (lastTime === undefined || startTime === undefined) {
+      startTime = now;
+      lastTime = now;
+    }
+    cb({
+      deltaSeconds: (now - lastTime) / 1000,
+      elapsedSeconds: (now - startTime) / 1000,
+    });
+    lastTime = now;
+  };
+
+  loop();
+
+  return () => {
+    if (frameId !== undefined) {
+      cancelAnimationFrame(frameId);
+    }
+  };
+}
+
 export function useFrame(cb: (ctx: FrameCtx) => void) {
   const latestCb = useRef(cb);
 
   useEffect(() => {
     latestCb.current = cb;
   }, [cb]);
 
-  useEffect(() => {
-    let frameId: number | undefined;
-    let startTime: number | undefined;
-    let lastTime: number | undefined;
-
-    const loop = () => {
-      frameId = requestAnimationFrame(loop);
-
-      const now = performance.now();
-      if (lastTime === undefined || startTime === undefined) {
-        startTime = now;
-        lastTime = now;
-      }
-      latestCb.current({
-        deltaSeconds: (now - lastTime) / 1000,
-        elapsedSeconds: (now - startTime) / 1000,
-      });
-      lastTime = now;
-    };
-
-    loop();
-
-    return () => {
-      if (frameId !== undefined) {
-        cancelAnimationFrame(frameId);
-      }
-    };
-  }, []);
+  useEffect(() => startFrameLoop((ctx) => latestCb.current(ctx)), []);
 }

packages/typegpu-react/src/react-native/index.ts

@@ -1,8 +1,14 @@
 import { WebGPUModule } from 'react-native-webgpu';
 
+import { registerTypegpuReactSerializables } from './serialization/register-serializables.ts';
+
 // Making sure the WebGPU module is installed before navigator.gpu is accessed
 WebGPUModule.install();
+// No-ops when react-native-worklets is not installed
+registerTypegpuReactSerializables();
 
 export * from '../shared-exports.ts';
 
-export { useConfigureContext } from './use-configure-context.ts';
+// Intentionally shadows the browser `useFrame`, this one can run the frame loop on the UI runtime
+export { useFrame } from './use-frame.ts';
+export { useConfigureContext } from './use-configure-worklet-context.ts';