typegpu/std/subgroup.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 { bool, i32, u32 } from "../data/numeric.js";
import { vec4u } from "../data/vector.js";
import { unify } from "../tgsl/conversion.js";
const errorMessage = 'Subgroup operations can only be used in the GPU context.';
export const subgroupAdd = dualImpl({
    name: 'subgroupAdd',
    signature: (arg) => ({ argTypes: [arg], returnType: arg }),
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [arg]) => stitch `subgroupAdd(${arg})`,
    sideEffects: false,
});
export const subgroupExclusiveAdd = dualImpl({
    name: 'subgroupExclusiveAdd',
    signature: (arg) => ({ argTypes: [arg], returnType: arg }),
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [arg]) => stitch `subgroupExclusiveAdd(${arg})`,
    sideEffects: false,
});
export const subgroupInclusiveAdd = dualImpl({
    name: 'subgroupInclusiveAdd',
    signature: (arg) => ({ argTypes: [arg], returnType: arg }),
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [arg]) => stitch `subgroupInclusiveAdd(${arg})`,
    sideEffects: false,
});
export const subgroupAll = dualImpl({
    name: 'subgroupAll',
    signature: { argTypes: [bool], returnType: bool },
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [e]) => stitch `subgroupAll(${e})`,
    sideEffects: false,
});
export const subgroupAnd = dualImpl({
    name: 'subgroupAnd',
    signature: (arg) => ({ argTypes: [arg], returnType: arg }),
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [e]) => stitch `subgroupAnd(${e})`,
    sideEffects: false,
});
export const subgroupAny = dualImpl({
    name: 'subgroupAny',
    signature: { argTypes: [bool], returnType: bool },
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [e]) => stitch `subgroupAny(${e})`,
    sideEffects: false,
});
export const subgroupBallot = dualImpl({
    name: 'subgroupBallot',
    signature: { argTypes: [bool], returnType: vec4u },
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [e]) => stitch `subgroupBallot(${e})`,
    sideEffects: false,
});
export const subgroupBroadcast = dualImpl({
    name: 'subgroupBroadcast',
    signature: (...args) => {
        const id = unify([args[1]], [i32, u32]);
        if (!id) {
            throw new Error(`subgroupBroadcast's second argument has to be compatible with i32 or u32. Got: ${args[1].type}`);
        }
        return { argTypes: [args[0], id[0]], returnType: args[0] };
    },
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [e, index]) => stitch `subgroupBroadcast(${e}, ${index})`,
    sideEffects: false,
});
export const subgroupBroadcastFirst = dualImpl({
    name: 'subgroupBroadcastFirst',
    signature: (arg) => ({ argTypes: [arg], returnType: arg }),
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [e]) => stitch `subgroupBroadcastFirst(${e})`,
    sideEffects: false,
});
export const subgroupElect = dualImpl({
    name: 'subgroupElect',
    signature: { argTypes: [], returnType: bool },
    normalImpl: errorMessage,
    codegenImpl: () => stitch `subgroupElect()`,
    sideEffects: false,
});
export const subgroupMax = dualImpl({
    name: 'subgroupMax',
    signature: (arg) => ({ argTypes: [arg], returnType: arg }),
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [arg]) => stitch `subgroupMax(${arg})`,
    sideEffects: false,
});
export const subgroupMin = dualImpl({
    name: 'subgroupMin',
    signature: (arg) => ({ argTypes: [arg], returnType: arg }),
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [arg]) => stitch `subgroupMin(${arg})`,
    sideEffects: false,
});
export const subgroupMul = dualImpl({
    name: 'subgroupMul',
    signature: (arg) => ({ argTypes: [arg], returnType: arg }),
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [arg]) => stitch `subgroupMul(${arg})`,
    sideEffects: false,
});
export const subgroupExclusiveMul = dualImpl({
    name: 'subgroupExclusiveMul',
    signature: (arg) => ({ argTypes: [arg], returnType: arg }),
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [arg]) => stitch `subgroupExclusiveMul(${arg})`,
    sideEffects: false,
});
export const subgroupInclusiveMul = dualImpl({
    name: 'subgroupInclusiveMul',
    signature: (arg) => ({ argTypes: [arg], returnType: arg }),
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [arg]) => stitch `subgroupInclusiveMul(${arg})`,
    sideEffects: false,
});
export const subgroupOr = dualImpl({
    name: 'subgroupOr',
    signature: (arg) => ({ argTypes: [arg], returnType: arg }),
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [e]) => stitch `subgroupOr(${e})`,
    sideEffects: false,
});
export const subgroupShuffle = dualImpl({
    name: 'subgroupShuffle',
    signature: (...args) => {
        const id = unify([args[1]], [i32, u32]);
        if (!id) {
            throw new Error(`subgroupShuffle's second argument has to be compatible with i32 or u32. Got: ${args[1].type}`);
        }
        return { argTypes: [args[0], id[0]], returnType: args[0] };
    },
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [e, index]) => stitch `subgroupShuffle(${e}, ${index})`,
    sideEffects: false,
});
export const subgroupShuffleDown = dualImpl({
    name: 'subgroupShuffleDown',
    signature: (...args) => {
        const delta = unify([args[1]], [u32]);
        if (!delta) {
            throw new Error(`subgroupShuffleDown's second argument has to be compatible with u32. Got: ${args[1].type}`);
        }
        return { argTypes: [args[0], delta[0]], returnType: args[0] };
    },
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [e, delta]) => stitch `subgroupShuffleDown(${e}, ${delta})`,
    sideEffects: false,
});
export const subgroupShuffleUp = dualImpl({
    name: 'subgroupShuffleUp',
    signature: (...args) => {
        const delta = unify([args[1]], [u32]);
        if (!delta) {
            throw new Error(`subgroupShuffleUp's second argument has to be compatible with u32. Got: ${args[1].type}`);
        }
        return { argTypes: [args[0], delta[0]], returnType: args[0] };
    },
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [e, delta]) => stitch `subgroupShuffleUp(${e}, ${delta})`,
    sideEffects: false,
});
export const subgroupShuffleXor = dualImpl({
    name: 'subgroupShuffleXor',
    signature: (...args) => {
        const mask = unify([args[1]], [u32]);
        if (!mask) {
            throw new Error(`subgroupShuffleXor's second argument has to be compatible with u32. Got: ${args[1].type}`);
        }
        return { argTypes: [args[0], mask[0]], returnType: args[0] };
    },
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [e, mask]) => stitch `subgroupShuffleXor(${e}, ${mask})`,
    sideEffects: false,
});
export const subgroupXor = dualImpl({
    name: 'subgroupXor',
    signature: (arg) => ({ argTypes: [arg], returnType: arg }),
    normalImpl: errorMessage,
    codegenImpl: (_ctx, [e]) => stitch `subgroupXor(${e})`,
    sideEffects: false,
});