chore(webgpu): Improve ComputePipeline implementation, add tests (#1979)

docs/api-reference/core/resources/compute-pass.md

@@ -6,6 +6,46 @@ WebGPU only
 
 A pass on which to run computations with compute pipelines.
 
+## Usage
+
+Create and run a compute shader that multiplies an array of numbers by 2.
+
+```ts
+const source = /*WGSL*/`\
+@group(0) @binding(0) var<storage, read_write> data: array<i32>;
+@compute @workgroup_size(1) fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+  let i = id.x;
+  data[i] = 2 * data[i];
+}`;
+
+const shader = webgpuDevice.createShader({source});
+const computePipeline = webgpuDevice.createComputePipeline({
+  shader,
+  shaderLayout: {
+    bindings: [{name: 'data', type: 'storage', location: 0}]
+  }
+});
+
+const workBuffer = webgpuDevice.createBuffer({
+  byteLength: 4,
+  usage: Buffer.STORAGE | Buffer.COPY_SRC | Buffer.COPY_DST,
+});
+
+workBuffer.write(new Int32Array([2]));
+
+computePipeline.setBindings({data: workBuffer});
+
+const computePass = webgpuDevice.beginComputePass({});
+computePass.setPipeline(computePipeline);
+computePass.dispatch(1);
+computePass.end();
+
+webgpuDevice.submit();
+
+const computedData = new Int32Array(await workBuffer.readAsync());
+// computedData[0] === 4
+```
+
 ## Types
 
 ### `ComputePassProps`

docs/api-reference/core/resources/compute-pipeline.md

@@ -1,16 +1,62 @@
 # ComputePipeline
 
 :::info
-`ComputePipeline` is only available on WebGPU. Note that you can still perform
-many GPU computations on WebGL using `RenderPipeline`.
+`ComputePipeline` is only available on WebGPU. Note on WebGL you can still perform
+many GPU computations on `RenderPipeline` using `TransformFeedback`.
 :::
 
 A `ComputePipeline` holds a compiled and linked compute shader.
 
+## Usage
+
+Create and run a compute shader that multiplies an array of numbers by 2.
+
+```ts
+const source = /*WGSL*/`\
+@group(0) @binding(0) var<storage, read_write> data: array<i32>;
+@compute @workgroup_size(1) fn main(@builtin(global_invocation_id) id: vec3<u32>) {
+  let i = id.x;
+  data[i] = 2 * data[i];
+}`;
+
+const shader = webgpuDevice.createShader({source});
+const computePipeline = webgpuDevice.createComputePipeline({
+  shader,
+  shaderLayout: {
+    bindings: [{name: 'data', type: 'storage', location: 0}]
+  }
+});
+
+const workBuffer = webgpuDevice.createBuffer({
+  byteLength: 4,
+  usage: Buffer.STORAGE | Buffer.COPY_SRC | Buffer.COPY_DST,
+});
+
+workBuffer.write(new Int32Array([2]));
+
+computePipeline.setBindings({data: workBuffer});
+
+const computePass = webgpuDevice.beginComputePass({});
+computePass.setPipeline(computePipeline);
+computePass.dispatch(1);
+computePass.end();
+
+webgpuDevice.submit();
+
+const computedData = new Int32Array(await workBuffer.readAsync());
+// computedData[0] === 4
+```
+
+## Types
+
+### `ComputePipelineProps`
+
+- `handle`: `unknown` - holds the underlying WebGPU object
+
 ## Members
 
 - `device`: `Device` - holds a reference to the device that created this resource
-- `handle`: `unknown` - holds the underlying WebGL or WebGPU object
+- `handle`: `unknown` - holds the underlying WebGPU object
 - `props`: `ComputePipelineProps` - holds a copy of the `ComputePipelineProps` used to create this `ComputePipeline`.
 
 ## Methods
@@ -28,5 +74,4 @@ const computePipeline = device.createComputePipeline({...})
 ```typescript
 destroy(): void
 ```
-
 Free up any GPU resources associated with this compute pipeline immediately (instead of waiting for garbage collection).

modules/core-tests/test/adapter/resources/compute-pipeline.spec.ts

