typegpu
Version:
A thin layer between JS and WebGPU/WGSL that improves development experience and allows for faster iteration.
293 lines (292 loc) • 12.3 kB
JavaScript
import { UnknownData } from "../data/dataTypes.js";
import { undecorate } from "../data/dataTypes.js";
import { derefSnippet, RefOperator } from "../data/ref.js";
import { schemaCallWrapperGPU } from "../data/schemaCallWrapper.js";
import { snip, withDataType } from "../data/snippet.js";
import { isMat, isPtr, isVec, isWgslStruct, } from "../data/wgslTypes.js";
import { invariant, WgslTypeError } from "../errors.js";
import { getName } from "../shared/meta.js";
import { safeStringify } from "../shared/stringify.js";
import { assertExhaustive } from "../shared/utilityTypes.js";
import { logger } from "../tgpuLogger.js";
import { accessStructProp } from "./accessStructProp.js";
const INFINITE_RANK = {
rank: Number.POSITIVE_INFINITY,
action: 'none',
};
function getAutoConversionRank(src, dest) {
const trueSrc = undecorate(src);
const trueDst = undecorate(dest);
if (trueSrc.type === trueDst.type) {
if (trueSrc.type === 'struct' && trueSrc !== trueDst) {
return { rank: 1, action: 'cast', targetType: dest };
}
return { rank: 0, action: 'none' };
}
if (trueSrc.type === 'abstractFloat') {
if (trueDst.type === 'f32')
return { rank: 1, action: 'none' };
if (trueDst.type === 'f16')
return { rank: 2, action: 'none' };
}
if (trueSrc.type === 'abstractInt') {
if (trueDst.type === 'i32')
return { rank: 3, action: 'none' };
if (trueDst.type === 'u32')
return { rank: 4, action: 'none' };
if (trueDst.type === 'abstractFloat')
return { rank: 5, action: 'none' };
if (trueDst.type === 'f32')
return { rank: 6, action: 'none' };
if (trueDst.type === 'f16')
return { rank: 7, action: 'none' };
}
if (isVec(trueSrc) &&
isVec(trueDst) &&
// Same length vectors
trueSrc.type[3] === trueDst.type[3]) {
return getAutoConversionRank(trueSrc.primitive, trueDst.primitive);
}
if (isMat(trueSrc) &&
isMat(trueDst) &&
// Same dimensions
trueSrc.type[3] === trueDst.type[3]) {
// Matrix conversion rank depends only on component type (always f32 for now)
return { rank: 0, action: 'none' };
}
return INFINITE_RANK;
}
function getImplicitConversionRank(src, dest) {
const trueSrc = undecorate(src);
const trueDst = undecorate(dest);
if (isPtr(trueSrc) &&
// Only dereferencing implicit pointers, otherwise we'd have a types mismatch between TS and WGSL
trueSrc.implicit &&
getAutoConversionRank(trueSrc.inner, trueDst).rank < Number.POSITIVE_INFINITY) {
return { rank: 0, action: 'deref' };
}
if (isPtr(trueDst) &&
getAutoConversionRank(trueSrc, trueDst.inner).rank < Number.POSITIVE_INFINITY) {
return { rank: 1, action: 'ref' };
}
const primitivePreference = {
f32: 0,
f16: 1,
i32: 2,
u32: 3,
bool: 4,
};
if (trueSrc.type in primitivePreference && trueDst.type in primitivePreference) {
const srcType = trueSrc.type;
const destType = trueDst.type;
if (srcType !== destType) {
const srcPref = primitivePreference[srcType];
const destPref = primitivePreference[destType];
const rank = destPref < srcPref ? 10 : 20;
return { rank: rank, action: 'cast', targetType: trueDst };
}
}
if ((trueSrc.type === 'u32' || trueSrc.type === 'i32') && trueDst.type === 'abstractFloat') {
// When one of the types is a float (abstract or not), we don't want to cast it to a non-float type,
// which would cause it to lose precision. We instead choose the common type to be f32.
return { rank: 1, action: 'cast', targetType: trueDst.concretized };
}
if (trueSrc.type === 'abstractFloat') {
if (trueDst.type === 'i32') {
return { rank: 2, action: 'cast', targetType: trueDst };
}
if (trueDst.type === 'u32') {
return { rank: 3, action: 'cast', targetType: trueDst };
}
}
return INFINITE_RANK;
}
function getConversionRank(src, dest, allowImplicit) {
const autoRank = getAutoConversionRank(src, dest);
if (autoRank.rank < Number.POSITIVE_INFINITY) {
return autoRank;
}
if (allowImplicit) {
return getImplicitConversionRank(src, dest);
}
return INFINITE_RANK;
}
function findBestType(types, uniqueTypes, allowImplicit) {
let bestResult;
for (const targetType of uniqueTypes) {
/**
* The type we end up converting to. Will be different than `targetType` if `targetType === abstractFloat`
*/
let destType = targetType;
const details = [];
let sum = 0;
for (const sourceType of types) {
const conversion = getConversionRank(sourceType, targetType, allowImplicit);
sum += conversion.rank;
if (conversion.rank === Number.POSITIVE_INFINITY) {
break;
}
details.push(conversion);
if (conversion.action === 'cast') {
destType = conversion.targetType;
}
}
if (sum < (bestResult?.sum ?? Number.POSITIVE_INFINITY)) {
bestResult = { type: destType, details, sum };
}
}
if (!bestResult) {
return undefined;
}
const actions = bestResult.details.map((detail, index) => ({
sourceIndex: index,
action: detail.action,
...(detail.action === 'cast' && {
targetType: detail.targetType,
}),
}));
return {
targetType: bestResult.type,
actions,
hasImplicitConversions: actions.some((action) => action.action === 'cast'),
};
}
export function getBestConversion(types, targetTypes) {
if (types.length === 0)
return undefined;
const uniqueTargetTypes = [...new Set((targetTypes || types).map(undecorate))];
const explicitResult = findBestType(types, uniqueTargetTypes, false);
if (explicitResult) {
return explicitResult;
}
const implicitResult = findBestType(types, uniqueTargetTypes, true);
if (implicitResult) {
return implicitResult;
}
return undefined;
}
function applyActionToSnippet(ctx, snippet, action, targetType) {
if (action.action === 'none') {
if (targetType === snippet.dataType) {
return snippet;
}
return withDataType(targetType, snippet);
}
switch (action.action) {
case 'ref':
return snip(new RefOperator(snippet, targetType), targetType, snippet.origin, snippet.possibleSideEffects);
case 'deref':
return derefSnippet(snippet);
case 'cast': {
if (isWgslStruct(snippet.dataType) && isWgslStruct(targetType)) {
const typeName = getName(snippet.dataType) ?? '<unnamed>';
const targetName = getName(targetType) ?? '<unnamed>';
// Struct to struct casting
if (snippet.possibleSideEffects) {
throw new Error(`Cannot resolve struct cast from '${typeName}' to '${targetName}'. Store the value to a variable first, then cast it.`);
}
const propSnips = Object.entries(targetType.propTypes).map(([key, dataType]) => {
const accessedProp = accessStructProp(snippet, key);
if (!accessedProp) {
throw new Error(`Cannot auto-convert struct '${typeName}' to '${targetName}' because the property '${key}' is missing.`);
}
const converted = convertToCommonType(ctx, [accessedProp], [dataType]);
if (!converted || !converted[0]) {
throw new Error(`Cannot auto-convert struct '${typeName}' to '${targetName}' because type '${getName(accessedProp.dataType) ?? '<unnamed>'}' is not convertible to '${getName(undecorate(dataType)) ?? '<unnamed>'}'.`);
}
return converted[0];
});
const targetSnippet = ctx.resolve(targetType).value;
return snip(`${targetSnippet}(${propSnips.map((snip) => snip.value).join(', ')})`, targetType, 'runtime', false);
}
// Casting means calling the schema with the snippet as an argument.
return schemaCallWrapperGPU(ctx, targetType, snippet);
}
default: {
assertExhaustive(action.action, 'applyActionToSnippet');
}
}
}
/**
* Unifies input types to a common type.
*/
export function unify(inTypes, restrictTo) {
if (inTypes.some((type) => type === UnknownData)) {
return undefined;
}
const conversion = getBestConversion(inTypes, restrictTo);
if (!conversion) {
return undefined;
}
return inTypes.map((type) => (isVec(type) || isMat(type) ? type : conversion.targetType));
}
/**
* Unifies input types to a common type.
* Unlike `unify`, it does not allow implicit conversions.
*/
export function unifyStrict(inTypes, restrictTo) {
if (inTypes.some((type) => type === UnknownData)) {
return undefined;
}
const uniqueTargetTypes = [...new Set((restrictTo || inTypes).map(undecorate))];
const conversion = findBestType(inTypes, uniqueTargetTypes, false);
if (!conversion) {
return undefined;
}
return inTypes.map((type) => (isVec(type) || isMat(type) ? type : conversion.targetType));
}
export function convertToCommonType(ctx, values, restrictTo, verbose = true) {
const types = values.map((value) => value.dataType);
if (types.some((type) => type === UnknownData)) {
return undefined;
}
// Calling convertToCommonType with an empty restrictTo array
// prevents any conversions from being made. If you intend to allow
// all conversions, pass undefined instead. If this was intended call
// the function conditionally since the result will always be undefined.
invariant(!(Array.isArray(restrictTo) && restrictTo.length === 0), "Internal error, expected 'restrictTo' to not be an empty array.");
const conversion = getBestConversion(types, restrictTo);
if (!conversion) {
return undefined;
}
if (verbose && conversion.hasImplicitConversions) {
logger.warn('implicit-conversion', `Implicit conversions from [\n${values
.map((v) => ` ${ctx.resolveSnippet(v).value}: ${safeStringify(v.dataType)}`)
.join(',\n')}\n] to ${conversion.targetType.type} are supported, but not recommended.
Consider using explicit conversions instead.`);
}
return values.map((value, index) => {
const action = conversion.actions[index];
invariant(action, 'Action should not be undefined');
return applyActionToSnippet(ctx, value, action, conversion.targetType);
});
}
export function tryConvertSnippet(ctx, snippet, targetDataTypes, verbose = true) {
const targets = Array.isArray(targetDataTypes) ? targetDataTypes : [targetDataTypes];
const { value, dataType, origin, possibleSideEffects } = snippet;
if (targets.length === 1) {
const target = targets[0];
if (target === dataType) {
return snip(value, target, origin, possibleSideEffects);
}
if (dataType === UnknownData) {
// Commit unknown to the expected type.
return ctx.resolveSnippet(snip(value, target, origin, possibleSideEffects));
}
}
const converted = convertToCommonType(ctx, [snippet], targets, verbose);
if (converted) {
return converted[0];
}
throw new WgslTypeError(`Cannot convert value of type '${String(dataType)}' to any of the target types: [${targets.map((t) => t.type).join(', ')}]`);
}
export function convertStructValues(ctx, structType, values) {
return Object.entries(structType.propTypes).map(([key, targetType]) => {
const val = values[key];
if (!val) {
throw new Error(`Missing property ${key}`);
}
const converted = convertToCommonType(ctx, [val], [targetType]);
return converted?.[0] ?? val;
});
}