TS to WGSL compiler proof-of-concept (#206)

apps/typegpu-docs/astro.config.mjs

@@ -20,6 +20,8 @@ export default defineConfig({
           '../../packages/typegpu/dist/macro/index.d.ts',
         'typegpu/dist/experimental/index.d.ts?raw':
           '../../packages/typegpu/dist/experimental/index.d.ts',
+        '@typegpu/jit/dist/index.d.ts?raw':
+          '../../packages/jit/dist/index.d.ts',
       }),
     ],
   },

apps/typegpu-docs/package.json

@@ -38,6 +38,7 @@
     "typedoc": "^0.26.5",
     "typedoc-plugin-markdown": "^4.2.3",
     "typegpu": "workspace:*",
+    "@typegpu/jit": "workspace:*",
     "typescript": "^5.5.4",
     "zod": "^3.23.8"
   },

apps/typegpu-docs/src/components/CodeEditor.tsx

@@ -1,4 +1,5 @@
 import Editor, { type Monaco } from '@monaco-editor/react';
+import typegpuJitDts from '@typegpu/jit/dist/index.d.ts?raw';
 import webgpuTypes from '@webgpu/types/dist/index.d.ts?raw';
 import { entries, map, pipe } from 'remeda';
 import typedBinary from 'typed-binary/dist/index.d.ts?raw';
@@ -55,6 +56,7 @@ function handleEditorWillMount(monaco: Monaco) {
   }
   tsDefaults.addExtraLib(toolkitTypes, 'example-toolkit.d.ts');
   tsDefaults.addExtraLib(typedBinary, 'typed-binary.d.ts');
+  tsDefaults.addExtraLib(typegpuJitDts, 'typegpu-jit.d.ts');
 
   tsDefaults.setCompilerOptions({
     ...tsCompilerOptions,
@@ -63,6 +65,7 @@ function handleEditorWillMount(monaco: Monaco) {
       'typegpu/experimental': ['typegpu/dist/experimental/index.d.ts'],
       'typegpu/data': ['typegpu/dist/data/index.d.ts'],
       'typegpu/macro': ['typegpu/dist/macro/index.d.ts'],
+      '@typegpu/jit': ['typegpu-jit.d.ts'],
     },
   });
 }

apps/typegpu-docs/src/content/examples/simple/gradient-tiles-tgsl.ts

@@ -0,0 +1,122 @@
+/*
+{
+  "title": "Gradient Tiles (TGSL)",
+  "category": "simple"
+}
+*/
+
+// -- Hooks into the example environment
+import {
+  addElement,
+  addSliderPlumParameter,
+  onCleanup,
+  onFrame,
+} from '@typegpu/example-toolkit';
+// --
+
+import { JitTranspiler } from '@typegpu/jit';
+import { f32, struct, u32, vec2f, vec4f } from 'typegpu/data';
+import tgpu, {
+  asUniform,
+  builtin,
+  createRuntime,
+  wgsl,
+} from 'typegpu/experimental';
+
+const xSpanPlum = addSliderPlumParameter('x span', 16, {
+  min: 1,
+  max: 16,
+  step: 1,
+});
+const ySpanPlum = addSliderPlumParameter('y span', 16, {
+  min: 1,
+  max: 16,
+  step: 1,
+});
+
+const spanPlum = wgsl.plum((get) => ({ x: get(xSpanPlum), y: get(ySpanPlum) }));
+const spanBuffer = tgpu
+  .createBuffer(struct({ x: u32, y: u32 }), spanPlum)
+  .$usage(tgpu.Uniform)
+  .$name('span');
+const spanUniform = asUniform(spanBuffer);
+
+const runtime = await createRuntime({ jitTranspiler: new JitTranspiler() });
+const canvas = await addElement('canvas', { aspectRatio: 1 });
+const context = canvas.getContext('webgpu') as GPUCanvasContext;
+const presentationFormat = navigator.gpu.getPreferredCanvasFormat();
+
+context.configure({
+  device: runtime.device,
+  format: presentationFormat,
+  alphaMode: 'premultiplied',
+});
+
+const VertexOutput = struct({ uv: vec2f });
+
+const vertexStep = tgpu
+  .fn([], VertexOutput)
+  .implement(() => {
+    const pos = [
+      vec2f(1, 1), // top-right
+      vec2f(-1, 1), // top-left
+      vec2f(1, -1), // bottom-right
+      vec2f(-1, -1), // bottom-left
+    ];
+
+    const uv = [
+      vec2f(1, 1), // top-right
+      vec2f(0, 1), // top-left
+      vec2f(1, 0), // bottom-right
+      vec2f(0, 0), // bottom-left
+    ];
+
+    return {
+      [builtin.position]: vec4f(pos[builtin.vertexIndex], 0.0, 1.0),
+      uv: uv[builtin.vertexIndex],
+    };
+  })
+  .$uses({});
+
+const fragmentStep = tgpu
+  .fn([VertexOutput])
+  .implement((input) => {
+    const span = spanUniform.value;
+    const red = std.floor(input.uv.x * f32(span.x)) / f32(span.x);
+    const green = std.floor(input.uv.y * f32(span.y)) / f32(span.y);
+    return vec4f(red, green, 0.5, 1.0);
+  })
+  .$uses({ spanUniform });
+
+onFrame(() => {
+  const textureView = context.getCurrentTexture().createView();
+
+  runtime.render(vertexStep, fragmentStep, {
+    colorAttachments: [
+      {
+        view: textureView,
+        clearValue: [0, 0, 0, 0],
+        loadOp: 'clear',
+        storeOp: 'store',
+      },
+    ],
+
+    target: [
+      {
+        format: presentationFormat,
+      },
+    ],
+
+    primitive: {
+      topology: 'triangle-strip',
+    },
+
+    vertexCount: 4,
+  });
+
+  runtime.flush();
+});
+
+onCleanup(() => {
+  runtime.dispose();
+});

apps/typegpu-docs/src/content/examples/simple/increment-tgsl.ts

@@ -0,0 +1,48 @@
+/*
+{
+  "title": "Increment (TGSL)",
+  "category": "simple"
+}
+*/
+
+import { addElement } from '@typegpu/example-toolkit';
+import { JitTranspiler } from '@typegpu/jit';
+import { vec2f } from 'typegpu/data';
+import { asMutable, createRuntime, tgpu } from 'typegpu/experimental';
+
+// ---
+addElement('button', {
+  label: 'Increment',
+  onClick: async () => {
+    const result = await doIncrement();
+    table.setMatrix([[result.x, result.y]]);
+  },
+});
+
+const table = await addElement('table', {
+  label: 'I am incremented on the GPU!',
+});
+table.setMatrix([[0]]);
+// ---
+
+const counterBuffer = tgpu
+  .createBuffer(vec2f, vec2f(0, 1))
+  .$usage(tgpu.Storage);
+const counter = asMutable(counterBuffer);
+
+const increment = tgpu
+  .procedure(() => {
+    const tmp = counter.value.x;
+    counter.value.x = counter.value.y;
+    counter.value.y += tmp;
+  })
+  .$uses({ counter });
+
+const runtime = await createRuntime({
+  jitTranspiler: new JitTranspiler(),
+});
+
+async function doIncrement() {
+  runtime.compute(increment);
+  return await runtime.readBuffer(counterBuffer);
+}