Skip to content

feat: array.length in TGSL #1097

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Mar 31, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -30,11 +30,10 @@ const TriangleData = d.struct({
});

const renderBindGroupLayout = tgpu.bindGroupLayout({
trianglePos: { uniform: d.arrayOf(TriangleData, triangleAmount) },
colorPalette: { uniform: d.vec3f },
});

const { trianglePos, colorPalette } = renderBindGroupLayout.bound;
const { colorPalette } = renderBindGroupLayout.bound;

const VertexOutput = {
position: d.builtin.position,
Expand All @@ -61,7 +60,6 @@ const mainVert = tgpu['~unstable']
return VertexOutput(pos, color);
}`)
.$uses({
trianglePos,
colorPalette,
getRotationFromVelocity,
rotate,
Expand Down Expand Up @@ -223,7 +221,7 @@ const mainCompute = tgpu['~unstable']
let alignmentCount = 0;
let cohesionCount = 0;

for (let i = d.u32(0); i < std.arrayLength(currentTrianglePos.value); i++) {
for (let i = d.u32(0); i < currentTrianglePos.value.length; i++) {
if (i === index) {
continue;
}
Expand Down Expand Up @@ -295,7 +293,6 @@ const computePipeline = root['~unstable']

const renderBindGroups = [0, 1].map((idx) =>
root.createBindGroup(renderBindGroupLayout, {
trianglePos: trianglePosBuffers[idx],
colorPalette: colorPaletteBuffer,
}),
);
Expand Down
25 changes: 22 additions & 3 deletions packages/typegpu/src/smol/wgslGenerator.ts
Original file line number Diff line number Diff line change
@@ -1,19 +1,20 @@
import type * as smol from 'tinyest';
import * as d from '../data';
import * as wgsl from '../data/wgslTypes';
import { abstractInt } from '../data/numeric.js';
import * as wgsl from '../data/wgslTypes.js';
import {
type ResolutionCtx,
type Resource,
UnknownData,
type Wgsl,
isWgsl,
} from '../types';
} from '../types.js';
import {
getTypeForIndexAccess,
getTypeForPropAccess,
getTypeFromWgsl,
numericLiteralToResource,
} from './generationHelpers';
} from './generationHelpers.js';

const parenthesizedOps = [
'==',
Expand Down Expand Up @@ -192,6 +193,24 @@ export function generateExpression(
dataType: getTypeForPropAccess(target.dataType as Wgsl, property),
};
}

if (wgsl.isWgslArray(target.dataType)) {
if (property === 'length') {
if (target.dataType.elementCount === 0) {
// Dynamically-sized array
return {
value: `arrayLength(&${ctx.resolve(target.value)})`,
dataType: d.u32,
};
}

return {
value: String(target.dataType.elementCount),
dataType: abstractInt,
};
}
}

// biome-ignore lint/suspicious/noExplicitAny: <sorry TypeScript>
const propValue = (target.value as any)[property];

Expand Down
116 changes: 91 additions & 25 deletions packages/typegpu/tests/array.test.ts
Original file line number Diff line number Diff line change
@@ -1,72 +1,74 @@
import { BufferReader, BufferWriter } from 'typed-binary';
import { describe, expect, it } from 'vitest';
import { arrayOf, sizeOf, vec3f, vec3u } from '../src/data';
import tgpu from '../src';
import * as d from '../src/data';
import { readData, writeData } from '../src/data/dataIO';
import { StrictNameRegistry } from '../src/nameRegistry';
import { resolve } from '../src/resolutionCtx';
import type { Infer } from '../src/shared/repr';
import { parse, parseResolved } from './utils/parseResolved';

describe('array', () => {
it('takes element alignment into account when measuring', () => {
const TestArray = arrayOf(vec3u, 3);
expect(sizeOf(TestArray)).toEqual(48);
const TestArray = d.arrayOf(d.vec3u, 3);
expect(d.sizeOf(TestArray)).toEqual(48);
});

it('aligns array elements when writing', () => {
const TestArray = arrayOf(vec3u, 3);
const buffer = new ArrayBuffer(sizeOf(TestArray));
const TestArray = d.arrayOf(d.vec3u, 3);
const buffer = new ArrayBuffer(d.sizeOf(TestArray));
const writer = new BufferWriter(buffer);

writeData(writer, TestArray, [
vec3u(1, 2, 3),
vec3u(4, 5, 6),
vec3u(7, 8, 9),
d.vec3u(1, 2, 3),
d.vec3u(4, 5, 6),
d.vec3u(7, 8, 9),
]);
expect([...new Uint32Array(buffer)]).toEqual([
1, 2, 3, 0, 4, 5, 6, 0, 7, 8, 9, 0,
]);
});

it('aligns array elements when reading', () => {
const TestArray = arrayOf(vec3u, 3);
const buffer = new ArrayBuffer(sizeOf(TestArray));
const TestArray = d.arrayOf(d.vec3u, 3);
const buffer = new ArrayBuffer(d.sizeOf(TestArray));
const reader = new BufferReader(buffer);

new Uint32Array(buffer).set([1, 2, 3, 0, 4, 5, 6, 0, 7, 8, 9, 0]);

expect(readData(reader, TestArray)).toEqual([
vec3u(1, 2, 3),
vec3u(4, 5, 6),
vec3u(7, 8, 9),
d.vec3u(1, 2, 3),
d.vec3u(4, 5, 6),
d.vec3u(7, 8, 9),
]);
});

it('encodes and decodes arrays properly', () => {
const TestArray = arrayOf(vec3f, 5);
const TestArray = d.arrayOf(d.vec3f, 5);

const buffer = new ArrayBuffer(sizeOf(TestArray));
const buffer = new ArrayBuffer(d.sizeOf(TestArray));

const value: Infer<typeof TestArray> = [
vec3f(1.5, 2, 3.5),
vec3f(),
vec3f(-1.5, 2, 3.5),
vec3f(1.5, -2, 3.5),
vec3f(1.5, 2, 15),
d.vec3f(1.5, 2, 3.5),
d.vec3f(),
d.vec3f(-1.5, 2, 3.5),
d.vec3f(1.5, -2, 3.5),
d.vec3f(1.5, 2, 15),
];

writeData(new BufferWriter(buffer), TestArray, value);
expect(readData(new BufferReader(buffer), TestArray)).toEqual(value);
});

it('throws when trying to read/write a runtime-sized array', () => {
const TestArray = arrayOf(vec3f, 0);
const TestArray = d.arrayOf(d.vec3f, 0);

expect(sizeOf(TestArray)).toBeNaN();
expect(d.sizeOf(TestArray)).toBeNaN();

expect(() =>
writeData(new BufferWriter(new ArrayBuffer(0)), TestArray, [
vec3f(),
vec3f(),
d.vec3f(),
d.vec3f(),
]),
).toThrow();

Expand All @@ -80,6 +82,70 @@ describe('array', () => {
});

it('throws when trying to nest runtime sized arrays', () => {
expect(() => arrayOf(arrayOf(vec3f, 0), 0)).toThrow();
expect(() => d.arrayOf(d.arrayOf(d.vec3f, 0), 0)).toThrow();
});
});

describe('array.length', () => {
it('works for dynamically-sized arrays in TGSL', () => {
const layout = tgpu.bindGroupLayout({
values: {
storage: (n: number) => d.arrayOf(d.f32, n),
access: 'mutable',
},
});

const foo = tgpu['~unstable'].fn([]).does(() => {
let acc = d.f32(1);
for (let i = 0; i < layout.bound.values.value.length; i++) {
layout.bound.values.value[i] = acc;
acc *= 2;
}
});

expect(parseResolved({ foo })).toEqual(
parse(/* wgsl */ `
@group(0) @binding(0) var <storage, read_write> values: array<f32>;

fn foo() {
var acc = f32(1);
for (var i = 0; (i < arrayLength(&values)); i++) {
values[i] = acc;
acc *= 2;
}
}
`),
);
});

it('works for statically-sized arrays in TGSL', () => {
const layout = tgpu.bindGroupLayout({
values: {
storage: d.arrayOf(d.f32, 128),
access: 'mutable',
},
});

const foo = tgpu['~unstable'].fn([]).does(() => {
let acc = d.f32(1);
for (let i = 0; i < layout.bound.values.value.length; i++) {
layout.bound.values.value[i] = acc;
acc *= 2;
}
});

expect(parseResolved({ foo })).toEqual(
parse(/* wgsl */ `
@group(0) @binding(0) var <storage, read_write> values: array<f32, 128>;

fn foo() {
var acc = f32(1);
for (var i = 0; (i < 128); i++) {
values[i] = acc;
acc *= 2;
}
}
`),
);
});
});