feat: WebGPU API (#7977)

Co-authored-by: Luca Casonato <lucacasonato@yahoo.com>

5 files changed, 284 insertions, 0 deletions

diff --git a/cli/tests/unit/unit_tests.ts b/cli/tests/unit/unit_tests.ts
index 8af28c1bb..2664a9ab0 100644
--- a/cli/tests/unit/unit_tests.ts
+++ b/cli/tests/unit/unit_tests.ts
@@ -76,3 +76,4 @@ import "./write_text_file_test.ts";
 import "./performance_test.ts";
 import "./version_test.ts";
 import "./websocket_test.ts";
+import "./webgpu_test.ts";
diff --git a/cli/tests/unit/webgpu_test.ts b/cli/tests/unit/webgpu_test.ts
new file mode 100644
index 000000000..628a3c8db
--- /dev/null
+++ b/cli/tests/unit/webgpu_test.ts
@@ -0,0 +1,225 @@
+// TODO(lucacasonato): remove when GPUBufferUsage and friends are added to dlint
+// deno-lint-ignore-file no-undef
+
+import { assert, assertEquals, unitTest } from "./test_util.ts";
+
+let isCI: boolean;
+try {
+  isCI = (Deno.env.get("CI")?.length ?? 0) > 0;
+} catch {
+  isCI = true;
+}
+
+// Skip this test on linux CI, because the vulkan emulator is not good enough
+// yet, and skip on macOS because these do not have virtual GPUs.
+unitTest({
+  perms: { read: true, env: true },
+  ignore: (Deno.build.os === "linux" || Deno.build.os === "darwin") && isCI,
+}, async function webgpuComputePass() {
+  const adapter = await navigator.gpu.requestAdapter();
+  assert(adapter);
+
+  const numbers = [1, 4, 3, 295];
+
+  const device = await adapter.requestDevice();
+  assert(device);
+
+  const shaderCode = await Deno.readTextFile(
+    "cli/tests/webgpu_computepass_shader.wgsl",
+  );
+
+  const shaderModule = device.createShaderModule({
+    code: shaderCode,
+  });
+
+  const size = new Uint32Array(numbers).byteLength;
+
+  const stagingBuffer = device.createBuffer({
+    size: size,
+    usage: 1 | 8,
+  });
+
+  const storageBuffer = device.createBuffer({
+    label: "Storage Buffer",
+    size: size,
+    usage: 0x80 | 8 | 4,
+    mappedAtCreation: true,
+  });
+
+  const buf = new Uint32Array(storageBuffer.getMappedRange());
+
+  buf.set(numbers);
+
+  storageBuffer.unmap();
+
+  const bindGroupLayout = device.createBindGroupLayout({
+    entries: [
+      {
+        binding: 0,
+        visibility: 4,
+        buffer: {
+          type: "storage",
+          minBindingSize: 4,
+        },
+      },
+    ],
+  });
+
+  const bindGroup = device.createBindGroup({
+    layout: bindGroupLayout,
+    entries: [
+      {
+        binding: 0,
+        resource: {
+          buffer: storageBuffer,
+        },
+      },
+    ],
+  });
+
+  const pipelineLayout = device.createPipelineLayout({
+    bindGroupLayouts: [bindGroupLayout],
+  });
+
+  const computePipeline = device.createComputePipeline({
+    layout: pipelineLayout,
+    compute: {
+      module: shaderModule,
+      entryPoint: "main",
+    },
+  });
+
+  const encoder = device.createCommandEncoder();
+
+  const computePass = encoder.beginComputePass();
+  computePass.setPipeline(computePipeline);
+  computePass.setBindGroup(0, bindGroup);
+  computePass.insertDebugMarker("compute collatz iterations");
+  computePass.dispatch(numbers.length);
+  computePass.endPass();
+
+  encoder.copyBufferToBuffer(storageBuffer, 0, stagingBuffer, 0, size);
+
+  device.queue.submit([encoder.finish()]);
+
+  await stagingBuffer.mapAsync(1);
+
+  const data = stagingBuffer.getMappedRange();
+
+  assertEquals(new Uint32Array(data), new Uint32Array([0, 2, 7, 55]));
+
+  stagingBuffer.unmap();
+
+  device.destroy();
+
+  // TODO(lucacasonato): webgpu spec should add a explicit destroy method for
+  // adapters.
+  const resources = Object.keys(Deno.resources());
+  Deno.close(Number(resources[resources.length - 1]));
+});
+
+// Skip this test on linux CI, because the vulkan emulator is not good enough
+// yet, and skip on macOS because these do not have virtual GPUs.
+unitTest({
+  perms: { read: true, env: true },
+  ignore: (Deno.build.os === "linux" || Deno.build.os === "darwin") && isCI,
+}, async function webgpuHelloTriangle() {
+  const adapter = await navigator.gpu.requestAdapter();
+  assert(adapter);
+
+  const device = await adapter.requestDevice();
+  assert(device);
+
+  const shaderCode = await Deno.readTextFile(
+    "cli/tests/webgpu_hellotriangle_shader.wgsl",
+  );
+
+  const shaderModule = device.createShaderModule({
+    code: shaderCode,
+  });
+
+  const pipelineLayout = device.createPipelineLayout({
+    bindGroupLayouts: [],
+  });
+
+  const renderPipeline = device.createRenderPipeline({
+    layout: pipelineLayout,
+    vertex: {
+      module: shaderModule,
+      entryPoint: "vs_main",
+    },
+    fragment: {
+      module: shaderModule,
+      entryPoint: "fs_main",
+      targets: [
+        {
+          format: "rgba8unorm-srgb",
+        },
+      ],
+    },
+  });
+
+  const dimensions = {
+    width: 200,
+    height: 200,
+  };
+  const unpaddedBytesPerRow = dimensions.width * 4;
+  const align = 256;
+  const paddedBytesPerRowPadding = (align - unpaddedBytesPerRow % align) %
+    align;
+  const paddedBytesPerRow = unpaddedBytesPerRow + paddedBytesPerRowPadding;
+
+  const outputBuffer = device.createBuffer({
+    label: "Capture",
+    size: paddedBytesPerRow * dimensions.height,
+    usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
+  });
+  const texture = device.createTexture({
+    label: "Capture",
+    size: dimensions,
+    format: "rgba8unorm-srgb",
+    usage: GPUTextureUsage.RENDER_ATTACHMENT | GPUTextureUsage.COPY_SRC,
+  });
+
+  const encoder = device.createCommandEncoder();
+  const renderPass = encoder.beginRenderPass({
+    colorAttachments: [
+      {
+        view: texture.createView(),
+        storeOp: "store",
+        loadValue: [0, 1, 0, 1],
+      },
+    ],
+  });
+  renderPass.setPipeline(renderPipeline);
+  renderPass.draw(3, 1);
+  renderPass.endPass();
+
+  encoder.copyTextureToBuffer(
+    {
+      texture,
+    },
+    {
+      buffer: outputBuffer,
+      bytesPerRow: paddedBytesPerRow,
+      rowsPerImage: 0,
+    },
+    dimensions,
+  );
+
+  device.queue.submit([encoder.finish()]);
+
+  await outputBuffer.mapAsync(1);
+  const data = new Uint8Array(outputBuffer.getMappedRange());
+
+  assertEquals(data, await Deno.readFile("cli/tests/webgpu_hellotriangle.out"));
+
+  outputBuffer.unmap();
+
+  device.destroy();
+
+  // TODO(lucacasonato): webgpu spec should add a explicit destroy method for
+  // adapters.
+  const resources = Object.keys(Deno.resources());
+  Deno.close(Number(resources[resources.length - 1]));
+});
diff --git a/cli/tests/webgpu_computepass_shader.wgsl b/cli/tests/webgpu_computepass_shader.wgsl
new file mode 100644
index 000000000..2433f3243
--- /dev/null
+++ b/cli/tests/webgpu_computepass_shader.wgsl
@@ -0,0 +1,39 @@
+[[builtin(global_invocation_id)]]
+var global_id: vec3<u32>;
+
+[[block]]
+struct PrimeIndices {
+    data: [[stride(4)]] array<u32>;
+}; // this is used as both input and output for convenience
+
+[[group(0), binding(0)]]
+var<storage> v_indices: [[access(read_write)]] PrimeIndices;
+
+// The Collatz Conjecture states that for any integer n:
+// If n is even, n = n/2
+// If n is odd, n = 3n+1
+// And repeat this process for each new n, you will always eventually reach 1.
+// Though the conjecture has not been proven, no counterexample has ever been found.
+// This function returns how many times this recurrence needs to be applied to reach 1.
+fn collatz_iterations(n_base: u32) -> u32{
+    var n: u32 = n_base;
+    var i: u32 = 0u;
+    loop {
+        if (n <= 1u) {
+            break;
+        }
+        if (n % 2u == 0u) {
+            n = n / 2u;
+        }
+        else {
+            n = 3u * n + 1u;
+        }
+        i = i + 1u;
+    }
+    return i;
+}
+
+[[stage(compute), workgroup_size(1)]]
+fn main() {
+    v_indices.data[global_id.x] = collatz_iterations(v_indices.data[global_id.x]);
+}
diff --git a/cli/tests/webgpu_hellotriangle.out b/cli/tests/webgpu_hellotriangle.out
new file mode 100644
index 000000000..91454dbfc
--- /dev/null
+++ b/cli/tests/webgpu_hellotriangle.out
diff --git a/cli/tests/webgpu_hellotriangle_shader.wgsl b/cli/tests/webgpu_hellotriangle_shader.wgsl
new file mode 100644
index 000000000..71934415b
--- /dev/null
+++ b/cli/tests/webgpu_hellotriangle_shader.wgsl
@@ -0,0 +1,19 @@
+[[builtin(vertex_index)]]
+var<in> in_vertex_index: u32;
+[[builtin(position)]]
+var<out> out_pos: vec4<f32>;
+
+[[stage(vertex)]]
+fn vs_main() {
+    var x: f32 = f32(i32(in_vertex_index) - 1);
+    var y: f32 = f32(i32(in_vertex_index & 1) * 2 - 1);
+    out_pos = vec4<f32>(x, y, 0.0, 1.0);
+}
+
+[[location(0)]]
+var<out> out_color: vec4<f32>;
+
+[[stage(fragment)]]
+fn fs_main() {
+    out_color = vec4<f32>(1.0, 0.0, 0.0, 1.0);
+}