typegpu
Version:
A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.
314 lines (313 loc) • 9.45 kB
JavaScript
import { dualImpl } from "../core/function/dualImpl.js";
import { bitcastF32toU32Impl, bitcastU32toF32Impl, bitcastU32toI32Impl, } from "../data/numberOps.js";
import { f16, f32, fromHalfBits, i32, toHalfBits, u32 } from "../data/numeric.js";
import { isVec } from "../data/wgslTypes.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 { unifyStrict } from "../tgsl/conversion.js";
import { SignatureNotSupportedError } from "../errors.js";
import { getName } from "../shared/meta.js";
import { comptime } from "../core/function/comptime.js";
import { coerceToSnippet } from "../tgsl/generationHelpers.js";
const u32AllowedSchemas = [u32, vec2u, vec3u, vec4u];
// TODO(#2731): Remove deprecated bitcasts. Remember about cpu implementations.
/**
* @deprecated Use 'std.bitcast' instead.
*/
export const bitcastU32toF32 = dualImpl({
name: 'bitcastU32toF32',
normalImpl: ((value) => {
if (typeof value === 'number') {
return bitcastU32toF32Impl(value);
}
return VectorOps.bitcastU32toF32[value.kind](value);
}),
codegenImpl: (ctx, [n], returnType) => {
return ctx.gen.emitCall('bitcast', [coerceToSnippet(returnType)], [n]);
},
signature: (...arg) => {
const uargs = unifyStrict(arg, u32AllowedSchemas);
if (!uargs) {
throw new SignatureNotSupportedError(arg, u32AllowedSchemas);
}
return {
argTypes: uargs,
returnType: isVec(uargs[0])
? uargs[0].type === 'vec2u'
? vec2f
: uargs[0].type === 'vec3u'
? vec3f
: vec4f
: f32,
};
},
sideEffects: false,
});
/**
* @deprecated Use 'std.bitcast' instead.
*/
export const bitcastU32toI32 = dualImpl({
name: 'bitcastU32toI32',
normalImpl: ((value) => {
if (typeof value === 'number') {
return bitcastU32toI32Impl(value);
}
return VectorOps.bitcastU32toI32[value.kind](value);
}),
codegenImpl: (ctx, [n], returnType) => {
return ctx.gen.emitCall('bitcast', [coerceToSnippet(returnType)], [n]);
},
signature: (...arg) => {
const uargs = unifyStrict(arg, u32AllowedSchemas);
if (!uargs) {
throw new SignatureNotSupportedError(arg, u32AllowedSchemas);
}
return {
argTypes: uargs,
returnType: isVec(uargs[0])
? uargs[0].type === 'vec2u'
? vec2i
: uargs[0].type === 'vec3u'
? vec3i
: vec4i
: i32,
};
},
sideEffects: false,
});
const f32AllowedSchemas = [f32, vec2f, vec3f, vec4f];
/**
* @deprecated Use 'std.bitcast' instead.
*/
export const bitcastF32toU32 = dualImpl({
name: 'bitcastF32toU32',
normalImpl: ((value) => {
if (typeof value === 'number') {
return bitcastF32toU32Impl(value);
}
return VectorOps.bitcastF32toU32[value.kind](value);
}),
codegenImpl: (ctx, [n], returnType) => {
return ctx.gen.emitCall('bitcast', [coerceToSnippet(returnType)], [n]);
},
signature: (...arg) => {
const uargs = unifyStrict(arg, f32AllowedSchemas);
if (!uargs) {
throw new SignatureNotSupportedError(arg, f32AllowedSchemas);
}
return {
argTypes: uargs,
returnType: isVec(uargs[0])
? uargs[0].type === 'vec2f'
? vec2u
: uargs[0].type === 'vec3f'
? vec3u
: vec4u
: u32,
};
},
sideEffects: false,
});
const bitcastAllowedSchemas = [
/* 2 bytes */
f16,
/* 4 bytes */
f32,
i32,
u32,
vec2h,
/* 6 bytes */
vec3h,
/* 8 bytes */
vec2f,
vec2i,
vec2u,
vec4h,
/* 12 bytes */
vec3f,
vec3i,
vec3u,
/* 16 bytes */
vec4f,
vec4i,
vec4u,
];
const buffer = new ArrayBuffer(16);
const bufViews = {
f32: new Float32Array(buffer),
u32: new Uint32Array(buffer),
i32: new Int32Array(buffer),
u16: new Uint16Array(buffer),
};
function writeToBuffer(item, target) {
if (typeof item === 'number') {
target[0] = item;
}
else {
for (let i = 0; i < item.length; i++) {
target[i] = item[i];
}
}
}
function writeFloat16ToBuffer(item, target) {
if (typeof item === 'number') {
target[0] = toHalfBits(item);
}
else {
for (let i = 0; i < item.length; i++) {
target[i] = toHalfBits(item[i]);
}
}
}
function readFromBuffer(buf, schema) {
const length = 'componentCount' in schema ? schema.componentCount : 1;
const items = [];
for (let i = 0; i < length; i++) {
items.push(buf[i]);
}
return schema(...items);
}
function readFloat16FromBuffer(buf, schema) {
const length = 'componentCount' in schema ? schema.componentCount : 1;
const items = [];
for (let i = 0; i < length; i++) {
items.push(fromHalfBits(buf[i]));
}
return schema(...items);
}
const getCpuBitcast = (inType, outType) => {
const writeToPrimitive = 'primitive' in inType ? inType.primitive : inType;
const readFromPrimitive = 'primitive' in outType ? outType.primitive : outType;
return (value) => {
if (writeToPrimitive.type === 'f16') {
writeFloat16ToBuffer(value, bufViews['u16']);
}
else {
writeToBuffer(value, bufViews[writeToPrimitive.type]);
}
if (readFromPrimitive.type === 'f16') {
return readFloat16FromBuffer(bufViews['u16'], outType);
}
return readFromBuffer(bufViews[readFromPrimitive.type], outType);
};
};
function bitcastFor(inType, outType) {
return dualImpl({
name: 'bitcast',
normalImpl: getCpuBitcast(inType, outType),
codegenImpl: (ctx, [n]) => ctx.gen.emitCall('bitcast', [coerceToSnippet(outType)], [n]),
signature: (arg) => {
const uarg = unifyStrict([arg], [inType]);
if (!uarg) {
throw new SignatureNotSupportedError([arg], [inType]);
}
return {
argTypes: uarg,
returnType: outType,
};
},
sideEffects: false,
});
}
const casts = {
/* 2 bytes */
f16: {
f16: bitcastFor(f16, f16),
},
/* 4 bytes */
f32: {
f32: bitcastFor(f32, f32),
i32: bitcastFor(f32, i32),
u32: bitcastFor(f32, u32),
vec2h: bitcastFor(f32, vec2h),
},
i32: {
f32: bitcastFor(i32, f32),
i32: bitcastFor(i32, i32),
u32: bitcastFor(i32, u32),
vec2h: bitcastFor(i32, vec2h),
},
u32: {
f32: bitcastFor(u32, f32),
i32: bitcastFor(u32, i32),
u32: bitcastFor(u32, u32),
vec2h: bitcastFor(u32, vec2h),
},
vec2h: {
f32: bitcastFor(vec2h, f32),
i32: bitcastFor(vec2h, i32),
u32: bitcastFor(vec2h, u32),
vec2h: bitcastFor(vec2h, vec2h),
},
/* 6 bytes */
vec3h: {
vec3h: bitcastFor(vec3h, vec3h),
},
/* 8 bytes */
vec2f: {
vec2f: bitcastFor(vec2f, vec2f),
vec2i: bitcastFor(vec2f, vec2i),
vec2u: bitcastFor(vec2f, vec2u),
vec4h: bitcastFor(vec2f, vec4h),
},
vec2i: {
vec2f: bitcastFor(vec2i, vec2f),
vec2i: bitcastFor(vec2i, vec2i),
vec2u: bitcastFor(vec2i, vec2u),
vec4h: bitcastFor(vec2i, vec4h),
},
vec2u: {
vec2f: bitcastFor(vec2u, vec2f),
vec2i: bitcastFor(vec2u, vec2i),
vec2u: bitcastFor(vec2u, vec2u),
vec4h: bitcastFor(vec2u, vec4h),
},
vec4h: {
vec2f: bitcastFor(vec4h, vec2f),
vec2i: bitcastFor(vec4h, vec2i),
vec2u: bitcastFor(vec4h, vec2u),
vec4h: bitcastFor(vec4h, vec4h),
},
/* 12 bytes */
vec3f: {
vec3f: bitcastFor(vec3f, vec3f),
vec3i: bitcastFor(vec3f, vec3i),
vec3u: bitcastFor(vec3f, vec3u),
},
vec3i: {
vec3f: bitcastFor(vec3i, vec3f),
vec3i: bitcastFor(vec3i, vec3i),
vec3u: bitcastFor(vec3i, vec3u),
},
vec3u: {
vec3f: bitcastFor(vec3u, vec3f),
vec3i: bitcastFor(vec3u, vec3i),
vec3u: bitcastFor(vec3u, vec3u),
},
/* 16 bytes */
vec4f: {
vec4f: bitcastFor(vec4f, vec4f),
vec4i: bitcastFor(vec4f, vec4i),
vec4u: bitcastFor(vec4f, vec4u),
},
vec4i: {
vec4f: bitcastFor(vec4i, vec4f),
vec4i: bitcastFor(vec4i, vec4i),
vec4u: bitcastFor(vec4i, vec4u),
},
vec4u: {
vec4f: bitcastFor(vec4u, vec4f),
vec4i: bitcastFor(vec4u, vec4i),
vec4u: bitcastFor(vec4u, vec4u),
},
};
function getBitcast(from, to) {
if ('type' in from && from.type in casts) {
const intermediate = casts[from.type];
if ('type' in to && to.type in intermediate) {
return intermediate[to.type];
}
}
throw new Error(`Incorrect bitcast from ${getName(from)} to ${getName(to)}.`);
}
export const bitcast = comptime(getBitcast);