@@ -0,0 +1,68 @@
+// luma.gl
+// SPDX-License-Identifier: MIT
+// Copyright (c) vis.gl contributors
+
+import test from 'tape-promise/tape';
+import {webgpuDevice, getTestDevices} from '@luma.gl/test-utils';
+import {ComputePipeline, Buffer} from '@luma.gl/core';
+
+const source = /* WGSL*/ `\
+@group(0) @binding(0) var<storage, read_write> data: array<i32>;
+
+@compute @workgroup_size(1) fn main(
+  @builtin(global_invocation_id) id: vec3<u32>
+) {
+  let i = id.x;
+  data[i] = 2 * data[i];
+}
+`;
+
+test.skip('ComputePipeline construct/delete', async t => {
+  await getTestDevices();
+  if (webgpuDevice) {
+    const shader = webgpuDevice.createShader({source});
+    const computePipeline = webgpuDevice.createComputePipeline({shader});
+    t.ok(computePipeline instanceof ComputePipeline, 'ComputePipeline construction successful');
+    computePipeline.destroy();
+    t.ok(computePipeline instanceof ComputePipeline, 'ComputePipeline delete successful');
+    computePipeline.destroy();
+    t.ok(computePipeline instanceof ComputePipeline, 'ComputePipeline repeated delete successful');
+  }
+  t.end();
+});
+
+test.skip('ComputePipeline compute', async t => {
+  await getTestDevices();
+  if (webgpuDevice) {
+    const shader = webgpuDevice.createShader({source});
+    const computePipeline = webgpuDevice.createComputePipeline({
+      shader,
+      shaderLayout: {
+        bindings: [{name: 'data', type: 'storage', location: 0}]
+      }
+    });
+
+    const workBuffer = webgpuDevice.createBuffer({
+      id: 'work buffer',
+      byteLength: 4,
+      usage: Buffer.STORAGE | Buffer.COPY_SRC | Buffer.COPY_DST
+    });
+
+    workBuffer.write(new Int32Array([2]));
+    const inputData = new Int32Array(await workBuffer.readAsync());
+    t.equal(inputData[0], 2, 'Input data is correct');
+
+    computePipeline.setBindings({data: workBuffer});
+
+    const computePass = webgpuDevice.beginComputePass({});
+    computePass.setPipeline(computePipeline);
+    computePass.dispatch(1);
+    computePass.end();
+
+    webgpuDevice.submit();
+
+    const computedData = new Int32Array(await workBuffer.readAsync());
+    t.equal(computedData[0], 4, 'Computed data is correct');
+  }
+  t.end();
+});

modules/core-tests/test/index.ts

@@ -22,10 +22,11 @@ import './adapter/texture-formats.spec';
 // Resources - TODO these tests only depend on Device and could move to API...
 import './adapter/resources/buffer.spec';
 import './adapter/resources/command-buffer.spec';
-import './adapter/resources/framebuffer.spec';
-import './adapter/resources/render-pipeline.spec';
 import './adapter/resources/shader.spec';
+import './adapter/resources/render-pipeline.spec';
+import './adapter/resources/compute-pipeline.spec';
 import './adapter/resources/sampler.spec';
 import './adapter/resources/texture.spec';
+import './adapter/resources/framebuffer.spec';
 import './adapter/resources/vertex-array.spec';
 import './adapter/resources/query-set.spec';

modules/core/src/adapter/resources/buffer.ts

