typegpu/core/commandEncoder/computePass.js

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

import { $internal } from "../../shared/symbols.js";
import { isGPUBuffer } from "../../types.js";
import { DISPATCH_INDIRECT_SIZE, resolveIndirectOffset } from "../pipeline/pipelineUtils.js";
import { ComputeDrawState, emitComputeDispatch, recordBindGroup, stampComputePipeline, } from "../pipeline/drawState.js";
import { unwrapTimestampWrites } from "./attachments.js";
// --------------
// Implementation
// --------------
export function INTERNAL_beginComputePass(encoder, descriptor) {
    const { rawEncoder, root } = encoder[$internal];
    const rawDescriptor = {};
    if (descriptor?.label !== undefined) {
        rawDescriptor.label = descriptor.label;
    }
    if (descriptor?.timestampWrites !== undefined) {
        rawDescriptor.timestampWrites = unwrapTimestampWrites(root, descriptor.timestampWrites);
    }
    return new TgpuComputePassImpl(root, rawEncoder.beginComputePass(rawDescriptor), encoder);
}
export function INTERNAL_adoptComputePass(root, rawPass) {
    const adopted = new TgpuComputePassImpl(root, rawPass, undefined);
    adopted[$internal].state.rawAccessed = true;
    return adopted;
}
class TgpuComputePassImpl {
    [$internal];
    resourceType = 'compute-pass';
    #root;
    constructor(root, rawPass, owner) {
        this.#root = root;
        this[$internal] = {
            rawPass,
            state: new ComputeDrawState(),
            owner,
            appliedVersion: undefined,
        };
    }
    #emit(emit) {
        const internals = this[$internal];
        const pipeline = internals.state.currentPipeline;
        if (!pipeline) {
            throw new Error('Cannot dispatch without a call to pass.setPipeline');
        }
        emitComputeDispatch(this.#root, internals, pipeline, emit);
    }
    setPipeline(pipeline) {
        stampComputePipeline(this[$internal].state, pipeline);
    }
    setBindGroup(first, bindGroup) {
        recordBindGroup(this[$internal].state, first, bindGroup);
    }
    dispatchWorkgroups(x, y, z) {
        this.#emit((rawPass) => rawPass.dispatchWorkgroups(x, y, z));
    }
    dispatchWorkgroupsIndirect(indirectBuffer, start) {
        const rawBuffer = isGPUBuffer(indirectBuffer) ? indirectBuffer : indirectBuffer.buffer;
        const offset = resolveIndirectOffset(indirectBuffer, start, DISPATCH_INDIRECT_SIZE, 'dispatchWorkgroupsIndirect');
        this.#emit((rawPass) => rawPass.dispatchWorkgroupsIndirect(rawBuffer, offset));
    }
    end() {
        this[$internal].rawPass.end();
    }
}