typegpu
Version:
A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.
892 lines (891 loc) • 32.5 kB
JavaScript
import { dualImpl, MissingCpuImplError } from "../core/function/dualImpl.js";
import { stitch } from "../core/resolve/stitch.js";
import { mat2x2f, mat3x3f, mat4x4f } from "../data/matrix.js";
import { clampScalar, smoothstepScalar } from "../data/numberOps.js";
import { abstractFloat, abstractInt, f16, f32, i32, u32 } from "../data/numeric.js";
import { abstruct } from "../data/struct.js";
import { vec2f, vec2h, vec2i, vec2u, vec3f, vec3h, vec3i, vec3u, vec4f, vec4h, vec4i, vec4u, } from "../data/vector.js";
import { VectorOps } from "../data/vectorOps.js";
import { generalizeFn, upCast } from "../data/generalizeFn.js";
import { isHalfPrecisionSchema, WORKAROUND_getSchema, } from "../data/wgslTypes.js";
import { SignatureNotSupportedError } from "../errors.js";
import { assertExhaustive } from "../shared/utilityTypes.js";
import { unify } from "../tgsl/conversion.js";
import { mul, sub } from "./operators.js";
// helpers
const unaryIdentitySignature = (arg) => {
return {
argTypes: [arg],
returnType: arg,
};
};
const variadicUnifySignature = (...args) => {
const uargs = unify(args) ?? args;
return {
argTypes: uargs,
returnType: uargs[0],
};
};
const unifyRestrictedSignature = (restrict) => (...args) => {
const uargs = unify(args, restrict);
if (!uargs) {
throw new SignatureNotSupportedError(args, restrict);
}
return {
argTypes: uargs,
returnType: uargs[0],
};
};
function variadicReduce(fn) {
return (fst, ...rest) => {
let acc = fst;
for (const r of rest) {
acc = fn(acc, r);
}
return acc;
};
}
function variadicStitch(wrapper) {
return (_ctx, [fst, ...rest]) => {
let acc = stitch `${fst}`;
for (const r of rest) {
acc = stitch `${wrapper}(${acc}, ${r})`;
}
return acc;
};
}
const anyFloatPrimitive = [f32, f16, abstractFloat];
const anyFloatVec = [vec2f, vec3f, vec4f, vec2h, vec3h, vec4h];
const anyFloat = [...anyFloatPrimitive, ...anyFloatVec];
const anyConcreteIntegerPrimitive = [i32, u32];
const anyConcreteIntegerVec = [vec2i, vec3i, vec4i, vec2u, vec3u, vec4u];
const anyConcreteInteger = [...anyConcreteIntegerPrimitive, ...anyConcreteIntegerVec];
function cpuAbs(value) {
return generalizeFn(Math.abs, [value]);
}
export const abs = dualImpl({
name: 'abs',
signature: unaryIdentitySignature,
normalImpl: cpuAbs,
codegenImpl: (_ctx, [value]) => stitch `abs(${value})`,
sideEffects: false,
});
function cpuAcos(value) {
return generalizeFn(Math.acos, [value]);
}
export const acos = dualImpl({
name: 'acos',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuAcos,
codegenImpl: (_ctx, [value]) => stitch `acos(${value})`,
sideEffects: false,
});
function cpuAcosh(value) {
return generalizeFn(Math.acosh, [value]);
}
export const acosh = dualImpl({
name: 'acosh',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuAcosh,
codegenImpl: (_ctx, [value]) => stitch `acosh(${value})`,
sideEffects: false,
});
function cpuAsin(value) {
return generalizeFn(Math.asin, [value]);
}
export const asin = dualImpl({
name: 'asin',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuAsin,
codegenImpl: (_ctx, [value]) => stitch `asin(${value})`,
sideEffects: false,
});
function cpuAsinh(value) {
return generalizeFn(Math.asinh, [value]);
}
export const asinh = dualImpl({
name: 'asinh',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuAsinh,
codegenImpl: (_ctx, [value]) => stitch `asinh(${value})`,
sideEffects: false,
});
function cpuAtan(value) {
return generalizeFn(Math.atan, [value]);
}
export const atan = dualImpl({
name: 'atan',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuAtan,
codegenImpl: (_ctx, [value]) => stitch `atan(${value})`,
sideEffects: false,
});
function cpuAtanh(value) {
return generalizeFn(Math.atanh, [value]);
}
export const atanh = dualImpl({
name: 'atanh',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuAtanh,
codegenImpl: (_ctx, [value]) => stitch `atanh(${value})`,
sideEffects: false,
});
function cpuAtan2(y, x) {
return generalizeFn(Math.atan2, [y, x]);
}
export const atan2 = dualImpl({
name: 'atan2',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuAtan2,
codegenImpl: (_ctx, [y, x]) => stitch `atan2(${y}, ${x})`,
sideEffects: false,
});
function cpuCeil(value) {
return generalizeFn(Math.ceil, [value]);
}
export const ceil = dualImpl({
name: 'ceil',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuCeil,
codegenImpl: (_ctx, [value]) => stitch `ceil(${value})`,
sideEffects: false,
});
function cpuClamp(value, low, high) {
return generalizeFn(clampScalar, [value, low, high]);
}
export const clamp = dualImpl({
name: 'clamp',
signature: variadicUnifySignature,
normalImpl: cpuClamp,
codegenImpl: (_ctx, [value, low, high]) => stitch `clamp(${value}, ${low}, ${high})`,
sideEffects: false,
});
function cpuCos(value) {
return generalizeFn(Math.cos, [value]);
}
export const cos = dualImpl({
name: 'cos',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuCos,
codegenImpl: (_ctx, [value]) => stitch `cos(${value})`,
sideEffects: false,
});
function cpuCosh(value) {
return generalizeFn(Math.cosh, [value]);
}
export const cosh = dualImpl({
name: 'cosh',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuCosh,
codegenImpl: (_ctx, [value]) => stitch `cosh(${value})`,
sideEffects: false,
});
function cpuCountLeadingZeros(_value) {
throw new Error('Unreachable code. The function is only used for the type.');
}
export const countLeadingZeros = dualImpl({
name: 'countLeadingZeros',
signature: unifyRestrictedSignature(anyConcreteInteger),
normalImpl: 'CPU implementation for countLeadingZeros not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [value]) => stitch `countLeadingZeros(${value})`,
sideEffects: false,
});
function cpuCountOneBits(_value) {
throw new Error('Unreachable code. The function is only used for the type.');
}
export const countOneBits = dualImpl({
name: 'countOneBits',
signature: unifyRestrictedSignature(anyConcreteInteger),
normalImpl: 'CPU implementation for countOneBits not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [value]) => stitch `countOneBits(${value})`,
sideEffects: false,
});
function cpuCountTrailingZeros(_value) {
throw new Error('Unreachable code. The function is only used for the type.');
}
export const countTrailingZeros = dualImpl({
name: 'countTrailingZeros',
signature: unifyRestrictedSignature(anyConcreteInteger),
normalImpl: 'CPU implementation for countTrailingZeros not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [value]) => stitch `countTrailingZeros(${value})`,
sideEffects: false,
});
export const cross = dualImpl({
name: 'cross',
signature: unifyRestrictedSignature([vec3f, vec3h]),
normalImpl: (a, b) => VectorOps.cross[a.kind](a, b),
codegenImpl: (_ctx, [a, b]) => stitch `cross(${a}, ${b})`,
sideEffects: false,
});
function cpuDegrees(value) {
if (typeof value === 'number') {
return ((value * 180) / Math.PI);
}
throw new MissingCpuImplError('CPU implementation for degrees on vectors not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues');
}
export const degrees = dualImpl({
name: 'degrees',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuDegrees,
codegenImpl: (_ctx, [value]) => stitch `degrees(${value})`,
sideEffects: false,
});
export const determinant = dualImpl({
name: 'determinant',
signature: (arg) => {
if (!(arg.type === 'mat2x2f' || arg.type === 'mat3x3f' || arg.type === 'mat4x4f')) {
throw new SignatureNotSupportedError([arg], [mat2x2f, mat3x3f, mat4x4f]);
}
return { argTypes: [arg], returnType: f32 };
},
normalImpl: 'CPU implementation for determinant not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [value]) => stitch `determinant(${value})`,
sideEffects: false,
});
function cpuDistance(a, b) {
if (typeof a === 'number' && typeof b === 'number') {
return Math.abs(a - b);
}
return length(sub(a, b));
}
export const distance = dualImpl({
name: 'distance',
signature: (...args) => {
const uargs = unify(args, anyFloat);
if (!uargs) {
throw new SignatureNotSupportedError(args, anyFloat);
}
return {
argTypes: uargs,
returnType: isHalfPrecisionSchema(uargs[0]) ? f16 : f32,
};
},
normalImpl: cpuDistance,
codegenImpl: (_ctx, [a, b]) => stitch `distance(${a}, ${b})`,
sideEffects: false,
});
export const dot = dualImpl({
name: 'dot',
signature: (...args) => ({
argTypes: args,
returnType: args[0].primitive,
}),
normalImpl: (lhs, rhs) => VectorOps.dot[lhs.kind](lhs, rhs),
codegenImpl: (_ctx, [lhs, rhs]) => stitch `dot(${lhs}, ${rhs})`,
sideEffects: false,
});
export const dot4U8Packed = dualImpl({
name: 'dot4U8Packed',
signature: { argTypes: [u32, u32], returnType: u32 },
normalImpl: 'CPU implementation for dot4U8Packed not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [e1, e2]) => stitch `dot4U8Packed(${e1}, ${e2})`,
sideEffects: false,
});
export const dot4I8Packed = dualImpl({
name: 'dot4I8Packed',
signature: { argTypes: [u32, u32], returnType: i32 },
normalImpl: 'CPU implementation for dot4I8Packed not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [e1, e2]) => stitch `dot4I8Packed(${e1}, ${e2})`,
sideEffects: false,
});
function cpuExp(value) {
return generalizeFn(Math.exp, [value]);
}
export const exp = dualImpl({
name: 'exp',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuExp,
codegenImpl: (_ctx, [value]) => stitch `exp(${value})`,
sideEffects: false,
});
function cpuExp2(value) {
return generalizeFn((val) => 2 ** val, [value]);
}
export const exp2 = dualImpl({
name: 'exp2',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuExp2,
codegenImpl: (_ctx, [value]) => stitch `exp2(${value})`,
sideEffects: false,
});
function cpuExtractBits(_e, _offset, _count) {
throw new Error('Unreachable code. The function is only used for the type.');
}
export const extractBits = dualImpl({
name: 'extractBits',
signature: (arg, _offset, _count) => {
const argRestricted = unify([arg], anyConcreteInteger)?.[0];
if (!argRestricted) {
throw new SignatureNotSupportedError([arg], anyConcreteInteger);
}
return {
argTypes: [argRestricted, u32, u32],
returnType: argRestricted,
};
},
normalImpl: 'CPU implementation for extractBits not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [e, offset, count]) => stitch `extractBits(${e}, ${offset}, ${count})`,
sideEffects: false,
});
export const faceForward = dualImpl({
name: 'faceForward',
signature: unifyRestrictedSignature(anyFloatVec),
normalImpl: 'CPU implementation for faceForward not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [e1, e2, e3]) => stitch `faceForward(${e1}, ${e2}, ${e3})`,
sideEffects: false,
});
function cpuFirstLeadingBit(_value) {
throw new Error('Unreachable code. The function is only used for the type.');
}
export const firstLeadingBit = dualImpl({
name: 'firstLeadingBit',
signature: unaryIdentitySignature,
normalImpl: 'CPU implementation for firstLeadingBit not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [value]) => stitch `firstLeadingBit(${value})`,
sideEffects: false,
});
function cpuFirstTrailingBit(_value) {
throw new Error('Unreachable code. The function is only used for the type.');
}
export const firstTrailingBit = dualImpl({
name: 'firstTrailingBit',
signature: unifyRestrictedSignature(anyConcreteInteger),
normalImpl: 'CPU implementation for firstTrailingBit not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [value]) => stitch `firstTrailingBit(${value})`,
sideEffects: false,
});
function cpuFloor(value) {
return generalizeFn(Math.floor, [value]);
}
export const floor = dualImpl({
name: 'floor',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuFloor,
codegenImpl: (_ctx, [arg]) => stitch `floor(${arg})`,
sideEffects: false,
});
function cpuFma(e1, e2, e3) {
if (typeof e1 === 'number') {
return (e1 * e2 + e3);
}
throw new MissingCpuImplError('CPU implementation for fma on vectors not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues');
}
export const fma = dualImpl({
name: 'fma',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuFma,
codegenImpl: (_ctx, [e1, e2, e3]) => stitch `fma(${e1}, ${e2}, ${e3})`,
sideEffects: false,
});
function cpuFract(value) {
return generalizeFn((value) => value - Math.floor(value), [value]);
}
export const fract = dualImpl({
name: 'fract',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuFract,
codegenImpl: (_ctx, [a]) => stitch `fract(${a})`,
sideEffects: false,
});
const FrexpResults = {
f32: abstruct({ fract: f32, exp: i32 }),
f16: abstruct({ fract: f16, exp: i32 }),
abstractFloat: abstruct({ fract: abstractFloat, exp: abstractInt }),
vec2f: abstruct({ fract: vec2f, exp: vec2i }),
vec3f: abstruct({ fract: vec3f, exp: vec3i }),
vec4f: abstruct({ fract: vec4f, exp: vec4i }),
vec2h: abstruct({ fract: vec2h, exp: vec2i }),
vec3h: abstruct({ fract: vec3h, exp: vec3i }),
vec4h: abstruct({ fract: vec4h, exp: vec4i }),
};
export const frexp = dualImpl({
name: 'frexp',
normalImpl: 'CPU implementation for frexp not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
signature: (value) => {
const returnType = FrexpResults[value.type];
if (!returnType) {
throw new SignatureNotSupportedError([value], anyFloat);
}
return { argTypes: [value], returnType };
},
codegenImpl: (_ctx, [value]) => stitch `frexp(${value})`,
sideEffects: false,
});
function cpuInsertBits(_e, _newbits, _offset, _count) {
throw new Error('Unreachable code. The function is only used for the type.');
}
export const insertBits = dualImpl({
name: 'insertBits',
signature: (e, newbits, _offset, _count) => {
const uargs = unify([e, newbits], anyConcreteInteger);
if (!uargs) {
throw new SignatureNotSupportedError([e, newbits], anyConcreteInteger);
}
return {
argTypes: [...uargs, u32, u32],
returnType: uargs[0],
};
},
normalImpl: 'CPU implementation for insertBits not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [e, newbits, offset, count]) => stitch `insertBits(${e}, ${newbits}, ${offset}, ${count})`,
sideEffects: false,
});
function cpuInverseSqrt(value) {
if (typeof value === 'number') {
return (1 / Math.sqrt(value));
}
throw new MissingCpuImplError('CPU implementation for inverseSqrt on vectors not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues');
}
export const inverseSqrt = dualImpl({
name: 'inverseSqrt',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuInverseSqrt,
codegenImpl: (_ctx, [value]) => stitch `inverseSqrt(${value})`,
sideEffects: false,
});
function cpuLdexp(_e1, _e2) {
throw new Error('Unreachable code. The function is only used for the type.');
}
export const ldexp = dualImpl({
name: 'ldexp',
signature: (e1, _e2) => {
switch (e1.type) {
case 'abstractFloat':
return { argTypes: [e1, abstractInt], returnType: e1 };
case 'f32':
case 'f16':
return { argTypes: [e1, i32], returnType: e1 };
case 'vec2f':
case 'vec2h':
return { argTypes: [e1, vec2i], returnType: e1 };
case 'vec3f':
case 'vec3h':
return { argTypes: [e1, vec3i], returnType: e1 };
case 'vec4f':
case 'vec4h':
return { argTypes: [e1, vec4i], returnType: e1 };
default:
throw new Error(`Unsupported data type for ldexp: ${e1.type}. Supported types are abstractFloat, f32, f16, vec2f, vec2h, vec3f, vec3h, vec4f, vec4h.`);
}
},
normalImpl: 'CPU implementation for ldexp not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [e1, e2]) => stitch `ldexp(${e1}, ${e2})`,
sideEffects: false,
});
function cpuLength(value) {
if (typeof value === 'number') {
return Math.abs(value);
}
return VectorOps.length[value.kind](value);
}
export const length = dualImpl({
name: 'length',
signature: (arg) => {
const uarg = unify([arg], anyFloat);
if (!uarg) {
throw new SignatureNotSupportedError([arg], anyFloat);
}
return {
argTypes: uarg,
returnType: isHalfPrecisionSchema(uarg[0]) ? f16 : f32,
};
},
normalImpl: cpuLength,
codegenImpl: (_ctx, [arg]) => stitch `length(${arg})`,
sideEffects: false,
});
function cpuLog(value) {
return generalizeFn(Math.log, [value]);
}
export const log = dualImpl({
name: 'log',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuLog,
codegenImpl: (_ctx, [value]) => stitch `log(${value})`,
sideEffects: false,
});
function cpuLog2(value) {
return generalizeFn(Math.log2, [value]);
}
export const log2 = dualImpl({
name: 'log2',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuLog2,
codegenImpl: (_ctx, [value]) => stitch `log2(${value})`,
sideEffects: false,
});
function cpuMax(a, b) {
return generalizeFn(Math.max, [a, b]);
}
export const max = dualImpl({
name: 'max',
signature: variadicUnifySignature,
normalImpl: variadicReduce(cpuMax),
codegenImpl: variadicStitch('max'),
sideEffects: false,
});
function cpuMin(a, b) {
return generalizeFn(Math.min, [a, b]);
}
export const min = dualImpl({
name: 'min',
signature: variadicUnifySignature,
normalImpl: variadicReduce(cpuMin),
codegenImpl: variadicStitch('min'),
sideEffects: false,
});
function cpuMix(e1, e2, e3) {
return generalizeFn((e1, e2, e3) => e1 * (1 - e3) + e2 * e3, [e1, ...upCast([e2, e3])]);
}
export const mix = dualImpl({
name: 'mix',
signature: (e1, e2, e3) => {
if (e1.type.startsWith('vec') && !e3.type.startsWith('vec')) {
const uarg = unify([e3], [e1.primitive]);
if (!uarg) {
throw new SignatureNotSupportedError([e3], [e1.primitive]);
}
return { argTypes: [e1, e2, uarg[0]], returnType: e1 };
}
const uargs = unify([e1, e2, e3], anyFloat);
if (!uargs) {
throw new SignatureNotSupportedError([e1, e2, e3], anyFloat);
}
return { argTypes: uargs, returnType: uargs[0] };
},
normalImpl: cpuMix,
codegenImpl: (_ctx, [e1, e2, e3]) => stitch `mix(${e1}, ${e2}, ${e3})`,
sideEffects: false,
});
const ModfResult = {
f32: abstruct({ fract: f32, whole: f32 }),
f16: abstruct({ fract: f16, whole: f16 }),
abstractFloat: abstruct({ fract: abstractFloat, whole: abstractFloat }),
vec2f: abstruct({ fract: vec2f, whole: vec2f }),
vec3f: abstruct({ fract: vec3f, whole: vec3f }),
vec4f: abstruct({ fract: vec4f, whole: vec4f }),
vec2h: abstruct({ fract: vec2h, whole: vec2h }),
vec3h: abstruct({ fract: vec3h, whole: vec3h }),
vec4h: abstruct({ fract: vec4h, whole: vec4h }),
};
function cpuModf(_value) {
throw new Error('Unreachable code. The function is only used for the type.');
}
export const modf = dualImpl({
name: 'modf',
signature: (e) => {
const returnType = ModfResult[e.type];
if (!returnType) {
throw new Error(`Unsupported data type for modf: ${e.type}. Supported types are f32, f16, abstractFloat, vec2f, vec3f, vec4f, vec2h, vec3h, vec4h.`);
}
return { argTypes: [e], returnType };
},
normalImpl: 'CPU implementation for modf not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [value]) => stitch `modf(${value})`,
sideEffects: false,
});
export const normalize = dualImpl({
name: 'normalize',
signature: unifyRestrictedSignature(anyFloatVec),
normalImpl: (v) => {
const len = length(v);
return generalizeFn((e) => e / len, [v]);
},
codegenImpl: (_ctx, [value]) => stitch `normalize(${value})`,
sideEffects: false,
});
function powCpu(base, exponent) {
return generalizeFn((a, b) => a ** b, [base, exponent]);
}
export const pow = dualImpl({
name: 'pow',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: powCpu,
codegenImpl: (_ctx, [lhs, rhs]) => stitch `pow(${lhs}, ${rhs})`,
sideEffects: false,
});
function cpuQuantizeToF16(_value) {
throw new Error('Unreachable code. The function is only used for the type.');
}
export const quantizeToF16 = dualImpl({
name: 'quantizeToF16',
signature: (arg) => {
const candidates = [vec2f, vec3f, vec4f, f32];
const uarg = unify([arg], candidates)?.[0];
if (!uarg) {
throw new SignatureNotSupportedError([arg], candidates);
}
return { argTypes: [uarg], returnType: uarg };
},
normalImpl: 'CPU implementation for quantizeToF16 not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [value]) => stitch `quantizeToF16(${value})`,
sideEffects: false,
});
function cpuRadians(value) {
if (typeof value === 'number') {
return ((value * Math.PI) / 180);
}
throw new MissingCpuImplError('CPU implementation for radians on vectors not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues');
}
export const radians = dualImpl({
name: 'radians',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuRadians,
codegenImpl: (_ctx, [value]) => stitch `radians(${value})`,
sideEffects: false,
});
export const reflect = dualImpl({
name: 'reflect',
signature: (...args) => {
const uargs = unify(args, anyFloatVec);
if (!uargs) {
throw new SignatureNotSupportedError(args, anyFloatVec);
}
return {
argTypes: uargs,
returnType: uargs[0],
};
},
normalImpl: (e1, e2) => sub(e1, mul(2 * dot(e2, e1), e2)),
codegenImpl: (_ctx, [e1, e2]) => stitch `reflect(${e1}, ${e2})`,
sideEffects: false,
});
export const refract = dualImpl({
name: 'refract',
normalImpl: 'CPU implementation for refract not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [e1, e2, e3]) => stitch `refract(${e1}, ${e2}, ${e3})`,
signature: (e1, e2, _e3) => ({
argTypes: [e1, e2, isHalfPrecisionSchema(e1) ? f16 : f32],
returnType: e1,
}),
sideEffects: false,
});
function cpuReverseBits(_value) {
throw new Error('Unreachable code. The function is only used for the type.');
}
export const reverseBits = dualImpl({
name: 'reverseBits',
signature: unifyRestrictedSignature(anyConcreteInteger),
normalImpl: 'CPU implementation for reverseBits not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [value]) => stitch `reverseBits(${value})`,
sideEffects: false,
});
function cpuRound(value) {
if (typeof value === 'number') {
const floor = Math.floor(value);
if (value === floor + 0.5) {
if (floor % 2 === 0) {
return floor;
}
return (floor + 1);
}
return Math.round(value);
}
throw new MissingCpuImplError('CPU implementation for round on vectors not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues');
}
export const round = dualImpl({
name: 'round',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuRound,
codegenImpl: (_ctx, [value]) => stitch `round(${value})`,
sideEffects: false,
});
function cpuSaturate(value) {
if (typeof value === 'number') {
return Math.max(0, Math.min(1, value));
}
throw new MissingCpuImplError('CPU implementation for saturate on vectors not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues');
}
export const saturate = dualImpl({
name: 'saturate',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuSaturate,
codegenImpl: (ctx, [value]) => ctx.gen.emitCall('saturate', [], [value]),
sideEffects: false,
});
function cpuSign(e) {
return generalizeFn(Math.sign, [e]);
}
export const sign = dualImpl({
name: 'sign',
signature: (arg) => {
const candidates = [...anyFloat, i32, vec2i, vec3i, vec4i];
const uarg = unify([arg], candidates)?.[0];
if (!uarg) {
throw new SignatureNotSupportedError([arg], candidates);
}
return { argTypes: [uarg], returnType: uarg };
},
normalImpl: cpuSign,
codegenImpl: (_ctx, [e]) => stitch `sign(${e})`,
sideEffects: false,
});
function cpuSin(value) {
return generalizeFn(Math.sin, [value]);
}
export const sin = dualImpl({
name: 'sin',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuSin,
codegenImpl: (_ctx, [value]) => stitch `sin(${value})`,
sideEffects: false,
});
function cpuSinh(value) {
return generalizeFn(Math.sinh, [value]);
}
export const sinh = dualImpl({
name: 'sinh',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuSinh,
codegenImpl: (_ctx, [value]) => stitch `sinh(${value})`,
sideEffects: false,
});
function cpuSmoothstep(edge0, edge1, x) {
return generalizeFn(smoothstepScalar, [edge0, edge1, x]);
}
export const smoothstep = dualImpl({
name: 'smoothstep',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuSmoothstep,
codegenImpl: (_ctx, [edge0, edge1, x]) => stitch `smoothstep(${edge0}, ${edge1}, ${x})`,
sideEffects: false,
});
function cpuSqrt(value) {
return generalizeFn(Math.sqrt, [value]);
}
export const sqrt = dualImpl({
name: 'sqrt',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuSqrt,
codegenImpl: (_ctx, [value]) => stitch `sqrt(${value})`,
sideEffects: false,
});
function cpuStep(edge, x) {
if (typeof edge === 'number') {
return (edge <= x ? 1.0 : 0.0);
}
throw new MissingCpuImplError('CPU implementation for step on vectors not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues');
}
export const step = dualImpl({
name: 'step',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuStep,
codegenImpl: (_ctx, [edge, x]) => stitch `step(${edge}, ${x})`,
sideEffects: false,
});
function cpuTan(value) {
if (typeof value === 'number') {
return Math.tan(value);
}
throw new MissingCpuImplError('CPU implementation for tan on vectors not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues');
}
export const tan = dualImpl({
name: 'tan',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuTan,
codegenImpl: (_ctx, [value]) => stitch `tan(${value})`,
sideEffects: false,
});
function cpuTanh(value) {
return generalizeFn(Math.tanh, [value]);
}
export const tanh = dualImpl({
name: 'tanh',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: cpuTanh,
codegenImpl: (_ctx, [value]) => stitch `tanh(${value})`,
sideEffects: false,
});
function cpuTranspose(value) {
const schema = WORKAROUND_getSchema(value);
// NOTE: This assumes all matrices are square
const transposed = schema();
const src = value.columns;
const dst = transposed.columns;
if (src.length === 2) {
dst[0][0] = src[0][0];
dst[0][1] = src[1][0];
dst[1][0] = src[0][1];
dst[1][1] = src[1][1];
}
else if (src.length === 3) {
dst[0][0] = src[0][0];
dst[0][1] = src[1][0];
dst[0][2] = src[2][0];
dst[1][0] = src[0][1];
dst[1][1] = src[1][1];
dst[1][2] = src[2][1];
// oxlint-disable-next-line typescript/no-non-null-assertion
const dst2 = dst[2];
dst2[0] = src[0][2];
dst2[1] = src[1][2];
dst2[2] = src[2][2];
}
else if (src.length === 4) {
dst[0][0] = src[0][0];
dst[0][1] = src[1][0];
dst[0][2] = src[2][0];
dst[0][3] = src[3][0];
dst[1][0] = src[0][1];
dst[1][1] = src[1][1];
dst[1][2] = src[2][1];
dst[1][3] = src[3][1];
// oxlint-disable-next-line typescript/no-non-null-assertion
const dst2 = dst[2];
dst2[0] = src[0][2];
dst2[1] = src[1][2];
dst2[2] = src[2][2];
dst2[3] = src[3][2];
// oxlint-disable-next-line typescript/no-non-null-assertion
const dst3 = dst[3];
dst3[0] = src[0][3];
dst3[1] = src[1][3];
dst3[2] = src[2][3];
dst3[3] = src[3][3];
}
else {
assertExhaustive(src, 'std/numeric.ts#cpuTranspose');
}
return transposed;
}
export const transpose = dualImpl({
name: 'transpose',
signature: unaryIdentitySignature,
normalImpl: cpuTranspose,
codegenImpl: (_ctx, [e]) => stitch `transpose(${e})`,
sideEffects: false,
});
function cpuTrunc(_value) {
throw new Error('Unreachable code. The function is only used for the type.');
}
export const trunc = dualImpl({
name: 'trunc',
signature: unifyRestrictedSignature(anyFloat),
normalImpl: 'CPU implementation for trunc not implemented yet. Please submit an issue at https://github.com/software-mansion/TypeGPU/issues',
codegenImpl: (_ctx, [value]) => stitch `trunc(${value})`,
sideEffects: false,
});
function cpuIntdiv(lhs, rhs) {
if (typeof lhs !== 'number' || typeof rhs !== 'number') {
throw new Error('std.intdiv called with invalid arguments.');
}
return Math.trunc(Math.trunc(lhs) / Math.trunc(rhs));
}
/**
* Performs integer division on the passed in scalars.
* Equivalent to `trunc(trunc(lhs) / trunc(rhs))`. Coerces both
* arguments to integers if they're floating point.
*/
export const intdiv = dualImpl({
name: 'intdiv',
signature: (lhs, rhs) => {
const unified = unify([lhs, rhs], [u32, i32]);
if (!unified) {
throw new SignatureNotSupportedError([lhs, rhs], [u32, i32, abstractInt]);
}
return { argTypes: unified, returnType: unified[0] };
},
normalImpl: cpuIntdiv,
codegenImpl: (ctx, [lhs, rhs]) => ctx.gen.emitBinaryOp(lhs, '/', rhs),
sideEffects: false,
});