typegpu/core/function/dualImpl.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 { setName } from "../../shared/meta.js";
import { $gpuCallable } from "../../shared/symbols.js";
import { tryConvertSnippet } from "../../tgsl/conversion.js";
import { concretize } from "../../tgsl/generationHelpers.js";
import { isKnownAtComptime, NormalState } from "../../types.js";
import { isPtr } from "../../data/wgslTypes.js";
export class MissingCpuImplError extends Error {
    constructor(message) {
        super(message);
        this.name = this.constructor.name;
    }
}
export function dualImpl(options) {
    const impl = ((...args) => {
        if (typeof options.normalImpl === 'string') {
            throw new MissingCpuImplError(options.normalImpl);
        }
        return options.normalImpl(...args);
    });
    if (options.name) {
        setName(impl, options.name);
    }
    impl.toString = () => options.name ?? '<unknown>';
    impl[$gpuCallable] = {
        get strictSignature() {
            return typeof options.signature !== 'function' ? options.signature : undefined;
        },
        call(ctx, args) {
            const { argTypes, returnType } = typeof options.signature === 'function'
                ? options.signature(...args.map((s) => {
                    // Dereference implicit pointers
                    if (isPtr(s.dataType) && s.dataType.implicit) {
                        return s.dataType.inner;
                    }
                    return s.dataType;
                }))
                : options.signature;
            const converted = args.map((s, idx) => {
                const argType = argTypes[idx];
                if (!argType) {
                    throw new Error('Function called with invalid arguments');
                }
                return tryConvertSnippet(ctx, s, argType, !options.ignoreImplicitCastWarning);
            });
            if (!options.noComptime &&
                converted.every((s) => isKnownAtComptime(s)) &&
                typeof options.normalImpl === 'function') {
                ctx.pushMode(new NormalState());
                try {
                    return snip(options.normalImpl(...converted.map((s) => s.value)), returnType, 
                    // Functions give up ownership of their return value
                    /* origin */ 'constant', options.sideEffects);
                }
                catch (e) {
                    // cpuImpl may in some cases be present but implemented only partially.
                    // In that case, if the MissingCpuImplError is thrown, we fallback to codegenImpl.
                    // If it is any other error, we just rethrow.
                    if (!(e instanceof MissingCpuImplError)) {
                        throw e;
                    }
                }
                finally {
                    ctx.popMode('normal');
                }
            }
            const possibleSideEffects = options.sideEffects || args.some((a) => a.possibleSideEffects);
            const concreteReturnType = concretize(returnType);
            return snip(options.codegenImpl(ctx, converted, concreteReturnType), concreteReturnType, 
            // Functions give up ownership of their return value
            /* origin */ 'runtime', possibleSideEffects);
        },
    };
    return impl;
}