Merge branch 'main' into chore/lower-node-tgpucli

Author: cieplypolar

packages/typegpu/src/builtin.ts

@@ -32,7 +32,9 @@ export type BuiltinSampleMask = Decorated<U32, [Builtin<'sample_mask'>]>;
 export type BuiltinLocalInvocationId = Decorated<Vec3u, [Builtin<'local_invocation_id'>]>;
 export type BuiltinLocalInvocationIndex = Decorated<U32, [Builtin<'local_invocation_index'>]>;
 export type BuiltinGlobalInvocationId = Decorated<Vec3u, [Builtin<'global_invocation_id'>]>;
+export type BuiltinGlobalInvocationIndex = Decorated<U32, [Builtin<'global_invocation_index'>]>;
 export type BuiltinWorkgroupId = Decorated<Vec3u, [Builtin<'workgroup_id'>]>;
+export type BuiltinWorkgroupIndex = Decorated<U32, [Builtin<'workgroup_index'>]>;
 export type BuiltinNumWorkgroups = Decorated<Vec3u, [Builtin<'num_workgroups'>]>;
 export type BuiltinSubgroupInvocationId = Decorated<U32, [Builtin<'subgroup_invocation_id'>]>;
 export type BuiltinSubgroupSize = Decorated<U32, [Builtin<'subgroup_size'>]>;
