Merge pull request #24 from blazejkustra/feat/shared-gpu-device

perf(ShaderView): share one GPUDevice across all ShaderViews

Author: blazejkustra

src/hooks/useWGPUSetup.tsx

@@ -5,8 +5,9 @@ import {
   type RNCanvasContext,
 } from 'react-native-webgpu';
 import { useCallback, useEffect, useRef, useState } from 'react';
-import { scheduleOnRuntime, type WorkletRuntime } from 'react-native-worklets';
+import { type WorkletRuntime } from 'react-native-worklets';
 import { BackgroundRuntime } from '../utils/backgroundRuntime';
+import { getSharedGPUDevice } from '../utils/gpuDevice';
 
 type GPUResources = {
   device: GPUDevice;
@@ -38,27 +39,18 @@ export function useWGPUSetup(): WGPUSetupResult {
   const configuredSizeRef = useRef<{ width: number; height: number } | null>(
     null
   );
-  // Kept so the device can be released on unmount.
-  const deviceRef = useRef<GPUDevice | null>(null);
 
   useEffect(() => {
     let cancelled = false;
 
     (async () => {
-      const adapter = await navigator.gpu.requestAdapter();
-      if (!adapter || cancelled) {
+      // Shared across every ShaderView — see getSharedGPUDevice. Returns null if
+      // no adapter is available.
+      const device = await getSharedGPUDevice();
+      if (!device || cancelled) {
         return;
       }
 
-      const device = await adapter.requestDevice();
-      if (cancelled) {
-        // Unmounted mid-init: the worklet never used this device, so it's safe
-        // to drop it right here on the JS thread.
-        device.destroy();
-        return;
-      }
-      deviceRef.current = device;
-
       const context = canvasRef.current!.getContext('webgpu')!;
       const canvas = context.canvas as CanvasWithSize;
       const dpr = PixelRatio.get();
@@ -82,19 +74,11 @@ export function useWGPUSetup(): WGPUSetupResult {
 
     return () => {
       cancelled = true;
-      const device = deviceRef.current;
-      deviceRef.current = null;
-      if (device) {
-        // Release the device (and all GPU resources created from it: pipeline,
-        // buffers, ...) on the render runtime — the thread that used it — after
-        // the render loop has been signalled to stop, so we don't race an
-        // in-flight frame. Any frame that does slip through hits a destroyed
-        // device and is swallowed by the render loop's try/catch.
-        scheduleOnRuntime(runtime, () => {
-          'worklet';
-          device.destroy();
-        });
-      }
+      // The device is shared across all ShaderViews and lives for the JS
+      // runtime's lifetime, so it must NOT be destroyed here. This view's own
+      // GPU resources (the uniform buffer) are torn down by the render loop when
+      // it stops; the pipeline/shader modules are released once the loop's
+      // worklet closure is garbage-collected.
     };
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, []);

src/utils/gpuDevice.ts

@@ -0,0 +1,38 @@
+let devicePromise: Promise<GPUDevice | null> | null = null;
+
+/**
+ * Lazily requests a single GPUDevice shared by every ShaderView.
+ *
+ * Requesting an adapter + device per view is expensive (real GPU/memory cost and
+ * slower init), and a screen that mounts several effects would otherwise spin up
+ * N devices on the one background runtime. The promise is cached for the lifetime
+ * of the JS runtime so every view awaits the same device.
+ *
+ * If the device is ever lost, the cache is cleared so the next ShaderView mount
+ * requests a fresh one instead of awaiting a dead device forever.
+ */
+export function getSharedGPUDevice(): Promise<GPUDevice | null> {
+  if (devicePromise) {
+    return devicePromise;
+  }
+
+  const promise = (async () => {
+    const adapter = await navigator.gpu.requestAdapter();
+    if (!adapter) {
+      return null;
+    }
+
+    const device = await adapter.requestDevice();
+    device.lost?.then(() => {
+      // Only clear if we're still the cached promise, so we don't clobber a
+      // newer device that a later mount may already have requested.
+      if (devicePromise === promise) {
+        devicePromise = null;
+      }
+    });
+    return device;
+  })();
+
+  devicePromise = promise;
+  return promise;
+}