typegpu/core/function/tgpuComputeFn.js

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

import { Void } from "../../data/wgslTypes.js";
import { getName, setName } from "../../shared/meta.js";
import { $getNameForward, $internal, $resolve } from "../../shared/symbols.js";
import { shaderStageSlot } from "../slot/internalSlots.js";
import { createFnCore } from "./fnCore.js";
import { separateBuiltins } from "./ioSchema.js";
import { stripTemplate } from "./templateUtils.js";
/**
 * Creates a shell of a typed entry function for the compute shader stage. Any function
 * that implements this shell can perform general-purpose computation.
 *
 * @param options.in
 *   Record with builtins used by the compute shader.
 * @param options.workgroupSize
 *   Size of blocks that the thread grid will be divided into (up to 3 dimensions).
 */
export function computeFn(options) {
    const shell = {
        argTypes: Object.values(options.in ?? {}),
        returnType: Void,
        workgroupSize: options.workgroupSize,
        entryPoint: 'compute',
    };
    const entryInput = separateBuiltins(options.in ?? {});
    const call = (arg, ...values) => createComputeFn(shell, options.workgroupSize, stripTemplate(arg, ...values), entryInput);
    return Object.assign(call, shell);
}
export function isTgpuComputeFn(value) {
    return value?.shell?.entryPoint === 'compute';
}
// --------------
// Implementation
// --------------
function createComputeFn(shell, workgroupSize, implementation, entryInput) {
    const core = createFnCore(implementation, 'compute', workgroupSize);
    const result = {
        shell,
        $uses(newExternals) {
            core.setExternals('userProvided', newExternals);
            return this;
        },
        [$internal]: true,
        [$getNameForward]: core,
        $name(newLabel) {
            setName(this, newLabel);
            return this;
        },
        [$resolve](ctx) {
            return ctx.withSlots([[shaderStageSlot, 'compute']], () => core.resolve(ctx, [], shell.returnType, entryInput));
        },
        toString() {
            return `computeFn:${getName(core) ?? '<unnamed>'}`;
        },
    };
    return result;
}