typegpu
Version:
A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.
231 lines (230 loc) • 9.12 kB
JavaScript
import { dualImpl } from "../core/function/dualImpl.js";
import { stitch } from "../core/resolve/stitch.js";
import { abstractFloat, f16, f32, u32 } from "../data/numeric.js";
import { vec2i, vec2u, vec3i, vec3u, vec4i, vec4u } from "../data/vector.js";
import { VectorOps } from "../data/vectorOps.js";
import { generalizeFn, upCast } from "../data/generalizeFn.js";
import { isFloat32VecInstance, isInteger32VecInstance, isMat, isMatInstance, isUint32VecInstance, isVec, isVecInstance, } from "../data/wgslTypes.js";
import { SignatureNotSupportedError } from "../errors.js";
import { unify } from "../tgsl/conversion.js";
const getPrimitive = (t) => ('primitive' in t ? t.primitive : t);
const makeBinarySignature = (opts) => (lhs, rhs) => {
const { restrict } = opts ?? {};
const fail = (msg) => {
if (restrict) {
throw new SignatureNotSupportedError([lhs, rhs], restrict);
}
throw new Error(`Cannot apply operator to ${lhs.type} and ${rhs.type}: ${msg}`);
};
if (opts?.noMat && (isMat(lhs) || isMat(rhs))) {
return fail('matrices not supported');
}
const lhsC = isVec(lhs) || isMat(lhs);
const rhsC = isVec(rhs) || isMat(rhs);
if (!lhsC && !rhsC) {
// scalar × scalar
const unified = unify([lhs, rhs], restrict);
if (!unified)
return fail('incompatible scalar types');
return { argTypes: unified, returnType: unified[0] };
}
if (lhsC && rhsC) {
// vec × mat or mat × vec
if (opts?.matVecProduct && isVec(lhs) !== isVec(rhs)) {
return { argTypes: [lhs, rhs], returnType: isVec(lhs) ? lhs : rhs };
}
// composite × composite (same kind)
if (lhs.type !== rhs.type)
return fail('operands must have the same type');
return { argTypes: [lhs, rhs], returnType: lhs };
}
// scalar × composite
const [scalar, composite] = lhsC ? [rhs, lhs] : [lhs, rhs];
const unified = unify([scalar], [getPrimitive(composite)]);
if (!unified) {
return fail(`scalar not convertible to ${getPrimitive(composite).type}`);
}
return {
argTypes: lhsC ? [lhs, unified[0]] : [unified[0], rhs],
returnType: composite,
};
};
const binaryArithmeticSignature = makeBinarySignature();
const binaryMulSignature = makeBinarySignature({ matVecProduct: true });
const binaryDivSignature = makeBinarySignature({
noMat: true,
restrict: [f32, f16, abstractFloat],
});
function cpuAdd(lhs, rhs) {
if (typeof lhs === 'number' && typeof rhs === 'number') {
return lhs + rhs; // default addition
}
if (typeof lhs === 'number' && isVecInstance(rhs)) {
return generalizeFn((e) => lhs + e, [rhs]); // mixed addition
}
if (isVecInstance(lhs) && typeof rhs === 'number') {
return generalizeFn((e) => e + rhs, [lhs]); // mixed addition
}
if ((isVecInstance(lhs) && isVecInstance(rhs)) || (isMatInstance(lhs) && isMatInstance(rhs))) {
return generalizeFn((a, b) => a + b, [lhs, rhs]); // component-wise addition
}
throw new Error('Add/Sub called with invalid arguments.');
}
export const add = dualImpl({
name: 'add',
signature: binaryArithmeticSignature,
normalImpl: cpuAdd,
codegenImpl: (ctx, [lhs, rhs]) => ctx.gen.emitBinaryOp(lhs, '+', rhs),
sideEffects: false,
});
function cpuSub(lhs, rhs) {
// while illegal on the wgsl side, we can do this in js
return cpuAdd(lhs, cpuMul(-1, rhs));
}
export const sub = dualImpl({
name: 'sub',
signature: binaryArithmeticSignature,
normalImpl: cpuSub,
codegenImpl: (ctx, [lhs, rhs]) => ctx.gen.emitBinaryOp(lhs, '-', rhs),
sideEffects: false,
});
function cpuMul(lhs, rhs) {
if (typeof lhs === 'number' && typeof rhs === 'number') {
return lhs * rhs; // default multiplication
}
if (typeof lhs === 'number' && (isVecInstance(rhs) || isMatInstance(rhs))) {
return generalizeFn((e) => lhs * e, [rhs]); // scale
}
if ((isVecInstance(lhs) || isMatInstance(lhs)) && typeof rhs === 'number') {
return generalizeFn((e) => e * rhs, [lhs]); // scale
}
if (isVecInstance(lhs) && isVecInstance(rhs)) {
return generalizeFn((a, b) => a * b, [lhs, rhs]); // component-wise
}
if (isFloat32VecInstance(lhs) && isMatInstance(rhs)) {
return VectorOps.mulVxM[rhs.kind](lhs, rhs); // row-vector-matrix
}
if (isMatInstance(lhs) && isFloat32VecInstance(rhs)) {
return VectorOps.mulMxV[lhs.kind](lhs, rhs); // matrix-column-vector
}
if (isMatInstance(lhs) && isMatInstance(rhs)) {
return VectorOps.mulMxM[lhs.kind](lhs, rhs); // matrix multiplication
}
throw new Error('Mul called with invalid arguments.');
}
export const mul = dualImpl({
name: 'mul',
signature: binaryMulSignature,
normalImpl: cpuMul,
codegenImpl: (ctx, [lhs, rhs]) => ctx.gen.emitBinaryOp(lhs, '*', rhs),
sideEffects: false,
});
function cpuDiv(lhs, rhs) {
return generalizeFn((a, b) => a / b, upCast([lhs, rhs]));
}
export const div = dualImpl({
name: 'div',
signature: binaryDivSignature,
normalImpl: cpuDiv,
codegenImpl: (ctx, [lhs, rhs]) => ctx.gen.emitBinaryOp(lhs, '/', rhs),
ignoreImplicitCastWarning: true,
sideEffects: false,
});
/**
* @privateRemarks
* Both JS and WGSL implementations use truncated definition of modulo
*/
export const mod = dualImpl({
name: 'mod',
signature: binaryDivSignature,
normalImpl: ((a, b) => {
return generalizeFn((a, b) => a % b, upCast([a, b]));
}),
codegenImpl: (ctx, [lhs, rhs]) => ctx.gen.emitBinaryOp(lhs, '%', rhs),
sideEffects: false,
});
function cpuNeg(value) {
return generalizeFn((value) => -value, [value]);
}
export const neg = dualImpl({
name: 'neg',
signature: (arg) => ({
argTypes: [arg],
returnType: arg,
}),
normalImpl: cpuNeg,
codegenImpl: (_ctx, [arg]) => stitch `-(${arg})`,
sideEffects: false,
});
const intVecToUnsignedVec = {
vec2i: vec2u,
vec2u: vec2u,
vec3i: vec3u,
vec3u: vec3u,
vec4i: vec4u,
vec4u: vec4u,
};
const bitShiftSignature = (lhs, rhs) => {
const lhsUnified = unify([lhs], [vec2i, vec3i, vec4i, vec2u, vec3u, vec4u])?.[0];
if (!lhsUnified || !isVec(lhsUnified)) {
throw new SignatureNotSupportedError([lhs], [vec2i, vec3i, vec4i, vec2u, vec3u, vec4u]);
}
const cc = lhsUnified.componentCount;
const vecU = cc === 2 ? vec2u : cc === 3 ? vec3u : vec4u;
const rhsUnified = unify([rhs], [u32, vecU])?.[0];
if (!rhsUnified) {
throw new SignatureNotSupportedError([rhs], [u32, vecU]);
}
return {
argTypes: [lhsUnified, rhsUnified],
returnType: lhsUnified,
};
};
function cpuBitShiftLeft(lhs, rhs) {
if (isInteger32VecInstance(lhs) && isUint32VecInstance(rhs) && lhs.length == rhs.length) {
return VectorOps.bitShiftLeft[lhs.kind](lhs, rhs);
}
if (isInteger32VecInstance(lhs) && typeof rhs === 'number') {
const rhsVec = intVecToUnsignedVec[lhs.kind](rhs);
return VectorOps.bitShiftLeft[lhs.kind](lhs, rhsVec);
}
throw new Error("'bitShiftLeft' called with invalid arguments, expected: left-hand side to be an integer vector, right-hand side to be a number or unsigned integer vector of the same arity as the left-hand side.");
}
export const bitShiftLeft = dualImpl({
name: 'bitShiftLeft',
signature: bitShiftSignature,
normalImpl: cpuBitShiftLeft,
codegenImpl: (ctx, [lhs, rhs]) => {
if (isVec(lhs.dataType) && !isVec(rhs.dataType)) {
const cc = lhs.dataType.componentCount;
const schema = cc === 2 ? vec2u : cc === 3 ? vec3u : vec4u;
return ctx.gen.emitBinaryOp(lhs, '<<', ctx.gen.typeInstantiation(schema, [rhs]));
}
return ctx.gen.emitBinaryOp(lhs, '<<', rhs);
},
sideEffects: false,
});
function cpuBitShiftRight(lhs, rhs) {
if (isInteger32VecInstance(lhs) && isUint32VecInstance(rhs) && lhs.length == rhs.length) {
return VectorOps.bitShiftRight[lhs.kind](lhs, rhs);
}
if (isInteger32VecInstance(lhs) && typeof rhs === 'number') {
const rhsVec = intVecToUnsignedVec[lhs.kind](rhs);
return VectorOps.bitShiftRight[lhs.kind](lhs, rhsVec);
}
throw new Error("'bitShiftRight' called with invalid arguments, expected: left-hand side to be an integer vector, right-hand side to be a number or unsigned integer vector of the same arity as the left-hand side.");
}
export const bitShiftRight = dualImpl({
name: 'bitShiftRight',
signature: bitShiftSignature,
normalImpl: cpuBitShiftRight,
codegenImpl: (ctx, [lhs, rhs]) => {
if (isVec(lhs.dataType) && !isVec(rhs.dataType)) {
const cc = lhs.dataType.componentCount;
const schema = cc === 2 ? vec2u : cc === 3 ? vec3u : vec4u;
return ctx.gen.emitBinaryOp(lhs, '>>', ctx.gen.typeInstantiation(schema, [rhs]));
}
return ctx.gen.emitBinaryOp(lhs, '>>', rhs);
},
sideEffects: false,
});