typegpu/std/atomic.js

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

import { dualImpl } from "../core/function/dualImpl.js";
import { stitch } from "../core/resolve/stitch.js";
import { i32, u32 } from "../data/numeric.js";
import { isAtomic, Void, } from "../data/wgslTypes.js";
import { safeStringify } from "../shared/stringify.js";
export const workgroupBarrier = dualImpl({
    name: 'workgroupBarrier',
    normalImpl: 'workgroupBarrier is a no-op outside of CODEGEN mode.',
    signature: { argTypes: [], returnType: Void },
    codegenImpl: () => 'workgroupBarrier()',
    sideEffects: true,
});
export const storageBarrier = dualImpl({
    name: 'storageBarrier',
    normalImpl: 'storageBarrier is a no-op outside of CODEGEN mode.',
    signature: { argTypes: [], returnType: Void },
    codegenImpl: () => 'storageBarrier()',
    sideEffects: true,
});
export const textureBarrier = dualImpl({
    name: 'textureBarrier',
    normalImpl: 'textureBarrier is a no-op outside of CODEGEN mode.',
    signature: { argTypes: [], returnType: Void },
    codegenImpl: () => 'textureBarrier()',
    sideEffects: true,
});
const atomicNormalError = 'Atomic operations are not supported outside of CODEGEN mode.';
export const atomicLoad = dualImpl({
    name: 'atomicLoad',
    normalImpl: atomicNormalError,
    signature: (a) => {
        if (!isAtomic(a)) {
            throw new Error(`Invalid atomic type: ${safeStringify(a)}`);
        }
        return { argTypes: [a], returnType: a.inner };
    },
    codegenImpl: (_ctx, [a]) => stitch `atomicLoad(&${a})`,
    sideEffects: true,
});
const atomicActionSignature = (a) => {
    if (!isAtomic(a)) {
        throw new Error(`Invalid atomic type: ${safeStringify(a)}`);
    }
    return {
        argTypes: [a, a.inner.type === 'u32' ? u32 : i32],
        returnType: Void,
    };
};
const atomicOpSignature = (a) => {
    if (!isAtomic(a)) {
        throw new Error(`Invalid atomic type: ${safeStringify(a)}`);
    }
    const paramType = a.inner.type === 'u32' ? u32 : i32;
    return {
        argTypes: [a, paramType],
        returnType: paramType,
    };
};
export const atomicStore = dualImpl({
    name: 'atomicStore',
    normalImpl: atomicNormalError,
    signature: atomicActionSignature,
    codegenImpl: (_ctx, [a, value]) => stitch `atomicStore(&${a}, ${value})`,
    sideEffects: true,
});
export const atomicAdd = dualImpl({
    name: 'atomicAdd',
    normalImpl: atomicNormalError,
    signature: atomicOpSignature,
    codegenImpl: (_ctx, [a, value]) => stitch `atomicAdd(&${a}, ${value})`,
    sideEffects: true,
});
export const atomicSub = dualImpl({
    name: 'atomicSub',
    normalImpl: atomicNormalError,
    signature: atomicOpSignature,
    codegenImpl: (_ctx, [a, value]) => stitch `atomicSub(&${a}, ${value})`,
    sideEffects: true,
});
export const atomicMax = dualImpl({
    name: 'atomicMax',
    normalImpl: atomicNormalError,
    signature: atomicOpSignature,
    codegenImpl: (_ctx, [a, value]) => stitch `atomicMax(&${a}, ${value})`,
    sideEffects: true,
});
export const atomicMin = dualImpl({
    name: 'atomicMin',
    normalImpl: atomicNormalError,
    signature: atomicOpSignature,
    codegenImpl: (_ctx, [a, value]) => stitch `atomicMin(&${a}, ${value})`,
    sideEffects: true,
});
export const atomicAnd = dualImpl({
    name: 'atomicAnd',
    normalImpl: atomicNormalError,
    signature: atomicOpSignature,
    codegenImpl: (_ctx, [a, value]) => stitch `atomicAnd(&${a}, ${value})`,
    sideEffects: true,
});
export const atomicOr = dualImpl({
    name: 'atomicOr',
    normalImpl: atomicNormalError,
    signature: atomicOpSignature,
    codegenImpl: (_ctx, [a, value]) => stitch `atomicOr(&${a}, ${value})`,
    sideEffects: true,
});
export const atomicXor = dualImpl({
    name: 'atomicXor',
    normalImpl: atomicNormalError,
    signature: atomicOpSignature,
    codegenImpl: (_ctx, [a, value]) => stitch `atomicXor(&${a}, ${value})`,
    sideEffects: true,
});