fix/feat: bit operation on vectors (#2276)

Author: cieplypolar

packages/typegpu/src/data/index.ts

@@ -36,6 +36,18 @@ assignInfixOperator(MatBase, 'add', Operator.plus);
 assignInfixOperator(MatBase, 'sub', Operator.minus);
 assignInfixOperator(MatBase, 'mul', Operator.star);
 
+// bitShift does not yet have tsover operator symbol
+{
+  // oxlint-disable-next-line typescript/no-explicit-any -- anything is possible
+  const proto = VecBase.prototype as any;
+  proto.bitShiftLeft = function (this: unknown, other: unknown) {
+    return (infixOperators.bitShiftLeft as (a: unknown, b: unknown) => unknown)(this, other);
+  };
+  proto.bitShiftRight = function (this: unknown, other: unknown) {
+    return (infixOperators.bitShiftRight as (a: unknown, b: unknown) => unknown)(this, other);
+  };
+}
+
 export { bool, f16, f32, i32, u16, u32 } from './numeric.ts';
 export {
   isAlignAttrib,

packages/typegpu/src/data/vectorOps.ts

@@ -117,6 +117,15 @@ const binaryComponentWise4i = (op: BinaryOp) => (a: wgsl.v4i, b: wgsl.v4i) =>
 const binaryComponentWise4u = (op: BinaryOp) => (a: wgsl.v4u, b: wgsl.v4u) =>
   vec4u(op(a.x, b.x), op(a.y, b.y), op(a.z, b.z), op(a.w, b.w));
 
+const binaryComponentWise2i2u = (op: BinaryOp) => (a: wgsl.v2i, b: wgsl.v2u) =>
+  vec2i(op(a.x, b.x), op(a.y, b.y));
+
+const binaryComponentWise3i3u = (op: BinaryOp) => (a: wgsl.v3i, b: wgsl.v3u) =>
+  vec3i(op(a.x, b.x), op(a.y, b.y), op(a.z, b.z));
+
+const binaryComponentWise4i4u = (op: BinaryOp) => (a: wgsl.v4i, b: wgsl.v4u) =>
+  vec4i(op(a.x, b.x), op(a.y, b.y), op(a.z, b.z), op(a.w, b.w));
+
 const binaryComponentWise2x2f = (op: BinaryOp) => (a: wgsl.m2x2f, b: wgsl.m2x2f) => {
   const a_ = a.columns as [wgsl.v2f, wgsl.v2f];
   const b_ = b.columns as [wgsl.v2f, wgsl.v2f];
@@ -1090,6 +1099,34 @@ export const VectorOps = {
     vec4h: unary4h(Math.tanh),
   } as Record<VecKind, <T extends vBase>(v: T) => T>,
 
+  bitShiftLeft: {
+    vec2i: binaryComponentWise2i2u((a, b) => a << b),
+    vec2u: binaryComponentWise2u((a, b) => a << b),
+
+    vec3i: binaryComponentWise3i3u((a, b) => a << b),
+    vec3u: binaryComponentWise3u((a, b) => a << b),
+
+    vec4i: binaryComponentWise4i4u((a, b) => a << b),
+    vec4u: binaryComponentWise4u((a, b) => a << b),
+  } as Record<
+    VecKind,
+    <T extends wgsl.AnyIntegerVecInstance, U extends wgsl.AnyUnsignedVecInstance>(a: T, b: U) => T
+  >,
+
+  bitShiftRight: {
+    vec2i: binaryComponentWise2i2u((a, b) => a >> b),
+    vec2u: binaryComponentWise2u((a, b) => a >> b),
+
+    vec3i: binaryComponentWise3i3u((a, b) => a >> b),
+    vec3u: binaryComponentWise3u((a, b) => a >> b),
+
+    vec4i: binaryComponentWise4i4u((a, b) => a >> b),
+    vec4u: binaryComponentWise4u((a, b) => a >> b),
+  } as Record<
+    VecKind,
+    <T extends wgsl.AnyIntegerVecInstance, U extends wgsl.AnyUnsignedVecInstance>(a: T, b: U) => T
+  >,
+
   bitcastU32toF32: {
     vec2u: (n: wgsl.v2u) => vec2f(bitcastU32toF32Impl(n.x), bitcastU32toF32Impl(n.y)),
     vec3u: (n: wgsl.v3u) =>

packages/typegpu/src/data/wgslTypes.ts

@@ -70,6 +70,17 @@ export interface vecInfixNotation<T extends vecBase> {
   [Symbol.operatorPercent](lhs: T | number, rhs: T | number): T;
 }
 
+export type vecIToVecU<T extends AnyIntegerVecInstance> = T extends v2i | v2u
+  ? v2u
+  : T extends v3i | v3u
+    ? v3u
+    : v4u;
+
+export interface vecBitShiftNotation<T extends AnyIntegerVecInstance> {
+  bitShiftLeft(rhs: vecIToVecU<T> | number): T;
+  bitShiftRight(rhs: vecIToVecU<T> | number): T;
+}
+
 /**
  * Matrix infix notation.
  *
@@ -228,7 +239,8 @@ export interface v2h extends Tuple2<number>, Swizzle2<v2h, v3h, v4h>, vecInfixNo
  * Interface representing its WGSL vector type counterpart: vec2i or vec2<i32>.
  * A vector with 2 elements of type i32
  */
-export interface v2i extends Tuple2<number>, Swizzle2<v2i, v3i, v4i>, vecInfixNotation<v2i> {
+export interface v2i
+  extends Tuple2<number>, Swizzle2<v2i, v3i, v4i>, vecInfixNotation<v2i>, vecBitShiftNotation<v2i> {
   readonly [$internal]: true;
   /** use to distinguish between vectors of the same size on the type level */
   readonly kind: 'vec2i';
@@ -242,7 +254,8 @@ export interface v2i extends Tuple2<number>, Swizzle2<v2i, v3i, v4i>, vecInfixNo
  * Interface representing its WGSL vector type counterpart: vec2u or vec2<u32>.
  * A vector with 2 elements of type u32
  */
-export interface v2u extends Tuple2<number>, Swizzle2<v2u, v3u, v4u>, vecInfixNotation<v2u> {
+export interface v2u
+  extends Tuple2<number>, Swizzle2<v2u, v3u, v4u>, vecInfixNotation<v2u>, vecBitShiftNotation<v2u> {
   readonly [$internal]: true;
   /** use to distinguish between vectors of the same size on the type level */
   readonly kind: 'vec2u';
@@ -302,7 +315,8 @@ export interface v3h extends Tuple3<number>, Swizzle3<v2h, v3h, v4h>, vecInfixNo
  * Interface representing its WGSL vector type counterpart: vec3i or vec3<i32>.
  * A vector with 3 elements of type i32
  */
-export interface v3i extends Tuple3<number>, Swizzle3<v2i, v3i, v4i>, vecInfixNotation<v3i> {
+export interface v3i
+  extends Tuple3<number>, Swizzle3<v2i, v3i, v4i>, vecInfixNotation<v3i>, vecBitShiftNotation<v3i> {
   readonly [$internal]: true;
   /** use to distinguish between vectors of the same size on the type level */
   readonly kind: 'vec3i';
@@ -318,7 +332,8 @@ export interface v3i extends Tuple3<number>, Swizzle3<v2i, v3i, v4i>, vecInfixNo
  * Interface representing its WGSL vector type counterpart: vec3u or vec3<u32>.
  * A vector with 3 elements of type u32
  */
-export interface v3u extends Tuple3<number>, Swizzle3<v2u, v3u, v4u>, vecInfixNotation<v3u> {
+export interface v3u
+  extends Tuple3<number>, Swizzle3<v2u, v3u, v4u>, vecInfixNotation<v3u>, vecBitShiftNotation<v3u> {
   readonly [$internal]: true;
   /** use to distinguish between vectors of the same size on the type level */
   readonly kind: 'vec3u';
@@ -386,7 +401,8 @@ export interface v4h extends Tuple4<number>, Swizzle4<v2h, v3h, v4h>, vecInfixNo
  * Interface representing its WGSL vector type counterpart: vec4i or vec4<i32>.
  * A vector with 4 elements of type i32
  */
-export interface v4i extends Tuple4<number>, Swizzle4<v2i, v3i, v4i>, vecInfixNotation<v4i> {
+export interface v4i
+  extends Tuple4<number>, Swizzle4<v2i, v3i, v4i>, vecInfixNotation<v4i>, vecBitShiftNotation<v4i> {
   readonly [$internal]: true;
   /** use to distinguish between vectors of the same size on the type level */
   readonly kind: 'vec4i';
@@ -404,7 +420,8 @@ export interface v4i extends Tuple4<number>, Swizzle4<v2i, v3i, v4i>, vecInfixNo
  * Interface representing its WGSL vector type counterpart: vec4u or vec4<u32>.
  * A vector with 4 elements of type u32
  */
-export interface v4u extends Tuple4<number>, Swizzle4<v2u, v3u, v4u>, vecInfixNotation<v4u> {
+export interface v4u
+  extends Tuple4<number>, Swizzle4<v2u, v3u, v4u>, vecInfixNotation<v4u>, vecBitShiftNotation<v4u> {
   readonly [$internal]: true;
   /** use to distinguish between vectors of the same size on the type level */
   readonly kind: 'vec4u';
@@ -1528,12 +1545,18 @@ export function isMat(value: unknown): value is Mat2x2f | Mat3x3f | Mat4x4f {
   return isMat2x2f(value) || isMat3x3f(value) || isMat4x4f(value);
 }
 
-export function isFloat32VecInstance(
-  element: number | AnyVecInstance | AnyMatInstance,
-): element is AnyFloat32VecInstance {
+export function isFloat32VecInstance(element: unknown): element is AnyFloat32VecInstance {
   return isVecInstance(element) && ['vec2f', 'vec3f', 'vec4f'].includes(element.kind);
 }
 
+export function isInteger32VecInstance(value: unknown): value is v2u | v2i | v3u | v3i | v4u | v4i {
+  return isVecInstance(value) && /[iu]$/.test(value.kind);
+}
+
+export function isUint32VecInstance(value: unknown): value is v2u | v3u | v4u {
+  return isVecInstance(value) && /[u]$/.test(value.kind);
+}
+
 export function isWgslData(value: unknown): value is AnyWgslData {
   return isMarkedInternal(value) && wgslTypeLiterals.includes((value as AnyWgslData)?.type);
 }

packages/typegpu/src/std/index.ts

@@ -70,7 +70,7 @@ export {
   trunc,
 } from './numeric.ts';
 
-export { add, div, mod, mul, neg, sub } from './operators.ts';
+export { add, bitShiftLeft, bitShiftRight, div, mod, mul, neg, sub } from './operators.ts';
 
 export { rotateX4, rotateY4, rotateZ4, scale4, translate4 } from './matrix.ts';
 

packages/typegpu/src/std/operators.ts

@@ -1,19 +1,23 @@
 import { dualImpl } from '../core/function/dualImpl.ts';
 import { stitch } from '../core/resolve/stitch.ts';
-import { abstractFloat, f16, f32 } from '../data/numeric.ts';
-import { vecTypeToConstructor } from '../data/vector.ts';
+import { abstractFloat, f16, f32, i32, u32 } from '../data/numeric.ts';
+import { vec2i, vec2u, vec3i, vec3u, vec4i, vec4u, vecTypeToConstructor } from '../data/vector.ts';
 import { VectorOps } from '../data/vectorOps.ts';
 import {
+  type AnyIntegerVecInstance,
   type AnyMatInstance,
   type AnyNumericVecInstance,
   type BaseData,
+  type mBaseForVec,
+  type vBaseForMat,
   isFloat32VecInstance,
+  isInteger32VecInstance,
   isMat,
   isMatInstance,
+  isUint32VecInstance,
   isVec,
   isVecInstance,
-  type mBaseForVec,
-  type vBaseForMat,
+  vecIToVecU,
 } from '../data/wgslTypes.ts';
 import { SignatureNotSupportedError } from '../errors.ts';
 import { unify } from '../tgsl/conversion.ts';
@@ -279,3 +283,111 @@ export const neg = dualImpl({
   normalImpl: cpuNeg,
   codegenImpl: (_ctx, [arg]) => stitch`-(${arg})`,
 });
+
+const anyConcreteInteger = [i32, u32, vec2i, vec3i, vec4i, vec2u, vec3u, vec4u] as BaseData[];
+
+const intVecToUnsignedVec = {
+  vec2i: vec2u,
+  vec2u: vec2u,
+  vec3i: vec3u,
+  vec3u: vec3u,
+  vec4i: vec4u,
+  vec4u: vec4u,
+} as const;
+
+const bitShiftSignature = (lhs: BaseData, rhs: BaseData) => {
+  const lhsUnified = unify([lhs], anyConcreteInteger)?.[0];
+  if (!lhsUnified) {
+    throw new SignatureNotSupportedError([lhs], anyConcreteInteger);
+  }
+
+  let rhsType: BaseData;
+  if (isVec(lhsUnified)) {
+    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]);
+    }
+    rhsType = rhsUnified;
+  } else {
+    rhsType = u32;
+  }
+
+  return {
+    argTypes: [lhsUnified, rhsType],
+    returnType: lhsUnified,
+  };
+};
+
+function cpuBitShiftLeft(lhs: number, rhs: number): number;
+function cpuBitShiftLeft<T extends AnyIntegerVecInstance>(lhs: T, rhs: number): T;
+function cpuBitShiftLeft<T extends AnyIntegerVecInstance>(lhs: T, rhs: vecIToVecU<T>): T;
+function cpuBitShiftLeft<T extends AnyIntegerVecInstance>(
+  lhs: number | AnyIntegerVecInstance,
+  rhs: number | vecIToVecU<T>,
+) {
+  if (typeof lhs === 'number' && typeof rhs === 'number') {
+    return 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 types: number or integer vector (rhs must be the same arity as lhs).',
+  );
+}
+
+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 stitch`(${lhs} << ${schema}(${rhs}))`;
+    }
+    return stitch`(${lhs} << ${rhs})`;
+  },
+});
+
+function cpuBitShiftRight(lhs: number, rhs: number): number;
+function cpuBitShiftRight<T extends AnyIntegerVecInstance>(lhs: T, rhs: number): T;
+function cpuBitShiftRight<T extends AnyIntegerVecInstance>(lhs: T, rhs: vecIToVecU<T>): T;
+function cpuBitShiftRight<T extends AnyIntegerVecInstance>(
+  lhs: number | AnyIntegerVecInstance,
+  rhs: number | vecIToVecU<T>,
+) {
+  if (typeof lhs === 'number' && typeof rhs === 'number') {
+    return 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 types: number or integer vector (rhs must be the same arity as lhs).',
+  );
+}
+
+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 stitch`(${lhs} >> ${schema}(${rhs}))`;
+    }
+    return stitch`(${lhs} >> ${rhs})`;
+  },
+});

packages/typegpu/src/tgsl/accessProp.ts

@@ -37,7 +37,7 @@ import {
   isWgslStruct,
 } from '../data/wgslTypes.ts';
 import { $gpuCallable } from '../shared/symbols.ts';
-import { add, div, mod, mul, sub } from '../std/operators.ts';
+import { add, bitShiftLeft, bitShiftRight, div, mod, mul, sub } from '../std/operators.ts';
 import { isKnownAtComptime } from '../types.ts';
 import { coerceToSnippet } from './generationHelpers.ts';
 
@@ -65,6 +65,8 @@ export const infixOperators = {
   mul,
   div,
   mod,
+  bitShiftLeft,
+  bitShiftRight,
 } as const;
 
 export type InfixOperator = keyof typeof infixOperators;

packages/typegpu/src/tgsl/wgslGenerator.ts

@@ -10,6 +10,7 @@ import {
   unptr,
 } from '../data/dataTypes.ts';
 import { bool, i32, u32 } from '../data/numeric.ts';
+import { vec2u, vec3u, vec4u } from '../data/vector.ts';
 import {
   fallthroughCopyOrigin,
   isEphemeralOrigin,
@@ -77,6 +78,8 @@ const parenthesizedOps = [
 
 const binaryLogicalOps = ['&&', '||', '==', '!=', '===', '!==', '<', '<=', '>', '>='];
 
+const bitShiftOps: string[] = ['<<', '>>', '<<=', '>>='];
+
 const OP_MAP = {
   //
   // binary
@@ -348,12 +351,28 @@ ${this.ctx.pre}}`;
         return codegen(this.ctx, [lhsExpr, rhsExpr]);
       }
 
-      const forcedType = exprType === NODE.assignmentExpr ? [lhsExpr.dataType] : undefined;
+      let convLhs: Snippet;
+      let convRhs: Snippet;
 
-      const [convLhs, convRhs] = convertToCommonType(this.ctx, [lhsExpr, rhsExpr], forcedType) ?? [
-        lhsExpr,
-        rhsExpr,
-      ];
+      if (bitShiftOps.includes(op)) {
+        // rhs must be u32 (or vecN<u32> for vector lhs)
+        let rhsTarget: wgsl.BaseData;
+        if (wgsl.isVec(lhsExpr.dataType)) {
+          const cc = lhsExpr.dataType.componentCount;
+          rhsTarget = cc === 2 ? vec2u : cc === 3 ? vec3u : vec4u;
+        } else {
+          rhsTarget = u32;
+        }
+        convRhs = tryConvertSnippet(this.ctx, rhsExpr, rhsTarget, false);
+        // if lhs is not an integer type, the browser will return a descriptive wgsl error
+        convLhs = lhsExpr;
+      } else {
+        const forcedType = exprType === NODE.assignmentExpr ? [lhsExpr.dataType] : undefined;
+        [convLhs, convRhs] = convertToCommonType(this.ctx, [lhsExpr, rhsExpr], forcedType) ?? [
+          lhsExpr,
+          rhsExpr,
+        ];
+      }
 
       const lhsStr = this.ctx.resolve(convLhs.value, convLhs.dataType).value;
       const rhsStr = this.ctx.resolve(convRhs.value, convRhs.dataType).value;

packages/typegpu/tests/examples/individual/fluid-with-atomics.test.ts

@@ -103,7 +103,7 @@ describe('fluid with atomics example', () => {
         }
         if (isWaterDrain(x, y)) {
           persistFlags(x, y);
-          updateCell(x, y, (3 << 24));
+          updateCell(x, y, (3 << 24u));
           return true;
         }
         if (((((y == 0u) || (y == (size.y - 1u))) || (x == 0u)) || (x == (size.x - 1u)))) {