subgroup: Add subgroup invocation id and subgroup size builtins

gfx-rs · Oct 1, 2023 · a80fa00 · a80fa00
1 parent 4f2a204
commit a80fa00
Show file tree

Hide file tree

Showing 8 changed files with 49 additions and 3 deletions.
diff --git a/src/back/glsl/mod.rs b/src/back/glsl/mod.rs
@@ -4171,6 +4171,9 @@ const fn glsl_built_in(
         Bi::WorkGroupId => "gl_WorkGroupID",
         Bi::WorkGroupSize => "gl_WorkGroupSize",
         Bi::NumWorkGroups => "gl_NumWorkGroups",
+        // subgroup
+        Bi::SubgroupInvocationId => "gl_SubgroupInvocationID",
+        Bi::SubgroupSize => "gl_SubgroupSize",
     }
 }
 

diff --git a/src/back/hlsl/conv.rs b/src/back/hlsl/conv.rs
@@ -163,9 +163,12 @@ impl crate::BuiltIn {
             // to this field will get replaced with references to `SPECIAL_CBUF_VAR`
             // in `Writer::write_expr`.
             Self::NumWorkGroups => "SV_GroupID",
-            Self::BaseInstance | Self::BaseVertex | Self::WorkGroupSize => {
-                return Err(Error::Unimplemented(format!("builtin {self:?}")))
-            }
+
+            Self::SubgroupInvocationId
+            | Self::SubgroupSize
+            | Self::BaseInstance
+            | Self::BaseVertex
+            | Self::WorkGroupSize => return Err(Error::Unimplemented(format!("builtin {self:?}"))),
             Self::PointSize | Self::ViewIndex | Self::PointCoord => {
                 return Err(Error::Custom(format!("Unsupported builtin {self:?}")))
             }

diff --git a/src/back/msl/mod.rs b/src/back/msl/mod.rs
@@ -411,6 +411,9 @@ impl ResolvedBinding {
                     Bi::WorkGroupId => "threadgroup_position_in_grid",
                     Bi::WorkGroupSize => "dispatch_threads_per_threadgroup",
                     Bi::NumWorkGroups => "threadgroups_per_grid",
+                    // subgroup
+                    Bi::SubgroupInvocationId => "simdgroup_index_in_threadgroup",
+                    Bi::SubgroupSize => "simdgroups_per_threadgroup",
                     Bi::CullDistance | Bi::ViewIndex => {
                         return Err(Error::UnsupportedBuiltIn(built_in))
                     }

diff --git a/src/back/spv/writer.rs b/src/back/spv/writer.rs
@@ -1546,6 +1546,24 @@ impl Writer {
                     Bi::WorkGroupId => BuiltIn::WorkgroupId,
                     Bi::WorkGroupSize => BuiltIn::WorkgroupSize,
                     Bi::NumWorkGroups => BuiltIn::NumWorkgroups,
+                    // Subgroup
+                    Bi::SubgroupInvocationId => {
+                        self.require_any(
+                            "`subgroup_invocation_id` built-in",
+                            &[spirv::Capability::GroupNonUniform],
+                        )?;
+                        BuiltIn::SubgroupLocalInvocationId
+                    }
+                    Bi::SubgroupSize => {
+                        self.require_any(
+                            "`subgroup_invocation_id` built-in",
+                            &[
+                                spirv::Capability::GroupNonUniform,
+                                spirv::Capability::SubgroupBallotKHR,
+                            ],
+                        )?;
+                        BuiltIn::SubgroupSize
+                    }
                 };
 
                 self.decorate(id, Decoration::BuiltIn, &[built_in as u32]);

diff --git a/src/back/wgsl/writer.rs b/src/back/wgsl/writer.rs
@@ -1778,6 +1778,8 @@ fn builtin_str(built_in: crate::BuiltIn) -> Result<&'static str, Error> {
         Bi::SampleMask => "sample_mask",
         Bi::PrimitiveIndex => "primitive_index",
         Bi::ViewIndex => "view_index",
+        Bi::SubgroupInvocationId => "subgroup_invocation_id",
+        Bi::SubgroupSize => "subgroup_size",
         Bi::BaseInstance
         | Bi::BaseVertex
         | Bi::ClipDistance

diff --git a/src/front/wgsl/parse/conv.rs b/src/front/wgsl/parse/conv.rs
@@ -34,6 +34,9 @@ pub fn map_built_in(word: &str, span: Span) -> Result<crate::BuiltIn, Error<'_>>
         "local_invocation_index" => crate::BuiltIn::LocalInvocationIndex,
         "workgroup_id" => crate::BuiltIn::WorkGroupId,
         "num_workgroups" => crate::BuiltIn::NumWorkGroups,
+        // subgroup
+        "subgroup_invocation_id" => crate::BuiltIn::SubgroupInvocationId,
+        "subgroup_size" => crate::BuiltIn::SubgroupSize,
         _ => return Err(Error::UnknownBuiltin(span)),
     })
 }

diff --git a/src/lib.rs b/src/lib.rs
@@ -435,6 +435,9 @@ pub enum BuiltIn {
     WorkGroupId,
     WorkGroupSize,
     NumWorkGroups,
+    // subgroup
+    SubgroupInvocationId,
+    SubgroupSize,
 }
 
 /// Number of bytes per scalar.

diff --git a/src/valid/interface.rs b/src/valid/interface.rs
@@ -299,6 +299,17 @@ impl VaryingContext<'_> {
                                 width,
                             },
                     ),
+                    Bi::SubgroupInvocationId | Bi::SubgroupSize => (
+                        match self.stage {
+                            St::Compute | St::Fragment => !self.output,
+                            St::Vertex => false,
+                        },
+                        *ty_inner
+                            == Ti::Scalar {
+                                kind: Sk::Uint,
+                                width,
+                            },
+                    ),
                 };
 
                 if !visible {