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..d44b395563 --- /dev/null +++ b/apps/typegpu-docs/src/content/docs/apis/buffer-views.mdx @@ -0,0 +1,207 @@ +--- +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.buffer(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.buffer(n)` works wherever a schema does, in storage and uniform buffers, workgroup variables and bind group layouts. +Bind group layouts also take `d.buffer` itself, for storage buffers of any size. +A `d.buffer` entry accepts buffers of any schema, so the CPU can write typed data that shaders view as bytes. + +```ts twoslash +import { tgpu, d } from 'typegpu'; + +const root = await tgpu.init(); +// ---cut--- +const packet = root.createMutable(d.buffer(1024)); +const lights = root.createUniform(d.buffer(512)); +const scratch = tgpu.workgroupVar(d.buffer(4096)); + +const layout = tgpu.bindGroupLayout({ + mesh: { storage: d.buffer }, +}); +``` + +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.buffer(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.buffer(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 view of a runtime-sized array, such as `d.arrayOf(d.f32)`, runs to the end of the buffer, and its length is however many elements fit. +When more data follows the array, `std.bufferArrayView(schema, buffer, byteOffset, byteSize)` ends it after `byteSize` bytes, so its length is `byteSize` divided by the array stride (the element size rounded up to its alignment, so 16 for `d.vec3f`). +This gives shaders arrays whose length is only known at runtime, such as a count stored in the buffer itself. A fixed-size `d.arrayOf(d.f32, n)` needs `n` when the shader is created. +`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 vertices and its indices 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.buffer }, + extent: { storage: d.vec2f, access: 'mutable' }, +}); + +const measure = tgpu.computeFn({ workgroupSize: [1] })(() => { + const header = std.bufferView(Header, layout.$.mesh, 0); + const verticesSize = header.vertexCount * d.sizeOf(d.vec2f); + const vertices = std.bufferArrayView( + d.arrayOf(d.vec2f), + layout.$.mesh, + d.sizeOf(Header), + verticesSize, + ); + const indices = std.bufferView(d.arrayOf(d.u32), layout.$.mesh, d.sizeOf(Header) + verticesSize); + + for (const index of indices) { + layout.$.extent = std.max(layout.$.extent, std.abs(vertices[index]!)); + } +}); +``` + +The indices follow the vertices, so `std.bufferArrayView` ends the vertices after `verticesSize` bytes. +The indices run to the end of the buffer, so `std.bufferView` covers them. + +## 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.buffer(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. +Clearing the sign bit through a `u32` view makes the `f32` at the same offset positive. + +```ts twoslash +import { tgpu, d, std } from 'typegpu'; + +const root = await tgpu.init(); +const packet = root.createMutable(d.buffer(1024)); +// ---cut--- +const absolute = tgpu.fn([d.u32], d.f32)((offset) => { + std.bufferView(d.u32, packet.$, offset).$ &= 0x7fffffff; + return std.bufferView(d.f32, packet.$, offset).$; +}); +``` + +## 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.buffer(1024)); + +const main = tgpu.computeFn({ + workgroupSize: [128], + in: { lid: d.builtin.localInvocationIndex }, +})(({ lid }) => { + const floats = std.bufferView(d.arrayOf(d.f32, 128), scratch.$, 0); + floats[lid] = d.f32(lid); + std.workgroupBarrier(); + const sum = floats[lid]! + floats[(lid + 1) % 128]!; + std.workgroupBarrier(); + + const pairs = std.bufferView(d.arrayOf(d.vec2f, 128), scratch.$, 0); + pairs[lid] = d.vec2f(sum, d.f32(lid)); +}); +``` + +## Passing buffers to functions + +Functions take buffers through pointers, such as `d.ptrStorage(d.buffer(), 'read-write')`. + +```ts twoslash +import { tgpu, d, std } from 'typegpu'; + +const root = await tgpu.init(); +const packet = root.createMutable(d.buffer(1024)); +// ---cut--- +const bump = tgpu.fn([d.ptrStorage(d.buffer(), '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/packed-shapes/index.html b/apps/typegpu-docs/src/examples/simple/packed-shapes/index.html new file mode 100644 index 0000000000..c72c5f1548 --- /dev/null +++ b/apps/typegpu-docs/src/examples/simple/packed-shapes/index.html @@ -0,0 +1,17 @@ + + + + diff --git a/apps/typegpu-docs/src/examples/simple/packed-shapes/index.ts b/apps/typegpu-docs/src/examples/simple/packed-shapes/index.ts new file mode 100644 index 0000000000..ba71b0808d --- /dev/null +++ b/apps/typegpu-docs/src/examples/simple/packed-shapes/index.ts @@ -0,0 +1,202 @@ +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 Motion = d.struct({ center: d.vec2f, orbit: d.f32, speed: d.f32, phase: d.f32 }); +const Circle = d.struct({ color: d.u32, radius: d.f32, motion: Motion }); +const Box = d.struct({ color: d.u32, halfSize: d.vec2f, motion: Motion }); +const Segment = d.struct({ color: d.u32, width: d.f32, halfLength: d.f32, motion: Motion }); +const Counts = d.struct({ circles: d.u32, boxes: d.u32, segments: d.u32 }); + +const Scene = (counts: d.Infer) => + d.struct({ + circles: d.arrayOf(Circle, counts.circles), + boxes: d.arrayOf(Box, counts.boxes), + segments: d.arrayOf(Segment, counts.segments), + }); + +const layout = tgpu.bindGroupLayout({ scene: { storage: d.buffer } }); +const counts = root.createUniform(Counts); +const time = root.createUniform(d.f32); +const smoothness = root.createUniform(d.f32, 0.12); + +const colors = [ + d.vec3f(1, 0.42, 0.42), + d.vec3f(1, 0.79, 0.34), + d.vec3f(0.28, 0.86, 0.98), + d.vec3f(0.11, 0.82, 0.63), + d.vec3f(0.64, 0.61, 1), + d.vec3f(1, 0.62, 0.95), +]; +const palette = tgpu.const(d.arrayOf(d.vec3f, colors.length), colors); +const background = tgpu.const(d.vec3f, d.vec3f(0.07, 0.07, 0.09)); + +const Surface = d.struct({ dist: d.f32, color: d.vec3f }); + +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 toLocal = (p: d.v2f, motion: d.Infer) => { + 'use gpu'; + const angle = time.$ * motion.speed + motion.phase; + const center = motion.center + d.vec2f(std.cos(angle), std.sin(angle)) * motion.orbit; + return rotate(p - center, angle); +}; + +const merge = (surface: d.Infer, dist: number, color: number) => { + 'use gpu'; + const k = std.max(smoothness.$, 0.0001); + const h = std.saturate(0.5 + (0.5 * (surface.dist - dist)) / k); + return Surface({ + dist: sdf.opSmoothUnion(surface.dist, dist, k), + color: std.mix(surface.color, palette.$[color], h), + }); +}; + +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); + + const circlesSize = counts.$.circles * d.sizeOf(Circle); + const boxesSize = counts.$.boxes * d.sizeOf(Box); + const circles = std.bufferArrayView(d.arrayOf(Circle), layout.$.scene, 0, circlesSize); + const boxes = std.bufferArrayView(d.arrayOf(Box), layout.$.scene, circlesSize, boxesSize); + const segments = std.bufferView(d.arrayOf(Segment), layout.$.scene, circlesSize + boxesSize); + + let surface = Surface({ dist: 1000, color: background.$ }); + for (const circle of circles) { + const dist = sdf.sdDisk(toLocal(p, circle.motion), circle.radius); + surface = merge(surface, dist, circle.color); + } + for (const box of boxes) { + const dist = sdf.sdRoundedBox2d(toLocal(p, box.motion), box.halfSize, 0.02); + surface = merge(surface, dist, box.color); + } + for (const segment of segments) { + const a = d.vec2f(-segment.halfLength, 0); + const b = d.vec2f(segment.halfLength, 0); + const dist = sdf.sdLine(toLocal(p, segment.motion), a, b) - segment.width; + surface = merge(surface, dist, segment.color); + } + + const fill = 1 - std.smoothstep(-pixel, pixel, surface.dist); + const shade = surface.color * (0.8 + 0.2 * std.saturate(-surface.dist * 12)); + const rings = background.$ * (0.85 + 0.15 * std.cos(surface.dist * 90)); + return d.vec4f(std.mix(rings, shade, fill), 1); +}); + +const pipeline = root.createRenderPipeline({ + vertex: common.fullScreenTriangle, + fragment, +}); + +const random = (min: number, max: number) => min + Math.random() * (max - min); +const randomColor = () => Math.floor(Math.random() * colors.length); +const randomMotion = () => ({ + center: d.vec2f(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), +}); + +const circle = () => ({ color: randomColor(), radius: random(0.06, 0.16), motion: randomMotion() }); + +const box = () => { + const size = random(0.06, 0.16); + return { color: randomColor(), halfSize: d.vec2f(size, size * 0.6), motion: randomMotion() }; +}; + +const segment = () => { + const size = random(0.06, 0.16); + return { + color: randomColor(), + width: size * 0.25, + halfLength: size * 1.5, + motion: randomMotion(), + }; +}; + +function createScene(shapeCount: number) { + const third = Math.floor(shapeCount / 3); + const sizes = { circles: shapeCount - 2 * third, boxes: third, segments: third }; + + const buffer = root + .createBuffer(Scene(sizes), { + circles: Array.from({ length: sizes.circles }, circle), + boxes: Array.from({ length: sizes.boxes }, box), + segments: Array.from({ length: sizes.segments }, segment), + }) + .$usage('storage'); + counts.write(sizes); + return { buffer, bindGroup: root.createBindGroup(layout, { scene: buffer }) }; +} + +let shapeCount = 24; +let scene = createScene(shapeCount); + +function regenerate() { + scene.buffer.destroy(); + scene = createScene(shapeCount); +} + +let frameId: number; + +function frame(timestamp: number) { + time.write(timestamp / 1000); + pipeline.with(scene.bindGroup).withColorAttachment({ view: context }).draw(3); + + frameId = requestAnimationFrame(frame); +} + +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; + regenerate(); + }, + }, + smoothness: supported && { + initial: 0.12, + min: 0, + max: 0.3, + step: 0.01, + onSliderChange: (value) => { + smoothness.write(value); + }, + }, + 'new scene': supported && { + onButtonClick: regenerate, + }, +}); + +export function onCleanup() { + cancelAnimationFrame(frameId); + root.destroy(); +} + +// #endregion diff --git a/apps/typegpu-docs/src/examples/simple/packed-shapes/meta.json b/apps/typegpu-docs/src/examples/simple/packed-shapes/meta.json new file mode 100644 index 0000000000..df3c3ec3e2 --- /dev/null +++ b/apps/typegpu-docs/src/examples/simple/packed-shapes/meta.json @@ -0,0 +1,7 @@ +{ + "title": "Packed Shapes", + "description": "Writes circles, boxes and segments into one typed buffer sized to the scene, and a fragment shader views each run of shapes 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/packed-shapes/thumbnail.png b/apps/typegpu-docs/src/examples/simple/packed-shapes/thumbnail.png new file mode 100644 index 0000000000..841efa6154 Binary files /dev/null and b/apps/typegpu-docs/src/examples/simple/packed-shapes/thumbnail.png differ diff --git a/apps/typegpu-docs/tests/individual-example-tests/dispatch.test.ts b/apps/typegpu-docs/tests/individual-example-tests/dispatch.test.ts index b661036a22..eeb5b180dd 100644 --- a/apps/typegpu-docs/tests/individual-example-tests/dispatch.test.ts +++ b/apps/typegpu-docs/tests/individual-example-tests/dispatch.test.ts @@ -112,11 +112,11 @@ describe('tgsl parsing test example', () => { @group(0) @binding(0) var sizeUniform: vec3u; - @group(1) @binding(0) var buffer: array; + @group(1) @binding(0) var buffer_1: array; fn wrappedCallback(_arg_0: u32, _arg_1: u32, _arg_2: u32) { - for (var i = 0u; (i < arrayLength((&buffer))); i++) { - buffer[i] *= 2u; + for (var i = 0u; (i < arrayLength((&buffer_1))); i++) { + buffer_1[i] *= 2u; } } 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..39e67035ef 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/buffer.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.buffer(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 aa6db6ef3b..2a6b8a3a8a 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 './buffer.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/buffer.ts b/packages/typegpu/src/data/buffer.ts new file mode 100644 index 0000000000..3f65540dd9 --- /dev/null +++ b/packages/typegpu/src/data/buffer.ts @@ -0,0 +1,101 @@ +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]: buffer; + 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.buffer(256), unless declared in a bind group layout with { storage: d.buffer }'; + // --- +} + +/** + * The contents of a `d.buffer`, as seen by shaders. + * Pass it to `std.bufferView`, `std.bufferArrayView` and `std.bufferLength`. + */ +export interface buffer { + readonly [$internal]: true; + readonly type: 'buffer'; +} + +/** + * Creates a schema for `byteSize` raw bytes, the WGSL `buffer` type, viewed in shaders + * with `std.bufferView` and `std.bufferArrayView`. Without a size it is runtime-sized, which only + * storage bindings allow, e.g. `{ storage: d.buffer }` in a bind group layout. + * + * Requires the `buffer_view` WGSL language extension, see `root.enabledWgslLanguageFeatures`. + * + * @example + * const packet = root.createBuffer(d.buffer(1024)).$usage('storage'); + */ +export function buffer(): WgslBuffer<0>; +export function buffer(byteSize: TSize): WgslBuffer; +export function buffer(byteSize = 0): WgslBuffer { + if (!Number.isInteger(byteSize) || byteSize < 0 || byteSize % 4 !== 0) { + throw new Error( + `Cannot create a buffer schema of ${byteSize} bytes, the size has to be a multiple of 4.`, + ); + } + 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]: buffer; + 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 ? 'buffer' : `buffer(${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..79f4195117 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 { buffer, type WgslBuffer } from './buffer.ts'; export { _ref as ref } from './ref.ts'; export { align, diff --git a/packages/typegpu/src/data/partialIO.ts b/packages/typegpu/src/data/partialIO.ts index df121a857b..52c74a8763 100644 --- a/packages/typegpu/src/data/partialIO.ts +++ b/packages/typegpu/src/data/partialIO.ts @@ -1,6 +1,7 @@ import { BufferWriter, getSystemEndianness } from 'typed-binary'; import { roundUp } from '../mathUtils.ts'; import { alignmentOf } from './alignmentOf.ts'; +import { isWgslBuffer } from './buffer.ts'; import { getCompiledWriter } from './compiledIO.ts'; import { writeData } from './dataIO.ts'; import { isDisarray, isUnstruct } from './dataTypes.ts'; @@ -82,6 +83,22 @@ export function getPatchInstructions( // Deferred until every range is validated const pendingWrites: (() => void)[] = []; + function collectRaw( + value: ArrayBuffer | ArrayBufferView, + offset: number, + maxSize: number, + padding?: number, + ) { + const bytes = ArrayBuffer.isView(value) + ? new Uint8Array(value.buffer, value.byteOffset, value.byteLength) + : new Uint8Array(value); + const copyLen = Math.min(bytes.byteLength, maxSize); + pendingWrites.push(() => { + new Uint8Array(buf, offset, copyLen).set(bytes.subarray(0, copyLen)); + }); + segments.push({ start: offset, end: offset + copyLen, padding }); + } + function collect(node: wgsl.BaseData, value: unknown, offset: number, padding?: number) { if (value === undefined || value === null) { return; @@ -108,13 +125,7 @@ export function getPatchInstructions( const elementPadding = elementSize - sizeOf(arrSchema.elementType); if (ArrayBuffer.isView(value)) { - const copyLen = Math.min(value.byteLength, arrSchema.elementCount * elementSize); - pendingWrites.push(() => { - new Uint8Array(buf, offset, copyLen).set( - new Uint8Array(value.buffer, value.byteOffset, copyLen), - ); - }); - segments.push({ start: offset, end: offset + copyLen, padding }); + collectRaw(value, offset, arrSchema.elementCount * elementSize, padding); return; } @@ -135,6 +146,11 @@ export function getPatchInstructions( return; } + if (isWgslBuffer(node)) { + collectRaw(value as ArrayBuffer | ArrayBufferView, offset, node.byteSize, padding); + return; + } + const leafSize = sizeOf(node); const compiledWriter = getCompiledWriter(node); pendingWrites.push(() => { diff --git a/packages/typegpu/src/data/ptr.ts b/packages/typegpu/src/data/ptr.ts index ed1e76eade..9844d9087f 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 './buffer.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..88160360dc 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 './buffer.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 37d358d879..1fc61b53af 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 './buffer.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 0ffee46a80..a7c9e5e727 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 './buffer.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..43c14a3738 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 './buffer.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/nameUtils.ts b/packages/typegpu/src/nameUtils.ts index 38f58f1d7e..c3fabec7ff 100644 --- a/packages/typegpu/src/nameUtils.ts +++ b/packages/typegpu/src/nameUtils.ts @@ -268,6 +268,11 @@ export const builtins = new Set([ 'select', // array 'arrayLength', + // buffer view + 'buffer', + 'bufferArrayView', + 'bufferLength', + 'bufferView', // numeric 'abs', 'acos', 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..7e289bb092 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/buffer.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.buffer(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..84185efc1d --- /dev/null +++ b/packages/typegpu/src/std/bufferView.ts @@ -0,0 +1,273 @@ +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 buffer, isWgslBuffer, type WgslBuffer } from '../data/buffer.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 +// ---------- + +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: buffer, + byteOffset: number, +) => ref>>; + +type BufferArrayViewFn = BaseData>( + schema: ValidateViewSchema, + buffer: buffer, + byteOffset: number, + byteSize: number, +) => ref>>; + +/** + * Reads the bytes of a `d.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.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.buffer` in bytes, mirroring WGSL's `bufferLength`. + */ +export const bufferLength = dualImpl<(buffer: buffer) => 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.buffer, got '${String(inner)}'. Declare the buffer with d.buffer(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/tgpuBindGroupLayout.ts b/packages/typegpu/src/tgpuBindGroupLayout.ts index e0946898b3..0fc7f62516 100644 --- a/packages/typegpu/src/tgpuBindGroupLayout.ts +++ b/packages/typegpu/src/tgpuBindGroupLayout.ts @@ -38,6 +38,7 @@ import { NotStorageError, type StorageFlag } from './extension.ts'; import type { TgpuNamable } from './shared/meta.ts'; import { getName, setName } from './shared/meta.ts'; import type { InferGPU, MemIdentity } from './shared/repr.ts'; +import type { WgslBuffer } from './data/buffer.ts'; import { safeStringify } from './shared/stringify.ts'; import type { TgpuDeviceOwningSoul, TgpuSoul } from './shared/soul.ts'; import { $gpuValueOf, $internal, $soul } from './shared/symbols.ts'; @@ -182,6 +183,12 @@ export interface TgpuBindGroupLayout< unwrap(unwrapper: Unwrapper): GPUBindGroupLayout; } +/** + * The schema of buffers an entry takes. Entries of `d.buffer` take buffers of any schema, + * since shaders view their bytes as any type. + */ +type BoundSchema = T extends WgslBuffer ? BaseData : MemIdentity; + export type LayoutEntryToInput = // Widest type TgpuLayoutEntry | null extends T @@ -198,15 +205,15 @@ export type LayoutEntryToInput = : // Strict type-checking T extends TgpuLayoutUniform ? - | (TgpuBuffer>> & UniformFlag) - | TgpuUniform>> + | (TgpuBuffer>> & UniformFlag) + | TgpuUniform>> | GPUBuffer : T extends TgpuLayoutStorage ? - | (TgpuBuffer>> & StorageFlag) + | (TgpuBuffer>> & StorageFlag) | (T extends { access: 'mutable' } - ? TgpuMutable>> - : TgpuReadonly>>) + ? TgpuMutable>> + : TgpuReadonly>>) | GPUBuffer : T extends TgpuLayoutSampler ? TgpuSampler | GPUSampler diff --git a/packages/typegpu/src/tgsl/conversion.ts b/packages/typegpu/src/tgsl/conversion.ts index 1a84ac3d7d..03a97e9d64 100644 --- a/packages/typegpu/src/tgsl/conversion.ts +++ b/packages/typegpu/src/tgsl/conversion.ts @@ -1,3 +1,4 @@ +import { isWgslBuffer, type WgslBuffer } from '../data/buffer.ts'; import { UnknownData } from '../data/dataTypes.ts'; import { undecorate } from '../data/dataTypes.ts'; import { derefSnippet, RefOperator } from '../data/ref.ts'; @@ -42,6 +43,17 @@ function getAutoConversionRank(src: BaseData, dest: BaseData): ConversionRankInf const trueSrc = undecorate(src); const trueDst = undecorate(dest); + if ( + isPtr(trueSrc) && + isPtr(trueDst) && + isWgslBuffer(trueSrc.inner) && + isWgslBuffer(trueDst.inner) + ) { + return fitsBufferParameter(trueSrc.inner, trueDst.inner) + ? { rank: 0, action: 'none' } + : INFINITE_RANK; + } + if (trueSrc.type === trueDst.type) { if (trueSrc.type === 'struct' && trueSrc !== trueDst) { return { rank: 1, action: 'cast', targetType: dest }; @@ -84,6 +96,10 @@ function getAutoConversionRank(src: BaseData, dest: BaseData): ConversionRankInf return INFINITE_RANK; } +function fitsBufferParameter(src: WgslBuffer, dest: WgslBuffer): boolean { + return dest.byteSize === 0 || (src.byteSize !== 0 && dest.byteSize <= src.byteSize); +} + function getImplicitConversionRank(src: BaseData, dest: BaseData): ConversionRankInfo { const trueSrc = undecorate(src); const trueDst = undecorate(dest); 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..8fe1914e8d 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/buffer.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.buffer(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..2bbc285eb6 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,9 @@ export interface ResolutionCtx { */ registerImmediate(immediate: TgpuImmediateVar): void; + /** Adds a `requires` directive, or throws when the root 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 { const mutable = buffer.as('mutable'); expect(tgpu.resolve([mutable])).toMatchInlineSnapshot( - `"@group(0) @binding(0) var buffer: f32;"`, + `"@group(0) @binding(0) var buffer_1: f32;"`, ); }); @@ -277,7 +277,7 @@ describe('root.createReadonly', () => { const readonly = buffer.as('readonly'); expect(tgpu.resolve([readonly])).toMatchInlineSnapshot( - `"@group(0) @binding(0) var buffer: f32;"`, + `"@group(0) @binding(0) var buffer_1: f32;"`, ); }); @@ -547,7 +547,7 @@ describe('root.createUniform', () => { const uniform = buffer.as('uniform'); expect(tgpu.resolve([uniform])).toMatchInlineSnapshot( - `"@group(0) @binding(0) var buffer: f32;"`, + `"@group(0) @binding(0) var buffer_1: f32;"`, ); }); diff --git a/packages/typegpu/tests/serializeDataSchema.test.ts b/packages/typegpu/tests/serializeDataSchema.test.ts index 5aef259cb0..c7859a50d5 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.buffer(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..ea27803ad6 --- /dev/null +++ b/packages/typegpu/tests/std/bufferView.test.ts @@ -0,0 +1,333 @@ +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.buffer }, + 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.buffer(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.buffer(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.buffer(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.buffer(512)); + const counters = root.createMutable(d.buffer(64)); + + const bump = tgpu.fn([d.ptrStorage(d.buffer(), '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_1: ptr, index: u32) { + (*bufferView(&(*buffer_1), (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('passes buffers only to parameters of the same or a smaller size', ({ root }) => { + const packet = root.createMutable(d.buffer(64)); + const smaller = tgpu.fn([d.ptrStorage(d.buffer(32), 'read-write')])((buffer) => { + std.bufferView(d.u32, buffer, 0).$ = 1; + }); + const larger = tgpu.fn([d.ptrStorage(d.buffer(128), 'read-write')])((buffer) => { + std.bufferView(d.u32, buffer, 0).$ = 1; + }); + + expect(() => tgpu.resolve([tgpu.fn([])(() => smaller(d.ref(packet.$)))])).not.toThrow(); + expect(() => tgpu.resolve([tgpu.fn([])(() => larger(d.ref(packet.$)))])) + .toThrowErrorMatchingInlineSnapshot(` + [Error: Resolution of the following tree failed: + - + - fn:: Cannot convert 'd.ref(packet.$)' of type 'ptr' to any of the target types: [ptr]] + `); + }); + + it('rejects views that the shader cannot hold', ({ root }) => { + const packet = root.createMutable(d.buffer(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.buffer schemas 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.buffer, got 'u32'. Declare the buffer with d.buffer(n).] + `); + }); + + it('rejects misaligned and out-of-bounds constant offsets', ({ root }) => { + const packet = root.createMutable(d.buffer(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.buffer(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.buffer(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.buffer(16)); + const page = root.createUniform(d.buffer(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)); diff --git a/packages/typegpu/tests/wgslBuffer.test.ts b/packages/typegpu/tests/wgslBuffer.test.ts new file mode 100644 index 0000000000..bf7057f47c --- /dev/null +++ b/packages/typegpu/tests/wgslBuffer.test.ts @@ -0,0 +1,175 @@ +import { describe, expect, expectTypeOf, vi } from 'vitest'; +import { d, std, tgpu } from 'typegpu'; +import { it } from 'typegpu-testing-utility'; + +describe('d.buffer', () => { + it('resolves to the WGSL buffer type and requires buffer_view', ({ root }) => { + const packet = root.createMutable(d.buffer(256)); + const page = root.createUniform(d.buffer(64)); + const scratch = tgpu.workgroupVar(d.buffer(1024)); + const layout = tgpu.bindGroupLayout({ mesh: { storage: d.buffer } }); + + const fill = tgpu.fn([d.ptrStorage(d.buffer(), '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.buffer(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 + - buffer(1024): d.buffer(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.buffer(16)); + const layout = tgpu.bindGroupLayout({ page: { uniform: d.buffer() } }); + 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.buffer(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.buffer needs a size in the 'uniform' address space, only storage allows runtime-sized buffers. Give it a size, e.g. d.buffer(256).] + `); + expect(() => d.struct({ data: d.buffer(16) })).toThrowErrorMatchingInlineSnapshot( + `[Error: Cannot use d.buffer(16) as struct property 'data', buffers can only be whole variables.]`, + ); + expect(() => d.arrayOf(d.buffer(16), 4)).toThrowErrorMatchingInlineSnapshot( + `[Error: Cannot create an array of d.buffer(16), buffers can only be whole variables.]`, + ); + expect(() => d.buffer(6)).toThrowErrorMatchingInlineSnapshot( + `[Error: Cannot create a buffer schema of 6 bytes, the size has to be a multiple of 4.]`, + ); + }); + + it('needs a size outside storage bind group layouts', ({ root }) => { + const layout = tgpu.bindGroupLayout({ mesh: { storage: d.buffer } }); + root.createBindGroup(layout, { mesh: root.createBuffer(d.buffer(512)).$usage('storage') }); + + if (false) { + // @ts-expect-error + root.createBuffer(d.buffer()); + // @ts-expect-error + root.createUniform(d.buffer()); + // @ts-expect-error + tgpu.workgroupVar(d.buffer()); + } + }); + + it('takes buffers of any schema in bind group layouts', ({ root }) => { + const layout = tgpu.bindGroupLayout({ + mesh: { storage: d.buffer }, + counts: { storage: d.buffer, access: 'mutable' }, + params: { uniform: d.buffer(16) }, + }); + + root.createBindGroup(layout, { + mesh: root.createBuffer(d.arrayOf(d.vec2f, 64)).$usage('storage'), + counts: root.createMutable(d.arrayOf(d.u32, 4)), + params: root.createUniform(d.vec4f), + }); + }); + + it('writes and reads raw bytes', async ({ root, device }) => { + const packet = root.createBuffer(d.buffer(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.buffer(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]); + + page.patch(new Uint32Array([5])); + const [, patchOffset, patch] = vi.mocked(device.queue.writeBuffer).mock.calls[1]!; + expect([patchOffset, (patch as ArrayBufferView).byteLength]).toStrictEqual([0, 4]); + + const read = await page.read(); + expectTypeOf(read).toEqualTypeOf(); + expect(read.byteLength).toBe(16); + }); + + it('does not shadow the buffer view built-ins', ({ root }) => { + const buffer = root.createMutable(d.buffer(64)); + const main = tgpu.computeFn({ workgroupSize: [1] })(() => { + const bufferView = std.bufferView(d.u32, buffer.$, 0); + bufferView.$ = std.bufferLength(buffer.$); + }); + + expect(tgpu.resolve([main])).toMatchInlineSnapshot(` + "requires buffer_view; + + @group(0) @binding(0) var buffer_1: buffer<64>; + + @compute @workgroup_size(1) fn main() { + let bufferView_1 = bufferView(&buffer_1, 0u); + (*bufferView_1) = bufferLength(&buffer_1); + }" + `); + }); + + it('measures and compares as raw bytes', () => { + expect(d.sizeOf(d.buffer(64))).toBe(64); + expect(d.sizeOf(d.buffer())).toBeNaN(); + expect(d.alignmentOf(d.buffer(64))).toBe(4); + expect(d.schemasEqualStructurally(d.buffer(64), d.buffer(64))).toBe(true); + expect(d.schemasEqualStructurally(d.buffer(64), d.buffer(32))).toBe(false); + }); +});