diff --git a/apps/typegpu-docs/astro.config.mjs b/apps/typegpu-docs/astro.config.mjs index 3e4a0ae245..3f7cf02a21 100644 --- a/apps/typegpu-docs/astro.config.mjs +++ b/apps/typegpu-docs/astro.config.mjs @@ -209,6 +209,11 @@ export default defineConfig({ label: 'Buffers', slug: 'apis/buffers', }, + { + label: 'Buffer Views', + slug: 'apis/buffer-views', + badge: { text: 'new' }, + }, { label: 'Textures', slug: 'apis/textures', diff --git a/apps/typegpu-docs/src/content/docs/apis/buffer-views.mdx b/apps/typegpu-docs/src/content/docs/apis/buffer-views.mdx new file mode 100644 index 0000000000..9cf66b2256 --- /dev/null +++ b/apps/typegpu-docs/src/content/docs/apis/buffer-views.mdx @@ -0,0 +1,193 @@ +--- +title: Buffer Views +description: A guide on how to store raw bytes in buffers and view them as typed data in shaders. +--- + +:::caution[Language extension] +Buffer views require the `buffer_view` WGSL language extension, which is not yet supported everywhere. +Check `root.enabledWgslLanguageFeatures.has('buffer_view')` and fall back to typed schemas, such as `d.struct` and `d.arrayOf`, when it is missing. +::: + +A `d.bytes(n)` schema describes `n` bytes with no type of their own, the WGSL `buffer` type. +Shaders read and write those bytes as any host-shareable type with `std.bufferView`, at offsets they compute at runtime. +One buffer can then hold differently typed data, arrays of varying length, or several views of the same memory. + +## Creating buffers + +`d.bytes(n)` works wherever a schema does, in storage and uniform buffers, workgroup variables and bind group layouts. +Bind group layouts also take `d.bytes` itself, for storage buffers of any size. + +```ts twoslash +import { tgpu, d } from 'typegpu'; + +const root = await tgpu.init(); +// ---cut--- +const packet = root.createMutable(d.bytes(1024)); +const lights = root.createUniform(d.bytes(512)); +const scratch = tgpu.workgroupVar(d.bytes(4096)); + +const layout = tgpu.bindGroupLayout({ + mesh: { storage: d.bytes }, +}); +``` + +On the CPU, the buffer takes an `ArrayBuffer` or a typed array, and reads back an `ArrayBuffer`. + +```ts twoslash +import { tgpu, d } from 'typegpu'; + +const root = await tgpu.init(); +// ---cut--- +const packet = root.createBuffer(d.bytes(16), new Uint32Array([1, 2, 3, 4])).$usage('storage'); + +packet.write(new Float32Array([0.5]), { startOffset: 8 }); + +const bytes = await packet.read(); +const values = new Uint32Array(bytes); +``` + +## Viewing bytes + +`std.bufferView(schema, buffer, byteOffset)` views the bytes at `byteOffset` as `schema` and returns a reference. +Access its members and elements directly, and read or write the whole value through `.$`. +The view can be written if the buffer can. + +```ts twoslash +import { tgpu, d, std } from 'typegpu'; + +const root = await tgpu.init(); +const packet = root.createMutable(d.bytes(1024)); +// ---cut--- +const Header = d.struct({ count: d.u32, scale: d.f32 }); + +const main = tgpu.computeFn({ workgroupSize: [1] })(() => { + const header = std.bufferView(Header, packet.$, 0); + header.count += 1; + std.bufferView(d.vec4f, packet.$, 16).$ = d.vec4f(header.scale); +}); +``` + +A runtime-sized schema, such as `d.arrayOf(d.f32)`, extends to the end of the buffer. +`std.bufferArrayView(schema, buffer, byteOffset, byteSize)` limits it to `byteSize` bytes, and `std.bufferLength(buffer)` returns the size of the whole buffer. + +Each view must start at a multiple of its schema's alignment and fit in the buffer. +TypeGPU checks both when the offset is known at resolution time. +Views cannot hold atomics or booleans, and buffers cannot be struct members or array elements. + +## Packing differently typed data + +A mesh can store a header, its indices and its vertices in one buffer, with offsets that depend on the header. + +```ts twoslash +import { tgpu, d, std } from 'typegpu'; + +const root = await tgpu.init(); +// ---cut--- +const Header = d.struct({ indexCount: d.u32, vertexCount: d.u32 }); + +const layout = tgpu.bindGroupLayout({ + mesh: { storage: d.bytes }, + extent: { storage: d.vec2f, access: 'mutable' }, +}); + +const measure = tgpu.computeFn({ workgroupSize: [1] })(() => { + const header = std.bufferView(Header, layout.$.mesh, 0); + const indicesSize = header.indexCount * 4; + const indices = std.bufferArrayView(d.arrayOf(d.u32), layout.$.mesh, 8, indicesSize); + const vertices = std.bufferView(d.arrayOf(d.vec2f), layout.$.mesh, 8 + indicesSize); + + for (const index of indices) { + layout.$.extent = std.max(layout.$.extent, std.abs(vertices[index]!)); + } +}); +``` + +## Arrays of varying length + +Records of different lengths can follow each other, each a count and its values. +`std.bufferArrayView` gives each record an array of the right length. + +```ts twoslash +import { tgpu, d, std } from 'typegpu'; + +const root = await tgpu.init(); +// ---cut--- +const records = root.createReadonly(d.bytes(4096)); +const totals = root.createMutable(d.arrayOf(d.f32, 64)); + +const sum = tgpu.computeFn({ + workgroupSize: [64], + in: { gid: d.builtin.globalInvocationId }, +})(({ gid }) => { + let offset = d.u32(0); + for (let i = d.u32(0); i < gid.x; i++) { + offset += 4 + std.bufferView(d.u32, records.$, offset).$ * 4; + } + + const count = std.bufferView(d.u32, records.$, offset).$; + const values = std.bufferArrayView(d.arrayOf(d.f32), records.$, offset + 4, count * 4); + + let total = d.f32(0); + for (const value of values) { + total += value; + } + totals.$[gid.x] = total; +}); +``` + +## Reinterpreting bytes + +Two views of the same bytes read them as different types. + +```ts twoslash +import { tgpu, d, std } from 'typegpu'; + +const root = await tgpu.init(); +const packet = root.createMutable(d.bytes(1024)); +// ---cut--- +const flipSign = tgpu.fn([d.u32])((offset) => { + const bits = std.bufferView(d.u32, packet.$, offset); + bits.$ ^= 0x80000000; +}); +``` + +## Reusing workgroup memory + +A workgroup buffer can hold different data in each phase of a computation. + +```ts twoslash +import { tgpu, d, std } from 'typegpu'; +// ---cut--- +const scratch = tgpu.workgroupVar(d.bytes(1024)); + +const main = tgpu.computeFn({ + workgroupSize: [128], + in: { lid: d.builtin.localInvocationIndex }, +})(({ lid }) => { + const floats = std.bufferView(d.arrayOf(d.f32, 256), scratch.$, 0); + floats[lid] = d.f32(lid); + std.workgroupBarrier(); + + const pairs = std.bufferView(d.arrayOf(d.vec2u, 128), scratch.$, 0); + pairs[lid] = d.vec2u(lid); +}); +``` + +## Passing buffers to functions + +Functions take buffers through pointers, such as `d.ptrStorage(d.bytes(), 'read-write')`. + +```ts twoslash +import { tgpu, d, std } from 'typegpu'; + +const root = await tgpu.init(); +const packet = root.createMutable(d.bytes(1024)); +// ---cut--- +const bump = tgpu.fn([d.ptrStorage(d.bytes(), 'read-write'), d.u32])((buffer, index) => { + std.bufferView(d.u32, buffer, index * 4).$ += 1; +}); + +const main = tgpu.computeFn({ workgroupSize: [1] })(() => { + bump(d.ref(packet.$), 0); +}); +``` diff --git a/apps/typegpu-docs/src/examples/simple/shape-stream/index.html b/apps/typegpu-docs/src/examples/simple/shape-stream/index.html new file mode 100644 index 0000000000..c72c5f1548 --- /dev/null +++ b/apps/typegpu-docs/src/examples/simple/shape-stream/index.html @@ -0,0 +1,17 @@ + + + + diff --git a/apps/typegpu-docs/src/examples/simple/shape-stream/index.ts b/apps/typegpu-docs/src/examples/simple/shape-stream/index.ts new file mode 100644 index 0000000000..e8e2f6e933 --- /dev/null +++ b/apps/typegpu-docs/src/examples/simple/shape-stream/index.ts @@ -0,0 +1,232 @@ +import { common, d, std, tgpu } from 'typegpu'; +import * as sdf from '@typegpu/sdf'; +import { defineControls } from '../../common/defineControls.ts'; + +const root = await tgpu.init(); +const canvas = document.querySelector('canvas') as HTMLCanvasElement; +const context = root.configureContext({ canvas, alphaMode: 'premultiplied' }); + +const supported = root.enabledWgslLanguageFeatures.has('buffer_view'); +(document.querySelector('#unsupported') as HTMLElement).hidden = supported; + +const Kind = { end: 0, blend: 1, circle: 2, box: 3, segment: 4 } as const; + +const Header = d.struct({ kind: d.u32, color: d.u32 }); +const Circle = d.struct({ center: d.vec2f, radius: d.f32 }); +const Box = d.struct({ center: d.vec2f, halfSize: d.vec2f, angle: d.f32 }); +const Segment = d.struct({ a: d.vec2f, b: d.vec2f, width: d.f32 }); + +const HEADER_SIZE = d.sizeOf(Header); +const CIRCLE_RECORD = HEADER_SIZE + d.sizeOf(Circle); +const BOX_RECORD = HEADER_SIZE + d.sizeOf(Box); +const SEGMENT_RECORD = HEADER_SIZE + d.sizeOf(Segment); + +const STREAM_BYTES = 4096; +const commands = root.createReadonly(d.bytes(STREAM_BYTES)); + +const rotate = (v: d.v2f, angle: number) => { + 'use gpu'; + const c = std.cos(angle); + const s = std.sin(angle); + return d.vec2f(c * v.x + s * v.y, c * v.y - s * v.x); +}; + +const shapeDistance = tgpu.fn( + [d.u32, d.u32, d.vec2f], + d.f32, +)((kind, at, p) => { + 'use gpu'; + if (kind === Kind.circle) { + const circle = std.bufferView(Circle, commands.$, at); + return sdf.sdDisk(p - circle.center, circle.radius); + } + if (kind === Kind.box) { + const box = std.bufferView(Box, commands.$, at); + return sdf.sdRoundedBox2d(rotate(p - box.center, box.angle), box.halfSize, 0.02); + } + const segment = std.bufferView(Segment, commands.$, at); + return sdf.sdLine(p, segment.a, segment.b) - segment.width; +}); + +const recordSize = tgpu.fn( + [d.u32], + d.u32, +)((kind) => { + if (kind === Kind.circle) { + return CIRCLE_RECORD; + } + if (kind === Kind.box) { + return BOX_RECORD; + } + if (kind === Kind.segment) { + return SEGMENT_RECORD; + } + return HEADER_SIZE; +}); + +const background = d.vec3f(0.07, 0.07, 0.09); + +const fragment = tgpu.fragmentFn({ + in: { uv: d.vec2f }, + out: d.vec4f, +})(({ uv }) => { + 'use gpu'; + const p = (uv * 2 - 1) * d.vec2f(1, -1); + const pixel = std.fwidth(p.x); + + let dist = d.f32(1000); + let color = d.vec3f(background); + let k = d.f32(0.0001); + let offset = d.u32(0); + const end = std.bufferLength(commands.$); + + while (offset < end) { + const header = std.bufferView(Header, commands.$, offset); + if (header.kind === Kind.end) { + break; + } + + if (header.kind === Kind.blend) { + k = std.max(std.bufferView(d.f32, commands.$, offset + 4).$, 0.0001); + } else { + const shape = shapeDistance(header.kind, offset + HEADER_SIZE, p); + const h = std.saturate(0.5 + (0.5 * (dist - shape)) / k); + dist = std.mix(dist, shape, h) - k * h * (1 - h); + color = std.mix(color, std.unpack4x8unorm(header.color).xyz, h); + } + offset += recordSize(header.kind); + } + + const fill = 1 - std.smoothstep(-pixel, pixel, dist); + const shade = color * (0.8 + 0.2 * std.saturate(-dist * 12)); + const rings = background * (0.85 + 0.15 * std.cos(dist * 90)); + return d.vec4f(std.mix(rings, shade, fill), 1); +}); + +const pipeline = root.createRenderPipeline({ + vertex: common.fullScreenTriangle, + fragment, +}); + +const palette = [0xff6b6b, 0xfeca57, 0x48dbfb, 0x1dd1a1, 0xa29bfe, 0xff9ff3].map( + (hex) => ((hex >> 16) | (hex & 0xff00) | ((hex & 0xff) << 16) | (0xff << 24)) >>> 0, +); + +interface Shape { + kind: typeof Kind.circle | typeof Kind.box | typeof Kind.segment; + color: number; + center: [number, number]; + orbit: number; + speed: number; + phase: number; + size: number; +} + +const random = (min: number, max: number) => min + Math.random() * (max - min); + +let shapes: Shape[] = []; +let shapeCount = 24; +let smoothness = 0.12; + +function generateScene() { + shapes = Array.from({ length: shapeCount }, () => ({ + kind: ([Kind.circle, Kind.box, Kind.segment] as const)[Math.floor(Math.random() * 3)], + color: palette[Math.floor(Math.random() * palette.length)], + center: [random(-0.7, 0.7), random(-0.7, 0.7)], + orbit: random(0.05, 0.25), + speed: random(0.2, 0.8) * (Math.random() < 0.5 ? -1 : 1), + phase: random(0, Math.PI * 2), + size: random(0.06, 0.16), + })); +} + +const encoded = new ArrayBuffer(STREAM_BYTES); +const view = new DataView(encoded); + +function encode(time: number) { + let offset = 0; + const header = (kind: number, word: number) => { + view.setUint32(offset, kind, true); + view.setUint32(offset + 4, word, true); + offset += HEADER_SIZE; + }; + const floats = (...values: number[]) => { + values.forEach((value, i) => view.setFloat32(offset + i * 4, value, true)); + }; + + shapes.forEach((shape, i) => { + if (i % 6 === 0) { + header(Kind.blend, 0); + view.setFloat32(offset - 4, smoothness * (0.5 + ((i / 6) % 3) * 0.5), true); + } + + const angle = time * shape.speed + shape.phase; + const x = shape.center[0] + Math.cos(angle) * shape.orbit; + const y = shape.center[1] + Math.sin(angle) * shape.orbit; + + header(shape.kind, shape.color); + if (shape.kind === Kind.circle) { + floats(x, y, shape.size); + offset += CIRCLE_RECORD - HEADER_SIZE; + } else if (shape.kind === Kind.box) { + floats(x, y, shape.size, shape.size * 0.6, angle); + offset += BOX_RECORD - HEADER_SIZE; + } else { + const dx = Math.cos(angle * 2) * shape.size * 1.5; + const dy = Math.sin(angle * 2) * shape.size * 1.5; + floats(x - dx, y - dy, x + dx, y + dy, shape.size * 0.25); + offset += SEGMENT_RECORD - HEADER_SIZE; + } + }); + + header(Kind.end, 0); +} + +let frameId: number; + +function frame(timestamp: number) { + encode(timestamp / 1000); + commands.write(encoded); + pipeline.withColorAttachment({ view: context }).draw(3); + + frameId = requestAnimationFrame(frame); +} + +generateScene(); +if (supported) { + frameId = requestAnimationFrame(frame); +} + +// #region Example controls and cleanup + +export const controls = defineControls({ + shapes: supported && { + initial: shapeCount, + min: 3, + max: 60, + step: 1, + onSliderChange: (value) => { + shapeCount = value; + generateScene(); + }, + }, + smoothness: supported && { + initial: smoothness, + min: 0, + max: 0.3, + step: 0.01, + onSliderChange: (value) => { + smoothness = value; + }, + }, + 'new scene': supported && { + onButtonClick: generateScene, + }, +}); + +export function onCleanup() { + cancelAnimationFrame(frameId); + root.destroy(); +} + +// #endregion diff --git a/apps/typegpu-docs/src/examples/simple/shape-stream/meta.json b/apps/typegpu-docs/src/examples/simple/shape-stream/meta.json new file mode 100644 index 0000000000..91433b0de2 --- /dev/null +++ b/apps/typegpu-docs/src/examples/simple/shape-stream/meta.json @@ -0,0 +1,7 @@ +{ + "title": "Shape Stream", + "description": "Packs circles, boxes, segments and blend settings into one byte buffer each frame, and a fragment shader walks it with buffer views to draw smoothly merging shapes.", + "category": "simple", + "tags": ["sdf", "buffer", "animation"], + "coolFactor": 5 +} diff --git a/apps/typegpu-docs/src/examples/simple/shape-stream/thumbnail.png b/apps/typegpu-docs/src/examples/simple/shape-stream/thumbnail.png new file mode 100644 index 0000000000..8dc3a707ae Binary files /dev/null and b/apps/typegpu-docs/src/examples/simple/shape-stream/thumbnail.png differ diff --git a/packages/typegpu-testing-utility/src/extendedIt.ts b/packages/typegpu-testing-utility/src/extendedIt.ts index b4c4755e51..81dc29e091 100644 --- a/packages/typegpu-testing-utility/src/extendedIt.ts +++ b/packages/typegpu-testing-utility/src/extendedIt.ts @@ -254,7 +254,7 @@ export const it = base __brand: 'GPU', requestAdapter: vi.fn(() => Promise.resolve(adapter)), getPreferredCanvasFormat: vi.fn(() => 'bgra8unorm'), - wgslLanguageFeatures: new Set(['immediate_address_space']), + wgslLanguageFeatures: new Set(['immediate_address_space', 'buffer_view']), }, mediaDevices: { getUserMedia: vi.fn(() => Promise.resolve()), diff --git a/packages/typegpu/src/core/resolve/resolveData.ts b/packages/typegpu/src/core/resolve/resolveData.ts index f5428fb819..bd7d43e9b5 100644 --- a/packages/typegpu/src/core/resolve/resolveData.ts +++ b/packages/typegpu/src/core/resolve/resolveData.ts @@ -252,6 +252,11 @@ export function resolveData(ctx: ResolutionCtx, data: AnyData): string { return `atomic<${resolveData(ctx, data.inner)}>`; } + if (data.type === 'buffer') { + ctx.requireLanguageFeature('buffer_view', `d.${data}`, 'typed schemas (d.struct, d.arrayOf)'); + return data.byteSize === 0 ? 'buffer' : `buffer<${data.byteSize}>`; + } + if (data.type === 'decorated') { return ctx.resolve(data.inner as AnyWgslData).value; } diff --git a/packages/typegpu/src/core/variable/tgpuVariable.ts b/packages/typegpu/src/core/variable/tgpuVariable.ts index bc2e856c91..d232b75c9a 100644 --- a/packages/typegpu/src/core/variable/tgpuVariable.ts +++ b/packages/typegpu/src/core/variable/tgpuVariable.ts @@ -1,3 +1,4 @@ +import type { WgslBuffer } from '../../data/bytes.ts'; import type { AnyData } from '../../data/dataTypes.ts'; import { snip } from '../../data/snippet.ts'; import type { BaseData } from '../../data/wgslTypes.ts'; @@ -63,9 +64,11 @@ export function privateVar( * @param dataType The schema of the held data's type */ export function workgroupVar( - dataType: TDataType, + dataType: TDataType extends WgslBuffer<0> + ? '(Error) Workgroup buffers need a size, e.g. d.bytes(1024)' + : TDataType, ): TgpuVar<'workgroup', TDataType> { - return new TgpuVarImpl('workgroup', dataType); + return new TgpuVarImpl('workgroup', dataType as TDataType); } export function isVariable(value: unknown): value is TgpuVar { diff --git a/packages/typegpu/src/data/alignmentOf.ts b/packages/typegpu/src/data/alignmentOf.ts index 15fe7df207..035a3d3d65 100644 --- a/packages/typegpu/src/data/alignmentOf.ts +++ b/packages/typegpu/src/data/alignmentOf.ts @@ -36,6 +36,7 @@ const knownAlignmentMap: Record = { mat3x3f: 16, mat4x4f: 16, atomic: 4, + buffer: 4, }; function computeAlignment(data: object): number { diff --git a/packages/typegpu/src/data/array.ts b/packages/typegpu/src/data/array.ts index 39fadaaa98..69437b798b 100644 --- a/packages/typegpu/src/data/array.ts +++ b/packages/typegpu/src/data/array.ts @@ -1,5 +1,6 @@ import { comptime, type TgpuComptime } from '../core/function/comptime.ts'; import { $internal } from '../shared/symbols.ts'; +import { isWgslBuffer } from './bytes.ts'; import { schemaCallWrapper } from './schemaCallWrapper.ts'; import { sizeOf } from './sizeOf.ts'; import type { AnyWgslData, Decorated, Location, WgslArray } from './wgslTypes.ts'; @@ -58,6 +59,11 @@ export const arrayOf: TgpuComptime = comptime((( elementType: AnyWgslData, elementCount?: number, ) => { + if (isWgslBuffer(elementType)) { + throw new Error( + `Cannot create an array of d.${elementType}, buffers can only be whole variables.`, + ); + } if (isDecorated(elementType) && !elementType.attribs.every(isLocationAttrib)) { throw new Error( 'Arrays cannot hold decorated types other than @location. Wrap it in a struct instead, e.g. d.arrayOf(d.struct({ value: d.align(16, d.u32) }), n).', diff --git a/packages/typegpu/src/data/bytes.ts b/packages/typegpu/src/data/bytes.ts new file mode 100644 index 0000000000..8337681ef9 --- /dev/null +++ b/packages/typegpu/src/data/bytes.ts @@ -0,0 +1,105 @@ +import { + $gpuRepr, + $inRepr, + $internal, + $invalidSchemaReason, + $memIdent, + $repr, + $validStorageSchema, + $validUniformSchema, + $validVertexSchema, + isMarkedInternal, +} from '../shared/symbols.ts'; +import type { BaseData } from './wgslTypes.ts'; + +// ---------- +// Public API +// ---------- + +/** + * Schema of the WGSL `buffer` type: `byteSize` bytes without a type of their own, + * read and written in shaders through `std.bufferView` and `std.bufferArrayView`. + * A `byteSize` of 0 stands for a runtime-sized `buffer`. + */ +export interface WgslBuffer extends BaseData { + readonly type: 'buffer'; + readonly byteSize: TSize; + + // Type-tokens, not available at runtime + readonly [$repr]: ArrayBuffer; + readonly [$inRepr]: ArrayBuffer | ArrayBufferView; + readonly [$gpuRepr]: bytes; + readonly [$memIdent]: WgslBuffer; + readonly [$validStorageSchema]: TSize extends 0 ? false : true; + readonly [$validUniformSchema]: TSize extends 0 ? false : true; + readonly [$validVertexSchema]: false; + readonly [$invalidSchemaReason]: 'Buffers need a size, e.g. d.bytes(256), unless declared in a bind group layout with { storage: d.bytes }'; + // --- +} + +/** + * The contents of a `d.bytes` buffer, as seen by shaders. + * Pass it to `std.bufferView`, `std.bufferArrayView` and `std.bufferLength`. + */ +export interface bytes { + readonly [$internal]: true; + readonly type: 'bytes'; +} + +/** + * Creates a schema for `byteSize` raw bytes, resolved to the WGSL `buffer` type. + * Shaders read and write them as any host-shareable type with `std.bufferView` + * and `std.bufferArrayView`. + * + * Without a size (or with 0), it describes a runtime-sized `buffer`, which only storage + * bindings allow. Bind group layouts take the function itself: `{ storage: d.bytes }`. + * + * Requires the `buffer_view` WGSL language extension, check `root.enabledWgslLanguageFeatures` + * for support. + * + * @example + * const packet = root.createBuffer(d.bytes(1024)).$usage('storage'); + * + * @param byteSize The size in bytes. + */ +export function bytes(): WgslBuffer<0>; +export function bytes(byteSize: TSize): WgslBuffer; +export function bytes(byteSize = 0): WgslBuffer { + if (!Number.isInteger(byteSize) || byteSize < 0) { + throw new Error(`Cannot create a buffer schema of ${byteSize} bytes.`); + } + return new WgslBufferImpl(byteSize); +} + +export function isWgslBuffer(schema: unknown): schema is WgslBuffer { + return isMarkedInternal(schema) && (schema as WgslBuffer).type === 'buffer'; +} + +// -------------- +// Implementation +// -------------- + +class WgslBufferImpl implements WgslBuffer { + readonly [$internal] = {}; + readonly type = 'buffer'; + readonly byteSize: TSize; + + // Type-tokens, not available at runtime + declare readonly [$repr]: ArrayBuffer; + declare readonly [$inRepr]: ArrayBuffer | ArrayBufferView; + declare readonly [$gpuRepr]: bytes; + declare readonly [$memIdent]: WgslBuffer; + declare readonly [$validStorageSchema]: TSize extends 0 ? false : true; + declare readonly [$validUniformSchema]: TSize extends 0 ? false : true; + declare readonly [$validVertexSchema]: false; + declare readonly [$invalidSchemaReason]: WgslBuffer[typeof $invalidSchemaReason]; + // --- + + constructor(byteSize: TSize) { + this.byteSize = byteSize; + } + + toString(): string { + return this.byteSize === 0 ? 'bytes' : `bytes(${this.byteSize})`; + } +} diff --git a/packages/typegpu/src/data/dataIO.ts b/packages/typegpu/src/data/dataIO.ts index 0ef7055df9..637d619d54 100644 --- a/packages/typegpu/src/data/dataIO.ts +++ b/packages/typegpu/src/data/dataIO.ts @@ -223,6 +223,15 @@ const dataWriters = { dataWriters[schema.inner.type]?.(output, schema, value); }, + buffer(output, schema, value) { + const beginning = output.currentByteOffset; + const bytes = ArrayBuffer.isView(value) + ? new Uint8Array(value.buffer, value.byteOffset, value.byteLength) + : new Uint8Array(value); + output.writeSlice(bytes.subarray(0, schema.byteSize)); + output.seekTo(beginning + schema.byteSize); + }, + decorated(output, schema, value) { const alignment = customAlignmentOf(schema); alignIO(output, alignment); @@ -674,6 +683,12 @@ const dataReaders = { return readData(input, schema.inner); }, + buffer(input, schema) { + const bytes = new Uint8Array(schema.byteSize); + input.readSlice(bytes, 0, schema.byteSize); + return bytes.buffer; + }, + decorated(input, schema: wgsl.Decorated) { const alignment = customAlignmentOf(schema); alignIO(input, alignment); diff --git a/packages/typegpu/src/data/index.ts b/packages/typegpu/src/data/index.ts index 7642412155..14246b6bcc 100644 --- a/packages/typegpu/src/data/index.ts +++ b/packages/typegpu/src/data/index.ts @@ -183,6 +183,7 @@ export { unstruct } from './unstruct.ts'; export { matToArray } from './matrix.ts'; export * from './vertexFormatData.ts'; export { atomic } from './atomic.ts'; +export { bytes, type WgslBuffer } from './bytes.ts'; export { _ref as ref } from './ref.ts'; export { align, diff --git a/packages/typegpu/src/data/ptr.ts b/packages/typegpu/src/data/ptr.ts index ed1e76eade..688b693adb 100644 --- a/packages/typegpu/src/data/ptr.ts +++ b/packages/typegpu/src/data/ptr.ts @@ -1,5 +1,6 @@ import { $internal } from '../shared/symbols.ts'; import { type Origin, type OriginToPtrParams, originToPtrParams } from './snippet.ts'; +import type { WgslBuffer } from './bytes.ts'; import type { Access, AddressSpace, BaseData, Ptr, StorableData } from './wgslTypes.ts'; export function ptrFn(inner: T): Ptr<'function', T, 'read-write'> { @@ -10,18 +11,22 @@ export function ptrPrivate(inner: T): Ptr<'private', T, return INTERNAL_createPtr('private', inner, 'read-write'); } -export function ptrWorkgroup(inner: T): Ptr<'workgroup', T, 'read-write'> { +export function ptrWorkgroup( + inner: T, +): Ptr<'workgroup', T, 'read-write'> { return INTERNAL_createPtr('workgroup', inner, 'read-write'); } -export function ptrStorage( - inner: T, - access: TAccess = 'read' as TAccess, -): Ptr<'storage', T, TAccess> { +export function ptrStorage< + T extends StorableData | WgslBuffer, + TAccess extends 'read' | 'read-write' = 'read', +>(inner: T, access: TAccess = 'read' as TAccess): Ptr<'storage', T, TAccess> { return INTERNAL_createPtr('storage', inner, access); } -export function ptrUniform(inner: T): Ptr<'uniform', T, 'read'> { +export function ptrUniform( + inner: T, +): Ptr<'uniform', T, 'read'> { return INTERNAL_createPtr('uniform', inner, 'read'); } diff --git a/packages/typegpu/src/data/schemaMemoryLayout.ts b/packages/typegpu/src/data/schemaMemoryLayout.ts index 691e4fd0c6..4c2f6c69a5 100644 --- a/packages/typegpu/src/data/schemaMemoryLayout.ts +++ b/packages/typegpu/src/data/schemaMemoryLayout.ts @@ -10,6 +10,7 @@ import { } from './dataTypes.ts'; import type { BaseData, WgslArray, WgslStruct, WgslTypeLiteral } from './wgslTypes.ts'; import { isDecorated, isWgslArray, isWgslStruct } from './wgslTypes.ts'; +import { isWgslBuffer } from './bytes.ts'; const knownSizesMap: Record = { bool: 4, @@ -243,6 +244,11 @@ function computeMemoryLayout(data: BaseData): SchemaMemoryLayout { return computeMLOfDisarray(data); } + if (isWgslBuffer(data)) { + const size = data.byteSize === 0 ? Number.NaN : data.byteSize; + return { size, isContiguous: true, longestContiguousPrefix: size }; + } + if (isDecorated(data) || isLooseDecorated(data)) { const size = getCustomSize(data); const undecoratedLayout = computeMemoryLayout(undecorate(data)); diff --git a/packages/typegpu/src/data/schemasEqual.ts b/packages/typegpu/src/data/schemasEqual.ts index 3123550527..ecd65fb943 100644 --- a/packages/typegpu/src/data/schemasEqual.ts +++ b/packages/typegpu/src/data/schemasEqual.ts @@ -3,6 +3,7 @@ import type { AnyAttribute } from './attributes.ts'; import { isDisarray, isLooseDecorated, isUnstruct } from './dataTypes.ts'; import type { AnyData } from './dataTypes.ts'; import { isWgslStorageTexture, isWgslTexture } from './texture.ts'; +import { isWgslBuffer } from './bytes.ts'; import { isAtomic, isDecorated, isPtr, isWgslArray, isWgslStruct } from './wgslTypes.ts'; /** @@ -108,6 +109,10 @@ function schemasEqual(a: AnyData, b: AnyData, structural: boolean): boolean { return schemasEqual(a.inner, b.inner, structural); } + if (isWgslBuffer(a) && isWgslBuffer(b)) { + return a.byteSize === b.byteSize; + } + if (isWgslTexture(a) && isWgslTexture(b)) { // Dimension and multisampling are already encoded in the type return schemasEqual(a.sampleType as AnyData, b.sampleType as AnyData, structural); diff --git a/packages/typegpu/src/data/struct.ts b/packages/typegpu/src/data/struct.ts index e231ec5878..bcf969eff9 100644 --- a/packages/typegpu/src/data/struct.ts +++ b/packages/typegpu/src/data/struct.ts @@ -1,6 +1,7 @@ import { validateProp } from '../nameUtils.ts'; import { getName, setName } from '../shared/meta.ts'; import { $internal } from '../shared/symbols.ts'; +import { isWgslBuffer } from './bytes.ts'; import { schemaCallWrapper } from './schemaCallWrapper.ts'; import type { AnyWgslData, BaseData, WgslStruct } from './wgslTypes.ts'; @@ -26,6 +27,13 @@ export function struct>( if (Object.keys(props).length === 0) { throw new Error('Cannot create struct schema with no properties.'); } + for (const [key, prop] of Object.entries(props)) { + if (isWgslBuffer(prop)) { + throw new Error( + `Cannot use d.${prop} as struct property '${key}', buffers can only be whole variables.`, + ); + } + } return INTERNAL_createStruct(props, false); } diff --git a/packages/typegpu/src/data/wgslTypes.ts b/packages/typegpu/src/data/wgslTypes.ts index f39bf80fb6..9f11c69ad1 100644 --- a/packages/typegpu/src/data/wgslTypes.ts +++ b/packages/typegpu/src/data/wgslTypes.ts @@ -37,6 +37,7 @@ import { $internal, isMarkedInternal } from '../shared/symbols.ts'; import type { Prettify, SwapNever } from '../shared/utilityTypes.ts'; import type { WgslExternalTexture, WgslStorageTexture, WgslTexture } from './texture.ts'; import type { WgslComparisonSampler, WgslSampler } from './sampler.ts'; +import type { WgslBuffer } from './bytes.ts'; import type { _ref as ref } from './ref.ts'; import type { DualFn } from '../types.ts'; @@ -1291,6 +1292,7 @@ export interface Ptr< // Type-tokens, not available at runtime readonly [$repr]: ref>; + readonly [$gpuRepr]: ref>; readonly [$invalidSchemaReason]: 'Pointers are not host-shareable'; // --- } @@ -1417,6 +1419,7 @@ export const wgslTypeLiterals = [ 'array', 'ptr', 'atomic', + 'buffer', 'decorated', 'abstractInt', 'abstractFloat', @@ -1540,6 +1543,7 @@ export type AnyWgslData = | Ptr | Atomic | Atomic + | WgslBuffer | Decorated | AbstractInt | AbstractFloat diff --git a/packages/typegpu/src/resolutionCtx.ts b/packages/typegpu/src/resolutionCtx.ts index c183d04cdd..1ba90a40c3 100644 --- a/packages/typegpu/src/resolutionCtx.ts +++ b/packages/typegpu/src/resolutionCtx.ts @@ -56,7 +56,7 @@ import type { Wgsl, } from './types.ts'; import { CodegenState, isSelfResolvable, NormalState, type FunctionArgument } from './types.ts'; -import type { WgslEnableExtension } from './wgslExtensions.ts'; +import type { WgslEnableExtension, WgslLanguageFeature } from './wgslExtensions.ts'; import { getName, hasTinyestMetadata, isNamable, setName } from './shared/meta.ts'; import { FuncParameterType } from 'tinyest'; import { accessProp } from './tgsl/accessProp.ts'; @@ -456,6 +456,8 @@ export class ResolutionCtxImpl implements ResolutionCtx { public expectedType: BaseData | undefined; #usedImmediate: TgpuImmediateVar | undefined; + readonly #root: ExperimentalTgpuRoot | undefined; + readonly requiredLanguageFeatures = new Set(); /** * A counter used to generate unique identifiers for globally-scoped definitions in the 'random' strategy. @@ -464,6 +466,7 @@ export class ResolutionCtxImpl implements ResolutionCtx { constructor(opts: ResolutionCtxImplOptions) { this.enableExtensions = opts.enableExtensions; + this.#root = opts.root; this.#logGenerator = opts.root ? new LogGeneratorImpl(opts.root) : new LogGeneratorNullImpl(); this.#namespaceInternal = opts.namespace[$internal]; this.gen = opts.shaderGenerator ?? new WgslGenerator(); @@ -485,6 +488,15 @@ export class ResolutionCtxImpl implements ResolutionCtx { return this.#usedImmediate; } + requireLanguageFeature(feature: WgslLanguageFeature, usage: string, fallback: string): void { + if (this.#root && !this.#root.enabledWgslLanguageFeatures.has(feature)) { + throw new Error( + `${usage} cannot be used, because the '${feature}' WGSL language extension is not supported in this environment. Check support with root.enabledWgslLanguageFeatures and fall back to ${fallback} when unavailable.`, + ); + } + this.requiredLanguageFeatures.add(feature); + } + isIdentifierBanned(name: string): boolean { return bannedTokens.has(name); } @@ -1172,7 +1184,7 @@ export interface ResolvedDeclaration { * @param code - The resolved code. * @param declarations - The module-scope declarations emitted by TypeGPU * during this resolution, in emission order. When resolving an array without a - * template, `code` equals `declarations.map((d) => d.code).join('\n\n')` (unless extensions are enabled or minification is enabled). + * template, `code` equals `declarations.map((d) => d.code).join('\n\n')` (unless extensions are enabled or required, or minification is enabled). * When resolving a template, the template itself is not included in `declarations`. * With a shared namespace, only declarations emitted by *this* resolution are * included (memoized ones are not re-emitted). @@ -1246,9 +1258,12 @@ export function resolve(item: Wgsl, options: ResolutionCtxImplOptions): Resoluti bindingReplacements.push([placeholder, String(idx)]); } - if (options.enableExtensions && options.enableExtensions.length > 0) { - const extensions = options.enableExtensions.map((ext) => `enable ${ext};`); - code = `${extensions.join('\n')}\n\n${code}`; + const directives = (options.enableExtensions ?? []).map((ext) => `enable ${ext};`); + if (ctx.requiredLanguageFeatures.size > 0) { + directives.push(`requires ${[...ctx.requiredLanguageFeatures].join(', ')};`); + } + if (directives.length > 0) { + code = `${directives.join('\n')}\n\n${code}`; } let declarations = ctx.declarations.map(({ name, code: declarationCode }) => ({ diff --git a/packages/typegpu/src/serial/schema.ts b/packages/typegpu/src/serial/schema.ts index 9543928af3..4617ac7888 100644 --- a/packages/typegpu/src/serial/schema.ts +++ b/packages/typegpu/src/serial/schema.ts @@ -1,5 +1,6 @@ import * as d from '../data/index.ts'; import { isWgslStorageTexture, isWgslTexture } from '../data/texture.ts'; +import { isWgslBuffer } from '../data/bytes.ts'; import { assertExhaustive } from '../shared/utilityTypes.ts'; type SerializedDataAttrib = @@ -17,6 +18,7 @@ export type SerializedDataSchema = | { type: 'struct'; props: [string, SerializedDataSchema][] } | { type: 'unstruct'; props: [string, SerializedDataSchema][] } | { type: 'atomic'; inner: SerializedDataSchema } + | { type: 'buffer'; byteSize: number } | { type: 'decorated'; inner: SerializedDataSchema; attribs: SerializedDataAttrib[] } | { type: 'sampled-texture'; kind: d.WgslTexture['type']; sampleType: SerializedDataSchema } | { @@ -180,6 +182,10 @@ export function serializeDataSchema(schema: d.BaseData): SerializedDataSchema { return { type: 'atomic', inner: serializeDataSchema(schema.inner as d.AnyData) }; } + if (isWgslBuffer(schema)) { + return { type: 'buffer', byteSize: schema.byteSize }; + } + if (d.isWgslArray(schema)) { return { type: 'array', @@ -239,6 +245,10 @@ export function deserializeDataSchema(schema: SerializedDataSchema): d.AnyData { return d.atomic(deserializeDataSchema(schema.inner) as d.U32 | d.I32) as d.AnyData; } + if (schema.type === 'buffer') { + return d.bytes(schema.byteSize); + } + if (schema.type === 'decorated') { let result = deserializeDataSchema(schema.inner); for (let i = schema.attribs.length - 1; i >= 0; i--) { diff --git a/packages/typegpu/src/std/bufferView.ts b/packages/typegpu/src/std/bufferView.ts new file mode 100644 index 0000000000..e01ef2b675 --- /dev/null +++ b/packages/typegpu/src/std/bufferView.ts @@ -0,0 +1,277 @@ +import { dualImpl, MissingCpuImplError } from '../core/function/dualImpl.ts'; +import { warnIfNotUniformAligned } from '../core/pipeline/webgpuLimitations.ts'; +import { stitch } from '../core/resolve/stitch.ts'; +import { alignmentOf } from '../data/alignmentOf.ts'; +import { type bytes, isWgslBuffer, type WgslBuffer } from '../data/bytes.ts'; +import { undecorate, unptr } from '../data/dataTypes.ts'; +import { u32 } from '../data/numeric.ts'; +import { offsetsForProps } from '../data/offsets.ts'; +import { createPtrFromOrigin } from '../data/ptr.ts'; +import { type _ref as ref, derefSnippet } from '../data/ref.ts'; +import { sizeOf } from '../data/sizeOf.ts'; +import { snip, type Snippet } from '../data/snippet.ts'; +import { + type Atomic, + type BaseData, + type Decorated, + isMat, + isVec, + isVecBool, + isWgslArray, + isWgslData, + isWgslStruct, + type StorableData, + type WgslArray, + type WgslStruct, +} from '../data/wgslTypes.ts'; +import { roundUp } from '../mathUtils.ts'; +import { setName } from '../shared/meta.ts'; +import type { ExtractInvalidSchemaError, InferGPU, IsValidStorageSchema } from '../shared/repr.ts'; +import { $gpuCallable } from '../shared/symbols.ts'; +import { tryConvertSnippet } from '../tgsl/conversion.ts'; +import { asComptime, type DualFn, type ResolutionCtx } from '../types.ts'; + +// ---------- +// Public API +// ---------- + +/** + * A schema, or a function from an element count to a schema, which stands for a + * runtime-sized schema (as in bind group layouts), e.g. `d.arrayOf(d.u32)`. + */ +type ViewSchema = BaseData | ((count: number) => BaseData); + +type UnwrapViewSchema = T extends BaseData + ? T + : T extends (count: number) => infer TSchema + ? TSchema extends BaseData + ? TSchema + : never + : never; + +type ContainsAtomic = T extends Atomic + ? true + : T extends Decorated + ? ContainsAtomic + : T extends WgslArray + ? ContainsAtomic + : T extends WgslStruct + ? true extends { [K in keyof TProps]: ContainsAtomic }[keyof TProps] + ? true + : false + : false; + +type ViewSchemaError = T extends WgslBuffer + ? '(Error) Buffers cannot be viewed as buffers' + : IsValidStorageSchema extends false + ? ExtractInvalidSchemaError + : ContainsAtomic extends true + ? '(Error) Buffer views cannot hold atomics' + : never; + +type ValidateViewSchema = [ViewSchemaError>] extends [never] + ? T + : ViewSchemaError>; + +type BufferViewFn = ( + schema: ValidateViewSchema, + buffer: bytes, + byteOffset: number, +) => ref>>; + +type BufferArrayViewFn = BaseData>( + schema: ValidateViewSchema, + buffer: bytes, + byteOffset: number, + byteSize: number, +) => ref>>; + +/** + * Reads the bytes of a `d.bytes` buffer at `byteOffset` as `schema`, mirroring WGSL's `bufferView`. + * Returns a reference: access its members and elements directly, and read or write the whole value with `.$`. + * The reference is mutable if the buffer is. + * + * A runtime-sized schema, such as `d.arrayOf(d.f32)`, extends to the end of the buffer. + * The offset must be a multiple of the schema's alignment. + * + * @example + * const header = std.bufferView(Header, packet.$, 0); + * header.count += 1; + * std.bufferView(d.u32, packet.$, 64).$ = 7; + */ +export const bufferView = createBufferViewFn('bufferView') as unknown as DualFn; + +/** + * Reads `byteSize` bytes of a `d.bytes` buffer at `byteOffset` as a runtime-sized `schema`, + * mirroring WGSL's `bufferArrayView`. Use it to bound arrays that don't extend to the end + * of the buffer: their `length` covers only `byteSize` bytes. + * + * @example + * const indices = std.bufferArrayView(d.arrayOf(d.u32), packet.$, 16, indexCount * 4); + */ +export const bufferArrayView = createBufferViewFn( + 'bufferArrayView', +) as unknown as DualFn; + +/** + * Returns the size of a `d.bytes` buffer in bytes, mirroring WGSL's `bufferLength`. + */ +export const bufferLength = dualImpl<(buffer: bytes) => number>({ + name: 'bufferLength', + normalImpl: 'bufferLength is not supported outside of CODEGEN mode.', + signature: (buffer) => { + assertBuffer('bufferLength', buffer); + return { argTypes: [buffer], returnType: u32 }; + }, + codegenImpl: (_ctx, [buffer]) => stitch`bufferLength(&${derefSnippet(buffer)})`, + sideEffects: false, +}); + +// -------------- +// Implementation +// -------------- + +function createBufferViewFn(name: 'bufferView' | 'bufferArrayView') { + const impl = () => { + throw new MissingCpuImplError(`${name} is not supported outside of CODEGEN mode.`); + }; + setName(impl, name); + impl.toString = () => name; + + return Object.assign(impl, { + [$gpuCallable]: { + call(ctx: ResolutionCtx, [schemaArg, bufferArg, offsetArg, sizeArg]: Snippet[]) { + if (!schemaArg || !bufferArg || !offsetArg || (name === 'bufferArrayView' && !sizeArg)) { + throw new Error(`Missing arguments to std.${name}.`); + } + + const schema = unwrapViewSchema(name, schemaArg); + const typeName = ctx.resolve(schema).value; + assertViewable(name, typeName, schema); + const buffer = derefSnippet(bufferArg); + assertBuffer(name, buffer.dataType); + const ptrType = createPtrFromOrigin(buffer.origin, schema as StorableData); + if (!ptrType) { + throw new Error( + `std.${name} needs a buffer in the storage, uniform or workgroup address space.`, + ); + } + + if (ptrType.addressSpace === 'uniform') { + warnIfNotUniformAligned(schema); + } + + const offset = tryConvertSnippet(ctx, offsetArg, u32); + const size = sizeArg && tryConvertSnippet(ctx, sizeArg, u32); + assertViewFits(name, typeName, schema, buffer.dataType, offset, size); + + return snip( + stitch`${name}<${typeName}>(&${buffer}, ${[offset, size]})`, + ptrType, + buffer.origin, + buffer.possibleSideEffects || offset.possibleSideEffects || !!size?.possibleSideEffects, + ); + }, + }, + }); +} + +function unwrapViewSchema(name: string, arg: Snippet): BaseData { + const value = asComptime(arg)?.value; + const schema = !isWgslData(value) && typeof value === 'function' ? value(0) : value; + if (!isWgslData(schema)) { + throw new Error(`std.${name} expects a schema as its first argument.`); + } + + const runtimeSized = Number.isNaN(sizeOf(schema)); + if (name === 'bufferArrayView' && !runtimeSized) { + throw new Error( + `std.bufferArrayView needs a runtime-sized schema, such as d.arrayOf(d.u32). Use std.bufferView for '${schema}'.`, + ); + } + return schema; +} + +function assertViewable(name: string, typeName: string, schema: BaseData): void { + const inner = undecorate(schema); + if (isWgslStruct(inner)) { + for (const prop of Object.values(inner.propTypes)) { + assertViewable(name, typeName, prop); + } + return; + } + if (isWgslArray(inner)) { + assertViewable(name, typeName, inner.elementType); + return; + } + if (inner.type === 'atomic') { + throw new Error( + `std.${name} cannot view '${typeName}', because WGSL buffer views cannot hold atomics.`, + ); + } + const hostShareable = + ['f32', 'f16', 'i32', 'u32'].includes(inner.type) || + (isVec(inner) && !isVecBool(inner)) || + isMat(inner); + if (!hostShareable) { + throw new Error( + `std.${name} cannot view '${typeName}', because '${inner.type}' is not host-shareable.`, + ); + } +} + +function assertBuffer(name: string, dataType: unknown): asserts dataType is WgslBuffer { + const inner = unptr(dataType as BaseData); + if (!isWgslBuffer(inner)) { + throw new Error( + `std.${name} expects a d.bytes buffer, got '${String(inner)}'. Declare the buffer with d.bytes(n).`, + ); + } +} + +/** + * The fewest bytes a view of `schema` covers: a runtime-sized array needs one element. + */ +function minSizeOf(schema: BaseData): number { + const size = sizeOf(schema); + if (!Number.isNaN(size)) { + return size; + } + if (isWgslArray(schema)) { + return roundUp(sizeOf(schema.elementType), alignmentOf(schema.elementType)); + } + if (isWgslStruct(schema)) { + const [key, lastProp] = Object.entries(schema.propTypes).pop() as [string, BaseData]; + return (offsetsForProps(schema)[key]?.offset ?? 0) + minSizeOf(lastProp); + } + return size; +} + +function assertViewFits( + name: string, + typeName: string, + schema: BaseData, + buffer: WgslBuffer, + offsetArg: Snippet, + sizeArg: Snippet | undefined, +): void { + const offset = asComptime(offsetArg)?.value; + if (typeof offset !== 'number') { + return; + } + + const alignment = alignmentOf(schema); + if (offset % alignment !== 0) { + throw new Error( + `std.${name} cannot view '${typeName}' at byte ${offset}, because the offset is not a multiple of its alignment (${alignment}).`, + ); + } + + const size = sizeArg && asComptime(sizeArg)?.value; + const viewSize = typeof size === 'number' ? Math.max(size, minSizeOf(schema)) : minSizeOf(schema); + if (buffer.byteSize > 0 && offset + viewSize > buffer.byteSize) { + throw new Error( + `std.${name} cannot view '${typeName}' at byte ${offset} of d.${buffer}, because the view needs ${viewSize} bytes and only ${Math.max(buffer.byteSize - offset, 0)} are left.`, + ); + } +} diff --git a/packages/typegpu/src/std/index.ts b/packages/typegpu/src/std/index.ts index d023752924..4066e987ec 100644 --- a/packages/typegpu/src/std/index.ts +++ b/packages/typegpu/src/std/index.ts @@ -142,6 +142,8 @@ export { export { arrayLength } from './array.ts'; +export { bufferArrayView, bufferLength, bufferView } from './bufferView.ts'; + // oxfmt-ignore export { pack2x16float, diff --git a/packages/typegpu/src/tgsl/forOfUtils.ts b/packages/typegpu/src/tgsl/forOfUtils.ts index eb30d74cda..1ded8e7b1b 100644 --- a/packages/typegpu/src/tgsl/forOfUtils.ts +++ b/packages/typegpu/src/tgsl/forOfUtils.ts @@ -63,7 +63,10 @@ export function getRangeSnippets( iterableSnippet: Snippet, unroll: boolean = false, ): { start: Snippet; end: Snippet; step: Snippet; comparison: '<' | '>' } { - const { value, dataType } = iterableSnippet; + const { value } = iterableSnippet; + const dataType = wgsl.isPtr(iterableSnippet.dataType) + ? iterableSnippet.dataType.inner + : iterableSnippet.dataType; if (isTgpuRange(value)) { const { start, end, step } = value; diff --git a/packages/typegpu/src/tgsl/wgslGenerator.ts b/packages/typegpu/src/tgsl/wgslGenerator.ts index 46bc0a665e..2327f4ded3 100644 --- a/packages/typegpu/src/tgsl/wgslGenerator.ts +++ b/packages/typegpu/src/tgsl/wgslGenerator.ts @@ -1,6 +1,7 @@ import * as tinyest from 'tinyest'; import { stitch } from '../core/resolve/stitch.ts'; import { arrayOf } from '../data/array.ts'; +import { isWgslBuffer, type WgslBuffer } from '../data/bytes.ts'; import { type AnyData, UnknownData, unptr } from '../data/dataTypes.ts'; import { bool, i32, u32 } from '../data/numeric.ts'; import { vec2u, vec3u, vec4u } from '../data/vector.ts'; @@ -1231,6 +1232,10 @@ The generated shader will evaluate them in declaration order.`, } public declareGlobalVar(options: VariableDefinitionOptions): ResolvedSnippet { + if (isWgslBuffer(options.dataType)) { + assertBufferScope(options.dataType, options.scope); + } + let pre = ''; if (options.group !== undefined) { @@ -2191,6 +2196,22 @@ ${stringifyNode(statement)}`); } } +function assertBufferScope(dataType: WgslBuffer, scope: VariableDefinitionOptions['scope']) { + if (scope === 'readonly' || scope === 'mutable') { + return; + } + if (scope !== 'uniform' && scope !== 'workgroup') { + throw new Error( + `d.${dataType} cannot be used in the '${scope}' address space, only in storage, uniform and workgroup variables.`, + ); + } + if (dataType.byteSize === 0) { + throw new Error( + `d.${dataType} needs a size in the '${scope}' address space, only storage allows runtime-sized buffers. Give it a size, e.g. d.bytes(256).`, + ); + } +} + const immutabilityReasons: Partial> = { constant: 'the left side is a constant', 'constant-immutable-def': 'the left side is a constant', diff --git a/packages/typegpu/src/types.ts b/packages/typegpu/src/types.ts index 3c66e1031f..2fdef32ff1 100644 --- a/packages/typegpu/src/types.ts +++ b/packages/typegpu/src/types.ts @@ -43,7 +43,7 @@ import { $resolve, } from './shared/symbols.ts'; import type { TgpuBindGroupLayout, TgpuLayoutEntry } from './tgpuBindGroupLayout.ts'; -import type { WgslEnableExtension } from './wgslExtensions.ts'; +import type { WgslEnableExtension, WgslLanguageFeature } from './wgslExtensions.ts'; import type { Infer } from './shared/repr.ts'; import type { ShaderGenerator } from './tgsl/shaderGenerator.ts'; import type { StorageFlag } from './extension.ts'; @@ -385,6 +385,15 @@ export interface ResolutionCtx { */ registerImmediate(immediate: TgpuImmediateVar): void; + /** + * Registers the use of a WGSL language extension, which adds a `requires` directive + * to the resolved code. + * @param usage What depends on the extension, named in the error. + * @param fallback What to use instead, suggested in the error. + * @throws When the root the code is resolved for does not support the extension. + */ + requireLanguageFeature(feature: WgslLanguageFeature, usage: string, fallback: string): void; + isIdentifierBanned(name: string): boolean; /** diff --git a/packages/typegpu/src/wgslExtensions.ts b/packages/typegpu/src/wgslExtensions.ts index b35f47a4fb..05c5c725af 100644 --- a/packages/typegpu/src/wgslExtensions.ts +++ b/packages/typegpu/src/wgslExtensions.ts @@ -14,3 +14,9 @@ export const wgslEnableExtensionToFeatureName: Record { + it('resolves to the WGSL buffer type and requires buffer_view', ({ root }) => { + const packet = root.createMutable(d.bytes(256)); + const page = root.createUniform(d.bytes(64)); + const scratch = tgpu.workgroupVar(d.bytes(1024)); + const layout = tgpu.bindGroupLayout({ mesh: { storage: d.bytes } }); + + const fill = tgpu.fn([d.ptrStorage(d.bytes(), 'read-write')])(() => {}); + + const main = tgpu.computeFn({ workgroupSize: [1] })(() => { + const mesh = layout.$.mesh; + const local = scratch.$; + fill(d.ref(packet.$)); + + if (false) { + expectTypeOf(mesh).toEqualTypeOf(); + expectTypeOf(page.$).toEqualTypeOf(); + } + }); + + expect(tgpu.resolve([main])).toMatchInlineSnapshot(` + "requires buffer_view; + + @group(1) @binding(0) var mesh_1: buffer; + + var scratch: buffer<1024>; + + fn fill(_arg_0: ptr) {} + + @group(0) @binding(0) var packet: buffer<256>; + + @compute @workgroup_size(1) fn main() { + let mesh = (&mesh_1); + let local = (&scratch); + fill((&packet)); + }" + `); + }); + + it('throws when the environment does not support buffer_view', ({ root, navigator }) => { + navigator.gpu.wgslLanguageFeatures.delete('buffer_view'); + const scratch = tgpu.workgroupVar(d.bytes(1024)); + const main = tgpu.computeFn({ workgroupSize: [1] })(() => { + const local = scratch.$; + }); + + expect(() => root.createComputePipeline({ compute: main }).dispatchWorkgroups(1)) + .toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - computePipelineCore + - computeFn:main + - d.ref(scratch.$) + - var:scratch.$ + - var:scratch + - bytes(1024): d.bytes(1024) cannot be used, because the 'buffer_view' WGSL language extension is not supported in this environment. Check support with root.enabledWgslLanguageFeatures and fall back to typed schemas (d.struct, d.arrayOf) when unavailable.] + `); + }); + + it('allows buffers only as whole storage, uniform and workgroup variables', () => { + const local = tgpu.privateVar(d.bytes(16)); + const layout = tgpu.bindGroupLayout({ page: { uniform: d.bytes() } }); + const readLocal = tgpu.fn([])(() => { + const bytes = local.$; + }); + const readPage = tgpu.fn([])(() => { + const bytes = layout.$.page; + }); + + expect(() => tgpu.resolve([readLocal])).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn:readLocal + - d.ref(local.$) + - var:local.$ + - var:local: d.bytes(16) cannot be used in the 'private' address space, only in storage, uniform and workgroup variables.] + `); + expect(() => tgpu.resolve([readPage])).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn:readPage + - d.ref(layout.$.page) + - uniform:page.$ + - uniform:page: d.bytes needs a size in the 'uniform' address space, only storage allows runtime-sized buffers. Give it a size, e.g. d.bytes(256).] + `); + expect(() => d.struct({ data: d.bytes(16) })).toThrowErrorMatchingInlineSnapshot( + `[Error: Cannot use d.bytes(16) as struct property 'data', buffers can only be whole variables.]`, + ); + expect(() => d.arrayOf(d.bytes(16), 4)).toThrowErrorMatchingInlineSnapshot( + `[Error: Cannot create an array of d.bytes(16), buffers can only be whole variables.]`, + ); + expect(() => d.bytes(1.5)).toThrowErrorMatchingInlineSnapshot( + `[Error: Cannot create a buffer schema of 1.5 bytes.]`, + ); + }); + + it('needs a size outside storage bind group layouts', ({ root }) => { + const layout = tgpu.bindGroupLayout({ mesh: { storage: d.bytes } }); + root.createBindGroup(layout, { mesh: root.createBuffer(d.bytes(512)).$usage('storage') }); + + if (false) { + // @ts-expect-error + root.createBuffer(d.bytes()); + // @ts-expect-error + root.createUniform(d.bytes()); + // @ts-expect-error + tgpu.workgroupVar(d.bytes()); + } + }); + + it('writes and reads raw bytes', async ({ root, device }) => { + const packet = root.createBuffer(d.bytes(16), new Uint32Array([1, 2, 3])); + const initial = vi.mocked(root.unwrap(packet).getMappedRange).mock.results[0]?.value; + expect([...new Uint32Array(initial)]).toStrictEqual([1, 2, 3, 0]); + + const page = root.createBuffer(d.bytes(16)); + page.write(new Uint16Array([7, 8]), { startOffset: 4 }); + const [, gpuOffset, data, dataOffset, size] = vi.mocked(device.queue.writeBuffer).mock + .calls[0]!; + expect([gpuOffset, dataOffset, size]).toStrictEqual([4, 4, 4]); + expect([...new Uint16Array(data as ArrayBuffer, 4, 2)]).toStrictEqual([7, 8]); + + const read = await page.read(); + expectTypeOf(read).toEqualTypeOf(); + expect(read.byteLength).toBe(16); + }); + + it('measures and compares as raw bytes', () => { + expect(d.sizeOf(d.bytes(64))).toBe(64); + expect(d.sizeOf(d.bytes())).toBeNaN(); + expect(d.alignmentOf(d.bytes(64))).toBe(4); + expect(d.schemasEqualStructurally(d.bytes(64), d.bytes(64))).toBe(true); + expect(d.schemasEqualStructurally(d.bytes(64), d.bytes(32))).toBe(false); + }); +}); diff --git a/packages/typegpu/tests/serializeDataSchema.test.ts b/packages/typegpu/tests/serializeDataSchema.test.ts index 5aef259cb0..5e9488e470 100644 --- a/packages/typegpu/tests/serializeDataSchema.test.ts +++ b/packages/typegpu/tests/serializeDataSchema.test.ts @@ -14,6 +14,7 @@ const schemas = [ ), d.disarrayOf(d.unstruct({ id: d.u32, packed: d.uint16x2 }), 3), d.atomic(d.u32), + d.bytes(64), d.align(16, d.size(32, d.struct({ value: d.vec2f }))), d.location(2, d.vec4f), d.interpolate('linear, sample', d.vec2f), @@ -93,6 +94,10 @@ describe('data schema serialization', () => { }, "type": "atomic", }, + { + "byteSize": 64, + "type": "buffer", + }, { "attribs": [ { diff --git a/packages/typegpu/tests/std/bufferView.test.ts b/packages/typegpu/tests/std/bufferView.test.ts new file mode 100644 index 0000000000..e5f55ed437 --- /dev/null +++ b/packages/typegpu/tests/std/bufferView.test.ts @@ -0,0 +1,315 @@ +import { describe, expect, expectTypeOf } from 'vitest'; +import { d, std, tgpu } from 'typegpu'; +import { it } from 'typegpu-testing-utility'; + +describe('buffer views', () => { + it('packs differently typed data at runtime offsets', () => { + const Header = d.struct({ indexCount: d.u32, vertexCount: d.u32 }); + const layout = tgpu.bindGroupLayout({ + mesh: { storage: d.bytes }, + sum: { storage: d.f32, access: 'mutable' }, + }); + + const main = tgpu.computeFn({ workgroupSize: [1] })(() => { + const header = std.bufferView(Header, layout.$.mesh, 0); + const indicesSize = header.indexCount * 4; + const indices = std.bufferArrayView(d.arrayOf(d.u32), layout.$.mesh, 8, indicesSize); + const vertices = std.bufferView(d.arrayOf(d.vec2f), layout.$.mesh, 8 + indicesSize); + + for (const index of indices) { + layout.$.sum += vertices[index]!.x; + } + layout.$.sum /= d.f32(std.bufferLength(layout.$.mesh)); + + if (false) { + expectTypeOf(header).toEqualTypeOf>(); + expectTypeOf(indices).toEqualTypeOf>(); + expectTypeOf(vertices).toEqualTypeOf>(); + expectTypeOf(std.bufferLength(layout.$.mesh)).toEqualTypeOf(); + } + }); + + expect(tgpu.resolve([main])).toMatchInlineSnapshot(` + "requires buffer_view; + + struct Header { + indexCount: u32, + vertexCount: u32, + } + + @group(0) @binding(0) var mesh: buffer; + + @group(0) @binding(1) var sum: f32; + + @compute @workgroup_size(1) fn main() { + let header = bufferView
(&mesh, 0u); + let indicesSize = ((*header).indexCount * 4u); + let indices = bufferArrayView>(&mesh, 8u, indicesSize); + let vertices = bufferView>(&mesh, (8u + indicesSize)); + for (var i = 0u; i < arrayLength(indices); i += 1u) { + let index = (*indices)[i]; + sum += (*vertices)[index].x; + } + sum /= f32(bufferLength(&mesh)); + }" + `); + }); + + it('reads ragged records', ({ root }) => { + const records = root.createReadonly(d.bytes(4096)); + const totals = root.createMutable(d.arrayOf(d.f32, 64)); + + const main = tgpu.computeFn({ + workgroupSize: [1], + in: { gid: d.builtin.globalInvocationId }, + })(({ gid }) => { + let offset = d.u32(0); + for (let i = d.u32(0); i < gid.x; i++) { + offset += 4 + std.bufferView(d.u32, records.$, offset).$ * 4; + } + const count = std.bufferView(d.u32, records.$, offset).$; + const values = std.bufferArrayView(d.arrayOf(d.f32), records.$, offset + 4, count * 4); + let total = d.f32(0); + for (const value of values) { + total += value; + } + totals.$[gid.x] = total; + }); + + expect(tgpu.resolve([main])).toMatchInlineSnapshot(` + "requires buffer_view; + + @group(0) @binding(0) var records: buffer<4096>; + + @group(0) @binding(1) var totals: array; + + @compute @workgroup_size(1) fn main(@builtin(global_invocation_id) gid: vec3u) { + var offset = 0u; + for (var i = 0u; (i < gid.x); i++) { + offset += (4u + ((*bufferView(&records, offset)) * 4u)); + } + let count = (*bufferView(&records, offset)); + let values = bufferArrayView>(&records, (offset + 4u), (count * 4u)); + var total = 0f; + for (var i = 0u; i < arrayLength(values); i += 1u) { + let value = (*values)[i]; + total += value; + } + totals[gid.x] = total; + }" + `); + }); + + it('reinterprets bytes as another type', ({ root }) => { + const packet = root.createMutable(d.bytes(16)); + + const negate = tgpu.fn([])(() => { + const bits = std.bufferView(d.u32, packet.$, 0); + bits.$ ^= 0x80000000; + std.bufferView(d.f32, packet.$, 4).$ = -std.bufferView(d.f32, packet.$, 4).$; + }); + + expect(tgpu.resolve([negate])).toMatchInlineSnapshot(` + "requires buffer_view; + + @group(0) @binding(0) var packet: buffer<16>; + + fn negate() { + let bits = bufferView(&packet, 0u); + (*bits) ^= 2147483648u; + (*bufferView(&packet, 4u)) = -((*bufferView(&packet, 4u))); + }" + `); + }); + + it('reuses workgroup memory across phases', () => { + const scratch = tgpu.workgroupVar(d.bytes(1024)); + + const main = tgpu.computeFn({ + workgroupSize: [128], + in: { lid: d.builtin.localInvocationIndex }, + })(({ lid }) => { + const floats = std.bufferView(d.arrayOf(d.f32, 256), scratch.$, 0); + floats[lid] = d.f32(lid); + std.workgroupBarrier(); + + const pairs = std.bufferView(d.arrayOf(d.vec2u, 128), scratch.$, 0); + pairs[lid] = d.vec2u(lid); + }); + + expect(tgpu.resolve([main])).toMatchInlineSnapshot(` + "requires buffer_view; + + var scratch: buffer<1024>; + + @compute @workgroup_size(128) fn main(@builtin(local_invocation_index) lid: u32) { + let floats = bufferView>(&scratch, 0u); + (*floats)[lid] = f32(lid); + workgroupBarrier(); + let pairs = bufferView>(&scratch, 0u); + (*pairs)[lid] = vec2u(lid); + }" + `); + }); + + it('views uniform buffers and pointer parameters', ({ root }) => { + const Light = d.struct({ position: d.vec3f, range: d.f32 }); + const lights = root.createUniform(d.bytes(512)); + const counters = root.createMutable(d.bytes(64)); + + const bump = tgpu.fn([d.ptrStorage(d.bytes(), 'read-write'), d.u32])((buffer, index) => { + std.bufferView(d.u32, buffer, index * 4).$ += 1; + }); + + const main = tgpu.fn( + [d.u32], + d.f32, + )((index) => { + bump(d.ref(counters.$), index); + return std.bufferView(Light, lights.$, index * 32).range; + }); + + expect(tgpu.resolve([main])).toMatchInlineSnapshot(` + "requires buffer_view; + + fn bump(buffer: ptr, index: u32) { + (*bufferView(&(*buffer), (index * 4u))) += 1u; + } + + @group(0) @binding(0) var counters: buffer<64>; + + struct Light { + position: vec3f, + range: f32, + } + + @group(0) @binding(1) var lights: buffer<512>; + + fn main(index: u32) -> f32 { + bump((&counters), index); + return (*bufferView(&lights, (index * 32u))).range; + }" + `); + }); + + it('rejects views that the shader cannot hold', ({ root }) => { + const packet = root.createMutable(d.bytes(16)); + const Counter = d.struct({ count: d.atomic(d.u32) }); + const resolveFn = (body: () => void) => () => tgpu.resolve([tgpu.fn([])(body)]); + + expect( + resolveFn(() => { + 'use gpu'; + // @ts-expect-error: buffer views cannot hold atomics + std.bufferView(Counter, packet.$, 0); + }), + ).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn: + - fn:bufferView: std.bufferView cannot view 'Counter', because WGSL buffer views cannot hold atomics.] + `); + expect( + resolveFn(() => { + 'use gpu'; + // @ts-expect-error: booleans are not host-shareable + std.bufferView(d.vec2b, packet.$, 0); + }), + ).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn: + - fn:bufferView: std.bufferView cannot view 'vec2', because 'vec2' is not host-shareable.] + `); + expect( + resolveFn(() => { + 'use gpu'; + // @ts-expect-error: std.bufferArrayView needs a runtime-sized schema + std.bufferArrayView(d.arrayOf(d.u32, 4), packet.$, 0, 16); + }), + ).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn: + - fn:bufferArrayView: std.bufferArrayView needs a runtime-sized schema, such as d.arrayOf(d.u32). Use std.bufferView for 'arrayOf(u32, 4)'.] + `); + expect( + resolveFn(() => { + 'use gpu'; + // @ts-expect-error: only d.bytes buffers can be viewed + std.bufferView(d.u32, d.u32(1), 0); + }), + ).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn: + - fn:bufferView: std.bufferView expects a d.bytes buffer, got 'u32'. Declare the buffer with d.bytes(n).] + `); + }); + + it('rejects misaligned and out-of-bounds constant offsets', ({ root }) => { + const packet = root.createMutable(d.bytes(64)); + const resolveFn = (body: () => void) => () => tgpu.resolve([tgpu.fn([])(body)]); + + expect( + resolveFn(() => { + 'use gpu'; + std.bufferView(d.vec4f, packet.$, 8); + }), + ).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn: + - fn:bufferView: std.bufferView cannot view 'vec4f' at byte 8, because the offset is not a multiple of its alignment (16).] + `); + expect( + resolveFn(() => { + 'use gpu'; + std.bufferView(d.mat2x2f, packet.$, 56); + }), + ).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn: + - fn:bufferView: std.bufferView cannot view 'mat2x2f' at byte 56 of d.bytes(64), because the view needs 16 bytes and only 8 are left.] + `); + expect( + resolveFn(() => { + 'use gpu'; + std.bufferArrayView(d.arrayOf(d.u32), packet.$, 32, 40); + }), + ).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn: + - fn:bufferArrayView: std.bufferArrayView cannot view 'array' at byte 32 of d.bytes(64), because the view needs 40 bytes and only 32 are left.] + `); + }); + + it('keeps views of read-only buffers read-only', ({ root }) => { + const packet = root.createReadonly(d.bytes(16)); + const page = root.createUniform(d.bytes(16)); + const resolveFn = (body: () => void) => () => tgpu.resolve([tgpu.fn([])(body)]); + + expect( + resolveFn(() => { + 'use gpu'; + std.bufferView(d.u32, packet.$, 0).$ = 1; + }), + ).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn:: 'std.bufferView(d.u32, packet.$, 0).$ = 1' is invalid, because readonly buffers cannot be mutated.] + `); + expect( + resolveFn(() => { + 'use gpu'; + std.bufferView(d.vec4f, page.$, 0).x = 1; + }), + ).toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn:: 'std.bufferView(d.vec4f, page.$, 0).x = 1' is invalid, because uniform buffers cannot be mutated.] + `); + }); +}); diff --git a/packages/typegpu/tests/tgsl/wgslGenerator.test.ts b/packages/typegpu/tests/tgsl/wgslGenerator.test.ts index aa4afd180e..9cd44a15bd 100644 --- a/packages/typegpu/tests/tgsl/wgslGenerator.test.ts +++ b/packages/typegpu/tests/tgsl/wgslGenerator.test.ts @@ -338,6 +338,54 @@ describe('WgslGenerator', () => { `); }); + it('creates correct code for "for ... of ..." statements using pointer iterables', () => { + const layout = tgpu.bindGroupLayout({ + arr: { storage: d.arrayOf(d.f32) }, + }); + + const sum = tgpu.fn( + [d.ptrStorage(d.arrayOf(d.f32, 4))], + d.f32, + )((values) => { + let res = d.f32(0); + for (const foo of values) { + res += foo; + } + return res; + }); + + const main = () => { + 'use gpu'; + const values = layout.$.arr; + let res = d.f32(0); + for (const foo of values) { + res += foo; + } + }; + + expect(tgpu.resolve([sum, main])).toMatchInlineSnapshot(` + "fn sum(values: ptr, read>) -> f32 { + var res = 0f; + for (var i = 0u; i < 4u; i += 1u) { + let foo = (*values)[i]; + res += foo; + } + return res; + } + + @group(0) @binding(0) var arr: array; + + fn main() { + let values = (&arr); + var res = 0f; + for (var i = 0u; i < arrayLength(values); i += 1u) { + let foo = (*values)[i]; + res += foo; + } + }" + `); + }); + it('creates correct code for "for ... of ..." statements using lazy and comptime iterables', () => { const comptimeVec = tgpu.comptime(() => d.vec2f(1, 2));