typegpu/core/pipeline/computePipeline.js

A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.

import { snip } from "../../data/snippet.js";
import { Void } from "../../data/wgslTypes.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 { INTERNAL_adoptCommandEncoder, INTERNAL_createCommandEncoder, } from "../commandEncoder/commandEncoder.js";
import { INTERNAL_adoptComputePass } from "../commandEncoder/computePass.js";
import { emitComputeDispatch, finalizeOwnEncoder } from "./drawState.js";
import { isGPUCommandEncoder, isGPUComputePassEncoder, isTgpuCommandEncoder, isTgpuComputePass, } from "./typeGuards.js";
import { isGPUBuffer } from "../../types.js";
import { wgslEnableExtensions, wgslEnableExtensionToFeatureName } from "../../wgslExtensions.js";
import { namespace } from "../resolve/namespace.js";
import { warnIfOverflow } from "./webgpuLimitations.js";
import { collectBindGroupPairs, DISPATCH_INDIRECT_SIZE, resolveIndirectOffset, restoreTimestampPriors, } from "./pipelineUtils.js";
import { invariant } from "../../errors.js";
import { createWithPerformanceCallback, createWithTimestampWrites, } from "./timeable.js";
import { nonTransferablePriorsOf } from "./priors.js";
import { NullPerformanceTracker, PerformanceTrackerImpl, } from "./performanceTracker.js";
import { logger } from "../../tgpuLogger.js";
export function INTERNAL_createComputePipeline(root, slotBindings, descriptor) {
    return new TgpuComputePipelineImpl(new ComputePipelineCore(root, slotBindings, descriptor), {});
}
export function INTERNAL_restoreComputePipeline(soul, ctx) {
    invariant(soul.raw, 'A compute pipeline soul is only complete once materialized.');
    const root = ctx.getRoot(soul.device);
    const core = ComputePipelineCore.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,
    });
    const pipeline = new TgpuComputePipelineImpl(core, {
        bindGroupLayoutMap: new Map(soul.bindGroups),
    });
    return restoreTimestampPriors(pipeline, soul);
}
class TgpuComputePipelineImpl {
    [$internal];
    [$soul];
    resourceType = 'compute-pipeline';
    [$getNameForward];
    constructor(core, priors) {
        this[$soul] = {
            type: 'compute-pipeline',
            device: core.root.device,
            raw: undefined,
            label: undefined,
        };
        this[$internal] = {
            core,
            priors,
            root: core.root,
            materialize: () => {
                const soul = this[$soul];
                if (!soul.raw) {
                    const memo = core.unwrap();
                    soul.raw = memo.pipeline;
                    soul.usedBindGroupLayouts = memo.usedBindGroupLayouts;
                    soul.bindGroups = collectBindGroupPairs(memo.usedBindGroupLayouts, memo.catchall, priors.bindGroupLayoutMap);
                    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 `computePipeline:${getName(this) ?? '<unnamed>'}`;
    }
    #withPriors(patch) {
        const { core, priors } = this[$internal];
        return new TgpuComputePipelineImpl(core, { ...priors, ...patch });
    }
    with(first, bindGroup) {
        const internals = this[$internal];
        if (isTgpuComputePass(first)) {
            return this.#withPriors({ pass: first, encoder: undefined });
        }
        if (isTgpuCommandEncoder(first)) {
            return this.#withPriors({ pass: undefined, encoder: first });
        }
        if (isGPUComputePassEncoder(first)) {
            return this.#withPriors({
                pass: INTERNAL_adoptComputePass(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, bindGroup];
            return this.#withPriors({
                bindGroupLayoutMap: new Map([
                    ...(internals.priors.bindGroupLayoutMap ?? []),
                    [layout, group],
                ]),
            });
        }
        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));
    }
    dispatchWorkgroups(x, y, z) {
        this.#execute((pass) => pass.dispatchWorkgroups(x, y, z));
    }
    dispatchWorkgroupsIndirect(indirectBuffer, start) {
        const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : indirectBuffer.buffer;
        const offset = resolveIndirectOffset(indirectBuffer, start, DISPATCH_INDIRECT_SIZE, 'dispatchWorkgroupsIndirect');
        this.#execute((pass) => pass.dispatchWorkgroupsIndirect(rawBuffer, offset));
    }
    initAsync() {
        return this[$internal].core.initAsync();
    }
    initSync() {
        this[$internal].core.initSync();
    }
    #execute(dispatch) {
        const { core, priors, root } = this[$internal];
        if (priors.pass) {
            emitComputeDispatch(root, priors.pass[$internal], this, dispatch);
            return;
        }
        const encoder = priors.encoder ?? INTERNAL_createCommandEncoder(root);
        const pass = encoder.beginComputePass({
            label: getName(core) ?? '<unnamed>',
            timestampWrites: priors.timestampWrites,
        });
        emitComputeDispatch(root, pass[$internal], this, dispatch, /* ownsPass */ true);
        pass.end();
        finalizeOwnEncoder(encoder, core, core.unwrap().logResources, priors);
    }
    $name(label) {
        setName(this, label);
        return this;
    }
}
class ComputePipelineCore {
    [$internal] = true;
    root;
    #performanceTracker;
    #initAsyncPromise;
    #memo;
    #slotBindings;
    #descriptor;
    #performanceCallbackQuerySet;
    constructor(root, slotBindings, descriptor) {
        this.root = root;
        this.#slotBindings = slotBindings;
        this.#descriptor = descriptor;
        this.#performanceTracker = PERF?.enabled
            ? new PerformanceTrackerImpl()
            : new NullPerformanceTracker();
    }
    static precompiled(root, memo) {
        const core = new ComputePipelineCore(root, [], undefined);
        core.#memo = memo;
        return core;
    }
    [$resolve](ctx) {
        const descriptor = this.#descriptor;
        if (!descriptor) {
            // Precompiled pipelines have nothing to contribute to the shader
            return snip('', Void, /* origin */ 'runtime');
        }
        return ctx.withSlots(this.#slotBindings, () => {
            ctx.resolve(descriptor.compute);
            return snip('', Void, /* origin */ 'runtime');
        });
    }
    toString() {
        return 'computePipelineCore';
    }
    get performanceCallbackQuerySet() {
        if (!this.root.enabledFeatures.has('timestamp-query')) {
            return undefined;
        }
        return (this.#performanceCallbackQuerySet ??= this.root.createQuerySet('timestamp', 2));
    }
    /**
     * @privateRemarks
     * This function cannot be a regular async function
     * because when called multiple times before the promise finishes,
     * we want it to return the same promise each time.
     */
    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.root.device;
            const { resolutionResult, module } = this.resolveAndCreateShaderModule();
            const { usedBindGroupLayouts, catchall, logResources } = resolutionResult;
            this.#initAsyncPromise = device
                .createComputePipelineAsync({
                label: getName(this) ?? '<unnamed>',
                layout: device.createPipelineLayout({
                    label: `${getName(this) ?? '<unnamed>'} - Pipeline Layout`,
                    bindGroupLayouts: usedBindGroupLayouts.map((l) => this.root.unwrap(l)),
                }),
                compute: { module },
            })
                .then((pipeline) => {
                this.#memo = { pipeline, usedBindGroupLayouts, catchall, logResources };
                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.root.device;
        const { resolutionResult, module } = this.resolveAndCreateShaderModule();
        const { usedBindGroupLayouts, catchall, logResources } = resolutionResult;
        this.#memo = {
            pipeline: device.createComputePipeline({
                label: getName(this) ?? '<unnamed>',
                layout: device.createPipelineLayout({
                    label: `${getName(this) ?? '<unnamed>'} - Pipeline Layout`,
                    bindGroupLayouts: usedBindGroupLayouts.map((l) => this.root.unwrap(l)),
                }),
                compute: { module },
            }),
            usedBindGroupLayouts,
            catchall,
            logResources,
        };
        this.#performanceTracker.measureCompile(device);
    }
    unwrap() {
        this.initSync();
        return this.#memo;
    }
    resolveAndCreateShaderModule() {
        const device = this.root.device;
        const enableExtensions = wgslEnableExtensions.filter((extension) => this.root.enabledFeatures.has(wgslEnableExtensionToFeatureName[extension]));
        // Resolving code
        const ns = namespace({ names: this.root.nameRegistrySetting });
        const resolutionResult = this.#performanceTracker.measureResolve(() => resolve(this, {
            namespace: ns,
            minify: this.root.minify,
            enableExtensions,
            shaderGenerator: this.root.shaderGeneratorClass
                ? new this.root.shaderGeneratorClass()
                : undefined,
            root: this.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,
        });
        return { resolutionResult, module };
    }
}