typegpu
Version:
A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.
337 lines (336 loc) • 12.4 kB
JavaScript
import { dualImpl } from "../core/function/dualImpl.js";
import { stitch } from "../core/resolve/stitch.js";
import { bool, f16, f32, i32, u32 } from "../data/numeric.js";
import { isSnippetNumeric, snip } from "../data/snippet.js";
import { vec2b, vec2f, vec2h, vec2i, vec2u, vec3b, vec3f, vec3h, vec3i, vec3u, vec4b, vec4f, vec4h, vec4i, vec4u, } from "../data/vector.js";
import { VectorOps } from "../data/vectorOps.js";
import { generalizeBoolFn, generalizeFn } from "../data/generalizeFn.js";
import { isBool, isVecBool, isVecBoolInstance, } from "../data/wgslTypes.js";
import { SignatureNotSupportedError } from "../errors.js";
import { unify } from "../tgsl/conversion.js";
import { cpuCopy } from "./copy.js";
function correspondingBooleanVectorSchema(dataType) {
if (dataType.type.includes('2')) {
return vec2b;
}
if (dataType.type.includes('3')) {
return vec3b;
}
return vec4b;
}
// comparison
/**
* Checks whether `lhs == rhs` on all components.
* Equivalent to `all(eq(lhs, rhs))`.
* @example
* allEq(vec2f(0.0, 1.0), vec2f(0.0, 2.0)) // returns false
* allEq(vec3u(0, 1, 2), vec3u(0, 1, 2)) // returns true
*/
export const allEq = dualImpl({
name: 'allEq',
signature: (...argTypes) => ({ argTypes, returnType: bool }),
normalImpl: (lhs, rhs) => cpuAll(cpuEq(lhs, rhs)),
codegenImpl: (_ctx, [lhs, rhs]) => stitch `all(${lhs} == ${rhs})`,
sideEffects: false,
});
const cpuEq = (lhs, rhs) => generalizeBoolFn((a, b) => a === b, [lhs, rhs]);
/**
* Checks **component-wise** whether `lhs == rhs`.
* This function does **not** return `bool`, for that use-case, wrap the result in `all`, or use `allEq`.
* @example
* eq(vec2f(0.0, 1.0), vec2f(0.0, 2.0)) // returns vec2b(true, false)
* eq(vec3u(0, 1, 2), vec3u(2, 1, 0)) // returns vec3b(false, true, false)
* all(eq(vec4i(4, 3, 2, 1), vec4i(4, 3, 2, 1))) // returns true
* allEq(vec4i(4, 3, 2, 1), vec4i(4, 3, 2, 1)) // returns true
*/
export const eq = dualImpl({
name: 'eq',
signature: (...argTypes) => ({
argTypes,
returnType: correspondingBooleanVectorSchema(argTypes[0]),
}),
normalImpl: cpuEq,
codegenImpl: (_ctx, [lhs, rhs]) => stitch `(${lhs} == ${rhs})`,
sideEffects: false,
});
/**
* Checks **component-wise** whether `lhs != rhs`.
* This function does **not** return `bool`, for that use-case, wrap the result in `any`.
* @example
* ne(vec2f(0.0, 1.0), vec2f(0.0, 2.0)) // returns vec2b(false, true)
* ne(vec3u(0, 1, 2), vec3u(2, 1, 0)) // returns vec3b(true, false, true)
* any(ne(vec4i(4, 3, 2, 1), vec4i(4, 2, 2, 1))) // returns true
*/
export const ne = dualImpl({
name: 'ne',
signature: (...argTypes) => ({
argTypes,
returnType: correspondingBooleanVectorSchema(argTypes[0]),
}),
normalImpl: (lhs, rhs) => cpuNot(cpuEq(lhs, rhs)),
codegenImpl: (_ctx, [lhs, rhs]) => stitch `(${lhs} != ${rhs})`,
sideEffects: false,
});
const cpuLt = (lhs, rhs) => generalizeBoolFn((a, b) => a < b, [lhs, rhs]);
/**
* Checks **component-wise** whether `lhs < rhs`.
* This function does **not** return `bool`, for that use-case, wrap the result in `all`.
* @example
* lt(vec2f(0.0, 0.0), vec2f(0.0, 1.0)) // returns vec2b(false, true)
* lt(vec3u(0, 1, 2), vec3u(2, 1, 0)) // returns vec3b(true, false, false)
* all(lt(vec4i(1, 2, 3, 4), vec4i(2, 3, 4, 5))) // returns true
*/
export const lt = dualImpl({
name: 'lt',
signature: (...argTypes) => ({
argTypes,
returnType: correspondingBooleanVectorSchema(argTypes[0]),
}),
normalImpl: cpuLt,
codegenImpl: (_ctx, [lhs, rhs]) => stitch `(${lhs} < ${rhs})`,
sideEffects: false,
});
/**
* Checks **component-wise** whether `lhs <= rhs`.
* This function does **not** return `bool`, for that use-case, wrap the result in `all`.
* @example
* le(vec2f(0.0, 0.0), vec2f(0.0, 1.0)) // returns vec2b(true, true)
* le(vec3u(0, 1, 2), vec3u(2, 1, 0)) // returns vec3b(true, true, false)
* all(le(vec4i(1, 2, 3, 4), vec4i(2, 3, 3, 5))) // returns true
*/
export const le = dualImpl({
name: 'le',
signature: (...argTypes) => ({
argTypes,
returnType: correspondingBooleanVectorSchema(argTypes[0]),
}),
normalImpl: (lhs, rhs) => cpuOr(cpuLt(lhs, rhs), cpuEq(lhs, rhs)),
codegenImpl: (_ctx, [lhs, rhs]) => stitch `(${lhs} <= ${rhs})`,
sideEffects: false,
});
/**
* Checks **component-wise** whether `lhs > rhs`.
* This function does **not** return `bool`, for that use-case, wrap the result in `all`.
* @example
* gt(vec2f(0.0, 0.0), vec2f(0.0, 1.0)) // returns vec2b(false, false)
* gt(vec3u(0, 1, 2), vec3u(2, 1, 0)) // returns vec3b(false, false, true)
* all(gt(vec4i(2, 3, 4, 5), vec4i(1, 2, 3, 4))) // returns true
*/
export const gt = dualImpl({
name: 'gt',
signature: (...argTypes) => ({
argTypes,
returnType: correspondingBooleanVectorSchema(argTypes[0]),
}),
normalImpl: (lhs, rhs) => cpuAnd(cpuNot(cpuLt(lhs, rhs)), cpuNot(cpuEq(lhs, rhs))),
codegenImpl: (_ctx, [lhs, rhs]) => stitch `(${lhs} > ${rhs})`,
sideEffects: false,
});
/**
* Checks **component-wise** whether `lhs >= rhs`.
* This function does **not** return `bool`, for that use-case, wrap the result in `all`.
* @example
* ge(vec2f(0.0, 0.0), vec2f(0.0, 1.0)) // returns vec2b(true, false)
* ge(vec3u(0, 1, 2), vec3u(2, 1, 0)) // returns vec3b(false, true, true)
* all(ge(vec4i(2, 2, 4, 5), vec4i(1, 2, 3, 4))) // returns true
*/
export const ge = dualImpl({
name: 'ge',
signature: (...argTypes) => ({
argTypes: argTypes,
returnType: correspondingBooleanVectorSchema(argTypes[0]),
}),
normalImpl: (lhs, rhs) => cpuNot(cpuLt(lhs, rhs)),
codegenImpl: (_ctx, [lhs, rhs]) => stitch `(${lhs} >= ${rhs})`,
sideEffects: false,
});
function cpuNot(value) {
if (typeof value === 'boolean') {
return !value;
}
if (!isVecBoolInstance(value)) {
throw new Error(`'std.not' requires a boolean or boolean vector.`);
}
switch (value.length) {
case 2:
return vec2b(cpuNot(value.x), cpuNot(value.y));
case 3:
return vec3b(cpuNot(value.x), cpuNot(value.y), cpuNot(value.z));
case 4:
return vec4b(cpuNot(value.x), cpuNot(value.y), cpuNot(value.z), cpuNot(value.w));
}
}
/**
* Returns the logical negation of the given value.
* For booleans returns `!value`.
* For boolean vectors, returns **component-wise** `!value`.
* @example
* not(true) // returns false
* not(vec3b(true, true, false)) // returns vec3b(false, false, true)
*/
export const not = dualImpl({
name: 'not',
signature: (arg) => {
if (!isBool(arg) && !isVecBool(arg)) {
throw new SignatureNotSupportedError([arg], [bool, vec2b, vec3b, vec4b]);
}
return {
argTypes: [arg],
returnType: arg,
};
},
normalImpl: cpuNot,
codegenImpl: (_ctx, [arg]) => stitch `!(${arg})`,
sideEffects: false,
});
const cpuOr = (lhs, rhs) => generalizeBoolFn((a, b) => a || b, [lhs, rhs]);
/**
* Returns **component-wise** logical `or` result.
* @example
* or(vec2b(false, true), vec2b(false, false)) // returns vec2b(false, true)
* or(vec3b(true, true, false), vec3b(false, true, false)) // returns vec3b(true, true, false)
*/
export const or = dualImpl({
name: 'or',
signature: (...argTypes) => ({ argTypes, returnType: argTypes[0] }),
normalImpl: cpuOr,
codegenImpl: (_ctx, [lhs, rhs]) => stitch `(${lhs} | ${rhs})`,
sideEffects: false,
});
const cpuAnd = (lhs, rhs) => cpuNot(cpuOr(cpuNot(lhs), cpuNot(rhs)));
/**
* Returns **component-wise** logical `and` result.
* @example
* and(vec2b(false, true), vec2b(true, true)) // returns vec2b(false, true)
* and(vec3b(true, true, false), vec3b(false, true, false)) // returns vec3b(false, true, false)
*/
export const and = dualImpl({
name: 'and',
signature: (...argTypes) => ({ argTypes, returnType: argTypes[0] }),
normalImpl: cpuAnd,
codegenImpl: (_ctx, [lhs, rhs]) => stitch `(${lhs} & ${rhs})`,
sideEffects: false,
});
// logical aggregation
const cpuAll = (value) => VectorOps.all[value.kind](value);
/**
* Returns `true` if each component of `value` is true.
* @example
* all(vec2b(false, true)) // returns false
* all(vec3b(true, true, true)) // returns true
*/
export const all = dualImpl({
name: 'all',
signature: (...argTypes) => ({ argTypes, returnType: bool }),
normalImpl: cpuAll,
codegenImpl: (_ctx, [value]) => stitch `all(${value})`,
sideEffects: false,
});
/**
* Returns `true` if any component of `value` is true.
* @example
* any(vec2b(false, true)) // returns true
* any(vec3b(false, false, false)) // returns false
*/
export const any = dualImpl({
name: 'any',
signature: (...argTypes) => ({ argTypes, returnType: bool }),
normalImpl: (value) => !cpuAll(cpuNot(value)),
codegenImpl: (_ctx, [arg]) => stitch `any(${arg})`,
sideEffects: false,
});
// other
/**
* Checks whether the given elements differ by at most the `precision` value.
* Checks all elements of `lhs` and `rhs` if arguments are vectors.
* @example
* isCloseTo(0, 0.1) // returns false
* isCloseTo(vec3f(0, 0, 0), vec3f(0.002, -0.009, 0)) // returns true
*
* @param {number} precision argument that specifies the maximum allowed difference, 0.01 by default.
*/
export const isCloseTo = dualImpl({
name: 'isCloseTo',
signature: (...args) => ({
argTypes: args,
returnType: bool,
}),
// CPU implementation
normalImpl: (lhs, rhs, precision = 0.01) => {
const componentResult = generalizeBoolFn((lhs, rhs) => Math.abs(lhs - rhs) < precision, [lhs, rhs]);
return typeof componentResult === 'boolean' ? componentResult : all(componentResult);
},
// GPU implementation
codegenImpl: (_ctx, [lhs, rhs, precision = snip(0.01, f32, /* origin */ 'constant', false)]) => {
if (isSnippetNumeric(lhs) && isSnippetNumeric(rhs)) {
return stitch `(abs(f32(${lhs}) - f32(${rhs})) <= ${precision})`;
}
if (!isSnippetNumeric(lhs) && !isSnippetNumeric(rhs)) {
// https://www.w3.org/TR/WGSL/#vector-multi-component:~:text=Binary%20arithmetic%20expressions%20with%20mixed%20scalar%20and%20vector%20operands
// (a-a)+prec creates a vector of a.length elements, all equal to prec
return stitch `all(abs(${lhs} - ${rhs}) <= (${lhs} - ${lhs}) + ${precision})`;
}
return 'false';
},
sideEffects: false,
});
function cpuSelect(f, t, cond) {
if (typeof cond === 'boolean') {
return cpuCopy(cond ? t : f);
}
// generalizeFn will handle this fine, it just has no mixed type overload.
return generalizeFn((f, t, c) => (c ? t : f), [f, t, cond]);
}
export const validSelectBranchTypes = [
f32,
f16,
i32,
u32,
bool,
vec2f,
vec3f,
vec4f,
vec2h,
vec3h,
vec4h,
vec2i,
vec3i,
vec4i,
vec2u,
vec3u,
vec4u,
vec2b,
vec3b,
vec4b,
];
/**
* Returns `t` if `cond` is `true`, and `f` otherwise.
* Component-wise if `cond` is a vector.
* @example
* select(1, 2, false) // returns 1
* select(1, 2, true) // returns 2
* select(vec2i(1, 2), vec2i(3, 4), true) // returns vec2i(3, 4)
* select(vec2i(1, 2), vec2i(3, 4), vec2b(false, true)) // returns vec2i(1, 4)
*/
export const select = dualImpl({
name: 'select',
signature: (f, t, cond) => {
const [uf, ut] = unify([f, t], validSelectBranchTypes) ?? [f, t];
return { argTypes: [uf, ut, cond], returnType: uf };
},
normalImpl: cpuSelect,
codegenImpl: (ctx, [f, t, cond]) => {
const result = ctx.gen.emitCall('select', [], [f, t, cond]);
if (!validSelectBranchTypes.includes(f.dataType) ||
!validSelectBranchTypes.includes(t.dataType)) {
throw new Error(`'${result}' is invalid, std.select requires both branches to be either scalars or vectors.`);
}
if (f.dataType !== t.dataType) {
const fStr = ctx.resolve(f.dataType);
const tStr = ctx.resolve(t.dataType);
throw new Error(`'${result}' is invalid, std.select requires both branches to be the same type, got [${fStr.value}, ${tStr.value}].`);
}
return result;
},
sideEffects: false,
});