fix(📀): accept shorthand syntax for buffer bindings (#320)

Author: wcandillon

apps/example/ios/Podfile.lock

@@ -1865,7 +1865,7 @@ PODS:
     - ReactCommon/turbomodule/core
     - SocketRocket
     - Yoga
-  - react-native-wgpu (0.5.3):
+  - react-native-wgpu (0.5.4):
     - boost
     - DoubleConversion
     - fast_float
@@ -2938,7 +2938,7 @@ SPEC CHECKSUMS:
   React-microtasksnativemodule: 75b6604b667d297292345302cc5bfb6b6aeccc1b
   react-native-safe-area-context: c00143b4823773bba23f2f19f85663ae89ceb460
   react-native-skia: 5bf2b2107cd7f2d806fd364f5e16b1c7554ed3cd
-  react-native-wgpu: 27d4c1aaa89ba015e8c02d5dbf8abeaa83c4d523
+  react-native-wgpu: 5528610fabc9eb435d4ee578b4d6f7c1e133bf56
   React-NativeModulesApple: 879fbdc5dcff7136abceb7880fe8a2022a1bd7c3
   React-oscompat: 93b5535ea7f7dff46aaee4f78309a70979bdde9d
   React-perflogger: 5536d2df3d18fe0920263466f7b46a56351c0510

apps/example/src/App.tsx

@@ -35,6 +35,7 @@ import { ComputeToys } from "./ComputeToys";
 import { Reanimated } from "./Reanimated";
 import { AsyncStarvation } from "./Diagnostics/AsyncStarvation";
 import { DeviceLostHang } from "./Diagnostics/DeviceLostHang";
+import { StorageBufferVertices } from "./StorageBufferVertices";
 
 // The two lines below are needed by three.js
 import "fast-text-encoding";
@@ -91,6 +92,10 @@ function App() {
           <Stack.Screen name="Reanimated" component={Reanimated} />
           <Stack.Screen name="AsyncStarvation" component={AsyncStarvation} />
           <Stack.Screen name="DeviceLostHang" component={DeviceLostHang} />
+          <Stack.Screen
+            name="StorageBufferVertices"
+            component={StorageBufferVertices}
+          />
         </Stack.Navigator>
       </NavigationContainer>
     </GestureHandlerRootView>

apps/example/src/Home.tsx

@@ -123,6 +123,10 @@ export const examples = [
     screen: "DeviceLostHang",
     title: "⚠️ Device Lost Hang",
   },
+  {
+    screen: "StorageBufferVertices",
+    title: "💾 Storage Buffer Vertices",
+  },
 ];
 
 const styles = StyleSheet.create({

apps/example/src/Route.ts

@@ -28,4 +28,5 @@ export type Routes = {
   Reanimated: undefined;
   AsyncStarvation: undefined;
   DeviceLostHang: undefined;
+  StorageBufferVertices: undefined;
 };

apps/example/src/StorageBufferVertices/StorageBufferVertices.tsx

@@ -0,0 +1,206 @@
+import React from "react";
+import { StyleSheet, View } from "react-native";
+import { Canvas } from "react-native-wgpu";
+
+import { useWebGPU } from "../components/useWebGPU";
+
+import { shaderCode } from "./shaders";
+
+const rand = (min?: number, max?: number) => {
+  if (min === undefined) {
+    min = 0;
+    max = 1;
+  } else if (max === undefined) {
+    max = min;
+    min = 0;
+  }
+  return min + Math.random() * (max - min);
+};
+
+function createCircleVertices({
+  radius = 1,
+  numSubdivisions = 24,
+  innerRadius = 0,
+  startAngle = 0,
+  endAngle = Math.PI * 2,
+} = {}) {
+  const numVertices = numSubdivisions * 3 * 2;
+  const vertexData = new Float32Array(numSubdivisions * 2 * 3 * 2);
+
+  let offset = 0;
+  const addVertex = (x: number, y: number) => {
+    vertexData[offset++] = x;
+    vertexData[offset++] = y;
+  };
+
+  for (let i = 0; i < numSubdivisions; ++i) {
+    const angle1 =
+      startAngle + ((i + 0) * (endAngle - startAngle)) / numSubdivisions;
+    const angle2 =
+      startAngle + ((i + 1) * (endAngle - startAngle)) / numSubdivisions;
+
+    const c1 = Math.cos(angle1);
+    const s1 = Math.sin(angle1);
+    const c2 = Math.cos(angle2);
+    const s2 = Math.sin(angle2);
+
+    // first triangle
+    addVertex(c1 * radius, s1 * radius);
+    addVertex(c2 * radius, s2 * radius);
+    addVertex(c1 * innerRadius, s1 * innerRadius);
+
+    // second triangle
+    addVertex(c1 * innerRadius, s1 * innerRadius);
+    addVertex(c2 * radius, s2 * radius);
+    addVertex(c2 * innerRadius, s2 * innerRadius);
+  }
+
+  return {
+    vertexData,
+    numVertices,
+  };
+}
+
+export function StorageBufferVertices() {
+  const ref = useWebGPU(({ context, device, presentationFormat, canvas }) => {
+    const module = device.createShaderModule({
+      code: shaderCode,
+    });
+
+    const pipeline = device.createRenderPipeline({
+      label: "storage buffer vertices",
+      layout: "auto",
+      vertex: {
+        module,
+      },
+      fragment: {
+        module,
+        targets: [{ format: presentationFormat }],
+      },
+    });
+
+    const kNumObjects = 100;
+    const objectInfos: { scale: number }[] = [];
+
+    // create 2 storage buffers
+    const staticUnitSize =
+      4 * 4 + // color is 4 32bit floats (4bytes each)
+      2 * 4 + // offset is 2 32bit floats (4bytes each)
+      2 * 4; // padding
+    const changingUnitSize = 2 * 4; // scale is 2 32bit floats (4bytes each)
+    const staticStorageBufferSize = staticUnitSize * kNumObjects;
+    const changingStorageBufferSize = changingUnitSize * kNumObjects;
+
+    const staticStorageBuffer = device.createBuffer({
+      label: "static storage for objects",
+      size: staticStorageBufferSize,
+      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+    });
+
+    const changingStorageBuffer = device.createBuffer({
+      label: "changing storage for objects",
+      size: changingStorageBufferSize,
+      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+    });
+
+    // offsets to the various uniform values in float32 indices
+    const kColorOffset = 0;
+    const kOffsetOffset = 4;
+    const kScaleOffset = 0;
+
+    const staticStorageValues = new Float32Array(staticStorageBufferSize / 4);
+    for (let i = 0; i < kNumObjects; ++i) {
+      const staticOffset = i * (staticUnitSize / 4);
+
+      // These are only set once so set them now
+      staticStorageValues.set(
+        [rand(), rand(), rand(), 1],
+        staticOffset + kColorOffset,
+      ); // set the color
+      staticStorageValues.set(
+        [rand(-0.9, 0.9), rand(-0.9, 0.9)],
+        staticOffset + kOffsetOffset,
+      ); // set the offset
+
+      objectInfos.push({
+        scale: rand(0.2, 0.5),
+      });
+    }
+    device.queue.writeBuffer(staticStorageBuffer, 0, staticStorageValues);
+
+    // a typed array we can use to update the changingStorageBuffer
+    const storageValues = new Float32Array(changingStorageBufferSize / 4);
+
+    // setup a storage buffer with vertex data
+    const { vertexData, numVertices } = createCircleVertices({
+      radius: 0.5,
+      innerRadius: 0.25,
+    });
+    const vertexStorageBuffer = device.createBuffer({
+      label: "storage buffer vertices",
+      size: vertexData.byteLength,
+      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+    });
+    device.queue.writeBuffer(vertexStorageBuffer, 0, vertexData);
+
+    const bindGroup = device.createBindGroup({
+      label: "bind group for objects",
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: staticStorageBuffer },
+        { binding: 1, resource: changingStorageBuffer },
+        { binding: 2, resource: vertexStorageBuffer },
+      ],
+    });
+
+    const renderPassDescriptor: GPURenderPassDescriptor = {
+      label: "our basic canvas renderPass",
+      colorAttachments: [
+        {
+          view: context.getCurrentTexture().createView(),
+          clearValue: [0.3, 0.3, 0.3, 1],
+          loadOp: "clear",
+          storeOp: "store",
+        },
+      ],
+    };
+
+    // Set the uniform values in our JavaScript side Float32Array
+    const aspect = canvas.width / canvas.height;
+
+    // set the scales for each object
+    objectInfos.forEach(({ scale }, ndx) => {
+      const offset = ndx * (changingUnitSize / 4);
+      storageValues.set([scale / aspect, scale], offset + kScaleOffset);
+    });
+    // upload all scales at once
+    device.queue.writeBuffer(changingStorageBuffer, 0, storageValues);
+
+    const encoder = device.createCommandEncoder();
+    const pass = encoder.beginRenderPass(renderPassDescriptor);
+    pass.setPipeline(pipeline);
+    pass.setBindGroup(0, bindGroup);
+    pass.draw(numVertices, kNumObjects);
+    pass.end();
+
+    const commandBuffer = encoder.finish();
+    device.queue.submit([commandBuffer]);
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    (context as any).present();
+  });
+
+  return (
+    <View style={style.container}>
+      <Canvas ref={ref} style={style.webgpu} />
+    </View>
+  );
+}
+
+const style = StyleSheet.create({
+  container: {
+    flex: 1,
+  },
+  webgpu: {
+    flex: 1,
+  },
+});

apps/example/src/StorageBufferVertices/index.ts

@@ -0,0 +1 @@
+export * from "./StorageBufferVertices";

apps/example/src/StorageBufferVertices/shaders.ts

@@ -0,0 +1,41 @@
+export const shaderCode = /* wgsl */ `
+  struct OurStruct {
+    color: vec4f,
+    offset: vec2f,
+  };
+
+  struct OtherStruct {
+    scale: vec2f,
+  };
+
+  struct Vertex {
+    position: vec2f,
+  };
+
+  struct VSOutput {
+    @builtin(position) position: vec4f,
+    @location(0) color: vec4f,
+  };
+
+  @group(0) @binding(0) var<storage, read> ourStructs: array<OurStruct>;
+  @group(0) @binding(1) var<storage, read> otherStructs: array<OtherStruct>;
+  @group(0) @binding(2) var<storage, read> pos: array<Vertex>;
+
+  @vertex fn vs(
+    @builtin(vertex_index) vertexIndex : u32,
+    @builtin(instance_index) instanceIndex: u32
+  ) -> VSOutput {
+    let otherStruct = otherStructs[instanceIndex];
+    let ourStruct = ourStructs[instanceIndex];
+
+    var vsOut: VSOutput;
+    vsOut.position = vec4f(
+        pos[vertexIndex].position * otherStruct.scale + ourStruct.offset, 0.0, 1.0);
+    vsOut.color = ourStruct.color;
+    return vsOut;
+  }
+
+  @fragment fn fs(vsOut: VSOutput) -> @location(0) vec4f {
+    return vsOut.color;
+  }
+`;

packages/webgpu/cpp/rnwgpu/api/descriptors/GPUBindGroupEntry.h

@@ -46,6 +46,12 @@ template <> struct JSIConverter<std::shared_ptr<rnwgpu::GPUBindGroupEntry>> {
           } else if (obj.hasNativeState<rnwgpu::GPUTextureView>(runtime)) {
             result->textureView =
                 obj.getNativeState<rnwgpu::GPUTextureView>(runtime);
+          } else if (obj.hasNativeState<rnwgpu::GPUBuffer>(runtime)) {
+            // Support passing GPUBuffer directly as resource (auto-wrap in
+            // GPUBufferBinding)
+            auto binding = std::make_shared<rnwgpu::GPUBufferBinding>();
+            binding->buffer = obj.getNativeState<rnwgpu::GPUBuffer>(runtime);
+            result->buffer = binding;
           } else {
             result->buffer = JSIConverter<
                 std::shared_ptr<rnwgpu::GPUBufferBinding>>::fromJSI(runtime,

packages/webgpu/package.json

@@ -1,6 +1,6 @@
 {
   "name": "react-native-wgpu",
-  "version": "0.5.4",
+  "version": "0.5.5",
   "description": "React Native WebGPU",
   "main": "lib/commonjs/index",
   "module": "lib/module/index",