@@ -64,7 +66,12 @@ export const builtin = {
   localInvocationId: defineBuiltin<BuiltinLocalInvocationId>(vec3u, 'local_invocation_id'),
   localInvocationIndex: defineBuiltin<BuiltinLocalInvocationIndex>(u32, 'local_invocation_index'),
   globalInvocationId: defineBuiltin<BuiltinGlobalInvocationId>(vec3u, 'global_invocation_id'),
+  globalInvocationIndex: defineBuiltin<BuiltinGlobalInvocationIndex>(
+    u32,
+    'global_invocation_index',
+  ),
   workgroupId: defineBuiltin<BuiltinWorkgroupId>(vec3u, 'workgroup_id'),
+  workgroupIndex: defineBuiltin<BuiltinWorkgroupIndex>(u32, 'workgroup_index'),
   numWorkgroups: defineBuiltin<BuiltinNumWorkgroups>(vec3u, 'num_workgroups'),
   subgroupInvocationId: defineBuiltin<BuiltinSubgroupInvocationId>(u32, 'subgroup_invocation_id'),
   subgroupSize: defineBuiltin<BuiltinSubgroupSize>(u32, 'subgroup_size'),
@@ -77,7 +84,9 @@ export type AnyComputeBuiltin =
   | BuiltinLocalInvocationId
   | BuiltinLocalInvocationIndex
   | BuiltinGlobalInvocationId
+  | BuiltinGlobalInvocationIndex
   | BuiltinWorkgroupId
+  | BuiltinWorkgroupIndex
   | BuiltinNumWorkgroups
   | BuiltinSubgroupInvocationId
   | BuiltinSubgroupSize

packages/typegpu/src/core/pipeline/computePipeline.ts

@@ -19,7 +19,7 @@ import { isGPUCommandEncoder, isGPUComputePassEncoder } from './typeGuards.ts';
 import { logDataFromGPU } from '../../tgsl/consoleLog/deserializers.ts';
 import type { LogResources } from '../../tgsl/consoleLog/types.ts';
 import { isGPUBuffer, type ResolutionCtx, type SelfResolvable } from '../../types.ts';
-import { wgslExtensions, wgslExtensionToFeatureName } from '../../wgslExtensions.ts';
+import { wgslEnableExtensions, wgslEnableExtensionToFeatureName } from '../../wgslExtensions.ts';
 import type { IORecord } from '../function/fnTypes.ts';
 import type { TgpuComputeFn } from '../function/tgpuComputeFn.ts';
 import { namespace } from '../resolve/namespace.ts';
@@ -355,8 +355,8 @@ class ComputePipelineCore implements SelfResolvable {
   public unwrap(): Memo {
     if (this._memo === undefined) {
       const device = this.root.device;
-      const enableExtensions = wgslExtensions.filter((extension) =>
-        this.root.enabledFeatures.has(wgslExtensionToFeatureName[extension]),
+      const enableExtensions = wgslEnableExtensions.filter((extension) =>
+        this.root.enabledFeatures.has(wgslEnableExtensionToFeatureName[extension]),
       );
 
       // Resolving code

packages/typegpu/src/core/pipeline/renderPipeline.ts

@@ -38,7 +38,7 @@ import { logDataFromGPU } from '../../tgsl/consoleLog/deserializers.ts';
 import type { LogResources } from '../../tgsl/consoleLog/types.ts';
 import type { ResolutionCtx, SelfResolvable } from '../../types.ts';
 import { isGPUBuffer } from '../../types.ts';
-import { wgslExtensions, wgslExtensionToFeatureName } from '../../wgslExtensions.ts';
+import { wgslEnableExtensions, wgslEnableExtensionToFeatureName } from '../../wgslExtensions.ts';
 import {
   type AnyAutoCustoms,
   AutoFragmentFn,
@@ -1073,8 +1073,8 @@ class RenderPipelineCore implements SelfResolvable {
 
     const { root, descriptor: tgpuDescriptor } = this.options;
     const device = root.device;
-    const enableExtensions = wgslExtensions.filter((extension) =>
-      root.enabledFeatures.has(wgslExtensionToFeatureName[extension]),
+    const enableExtensions = wgslEnableExtensions.filter((extension) =>
+      root.enabledFeatures.has(wgslEnableExtensionToFeatureName[extension]),
     );
 
     // Resolving code

packages/typegpu/src/core/resolve/tgpuResolve.ts

@@ -5,7 +5,7 @@ import { $internal, $resolve } from '../../shared/symbols.ts';
 import { isBindGroupLayout } from '../../tgpuBindGroupLayout.ts';
 import type { ShaderGenerator } from '../../tgsl/shaderGenerator.ts';
 import type { ResolvableObject, SelfResolvable, Wgsl } from '../../types.ts';
-import type { WgslExtension } from '../../wgslExtensions.ts';
+import type { WgslEnableExtension } from '../../wgslExtensions.ts';
 import { isPipeline } from '../pipeline/typeGuards.ts';
 import type { Configurable, ExperimentalTgpuRoot } from '../root/rootTypes.ts';
 import { applyExternals, replaceExternalsInWgsl } from './externals.ts';
@@ -32,7 +32,7 @@ export interface TgpuResolveOptions {
   /**
    * List of WGSL shader extensions to enable.
    */
-  enableExtensions?: WgslExtension[] | undefined;
+  enableExtensions?: WgslEnableExtension[] | undefined;
   /**
    * **NOTE: This is an unstable API and may change in the future.**
    *

packages/typegpu/src/data/attributes.ts

@@ -75,7 +75,9 @@ export const builtinNames = [
   'local_invocation_id',
   'local_invocation_index',
   'global_invocation_id',
+  'global_invocation_index',
   'workgroup_id',
+  'workgroup_index',
   'num_workgroups',
   'subgroup_invocation_id',
   'subgroup_size',

packages/typegpu/src/resolutionCtx.ts

@@ -50,7 +50,7 @@ import type {
   Wgsl,
 } from './types.ts';
 import { CodegenState, isSelfResolvable, NormalState } from './types.ts';
-import type { WgslExtension } from './wgslExtensions.ts';
+import type { WgslEnableExtension } from './wgslExtensions.ts';
 import { getName, hasTinyestMetadata, setName } from './shared/meta.ts';
 import { FuncParameterType } from 'tinyest';
 import { accessProp } from './tgsl/accessProp.ts';
@@ -73,7 +73,7 @@ import { validateIdentifier, sanitizePrimer } from './nameUtils.ts';
 const CATCHALL_BIND_GROUP_IDX_MARKER = '#CATCHALL#';
 
 export type ResolutionCtxImplOptions = {
-  readonly enableExtensions?: WgslExtension[] | undefined;
+  readonly enableExtensions?: WgslEnableExtension[] | undefined;
   readonly shaderGenerator?: ShaderGenerator | undefined;
   readonly config?: ((cfg: Configurable) => Configurable) | undefined;
   readonly root?: ExperimentalTgpuRoot | undefined;
@@ -400,7 +400,7 @@ export class ResolutionCtxImpl implements ResolutionCtx {
   public readonly fixedBindings: FixedBindingConfig[] = [];
   // --
 
-  public readonly enableExtensions: WgslExtension[] | undefined;
+  public readonly enableExtensions: WgslEnableExtension[] | undefined;
   public expectedType: BaseData | undefined;
 
   /**

packages/typegpu/src/std/extensions.ts

@@ -1,16 +1,16 @@
 import { comptime } from '../core/function/comptime.ts';
 import { getResolutionCtx } from '../execMode.ts';
-import { type WgslExtension, wgslExtensions } from '../wgslExtensions.ts';
+import { type WgslEnableExtension, wgslEnableExtensions } from '../wgslExtensions.ts';
 
-export const extensionEnabled = comptime((extensionName: WgslExtension): boolean => {
+export const extensionEnabled = comptime((extensionName: WgslEnableExtension): boolean => {
   const resolutionCtx = getResolutionCtx();
   if (!resolutionCtx) {
     throw new Error(
       "Functions using `extensionEnabled` cannot be called directly. Either generate WGSL from them, or use tgpu['~unstable'].simulate(...)",
     );
   }
 
-  if (typeof extensionName !== 'string' || !wgslExtensions.includes(extensionName)) {
+  if (typeof extensionName !== 'string' || !wgslEnableExtensions.includes(extensionName)) {
     throw new Error(
       `extensionEnabled has to be called with a string literal representing a valid WGSL extension name. Got: '${extensionName}'`,
     );

packages/typegpu/src/types.ts

@@ -45,7 +45,7 @@ import {
   $resolve,
 } from './shared/symbols.ts';
 import type { TgpuBindGroupLayout, TgpuLayoutEntry } from './tgpuBindGroupLayout.ts';
-import type { WgslExtension } from './wgslExtensions.ts';
+import type { WgslEnableExtension } from './wgslExtensions.ts';
 import type { Infer } from './shared/repr.ts';
 import { ShaderGenerator } from './tgsl/shaderGenerator.ts';
 
@@ -275,7 +275,7 @@ export interface ResolutionCtx {
   };
 
   readonly mode: ExecState;
-  readonly enableExtensions: WgslExtension[] | undefined;
+  readonly enableExtensions: WgslEnableExtension[] | undefined;
   readonly gen: ShaderGenerator;
 
   addDeclaration(declaration: string): void;

packages/typegpu/src/wgslExtensions.ts

@@ -1,16 +1,16 @@
-export const wgslExtensions = [
+export const wgslEnableExtensions = [
   'f16',
   'clip_distances',
   'dual_source_blending',
   'subgroups',
   'primitive_index',
 ] as const;
-export type WgslExtension = (typeof wgslExtensions)[number];
+export type WgslEnableExtension = (typeof wgslEnableExtensions)[number];
 
-export const wgslExtensionToFeatureName: Record<WgslExtension, GPUFeatureName> = {
+export const wgslEnableExtensionToFeatureName: Record<WgslEnableExtension, GPUFeatureName> = {
   f16: 'shader-f16',
   clip_distances: 'clip-distances',
   dual_source_blending: 'dual-source-blending',
   subgroups: 'subgroups',
-  primitive_index: 'primitive-index' as GPUFeatureName,
+  primitive_index: 'primitive-index',
 };

packages/typegpu/tests/tgslFn.test.ts

@@ -302,6 +302,24 @@ describe('TGSL tgpu.fn function', () => {
     `);
   });
 
+  it('resolves linear compute builtins', () => {
+    const computeFn = tgpu.computeFn({
+      in: {
+        globalIndex: builtin.globalInvocationIndex,
+        workgroupIndex: builtin.workgroupIndex,
+      },
+      workgroupSize: [24],
+    })((input) => {
+      const index = input.globalIndex + input.workgroupIndex;
+    });
+
+    expect(tgpu.resolve([computeFn])).toMatchInlineSnapshot(`
+      "@compute @workgroup_size(24) fn computeFn(@builtin(global_invocation_index) globalIndex: u32, @builtin(workgroup_index) workgroupIndex: u32) {
+        let index = (globalIndex + workgroupIndex);
+      }"
+    `);
+  });
+
   it('allows destructuring the input argument in computeFn', () => {
     const computeFn = tgpu
       .computeFn({