feat: Texture std expansion and generator adjustments (#1189)

Author: reczkok

packages/typegpu/src/core/root/init.ts

@@ -17,7 +17,7 @@ import type { Infer } from '../../shared/repr.ts';
 import { $internal } from '../../shared/symbols.ts';
 import type { AnyVertexAttribs } from '../../shared/vertexFormat.ts';
 import type {
-  LayoutEntryToInput,
+  ExtractBindGroupInputFromLayout,
   TgpuBindGroup,
   TgpuBindGroupLayout,
   TgpuLayoutEntry,
@@ -289,9 +289,7 @@ class TgpuRootImpl
     >,
   >(
     layout: TgpuBindGroupLayout<Entries>,
-    entries: {
-      [K in keyof Entries]: LayoutEntryToInput<Entries[K]>;
-    },
+    entries: ExtractBindGroupInputFromLayout<Entries>,
   ) {
     return new TgpuBindGroupImpl(layout, entries);
   }

packages/typegpu/src/data/vectorOps.ts

@@ -829,6 +829,32 @@ export const VectorOps = {
     <T extends wgsl.AnyFloatVecInstance | number>(a: T, b: T) => T
   >,
 
+  sqrt: {
+    vec2f: unary2f(Math.sqrt),
+    vec2h: unary2h(Math.sqrt),
+
+    vec3f: unary3f(Math.sqrt),
+    vec3h: unary3h(Math.sqrt),
+
+    vec4f: unary4f(Math.sqrt),
+    vec4h: unary4h(Math.sqrt),
+  } as Record<VecKind, <T extends vBase>(v: T) => T>,
+
+  div: {
+    vec2f: binaryComponentWise2f((a, b) => a / b),
+    vec2h: binaryComponentWise2h((a, b) => a / b),
+    vec2i: binaryComponentWise2i((a, b) => a / b),
+    vec2u: binaryComponentWise2u((a, b) => a / b),
+    vec3f: binaryComponentWise3f((a, b) => a / b),
+    vec3h: binaryComponentWise3h((a, b) => a / b),
+    vec3i: binaryComponentWise3i((a, b) => a / b),
+    vec3u: binaryComponentWise3u((a, b) => a / b),
+    vec4f: binaryComponentWise4f((a, b) => a / b),
+    vec4h: binaryComponentWise4h((a, b) => a / b),
+    vec4i: binaryComponentWise4i((a, b) => a / b),
+    vec4u: binaryComponentWise4u((a, b) => a / b),
+  } as Record<VecKind, <T extends vBase>(a: T, b: T) => T>,
+
   mix: {
     vec2f: (e1: wgsl.v2f, e2: wgsl.v2f, e3: wgsl.v2f | number) => {
       if (typeof e3 === 'number') {

packages/typegpu/src/index.ts

@@ -153,6 +153,7 @@ export type {
   BindLayoutEntry,
   LayoutEntryToInput,
   TgpuBindGroup,
+  ExtractBindGroupInputFromLayout,
 } from './tgpuBindGroupLayout.ts';
 export type {
   TgpuFn,

packages/typegpu/src/std/index.ts

@@ -28,6 +28,8 @@ export {
   pow,
   reflect,
   neg,
+  sqrt,
+  div,
 } from './numeric.js';
 
 export {

packages/typegpu/src/std/numeric.ts

@@ -456,3 +456,40 @@ export const neg = createDualImpl(
   // GPU implementation
   (value) => ({ value: `-(${value.value})`, dataType: value.dataType }),
 );
+
+export const sqrt = createDualImpl(
+  // CPU implementation
+  <T extends AnyFloatVecInstance | number>(value: T): T => {
+    if (typeof value === 'number') {
+      return Math.sqrt(value) as T;
+    }
+    return VectorOps.sqrt[value.kind](value) as T;
+  },
+  // GPU implementation
+  (value) => ({ value: `sqrt(${value.value})`, dataType: value.dataType }),
+);
+
+export const div = createDualImpl(
+  // CPU implementation
+  <T extends AnyNumericVecInstance | number>(lhs: T, rhs: T | number): T => {
+    if (typeof lhs === 'number' && typeof rhs === 'number') {
+      return (lhs / rhs) as T;
+    }
+    if (typeof rhs === 'number') {
+      return VectorOps.mulSxV[(lhs as AnyNumericVecInstance).kind](
+        1 / rhs,
+        lhs as AnyNumericVecInstance,
+      ) as T;
+    }
+    // Vector / Vector case
+    return VectorOps.div[(lhs as AnyNumericVecInstance).kind](
+      lhs as AnyNumericVecInstance,
+      rhs as AnyNumericVecInstance,
+    ) as T;
+  },
+  // GPU implementation
+  (lhs, rhs) => ({
+    value: `(${lhs.value} / ${rhs.value})`,
+    dataType: lhs.dataType,
+  }),
+);

packages/typegpu/src/tgpuBindGroupLayout.ts

@@ -341,6 +341,12 @@ export type InferLayoutEntry<T extends TgpuLayoutEntry | null> =
               ? StorageTextureUsageForEntry<T>
               : never;
 
+export type ExtractBindGroupInputFromLayout<
+  T extends Record<string, TgpuLayoutEntry | null>,
+> = {
+  [K in keyof T]: LayoutEntryToInput<T[K]>;
+};
+
 export type TgpuBindGroup<
   Entries extends Record<string, TgpuLayoutEntry | null> = Record<
     string,

packages/typegpu/src/tgsl/generationHelpers.ts

@@ -40,6 +40,7 @@ import {
   UnknownData,
   type Wgsl,
   hasInternalDataType,
+  isSelfResolvable,
 } from '../types.ts';
 
 const swizzleableTypes = [
@@ -245,7 +246,10 @@ export function getTypeFromWgsl(resource: Wgsl): AnyData | UnknownData {
     }
   }
 
-  return isWgslData(resource) ? resource : UnknownData;
+  // TODO: do not treat self-resolvable as wgsl data (when we have proper texture schemas)
+  return isWgslData(resource) || isSelfResolvable(resource)
+    ? (resource as unknown as AnyData)
+    : UnknownData;
 }
 
 export function numericLiteralToSnippet(value: string): Snippet | undefined {

packages/typegpu/tests/std/numeric/div.test.ts

@@ -0,0 +1,78 @@
+import { describe, expect, it } from 'vitest';
+import * as d from '../../../src/data/index.ts';
+import { div } from '../../../src/std/index.ts';
+
+describe('div', () => {
+  it('computes quotient of a d.vec2f and a number', () => {
+    expect(div(d.vec2f(0, 0), 17)).toEqual(d.vec2f(0, 0));
+    expect(div(d.vec2f(0.24, 0), 0.4).x).toBeCloseTo(0.6);
+    expect(div(d.vec2f(0, 0), 5)).toEqual(d.vec2f());
+  });
+
+  it('computes quotient of a d.vec2u and a number', () => {
+    expect(div(d.vec2u(0, 0), 2)).toEqual(d.vec2u(0, 0));
+    expect(div(d.vec2u(3, 3), 3)).toEqual(d.vec2u(1));
+  });
+
+  it('computes quotient of a d.vec2i and a number', () => {
+    expect(div(d.vec2i(0, 0), 9)).toEqual(d.vec2i(0, 0));
+    expect(div(d.vec2i(-3, -3), -3)).toEqual(d.vec2i(1));
+    expect(div(d.vec2i(0, 0), 5)).toEqual(d.vec2i());
+  });
+
+  it('computes quotient of a d.vec3f and a number', () => {
+    expect(div(d.vec3f(-3, -4, -6), 2)).toEqual(d.vec3f(-1.5, -2, -3));
+    expect(div(d.vec3f(-2, -2, -2), -2)).toEqual(d.vec3f(1));
+    expect(div(d.vec3f(0, 0, 0), 5)).toEqual(d.vec3f());
+  });
+
+  it('computes quotient of a d.vec3u and a number', () => {
+    expect(div(d.vec3u(2, 2, 2), 2)).toEqual(d.vec3u(1));
+    expect(div(d.vec3u(0, 0, 0), 5)).toEqual(d.vec3u());
+  });
+
+  it('computes quotient of a d.vec3i and a number', () => {
+    expect(div(d.vec3i(-1, -2, -3), 1)).toEqual(d.vec3i(-1, -2, -3));
+    expect(div(d.vec3i(-4, -6, -8), -2)).toEqual(d.vec3i(2, 3, 4));
+    expect(div(d.vec3i(2, 2, 2), 2)).toEqual(d.vec3i(1));
+  });
+
+  it('computes quotient of a d.vec4f and a number', () => {
+    expect(div(d.vec4f(1.5, 2, 3, 4), 2)).toEqual(d.vec4f(0.75, 1, 1.5, 2));
+    expect(div(d.vec4f(4, 7, 8, 10), 2)).toEqual(d.vec4f(2, 3.5, 4, 5));
+    expect(div(d.vec4f(0.09, 0.09, 0.09, 0.09), 0.3)).toEqual(d.vec4f(0.3));
+  });
+
+  it('computes quotient of a d.vec4u and a number', () => {
+    expect(div(d.vec4u(2, 2, 2, 2), 2)).toEqual(d.vec4u(1));
+    expect(div(d.vec4u(16, 16, 16, 16), 8)).toEqual(d.vec4u(2));
+  });
+
+  it('computes quotient of a d.vec4i and a number', () => {
+    expect(div(d.vec4i(1, 2, 3, 4), -1)).toEqual(d.vec4i(-1, -2, -3, -4));
+    expect(div(d.vec4i(0, 0, 0, 0), 1)).toEqual(d.vec4i());
+    expect(div(d.vec4i(16, 16, 16, 16), 8)).toEqual(d.vec4i(2));
+  });
+
+  it('computes quotient of a d.vec4f and a d.vec4f', () => {
+    expect(div(d.vec4f(1.5, 2, 3, 4), d.vec4f(2, 2, 2, 2))).toEqual(
+      d.vec4f(0.75, 1, 1.5, 2),
+    );
+    expect(div(d.vec4f(0.09, 0.09, 0.09, 0.09), d.vec4f(0.3))).toEqual(
+      d.vec4f(0.3),
+    );
+  });
+
+  it('computes quotient of a d.vec4u and a d.vec4u', () => {
+    expect(div(d.vec4u(2, 2, 2, 2), d.vec4u(2))).toEqual(d.vec4u(1));
+    expect(div(d.vec4u(16, 16, 16, 16), d.vec4u(8))).toEqual(d.vec4u(2));
+  });
+
+  it('computes quotient of a d.vec4i and a d.vec4i', () => {
+    expect(div(d.vec4i(1, 2, 3, 4), d.vec4i(-1))).toEqual(
+      d.vec4i(-1, -2, -3, -4),
+    );
+    expect(div(d.vec4i(0, 0, 0, 0), d.vec4i(1))).toEqual(d.vec4i());
+    expect(div(d.vec4i(16, 16, 16, 16), d.vec4i(8))).toEqual(d.vec4i(2));
+  });
+});