typegpu/core/pipeline/timeable.js

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

import { isQuerySet } from "../querySet/querySet.js";
import { $internal } from "../../shared/symbols.js";
import { logger } from "../../tgpuLogger.js";
export function createWithPerformanceCallback(currentPriors, callback, querySet) {
    if (!currentPriors.timestampWrites) {
        return {
            ...currentPriors,
            performanceCallback: callback,
            timestampWrites: {
                querySet,
                beginningOfPassWriteIndex: 0,
                endOfPassWriteIndex: 1,
            },
        };
    }
    return {
        ...currentPriors,
        performanceCallback: callback,
    };
}
export function createWithTimestampWrites(currentPriors, options, root) {
    if (!root.enabledFeatures.has('timestamp-query')) {
        throw new Error('Timestamp writes require the "timestamp-query" feature to be enabled on GPU device.');
    }
    const timestampWrites = {
        querySet: options.querySet,
    };
    if (options.beginningOfPassWriteIndex !== undefined) {
        timestampWrites.beginningOfPassWriteIndex = options.beginningOfPassWriteIndex;
    }
    if (options.endOfPassWriteIndex !== undefined) {
        timestampWrites.endOfPassWriteIndex = options.endOfPassWriteIndex;
    }
    return {
        ...currentPriors,
        timestampWrites,
    };
}
const pendingTimestampReads = new WeakMap();
async function readTimestamps(querySet, registrations) {
    if (!querySet.available) {
        return;
    }
    const result = await querySet.read();
    for (const { priors, callback } of registrations) {
        const start = result[priors.timestampWrites?.beginningOfPassWriteIndex ?? 0];
        const end = result[priors.timestampWrites?.endOfPassWriteIndex ?? 1];
        if (start === undefined || end === undefined) {
            throw new Error('QuerySet did not return valid timestamps.');
        }
        await callback(start, end);
    }
}
/** Returns false when the encoder is one we cannot defer work to, meaning the callback never fires */
export function queueTimestampResolve(encoder, priors) {
    const querySet = priors.timestampWrites?.querySet;
    const callback = priors.performanceCallback;
    if (!querySet) {
        throw new Error('Cannot dispatch workgroups with performance callback without a query set.');
    }
    if (!isQuerySet(querySet)) {
        throw new Error('Performance callback with raw GPUQuerySet is not supported. Use TgpuQuerySet instead.');
    }
    const internals = encoder[$internal];
    if (internals.adopted) {
        return false;
    }
    const { root } = internals;
    // recorded at submission time to capture the last pass written into this encoder
    internals.beforeFinish.set(querySet, (rawEncoder) => {
        rawEncoder.resolveQuerySet(root.unwrap(querySet), 0, querySet.count, querySet[$internal].resolveBuffer, 0);
    });
    let byQuerySet = pendingTimestampReads.get(encoder);
    if (!byQuerySet) {
        byQuerySet = new Map();
        pendingTimestampReads.set(encoder, byQuerySet);
    }
    let registrations = byQuerySet.get(querySet);
    if (!registrations) {
        const regs = [];
        registrations = regs;
        byQuerySet.set(querySet, regs);
        internals.afterSubmit.set(querySet, () => {
            byQuerySet.delete(querySet);
            void readTimestamps(querySet, regs);
        });
    }
    if (registrations.some((reg) => reg.priors === priors)) {
        logger.warnOnce('suspicious', querySet, 'repeated-timed-execution', 'Repeated executions of a timed pipeline within one command encoder write to the same query set indices, so the performance callback reports only the last execution.');
    }
    else {
        registrations.push({ priors, callback });
    }
    return true;
}