typegpu
Version:
A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.
321 lines (320 loc) • 10.4 kB
JavaScript
import { $internal } from "../shared/symbols.js";
import { isBool, } from "./wgslTypes.js";
import { callableSchema } from "../core/function/createCallableSchema.js";
import { FiniteMathAssumptionError, SignatureNotSupportedError, WgslTypeError } from "../errors.js";
import { isSnippetNumeric } from "./snippet.js";
import { UnknownData } from "./dataTypes.js";
const boolCast = callableSchema({
name: 'bool',
schema: () => bool,
argTypes: (arg) => (arg ? [arg] : []),
normalImpl(v = false) {
if (typeof v === 'boolean') {
return v;
}
if (typeof v === 'number') {
if (!Number.isFinite(v)) {
throw new FiniteMathAssumptionError(v, bool);
}
return Boolean(v);
}
throw new Error(`Invalid argument type for 'd.bool'. Got '${typeof v}', expected number or boolean`);
},
codegenImpl: (ctx, [v]) => {
// check if zero arguments were passed
if (v === undefined) {
return ctx.gen.typeInstantiation(bool, []);
}
if (isBool(v.dataType) || isSnippetNumeric(v)) {
return ctx.gen.typeInstantiation(bool, [v]);
}
if (v.dataType === UnknownData) {
throw new WgslTypeError("Unknown argument type for 'd.bool'.");
}
throw new SignatureNotSupportedError([v.dataType], [bool, u32, i32, f32, f16]);
},
});
/**
* A schema that represents a boolean value. (equivalent to `bool` in WGSL)
*
* Can also be called to cast a value to a bool in accordance with WGSL casting rules.
*
* @example
* const value = bool(); // false
* @example
* const value = bool(0); // false
* @example
* const value = bool(-0); // false
* @example
* const value = bool(21.37); // true
*/
export const bool = Object.assign(boolCast, {
[$internal]: {},
type: 'bool',
});
const u32Cast = callableSchema({
name: 'u32',
schema: () => u32,
argTypes: (arg) => (arg ? [arg] : []),
normalImpl(v) {
if (v === undefined) {
return 0;
}
if (typeof v === 'boolean') {
return v ? 1 : 0;
}
if (!Number.isFinite(v)) {
throw new FiniteMathAssumptionError(v, u32);
}
if (!Number.isInteger(v)) {
const truncated = Math.trunc(v);
if (truncated < 0) {
return 0;
}
if (truncated > 0xffffffff) {
return 0xffffffff;
}
return truncated;
}
// Integer input: treat as bit reinterpretation (i32 -> u32)
return (v & 0xffffffff) >>> 0;
},
codegenImpl: (ctx, [v]) => ctx.gen.typeInstantiation(u32, v ? [v] : []),
});
/**
* A schema that represents an unsigned 32-bit integer value. (equivalent to `u32` in WGSL)
*
* Can also be called to cast a value to an u32 in accordance with WGSL casting rules.
*
* @example
* const value = u32(); // 0
* @example
* const value = u32(7); // 7
* @example
* const value = u32(3.14); // 3
* @example
* const value = u32(-1); // 4294967295
* @example
* const value = u32(-3.1); // 0
*/
export const u32 = Object.assign(u32Cast, {
[$internal]: {},
type: 'u32',
});
const i32Cast = callableSchema({
name: 'i32',
schema: () => i32,
argTypes: (arg) => (arg ? [arg] : []),
normalImpl(v) {
if (v === undefined) {
return 0;
}
if (typeof v === 'boolean') {
return v ? 1 : 0;
}
if (!Number.isFinite(v)) {
throw new FiniteMathAssumptionError(v, i32);
}
return v | 0;
},
codegenImpl: (ctx, [v]) => ctx.gen.typeInstantiation(i32, v ? [v] : []),
});
export const u16 = {
[$internal]: {},
type: 'u16',
};
/**
* A schema that represents a signed 32-bit integer value. (equivalent to `i32` in WGSL)
*
* Can also be called to cast a value to an i32 in accordance with WGSL casting rules.
*
* @example
* const value = i32(); // 0
* @example
* const value = i32(3.14); // 3
* @example
* const value = i32(-3.9); // -3
* @example
* const value = i32(10000000000) // 1410065408
*/
export const i32 = Object.assign(i32Cast, {
[$internal]: {},
type: 'i32',
});
const f32Cast = callableSchema({
name: 'f32',
schema: () => f32,
argTypes: (arg) => (arg ? [arg] : []),
normalImpl(v) {
if (v === undefined) {
return 0;
}
if (typeof v === 'boolean') {
return v ? 1 : 0;
}
if (!Number.isFinite(v)) {
throw new FiniteMathAssumptionError(v, f32);
}
return Math.fround(v);
},
codegenImpl: (ctx, [v]) => ctx.gen.typeInstantiation(f32, v ? [v] : []),
});
/**
* A schema that represents a 32-bit float value. (equivalent to `f32` in WGSL)
*
* Can also be called to cast a value to an f32.
*
* @example
* const value = f32(); // 0
* @example
* const value = f32(1.23); // 1.23
* @example
* const value = f32(true); // 1
*/
export const f32 = Object.assign(f32Cast, {
[$internal]: {},
type: 'f32',
});
// helpers for floating point conversion
const buf32 = new ArrayBuffer(4);
const f32arr = new Float32Array(buf32);
const u32arr = new Uint32Array(buf32);
/**
* Convert a JavaScript number (treated as float32) to **binary16** bit pattern.
* @param x 32-bit floating-point value
* @returns 16-bit half-precision encoding (stored in a JS number)
*/
export function toHalfBits(x) {
f32arr[0] = x; // Write value; shared buffer now contains raw bits.
const bits = u32arr[0]; // Read those bits as unsigned int.
// 1. Extract sign, exponent, and mantissa from the 32‑bit layout.
const sign = (bits >>> 31) & 0x1; // Bit 31 is the sign.
let exp = (bits >>> 23) & 0xff; // Bits 30‑23 form the biased exponent.
const mant = bits & 0x7fffff; // Bits 22‑0 are the significand.
// 2. Handle special values (NaN, ±∞) before re‑biasing.
if (exp === 0xff) {
// Preserve the quiet‑NaN bit if mant≠0; otherwise this is ±∞.
return (sign << 15) | 0x7c00 | (mant ? 0x0200 : 0);
}
// 3. Re‑bias the exponent from 127 → 15 (binary32 → binary16).
exp = exp - 127 + 15;
// 4. Underflow: exponent ≤ 0 yields sub‑normals or signed zero.
if (exp <= 0) {
// Below the smallest representable subnormal magnitude, round to ±0.
if (exp < -10) {
return sign << 15;
}
// Produce a sub‑normal: prepend the hidden 1, then round to nearest,
// ties to even. `shift` is the number of low bits dropped from the
// 24‑bit significand; the bit just below it is the rounding bit and
// everything under that forms the sticky bit.
const full = mant | 0x800000; // 24-bit significand incl. the implicit 1.
const shift = 14 - exp; // in [14, 24]
const roundBit = (full >>> (shift - 1)) & 1;
const sticky = full & ((1 << (shift - 1)) - 1) ? 1 : 0;
let half = full >>> shift;
if (roundBit & (sticky | (half & 1))) {
half += 1; // A carry here promotes to the smallest normal — that's fine,
// the bit pattern (exp field 1, mant 0) is exactly 2^-14.
}
return (sign << 15) | half;
}
// 5. Overflow: if the biased exponent is 31 (0x1f) or higher, the number
// cannot be represented in half precision, so we return ±∞.
if (exp >= 0x1f) {
return (sign << 15) | 0x7c00; // ±∞
}
// 6. Normalised number: round mantissa to nearest, ties to even, then pack.
const roundBit = (mant >>> 12) & 1;
const sticky = mant & 0xfff ? 1 : 0;
let half = mant >>> 13;
if (roundBit & (sticky | (half & 1))) {
half += 1;
}
if (half === 0x400) {
// The carry propagated out of the 10‑bit mantissa; it overflowed.
half = 0; // Rounded up to 1.0 × 2^(exp+1).
++exp; // Increment exponent (may overflow to ±∞).
if (exp >= 0x1f) {
return (sign << 15) | 0x7c00;
}
}
return (sign << 15) | (exp << 10) | half;
}
/**
* Convert a **binary16** encoded bit pattern back to JavaScript number.
* @param h 16-bit half-precision bits
* @returns JavaScript number (64-bit float) with same numerical value
*/
export function fromHalfBits(h) {
const sign = h & 0x8000 ? -1 : 1; // Sign multiplier (preserves −0).
const exp = (h >> 10) & 0x1f; // 5‑bit exponent.
const mant = h & 0x03ff; // 10‑bit significand.
// 1. Zero and sub‑normals.
if (exp === 0) {
// oxlint-disable-next-line oxc/erasing-op -- negative zero exists
return mant ? sign * mant * 2 ** -24 : sign * 0;
}
// 2. Special cases (exp == 31).
if (exp === 0x1f) {
return mant ? Number.NaN : sign === 1 ? Number.POSITIVE_INFINITY : Number.NEGATIVE_INFINITY;
}
// 3. Normalised numbers.
return sign * (1 + mant / 1024) * 2 ** (exp - 15);
}
function roundToF16(x) {
return fromHalfBits(toHalfBits(x));
}
const f16Cast = callableSchema({
name: 'f16',
schema: () => f16,
argTypes: (arg) => (arg ? [arg] : []),
normalImpl(v) {
if (v === undefined) {
return 0;
}
if (typeof v === 'boolean') {
return v ? 1 : 0;
}
if (!Number.isFinite(v)) {
throw new FiniteMathAssumptionError(v, f16);
}
return roundToF16(v);
},
// TODO: make usage of f16() in GPU mode check for feature availability and throw if not available
codegenImpl: (ctx, [v]) => ctx.gen.typeInstantiation(f16, v ? [v] : []),
});
/**
* A schema that represents a 16-bit float value. (equivalent to `f16` in WGSL)
*
* Can also be called to cast a value to an f16.
*
* @example
* const value = f16(); // 0
* @example
* const value = f32(1.23); // 1.23
* @example
* const value = f16(true); // 1
* @example
* const value = f16(21877.5); // 21872
*/
export const f16 = Object.assign(f16Cast, {
[$internal]: {},
type: 'f16',
});
export const abstractInt = {
[$internal]: {},
type: 'abstractInt',
toString() {
return 'abstractInt';
},
concretized: i32,
};
export const abstractFloat = {
[$internal]: {},
type: 'abstractFloat',
toString() {
return 'abstractFloat';
},
concretized: f32,
};