typegpu
Version:
A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.
365 lines (364 loc) • 14.8 kB
JavaScript
import { builtin } from "../../builtin.js";
import { INTERNAL_createQuerySet } from "../querySet/querySet.js";
import { WeakMemo } from "../../memo.js";
import { clearTextureUtilsCache } from "../texture/textureUtils.js";
import { $getNameForward, $internal, $soul } from "../../shared/symbols.js";
import { isBindGroup, isBindGroupLayout, TgpuBindGroupImpl } from "../../tgpuBindGroupLayout.js";
import { INTERNAL_createBuffer } from "../buffer/buffer.js";
import { isBufferBinding, TgpuBufferBindingImpl, } from "../buffer/bufferBinding.js";
import { computeFn } from "../function/tgpuComputeFn.js";
import { fn } from "../function/tgpuFn.js";
import { INTERNAL_createComputePipeline, } from "../pipeline/computePipeline.js";
import { INTERNAL_createRenderPipeline, } from "../pipeline/renderPipeline.js";
import { isTgpuCommandEncoder, isTgpuComputePass, isTgpuRenderCommands, } from "../pipeline/typeGuards.js";
import { INTERNAL_createCommandEncoder, } from "../commandEncoder/commandEncoder.js";
import { INTERNAL_createRenderBundleEncoder, } from "../commandEncoder/renderPass.js";
import { INTERNAL_createComparisonSampler, INTERNAL_createSampler, isComparisonSampler, isSampler, } from "../sampler/sampler.js";
import { isAccessor, isMutableAccessor, } from "../slot/slotTypes.js";
import { INTERNAL_createTexture, isTextureView, } from "../texture/texture.js";
import { isVertexLayout } from "../vertexLayout/vertexLayout.js";
import { ConfigurableImpl } from "./configurableImpl.js";
import { vec3f, vec3u } from "../../data/vector.js";
import { u32 } from "../../data/numeric.js";
import { ceil } from "../../std/numeric.js";
import { allEq } from "../../std/boolean.js";
import { getName, setName } from "../../shared/meta.js";
import { logger } from "../../tgpuLogger.js";
import { safeStringify } from "../../shared/stringify.js";
/**
* Changes the given array to a vec of 3 numbers, filling missing values with 1.
*/
function toVec3(arr) {
if (arr.includes(0)) {
throw new Error('Size and workgroupSize cannot contain zeroes.');
}
return vec3u(arr[0] ?? 1, arr[1] ?? 1, arr[2] ?? 1);
}
const workgroupSizeConfigs = [
vec3u(1, 1, 1),
vec3u(256, 1, 1),
vec3u(16, 16, 1),
vec3u(8, 8, 4),
];
export class TgpuGuardedComputePipelineImpl {
resourceType = 'guarded-compute-pipeline';
[$soul];
#root;
#lastSize;
constructor(root, pipeline, sizeUniform, workgroupSize) {
this.#root = root;
this.#lastSize = vec3u();
this[$soul] = {
type: 'guarded-compute-pipeline',
device: root.device,
pipeline,
sizeUniform,
workgroupSize,
label: undefined,
};
}
with(bindGroup) {
if (!isBindGroup(bindGroup)) {
throw new Error('Guarded pipelines only accept bind groups in .with(). To record into passes or encoders, use a regular compute pipeline.');
}
return new TgpuGuardedComputePipelineImpl(this.#root, this[$soul].pipeline.with(bindGroup), this[$soul].sizeUniform, this[$soul].workgroupSize);
}
withPerformanceCallback(callback) {
return new TgpuGuardedComputePipelineImpl(this.#root, this[$soul].pipeline.withPerformanceCallback(callback), this[$soul].sizeUniform, this[$soul].workgroupSize);
}
withTimestampWrites(options) {
return new TgpuGuardedComputePipelineImpl(this.#root, this[$soul].pipeline.withTimestampWrites(options), this[$soul].sizeUniform, this[$soul].workgroupSize);
}
dispatchThreads(...threads) {
const sanitizedSize = toVec3(threads);
const workgroupCount = ceil(vec3f(sanitizedSize).div(vec3f(this[$soul].workgroupSize)));
if (!allEq(sanitizedSize, this.#lastSize)) {
// Only updating the size if it has changed from the last
// invocation. This removes the need for flushing.
this.#lastSize = sanitizedSize;
this[$soul].sizeUniform.write(sanitizedSize);
}
this[$soul].pipeline.dispatchWorkgroups(workgroupCount.x, workgroupCount.y, workgroupCount.z);
}
initAsync() {
return this[$soul].pipeline.initAsync();
}
initSync() {
this[$soul].pipeline.initSync();
}
get pipeline() {
return this[$soul].pipeline;
}
get sizeUniform() {
return this[$soul].sizeUniform;
}
[$internal] = true;
get [$getNameForward]() {
return this[$soul].pipeline;
}
$name(label) {
setName(this, label);
return this;
}
}
export function INTERNAL_restoreRoot(soul, ctx) {
return ctx.getRoot(soul.device);
}
export function INTERNAL_restoreGuardedComputePipeline(soul, ctx) {
return new TgpuGuardedComputePipelineImpl(ctx.getRoot(soul.device), soul.pipeline, soul.sizeUniform, soul.workgroupSize);
}
class WithBindingImpl {
#getRoot;
#slotBindings;
constructor(getRoot, slotBindings) {
this.#getRoot = getRoot;
this.#slotBindings = slotBindings;
}
with(slot, value) {
return new WithBindingImpl(this.#getRoot, [
...this.#slotBindings,
[isAccessor(slot) || isMutableAccessor(slot) ? slot.slot : slot, value],
]);
}
createComputePipeline(descriptor) {
return INTERNAL_createComputePipeline(this.#getRoot(), this.#slotBindings, descriptor);
}
createRenderPipeline(descriptor) {
return INTERNAL_createRenderPipeline({
root: this.#getRoot(),
slotBindings: this.#slotBindings,
descriptor,
});
}
createGuardedComputePipeline(callback) {
const root = this.#getRoot();
if (callback.length >= 4) {
throw new Error('Guarded compute callback only supports up to three dimensions.');
}
const workgroupSize = workgroupSizeConfigs[callback.length];
const wrappedCallback = fn([u32, u32, u32])(callback);
if (getName(wrappedCallback) === undefined) {
wrappedCallback.$name('wrappedCallback');
}
const sizeUniform = root.createUniform(vec3u);
// WGSL instead of JS because we do not run unplugin
// before shipping the typegpu package
const mainCompute = computeFn({
workgroupSize,
in: { id: builtin.globalInvocationId },
}) `{
if (any(in.id >= sizeUniform)) {
return;
}
wrappedCallback(in.id.x, in.id.y, in.id.z);
}`.$uses({ sizeUniform, wrappedCallback });
// NOTE: in certain setups, unplugin can run on package typegpu, so we have to avoid auto-naming triggering here
const pipeline = (() => this.createComputePipeline({ compute: mainCompute }))();
return new TgpuGuardedComputePipelineImpl(root, pipeline, sizeUniform, workgroupSize);
}
pipe(transform) {
const newCfg = transform(new ConfigurableImpl([]));
return new WithBindingImpl(this.#getRoot, [...this.#slotBindings, ...newCfg.bindings]);
}
}
/**
* Holds all data that is necessary to facilitate CPU and GPU communication.
* Programs that share a root can interact via GPU buffers.
*/
class TgpuRootImpl extends WithBindingImpl {
'~unstable';
resourceType = 'root';
[$soul];
device;
nameRegistrySetting;
shaderGeneratorClass;
#unwrappedBindGroupLayouts = new WeakMemo((key) => key.unwrap(this));
#unwrappedBindGroups = new WeakMemo((key) => key.unwrap(this));
#ownDevice;
[$internal];
constructor(device, nameRegistrySetting, minify, ownDevice, logOptions, shaderGeneratorClass) {
super(() => this, []);
this.device = device;
this.nameRegistrySetting = nameRegistrySetting;
this.#ownDevice = ownDevice;
this.shaderGeneratorClass = shaderGeneratorClass;
this['~unstable'] = this;
this[$soul] = {
type: 'root',
device,
nameRegistrySetting,
logOptions,
minify,
nonTransferablePriors: shaderGeneratorClass ? ['shaderGeneratorClass'] : undefined,
label: undefined,
};
this[$internal] = {
logOptions,
};
}
configureContext(options) {
const context = options.canvas.getContext('webgpu');
if (!context) {
throw new Error("Unable to initialize 'webgpu' context on the provided canvas.");
}
context.configure({
...options,
device: this.device,
format: options.format ?? navigator.gpu.getPreferredCanvasFormat(),
});
return context;
}
get enabledFeatures() {
return new Set(this.device.features);
}
get minify() {
return this[$soul].minify;
}
createBuffer(typeSchema, initialOrBuffer) {
return INTERNAL_createBuffer(this, typeSchema, initialOrBuffer);
}
createUniform(typeSchema, initialOrBuffer) {
const buffer = INTERNAL_createBuffer(this, typeSchema, initialOrBuffer)
// oxlint-disable-next-line typescript/no-explicit-any -- i'm sure it's fine
.$usage('uniform');
return new TgpuBufferBindingImpl('uniform', buffer);
}
createMutable(typeSchema, initialOrBuffer) {
const buffer = INTERNAL_createBuffer(this, typeSchema, initialOrBuffer)
// oxlint-disable-next-line typescript/no-explicit-any -- i'm sure it's fine
.$usage('storage');
return new TgpuBufferBindingImpl('mutable', buffer);
}
createReadonly(typeSchema, initialOrBuffer) {
const buffer = INTERNAL_createBuffer(this, typeSchema, initialOrBuffer)
// oxlint-disable-next-line typescript/no-explicit-any -- i'm sure it's fine
.$usage('storage');
return new TgpuBufferBindingImpl('readonly', buffer);
}
createQuerySet(type, count, rawQuerySet) {
return INTERNAL_createQuerySet(this, type, count, rawQuerySet);
}
createBindGroup(layout, entries) {
return new TgpuBindGroupImpl(this, layout, entries);
}
destroy() {
clearTextureUtilsCache(this.device);
if (this.#ownDevice) {
this.device.destroy();
}
}
createTexture(props) {
const texture = INTERNAL_createTexture(props, this);
// oxlint-disable-next-line typescript/no-explicit-any -- too much type wrangling
return texture;
}
createSampler(props) {
return INTERNAL_createSampler(props, this);
}
createComparisonSampler(props) {
return INTERNAL_createComparisonSampler(props, this);
}
unwrap(resource) {
if (isTgpuCommandEncoder(resource)) {
return resource[$internal].rawEncoder;
}
if (isTgpuRenderCommands(resource) || isTgpuComputePass(resource)) {
resource[$internal].state.rawAccessed = true;
return resource[$internal].rawPass;
}
if (isBindGroupLayout(resource)) {
return this.#unwrappedBindGroupLayouts.getOrMake(resource);
}
if (isBindGroup(resource)) {
return this.#unwrappedBindGroups.getOrMake(resource);
}
if (isBufferBinding(resource)) {
return resource.buffer.buffer;
}
if (isTextureView(resource)) {
if (!resource[$internal].unwrap) {
throw new Error('Cannot unwrap laid-out texture view as it has no underlying resource.');
}
return resource[$internal].unwrap();
}
if (isVertexLayout(resource)) {
return resource.vertexLayout;
}
const internals = resource[$internal];
if (internals?.materialize) {
return internals.materialize();
}
if (isSampler(resource) || isComparisonSampler(resource)) {
throw new Error('Cannot unwrap laid-out sampler.');
}
throw new Error(`Unknown resource type: ${safeStringify(resource)}`);
}
createCommandEncoder(descriptor) {
return INTERNAL_createCommandEncoder(this, descriptor);
}
createRenderBundleEncoder(descriptor) {
return INTERNAL_createRenderBundleEncoder(this, descriptor);
}
}
/**
* Requests a new GPU device and creates a root around it.
* If a specific device should be used instead, use @see initFromDevice.
*
* @example
* When given no options, the function will ask the browser for a suitable GPU device.
* ```ts
* const root = await tgpu.init();
* ```
*
* @example
* If there are specific options that should be used when requesting a device, you can pass those in.
* ```ts
* const adapterOptions: GPURequestAdapterOptions = ...;
* const deviceDescriptor: GPUDeviceDescriptor = ...;
* const root = await tgpu.init({ adapter: adapterOptions, device: deviceDescriptor });
* ```
*/
export async function init(options) {
const { adapter: adapterOpt, device: deviceOpt, unstable_names: names = 'strict', unstable_minify: minify = false, unstable_logOptions: logOptions, unstable_shaderGeneratorClass: shaderGeneratorClass, } = options ?? {};
const { optionalFeatures, ...deviceDescriptor } = deviceOpt ?? {};
if (!navigator.gpu) {
throw new Error('WebGPU is not supported by this browser.');
}
const adapter = await navigator.gpu.requestAdapter(adapterOpt);
if (!adapter) {
throw new Error('Could not find a compatible GPU');
}
const availableFeatures = [];
for (const feature of deviceDescriptor.requiredFeatures ?? []) {
if (!adapter.features.has(feature)) {
throw new Error(`Requested feature "${feature}" is not supported by the adapter.`);
}
availableFeatures.push(feature);
}
for (const feature of optionalFeatures ?? []) {
if (adapter.features.has(feature)) {
availableFeatures.push(feature);
}
else {
logger.warn('webgpu-feature-missing', `Optional feature "${feature}" is not supported by the adapter.`);
}
}
const device = await adapter.requestDevice({
...deviceDescriptor,
requiredFeatures: availableFeatures,
});
return new TgpuRootImpl(device, names, minify, true, logOptions ?? {}, shaderGeneratorClass);
}
/**
* Creates a root from the given device, instead of requesting it like @see init.
*
* @example
* ```ts
* const device: GPUDevice = ...;
* const root = tgpu.initFromDevice({ device });
* ```
*/
export function initFromDevice(options) {
const { device, unstable_names: names = 'strict', unstable_minify: minify = false, unstable_logOptions: logOptions, unstable_shaderGeneratorClass: shaderGeneratorClass, } = options ?? {};
return new TgpuRootImpl(device, names, minify, false, logOptions ?? {}, shaderGeneratorClass);
}