@@ -82,7 +82,7 @@ export abstract class Buffer extends Resource<BufferProps> {
 
   /** Write data to buffer */
   abstract write(data: ArrayBufferView, byteOffset?: number): void;
-  /** Read data asynchronoursly */
+  /** Read data asynchronously */
   abstract readAsync(byteOffset?: number, byteLength?: number): Promise<Uint8Array>;
 
   /** Read data synchronously. @note WebGL2 only */

modules/core/src/adapter/resources/compute-pipeline.ts

@@ -3,7 +3,7 @@
 // Copyright (c) vis.gl contributors
 
 import {Resource, ResourceProps} from './resource';
-import {BindingDeclaration} from '../types/shader-layout';
+import type {ComputeShaderLayout} from '../types/shader-layout';
 import type {Device} from '../device';
 import type {Shader} from './shader';
 
@@ -12,10 +12,14 @@ import type {Shader} from './shader';
  */
 export type ComputePipelineProps = ResourceProps & {
   handle?: unknown;
-  cs: Shader;
-  csEntryPoint?: string;
-  csConstants?: Record<string, number>; // WGSL only
-  shaderLayout?: BindingDeclaration[];
+  /** Compiled shader object */
+  shader: Shader;
+  /** The entry point, defaults to main */
+  entryPoint?: string;
+  /** These are WGSL constant values - different from GLSL defines in that shader does not need to be recompiled */
+  constants?: Record<string, number>;
+  /** Describes the attributes and bindings exposed by the pipeline shader(s). */
+  shaderLayout?: ComputeShaderLayout | null;
 };
 
 /**
@@ -24,10 +28,10 @@ export type ComputePipelineProps = ResourceProps & {
 export abstract class ComputePipeline extends Resource<ComputePipelineProps> {
   static override defaultProps: Required<ComputePipelineProps> = {
     ...Resource.defaultProps,
-    cs: undefined,
-    csEntryPoint: undefined,
-    csConstants: {},
-    shaderLayout: []
+    shader: undefined,
+    entryPoint: undefined,
+    constants: {},
+    shaderLayout: undefined
   };
 
   override get [Symbol.toStringTag](): string {

modules/core/src/adapter/resources/shader.ts

@@ -16,8 +16,8 @@ import {getShaderInfo} from '../../lib/compiler-log/get-shader-info';
 export type ShaderProps = ResourceProps & {
   /** Shader language (defaults to auto) */
   language?: 'glsl' | 'wgsl' | 'auto';
-  /** Which stage are we compiling? Required by WebGL and GLSL transpiler */
-  stage: 'vertex' | 'fragment' | 'compute';
+  /** Which stage are we compiling? Required for GLSL. Ignored for WGSL. */
+  stage?: 'vertex' | 'fragment' | 'compute';
   /** Shader source code */
   source: string;
   /** Optional shader source map (WebGPU only) */
@@ -36,7 +36,7 @@ export abstract class Shader extends Resource<ShaderProps> {
   static override defaultProps: Required<ShaderProps> = {
     ...Resource.defaultProps,
     language: 'auto',
-    stage: 'vertex',
+    stage: undefined,
     source: '',
     sourceMap: null,
     entryPoint: 'main',

modules/core/src/adapter/types/shader-layout.ts

@@ -35,14 +35,19 @@ import type {TextureView} from '../resources/texture-view';
 export type ShaderLayout = {
   /** All attributes, their locations, and basic type information. Also an auto-deduced step mode */
   attributes: AttributeDeclaration[];
-  /** All bidning points (textures, samplers, uniform buffers) with their locations and type */
+  /** All binding points (textures, samplers, uniform buffers) with their locations and type */
   bindings: BindingDeclaration[];
   /** WebGL only (WebGPU use bindings and uniform buffers) */
   uniforms?: any[];
   /** WebGL2 only (WebGPU use compute shaders) */
   varyings?: VaryingBinding[];
 };
 
+export type ComputeShaderLayout = {
+  /** All binding points (textures, samplers, uniform buffers) with their locations and type */
+  bindings: BindingDeclaration[];
+};
+
 /**
  * Declares one for attributes
  */
@@ -88,7 +93,7 @@ type Binding = {
 /** ShaderLayout for bindings */
 export type BindingDeclaration =
   | UniformBufferBindingLayout
-  | BufferBindingLayout
+  | StorageBufferBindingLayout
   | TextureBindingLayout
   | SamplerBindingLayout
   | StorageTextureBindingLayout;
@@ -112,8 +117,8 @@ export type UniformInfo = {
   byteStride: number;
 };
 
-export type BufferBindingLayout = {
-  type: 'uniform' | 'storage' | 'read-only-storage';
+export type StorageBufferBindingLayout = {
+  type: 'storage' | 'read-only-storage';
   name: string;
   location: number;
   visibility?: number;

modules/core/src/index.ts

@@ -109,6 +109,7 @@ export type {ColorAttachment, DepthStencilAttachment} from './adapter/types/type
 
 export type {
   ShaderLayout,
+  ComputeShaderLayout,
   AttributeDeclaration,
   BindingDeclaration,
   Binding