feat: Add primitive_index, subgroup_id and num_subgroups builtins (#2028)

Author: iwoplaza

packages/typegpu/src/builtin.ts

@@ -23,13 +23,17 @@ import { $internal } from './shared/symbols.ts';
 
 export type BuiltinVertexIndex = Decorated<U32, [Builtin<'vertex_index'>]>;
 export type BuiltinInstanceIndex = Decorated<U32, [Builtin<'instance_index'>]>;
-export type BuiltinPosition = Decorated<Vec4f, [Builtin<'position'>]>;
 export type BuiltinClipDistances = Decorated<
   WgslArray<U32>,
   [Builtin<'clip_distances'>]
 >;
+export type BuiltinPosition = Decorated<Vec4f, [Builtin<'position'>]>;
 export type BuiltinFrontFacing = Decorated<Bool, [Builtin<'front_facing'>]>;
 export type BuiltinFragDepth = Decorated<F32, [Builtin<'frag_depth'>]>;
+export type BuiltinPrimitiveIndex = Decorated<
+  U32,
+  [Builtin<'primitive_index'>]
+>;
 export type BuiltinSampleIndex = Decorated<U32, [Builtin<'sample_index'>]>;
 export type BuiltinSampleMask = Decorated<U32, [Builtin<'sample_mask'>]>;
 export type BuiltinLocalInvocationId = Decorated<
@@ -54,6 +58,8 @@ export type BuiltinSubgroupInvocationId = Decorated<
   [Builtin<'subgroup_invocation_id'>]
 >;
 export type BuiltinSubgroupSize = Decorated<U32, [Builtin<'subgroup_size'>]>;
+export type BuiltinSubgroupId = Decorated<U32, [Builtin<'subgroup_id'>]>;
+export type BuiltinNumSubgroups = Decorated<U32, [Builtin<'num_subgroups'>]>;
 
 function defineBuiltin<T extends Decorated | LooseDecorated>(
   dataType: AnyWgslData,
@@ -70,13 +76,14 @@ function defineBuiltin<T extends Decorated | LooseDecorated>(
 export const builtin = {
   vertexIndex: defineBuiltin<BuiltinVertexIndex>(u32, 'vertex_index'),
   instanceIndex: defineBuiltin<BuiltinInstanceIndex>(u32, 'instance_index'),
-  position: defineBuiltin<BuiltinPosition>(vec4f, 'position'),
   clipDistances: defineBuiltin<BuiltinClipDistances>(
     arrayOf(u32, 8),
     'clip_distances',
   ),
+  position: defineBuiltin<BuiltinPosition>(vec4f, 'position'),
   frontFacing: defineBuiltin<BuiltinFrontFacing>(bool, 'front_facing'),
   fragDepth: defineBuiltin<BuiltinFragDepth>(f32, 'frag_depth'),
+  primitiveIndex: defineBuiltin<BuiltinPrimitiveIndex>(u32, 'primitive_index'),
   sampleIndex: defineBuiltin<BuiltinSampleIndex>(u32, 'sample_index'),
   sampleMask: defineBuiltin<BuiltinSampleMask>(u32, 'sample_mask'),
   localInvocationId: defineBuiltin<BuiltinLocalInvocationId>(
@@ -98,6 +105,8 @@ export const builtin = {
     'subgroup_invocation_id',
   ),
   subgroupSize: defineBuiltin<BuiltinSubgroupSize>(u32, 'subgroup_size'),
+  subgroupId: defineBuiltin<BuiltinSubgroupId>(u32, 'subgroup_id'),
+  numSubgroups: defineBuiltin<BuiltinNumSubgroups>(u32, 'num_subgroups'),
 } as const;
 
 export type AnyBuiltin = (typeof builtin)[keyof typeof builtin];
@@ -108,12 +117,15 @@ export type AnyComputeBuiltin =
   | BuiltinWorkgroupId
   | BuiltinNumWorkgroups
   | BuiltinSubgroupInvocationId
-  | BuiltinSubgroupSize;
+  | BuiltinSubgroupSize
+  | BuiltinSubgroupId
+  | BuiltinNumSubgroups;
 export type AnyVertexInputBuiltin = BuiltinVertexIndex | BuiltinInstanceIndex;
 export type AnyVertexOutputBuiltin = BuiltinClipDistances | BuiltinPosition;
 export type AnyFragmentInputBuiltin =
   | BuiltinPosition
   | BuiltinFrontFacing
+  | BuiltinPrimitiveIndex
   | BuiltinSampleIndex
   | BuiltinSampleMask
   | BuiltinSubgroupInvocationId

packages/typegpu/src/data/attributes.ts

@@ -60,10 +60,11 @@ import {
 export const builtinNames = [
   'vertex_index',
   'instance_index',
-  'position',
   'clip_distances',
+  'position',
   'front_facing',
   'frag_depth',
+  'primitive_index',
   'sample_index',
   'sample_mask',
   'fragment',
@@ -74,6 +75,8 @@ export const builtinNames = [
   'num_workgroups',
   'subgroup_invocation_id',
   'subgroup_size',
+  'subgroup_id',
+  'num_subgroups',
 ] as const;
 
 export type BuiltinName = (typeof builtinNames)[number];