typegpu
Version:
A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.
499 lines (498 loc) • 22.2 kB
JavaScript
import { isBuiltin } from "../../data/attributes.js";
import { getCustomLocation } from "../../data/dataTypes.js";
import { sizeOf } from "../../data/sizeOf.js";
import { snip } from "../../data/snippet.js";
import { formatToWGSLType } from "../../data/vertexFormatData.js";
import { isWgslData, Void, } from "../../data/wgslTypes.js";
import { invariant } from "../../errors.js";
import { resolve } from "../../resolutionCtx.js";
import { getName, PERF, setName } from "../../shared/meta.js";
import { $getNameForward, $internal, $resolve, $soul } from "../../shared/symbols.js";
import { isBindGroup, isBindGroupLayout, } from "../../tgpuBindGroupLayout.js";
import { isGPUBuffer } from "../../types.js";
import { wgslEnableExtensions, wgslEnableExtensionToFeatureName } from "../../wgslExtensions.js";
import { AutoFragmentFn, AutoVertexFn, } from "../function/autoIO.js";
import { namespace } from "../resolve/namespace.js";
import { connectAttributesToShader } from "../vertexLayout/connectAttributesToShader.js";
import { isVertexLayout } from "../vertexLayout/vertexLayout.js";
import { connectAttachmentToShader } from "./connectAttachmentToShader.js";
import { connectTargetsToShader } from "./connectTargetsToShader.js";
import { INTERNAL_adoptCommandEncoder, INTERNAL_createCommandEncoder, } from "../commandEncoder/commandEncoder.js";
import { INTERNAL_adoptRenderCommands, } from "../commandEncoder/renderPass.js";
import { emitRenderDraw, finalizeOwnEncoder, requireIndexBuffer } from "./drawState.js";
import { isGPUCommandEncoder, isGPURenderBundleEncoder, isGPURenderPassEncoder, isTgpuCommandEncoder, isTgpuRenderCommands, } from "./typeGuards.js";
import { createWithPerformanceCallback, createWithTimestampWrites, } from "./timeable.js";
import { nonTransferablePriorsOf } from "./priors.js";
import {} from "../../data/offsetUtils.js";
import { warnIfOverflow } from "./webgpuLimitations.js";
import { collectBindGroupPairs, collectVertexBufferPairs, DRAW_INDEXED_INDIRECT_SIZE, DRAW_INDIRECT_SIZE, resolveIndirectOffset, restoreTimestampPriors, } from "./pipelineUtils.js";
import { NullPerformanceTracker, PerformanceTrackerImpl, } from "./performanceTracker.js";
import { logger } from "../../tgpuLogger.js";
export function INTERNAL_createRenderPipeline(options) {
return new TgpuRenderPipelineImpl(new RenderPipelineCore(options), {});
}
export function INTERNAL_restoreRenderPipeline(soul, ctx) {
invariant(soul.raw, 'A render pipeline soul is only complete once materialized.');
const root = ctx.getRoot(soul.device);
const core = RenderPipelineCore.precompiled(root, {
pipeline: soul.raw,
usedBindGroupLayouts: soul.usedBindGroupLayouts ?? [],
// The catchall group is already one of `bindGroups`, keyed by the layout it was resolved with
catchall: undefined,
logResources: undefined,
usedVertexLayouts: soul.usedVertexLayouts ?? [],
fragmentOut: soul.fragmentOut,
});
const pipeline = new TgpuRenderPipelineImpl(core, {
bindGroupLayoutMap: new Map(soul.bindGroups),
vertexLayoutMap: new Map(soul.vertexBuffers),
indexBuffer: soul.indexBuffer,
stencilReference: soul.stencilReference,
});
return restoreTimestampPriors(pipeline, soul);
}
class TgpuRenderPipelineImpl {
[$internal];
[$soul];
resourceType = 'render-pipeline';
[$getNameForward];
constructor(core, priors) {
this[$soul] = {
type: 'render-pipeline',
device: core.options.root.device,
raw: undefined,
label: undefined,
};
this[$internal] = {
core,
priors,
root: core.options.root,
materialize: () => {
const soul = this[$soul];
if (!soul.raw) {
const memo = core.unwrap();
soul.raw = memo.pipeline;
soul.usedBindGroupLayouts = memo.usedBindGroupLayouts;
soul.usedVertexLayouts = memo.usedVertexLayouts;
soul.fragmentOut = memo.fragmentOut;
soul.bindGroups = collectBindGroupPairs(memo.usedBindGroupLayouts, memo.catchall, priors.bindGroupLayoutMap);
soul.vertexBuffers = collectVertexBufferPairs(memo.usedVertexLayouts, priors.vertexLayoutMap);
soul.indexBuffer = priors.indexBuffer;
soul.stencilReference = priors.stencilReference;
soul.timestampWrites = priors.timestampWrites;
soul.performanceCallback = priors.performanceCallback;
soul.nonTransferablePriors = nonTransferablePriorsOf(priors);
}
return soul.raw;
},
};
this[$getNameForward] = core;
}
[$resolve](ctx) {
return ctx.resolve(this[$internal].core);
}
toString() {
return `renderPipeline:${getName(this) ?? '<unnamed>'}`;
}
$name(label) {
setName(this, label);
return this;
}
#withPriors(patch) {
const { core, priors } = this[$internal];
return new TgpuRenderPipelineImpl(core, { ...priors, ...patch });
}
with(first, resource) {
const internals = this[$internal];
if (isTgpuRenderCommands(first)) {
return this.#withPriors({ pass: first, encoder: undefined });
}
if (isTgpuCommandEncoder(first)) {
return this.#withPriors({ pass: undefined, encoder: first });
}
if (isGPURenderPassEncoder(first) || isGPURenderBundleEncoder(first)) {
return this.#withPriors({
pass: INTERNAL_adoptRenderCommands(internals.root, first),
encoder: undefined,
});
}
if (isGPUCommandEncoder(first)) {
return this.#withPriors({
pass: undefined,
encoder: INTERNAL_adoptCommandEncoder(internals.root, first),
});
}
if (isBindGroup(first) || isBindGroupLayout(first)) {
const [layout, group] = isBindGroup(first)
? [first.layout, first]
: [first, resource];
return this.#withPriors({
bindGroupLayoutMap: new Map([
...(internals.priors.bindGroupLayoutMap ?? []),
[layout, group],
]),
});
}
if (isVertexLayout(first)) {
return this.#withPriors({
vertexLayoutMap: new Map([
...(internals.priors.vertexLayoutMap ?? []),
[first, resource],
]),
});
}
throw new Error('Unsupported value passed into .with()');
}
withPerformanceCallback(callback) {
const internals = this[$internal];
if (internals.priors.timestampWrites) {
return this.#withPriors({ performanceCallback: callback });
}
const querySet = internals.core.performanceCallbackQuerySet;
if (!querySet) {
logger.warn('webgpu-feature-missing', 'Performance callback cannot be used because the timestamp-query feature is not enabled on the root.');
return this;
}
return this.#withPriors(createWithPerformanceCallback(internals.priors, callback, querySet));
}
withTimestampWrites(options) {
const internals = this[$internal];
return this.#withPriors(createWithTimestampWrites(internals.priors, options, internals.root));
}
withColorAttachment(attachment) {
return this.#withPriors({ colorAttachment: attachment });
}
withDepthStencilAttachment(attachment) {
return this.#withPriors({ depthStencilAttachment: attachment });
}
withStencilReference(reference) {
return this.#withPriors({ stencilReference: reference });
}
withIndexBuffer(buffer, indexFormatOrOffset, offsetElementsOrSizeBytes, sizeElementsOrUndefined) {
if (isGPUBuffer(buffer)) {
if (typeof indexFormatOrOffset !== 'string') {
throw new Error('If a GPUBuffer is passed, indexFormat must be provided.');
}
return this.#withPriors({
indexBuffer: {
buffer,
indexFormat: indexFormatOrOffset,
offsetBytes: offsetElementsOrSizeBytes,
sizeBytes: sizeElementsOrUndefined,
},
});
}
const dataTypeToIndexFormat = {
u32: 'uint32',
u16: 'uint16',
};
const elementType = buffer.dataType.elementType;
return this.#withPriors({
indexBuffer: {
buffer,
indexFormat: dataTypeToIndexFormat[elementType.type],
offsetBytes: indexFormatOrOffset !== undefined
? indexFormatOrOffset * sizeOf(elementType)
: undefined,
sizeBytes: sizeElementsOrUndefined !== undefined
? sizeElementsOrUndefined * sizeOf(elementType)
: undefined,
},
});
}
initAsync() {
return this[$internal].core.initAsync();
}
initSync() {
this[$internal].core.initSync();
}
get hasIndexBuffer() {
return this[$internal].priors.indexBuffer !== undefined;
}
#ownPassDescriptor() {
const { core, priors } = this[$internal];
const { fragmentOut } = core.unwrap();
return {
label: getName(core) ?? '<unnamed>',
colorAttachments: fragmentOut
? connectAttachmentToShader(fragmentOut, priors.colorAttachment ?? {})
: [],
depthStencilAttachment: priors.depthStencilAttachment,
timestampWrites: priors.timestampWrites,
};
}
#execute(usesIndexBuffer, emit) {
const { core, priors, root } = this[$internal];
if (priors.pass) {
emitRenderDraw(root, priors.pass[$internal], this, usesIndexBuffer, emit);
return;
}
// checked up front so a rejected draw never leaves a half-recorded pass behind
if (usesIndexBuffer) {
requireIndexBuffer(priors.indexBuffer);
}
const encoder = priors.encoder ?? INTERNAL_createCommandEncoder(root);
const pass = encoder.beginRenderPass(this.#ownPassDescriptor());
emitRenderDraw(root, pass[$internal], this, usesIndexBuffer, emit, /* ownsPass */ true);
pass.end();
finalizeOwnEncoder(encoder, core, core.unwrap().logResources, priors);
}
draw(vertexCount, instanceCount, firstVertex, firstInstance) {
this.#execute(false, (rawPass) => rawPass.draw(vertexCount, instanceCount, firstVertex, firstInstance));
}
drawIndexed(indexCount, instanceCount, firstIndex, baseVertex, firstInstance) {
this.#execute(true, (rawPass) => rawPass.drawIndexed(indexCount, instanceCount, firstIndex, baseVertex, firstInstance));
}
drawIndirect(indirectBuffer, indirectOffset) {
const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : indirectBuffer.buffer;
const offset = resolveIndirectOffset(indirectBuffer, indirectOffset, DRAW_INDIRECT_SIZE, 'drawIndirect');
this.#execute(false, (rawPass) => rawPass.drawIndirect(rawBuffer, offset));
}
drawIndexedIndirect(indirectBuffer, indirectOffset) {
const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : indirectBuffer.buffer;
const offset = resolveIndirectOffset(indirectBuffer, indirectOffset, DRAW_INDEXED_INDIRECT_SIZE, 'drawIndexedIndirect');
this.#execute(true, (rawPass) => rawPass.drawIndexedIndirect(rawBuffer, offset));
}
}
class RenderPipelineCore {
[$internal] = true;
options;
#performanceTracker;
#initAsyncPromise;
#memo;
#latestAutoVertexIn;
#latestAutoFragmentOut;
#performanceCallbackQuerySet;
constructor(options) {
this.options = options;
this.#performanceTracker = PERF?.enabled
? new PerformanceTrackerImpl()
: new NullPerformanceTracker();
}
static precompiled(root, memo) {
const core = new RenderPipelineCore({
root,
slotBindings: [],
descriptor: undefined,
});
core.#memo = memo;
return core;
}
[$resolve](ctx) {
const { slotBindings } = this.options;
const { vertex, fragment, attribs = {} } = this.options.descriptor ?? {};
if (!vertex) {
return snip('', Void, /* origin */ 'runtime');
}
this.#latestAutoVertexIn = undefined;
this.#latestAutoFragmentOut = undefined;
const locations = matchUpVaryingLocations(vertex?.shell?.out, fragment?.shell?.in, getName(vertex) ?? '<unnamed>', getName(fragment) ?? '<unnamed>');
return ctx.withVaryingLocations(locations, () => ctx.withSlots(slotBindings, () => {
let vertexOut;
if (typeof vertex === 'function') {
const defaultAttribData = Object.fromEntries(Object.entries(attribs).map(([key, value]) => [
key,
formatToWGSLType[value.format],
]));
const autoFn = new AutoVertexFn(vertex, defaultAttribData, locations);
ctx.resolve(autoFn);
this.#latestAutoVertexIn = autoFn.autoIn.completeStruct.propTypes;
vertexOut = autoFn.autoOut.completeStruct;
}
else {
vertexOut = ctx.resolve(vertex).dataType;
}
if (fragment) {
if (typeof fragment === 'function') {
const varyings = Object.fromEntries(Object.entries(vertexOut.propTypes).filter(([, dataType]) => !isBuiltin(dataType)));
const fragOut = ctx.resolve(new AutoFragmentFn(fragment, varyings, locations));
this.#latestAutoFragmentOut = fragOut.dataType;
}
else {
ctx.resolve(fragment);
}
}
return snip('', Void, /* origin */ 'runtime');
}));
}
toString() {
return 'renderPipelineCore';
}
get performanceCallbackQuerySet() {
if (!this.options.root.enabledFeatures.has('timestamp-query')) {
return undefined;
}
return (this.#performanceCallbackQuerySet ??= this.options.root.createQuerySet('timestamp', 2));
}
initAsync() {
if (this.#memo !== undefined) {
// the pipeline was already resolved & compiled
return Promise.resolve();
}
if (this.#initAsyncPromise === undefined) {
// the pipeline did not start resolution & compilation
const device = this.options.root.device;
const { resolutionResult, descriptor, connectedAttribs, fragmentOut } = this.resolveAndCreateShaderModule();
const { usedBindGroupLayouts, catchall, logResources } = resolutionResult;
this.#initAsyncPromise = device
.createRenderPipelineAsync(descriptor)
.then((pipeline) => {
this.#memo = {
pipeline,
usedBindGroupLayouts,
catchall,
logResources,
usedVertexLayouts: connectedAttribs.usedVertexLayouts,
fragmentOut,
};
this.#performanceTracker.measureCompile(device);
})
.finally(() => {
this.#initAsyncPromise = undefined;
});
}
return this.#initAsyncPromise;
}
initSync() {
if (this.#memo !== undefined) {
return;
}
if (this.#initAsyncPromise !== undefined) {
throw new Error("'pipeline.initAsync()' was called and is not yet resolved.");
}
const device = this.options.root.device;
const { resolutionResult, descriptor, connectedAttribs, fragmentOut } = this.resolveAndCreateShaderModule();
const { usedBindGroupLayouts, catchall, logResources } = resolutionResult;
this.#memo = {
pipeline: device.createRenderPipeline(descriptor),
usedBindGroupLayouts,
catchall,
logResources,
usedVertexLayouts: connectedAttribs.usedVertexLayouts,
fragmentOut,
};
this.#performanceTracker.measureCompile(device);
}
unwrap() {
this.initSync();
return this.#memo;
}
resolveAndCreateShaderModule() {
const { root, descriptor: tgpuDescriptor } = this.options;
if (!tgpuDescriptor) {
throw new Error('Precompiled pipelines are never resolved again.');
}
const device = root.device;
const enableExtensions = wgslEnableExtensions.filter((extension) => root.enabledFeatures.has(wgslEnableExtensionToFeatureName[extension]));
// Resolving code
const ns = namespace({ names: root.nameRegistrySetting });
const resolutionResult = this.#performanceTracker.measureResolve(() => resolve(this, {
namespace: ns,
minify: root.minify,
enableExtensions,
shaderGenerator: root.shaderGeneratorClass ? new root.shaderGeneratorClass() : undefined,
root,
}));
const { code, usedBindGroupLayouts, catchall } = resolutionResult;
if (catchall !== undefined) {
usedBindGroupLayouts[catchall[0]]?.$name(`${getName(this) ?? '<unnamed>'} - Automatic Bind Group & Layout`);
}
warnIfOverflow(usedBindGroupLayouts, device.limits);
const module = device.createShaderModule({
label: `${getName(this) ?? '<unnamed>'} - Shader`,
code,
});
const { vertex, fragment, attribs = {}, targets } = tgpuDescriptor;
const connectedAttribs = connectAttributesToShader(vertex?.shell?.in ?? this.#latestAutoVertexIn ?? {}, attribs);
// If the fragment output is a single builtin, then this.#latestAutoFragmentOut doesn't
// retain that information, that's why we use that here and fallback to this.#latestAutoFragmentOut.
// One other reason is that if the shelled fragment has already been resolved in this namespace,
// then this.#latestAutoFragmentOut will be undefined.
const fragmentOut = fragment?.shell?.returnType ?? this.#latestAutoFragmentOut;
const connectedTargets = fragmentOut ? connectTargetsToShader(fragmentOut, targets) : [null];
const descriptor = {
layout: device.createPipelineLayout({
label: `${getName(this) ?? '<unnamed>'} - Pipeline Layout`,
bindGroupLayouts: usedBindGroupLayouts.map((l) => root.unwrap(l)),
}),
vertex: {
module,
buffers: connectedAttribs.bufferDefinitions,
},
};
const label = getName(this);
if (label !== undefined) {
descriptor.label = label;
}
if (tgpuDescriptor.fragment) {
descriptor.fragment = {
module,
targets: connectedTargets,
};
}
if (tgpuDescriptor.primitive) {
if (isWgslData(tgpuDescriptor.primitive.stripIndexFormat)) {
descriptor.primitive = {
...tgpuDescriptor.primitive,
stripIndexFormat: {
u32: 'uint32',
u16: 'uint16',
}[tgpuDescriptor.primitive.stripIndexFormat.type],
};
}
else {
descriptor.primitive = tgpuDescriptor.primitive;
}
}
if (tgpuDescriptor.depthStencil) {
descriptor.depthStencil = tgpuDescriptor.depthStencil;
}
if (tgpuDescriptor.multisample) {
descriptor.multisample = tgpuDescriptor.multisample;
}
return { resolutionResult, descriptor, connectedAttribs, fragmentOut };
}
}
/**
* Assumes vertexOut and fragmentIn are matching when it comes to the keys, that is fragmentIn's keyset is a subset of vertexOut's
* Logs a warning, when they don't match in terms of custom locations
*/
export function matchUpVaryingLocations(vertexOut = {}, fragmentIn = {}, vertexFnName, fragmentFnName) {
const locations = {};
const usedLocations = new Set();
function saveLocation(key, location) {
locations[key] = location;
usedLocations.add(location);
}
// respect custom locations and pair up vertex and fragment varying with the same key
for (const [key, value] of Object.entries(vertexOut)) {
const customLocation = getCustomLocation(value);
if (customLocation !== undefined) {
saveLocation(key, customLocation);
}
}
for (const [key, value] of Object.entries(fragmentIn)) {
const customLocation = getCustomLocation(value);
if (customLocation === undefined) {
continue;
}
if (locations[key] === undefined) {
saveLocation(key, customLocation);
}
else if (locations[key] !== customLocation) {
logger.warn('locations-mismatched', `Mismatched location between vertexFn (${vertexFnName}) output (${locations[key]}) and fragmentFn (${fragmentFnName}) input (${customLocation}) for the key "${key}", using the location set on vertex output.`);
}
}
// automatically assign remaining locations to the rest
let nextLocation = 0;
for (const key of Object.keys(vertexOut ?? {})) {
if (isBuiltin(vertexOut[key]) || locations[key] !== undefined) {
continue;
}
while (usedLocations.has(nextLocation)) {
nextLocation++;
}
saveLocation(key, nextLocation);
}
return locations;
}