hlsl-ir: implement support for subgroup_invocation_id and subgroup_size Bug: b/363199902 Bug: b/42251045 Change-Id: I839c58a2743a59a4b07ced6a84b19317bf7a8315 Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/205419 Reviewed-by: James Price <jrprice@google.com> Commit-Queue: Antonio Maiorano <amaiorano@google.com>

diff --git a/src/tint/lang/core/ir/transform/shader_io.cc b/src/tint/lang/core/ir/transform/shader_io.cc
index 4fd3502..20c295e 100644
--- a/src/tint/lang/core/ir/transform/shader_io.cc
+++ b/src/tint/lang/core/ir/transform/shader_io.cc

@@ -106,8 +106,8 @@
         TINT_DEFER(backend = nullptr);
 
         // Process the parameters and return value to prepare for building a wrapper function.
-        GatherInputs();
-        GatherOutput();
+        GatherInputs();  // Calls backend->AddInput() for each input
+        GatherOutput();  // Calls backend->AddOutput() for each output
 
         // Add an output for the vertex point size if needed.
         std::optional<uint32_t> vertex_point_size_index;
@@ -166,14 +166,17 @@
                     auto attributes = member->Attributes();
                     if (attributes.interpolation &&
                         ep->Stage() != Function::PipelineStage::kFragment) {
+                        // Strip interpolation on non-fragment inputs
                         attributes.interpolation = {};
                     }
-                    backend->AddInput(ir.symbols.Register(name), member->Type(), attributes);
+                    backend->AddInput(ir.symbols.Register(name), member->Type(),
+                                      std::move(attributes));
                 }
             } else {
                 // Pull out the IO attributes and remove them from the parameter.
                 auto attributes = param->Attributes();
                 if (attributes.interpolation && ep->Stage() != Function::PipelineStage::kFragment) {
+                    // Strip interpolation on non-fragment inputs
                     attributes.interpolation = {};
                 }
                 param->SetAttributes({});
@@ -195,6 +198,7 @@
                 auto name = str->Name().Name() + "_" + member->Name().Name();
                 auto attributes = member->Attributes();
                 if (attributes.interpolation && ep->Stage() != Function::PipelineStage::kVertex) {
+                    // Strip interpolation on non-vertex outputs
                     attributes.interpolation = {};
                 }
                 backend->AddOutput(ir.symbols.Register(name), member->Type(), attributes);
@@ -203,6 +207,7 @@
             // Pull out the IO attributes and remove them from the original function.
             auto attributes = ep->ReturnAttributes();
             if (attributes.interpolation && ep->Stage() != Function::PipelineStage::kVertex) {
+                // Strip interpolation on non-vertex outputs
                 attributes.interpolation = {};
             }
             ep->SetReturnAttributes({});

diff --git a/src/tint/lang/hlsl/builtin_fn.cc b/src/tint/lang/hlsl/builtin_fn.cc
index 67edf50c..197dc17 100644
--- a/src/tint/lang/hlsl/builtin_fn.cc
+++ b/src/tint/lang/hlsl/builtin_fn.cc

@@ -90,6 +90,8 @@
             return "unpack_u8u32";
         case BuiltinFn::kWaveGetLaneIndex:
             return "WaveGetLaneIndex";
+        case BuiltinFn::kWaveGetLaneCount:
+            return "WaveGetLaneCount";
         case BuiltinFn::kWaveReadLaneAt:
             return "WaveReadLaneAt";
         case BuiltinFn::kLoad:

diff --git a/src/tint/lang/hlsl/builtin_fn.h b/src/tint/lang/hlsl/builtin_fn.h
index 253a766..86c48aa 100644
--- a/src/tint/lang/hlsl/builtin_fn.h
+++ b/src/tint/lang/hlsl/builtin_fn.h

@@ -71,6 +71,7 @@
     kUnpackS8S32,
     kUnpackU8U32,
     kWaveGetLaneIndex,
+    kWaveGetLaneCount,
     kWaveReadLaneAt,
     kLoad,
     kLoad2,

diff --git a/src/tint/lang/hlsl/hlsl.def b/src/tint/lang/hlsl/hlsl.def
index eb9795d..d80b9ba 100644
--- a/src/tint/lang/hlsl/hlsl.def
+++ b/src/tint/lang/hlsl/hlsl.def

@@ -236,6 +236,7 @@
 fn unpack_u8u32(uint8_t4_packed) -> vec4<u32>
 
 fn WaveGetLaneIndex() -> u32
+fn WaveGetLaneCount() -> u32
 implicit(T: fiu32_f16) fn WaveReadLaneAt(T, u32) -> T
 implicit(N: num, T: fiu32_f16) fn WaveReadLaneAt(vec<N, T>, u32) -> vec<N, T>
 

diff --git a/src/tint/lang/hlsl/intrinsic/data.cc b/src/tint/lang/hlsl/intrinsic/data.cc
index e3cf060..f2c9bb9 100644
--- a/src/tint/lang/hlsl/intrinsic/data.cc
+++ b/src/tint/lang/hlsl/intrinsic/data.cc

@@ -5657,13 +5657,19 @@
   },
   {
     /* [24] */
+    /* fn WaveGetLaneCount() -> u32 */
+    /* num overloads */ 1,
+    /* overloads */ OverloadIndex(156),
+  },
+  {
+    /* [25] */
     /* fn WaveReadLaneAt[T : fiu32_f16](T, u32) -> T */
     /* fn WaveReadLaneAt[N : num, T : fiu32_f16](vec<N, T>, u32) -> vec<N, T> */
     /* num overloads */ 2,
     /* overloads */ OverloadIndex(148),
   },
   {
-    /* [25] */
+    /* [26] */
     /* fn Load(byte_address_buffer<readable>, offset: u32) -> u32 */
     /* fn Load[T : fiu32](texture: texture_1d<T>, location: vec2<i32>) -> vec4<T> */
     /* fn Load[T : fiu32](texture: texture_2d<T>, location: vec3<i32>) -> vec4<T> */
@@ -5689,97 +5695,97 @@
     /* overloads */ OverloadIndex(27),
   },
   {
-    /* [26] */
+    /* [27] */
     /* fn Load2(byte_address_buffer<readable>, offset: u32) -> vec2<u32> */
     /* num overloads */ 1,
     /* overloads */ OverloadIndex(157),
   },
   {
-    /* [27] */
+    /* [28] */
     /* fn Load3(byte_address_buffer<readable>, offset: u32) -> vec3<u32> */
     /* num overloads */ 1,
     /* overloads */ OverloadIndex(158),
   },
   {
-    /* [28] */
+    /* [29] */
     /* fn Load4(byte_address_buffer<readable>, offset: u32) -> vec4<u32> */
     /* num overloads */ 1,
     /* overloads */ OverloadIndex(159),
   },
   {
-    /* [29] */
+    /* [30] */
     /* fn LoadF16(byte_address_buffer<readable>, offset: u32) -> f16 */
     /* num overloads */ 1,
     /* overloads */ OverloadIndex(160),
   },
   {
-    /* [30] */
+    /* [31] */
     /* fn Load2F16(byte_address_buffer<readable>, offset: u32) -> vec2<f16> */
     /* num overloads */ 1,
     /* overloads */ OverloadIndex(161),
   },
   {
-    /* [31] */
+    /* [32] */
     /* fn Load3F16(byte_address_buffer<readable>, offset: u32) -> vec3<f16> */
     /* num overloads */ 1,
     /* overloads */ OverloadIndex(162),
   },
   {
-    /* [32] */
+    /* [33] */
     /* fn Load4F16(byte_address_buffer<readable>, offset: u32) -> vec4<f16> */
     /* num overloads */ 1,
     /* overloads */ OverloadIndex(163),
   },
   {
-    /* [33] */
+    /* [34] */
     /* fn Store(byte_address_buffer<writable>, offset: u32, value: u32) */
     /* num overloads */ 1,
     /* overloads */ OverloadIndex(164),
   },
   {
-    /* [34] */
+    /* [35] */
     /* fn Store2(byte_address_buffer<writable>, offset: u32, value: vec2<u32>) */
     /* num overloads */ 1,
     /* overloads */ OverloadIndex(165),
   },
   {
-    /* [35] */
+    /* [36] */
     /* fn Store3(byte_address_buffer<writable>, offset: u32, value: vec3<u32>) */
     /* num overloads */ 1,
     /* overloads */ OverloadIndex(166),
   },
   {
-    /* [36] */
+    /* [37] */
     /* fn Store4(byte_address_buffer<writable>, offset: u32, value: vec4<u32>) */
     /* num overloads */ 1,
     /* overloads */ OverloadIndex(167),
   },
   {
-    /* [37] */
+    /* [38] */
     /* fn StoreF16(byte_address_buffer<writable>, offset: u32, value: f16) */
     /* num overloads */ 1,
     /* overloads */ OverloadIndex(168),
   },
   {
-    /* [38] */
+    /* [39] */
     /* fn Store2F16(byte_address_buffer<writable>, offset: u32, value: vec2<f16>) */
     /* num overloads */ 1,
     /* overloads */ OverloadIndex(169),
   },
   {
-    /* [39] */
+    /* [40] */
     /* fn Store3F16(byte_address_buffer<writable>, offset: u32, value: vec3<f16>) */
     /* num overloads */ 1,
     /* overloads */ OverloadIndex(170),
   },
   {
-    /* [40] */
+    /* [41] */
     /* fn Store4F16(byte_address_buffer<writable>, offset: u32, value: vec4<f16>) */
     /* num overloads */ 1,
     /* overloads */ OverloadIndex(171),
   },
   {
-    /* [41] */
+    /* [42] */
     /* fn GatherCmp(texture: texture_depth_2d, sampler: sampler_comparison, coords: vec2<f32>, depth_ref: f32) -> vec4<f32> */
     /* fn GatherCmp(texture: texture_depth_2d, sampler: sampler_comparison, coords: vec2<f32>, depth_ref: f32, offset: vec2<i32>) -> vec4<f32> */
     /* fn GatherCmp(texture: texture_depth_2d_array, sampler: sampler_comparison, coords: vec3<f32>, depth_ref: f32) -> vec4<f32> */
@@ -5790,7 +5796,7 @@
     /* overloads */ OverloadIndex(105),
   },
   {
-    /* [42] */
+    /* [43] */
     /* fn Gather(texture: texture_depth_2d, sampler: sampler, coords: vec2<f32>) -> vec4<f32> */
     /* fn Gather(texture: texture_depth_2d, sampler: sampler, coords: vec2<f32>, offset: vec2<i32>) -> vec4<f32> */
     /* fn Gather(texture: texture_depth_2d_array, sampler: sampler, coords: vec3<f32>) -> vec4<f32> */
@@ -5801,7 +5807,7 @@
     /* overloads */ OverloadIndex(111),
   },
   {
-    /* [43] */
+    /* [44] */
     /* fn GatherAlpha[T : fiu32](texture: texture_2d<T>, sampler: sampler, coords: vec2<f32>) -> vec4<T> */
     /* fn GatherAlpha[T : fiu32](texture: texture_2d<T>, sampler: sampler, coords: vec2<f32>, offset: vec2<i32>) -> vec4<T> */
     /* fn GatherAlpha[T : fiu32](texture: texture_2d_array<T>, sampler: sampler, coords: vec3<f32>) -> vec4<T> */
@@ -5812,7 +5818,7 @@
     /* overloads */ OverloadIndex(117),
   },
   {
-    /* [44] */
+    /* [45] */
     /* fn GatherBlue[T : fiu32](texture: texture_2d<T>, sampler: sampler, coords: vec2<f32>) -> vec4<T> */
     /* fn GatherBlue[T : fiu32](texture: texture_2d<T>, sampler: sampler, coords: vec2<f32>, offset: vec2<i32>) -> vec4<T> */
     /* fn GatherBlue[T : fiu32](texture: texture_2d_array<T>, sampler: sampler, coords: vec3<f32>) -> vec4<T> */
@@ -5823,7 +5829,7 @@
     /* overloads */ OverloadIndex(117),
   },
   {
-    /* [45] */
+    /* [46] */
     /* fn GatherGreen[T : fiu32](texture: texture_2d<T>, sampler: sampler, coords: vec2<f32>) -> vec4<T> */
     /* fn GatherGreen[T : fiu32](texture: texture_2d<T>, sampler: sampler, coords: vec2<f32>, offset: vec2<i32>) -> vec4<T> */
     /* fn GatherGreen[T : fiu32](texture: texture_2d_array<T>, sampler: sampler, coords: vec3<f32>) -> vec4<T> */
@@ -5834,7 +5840,7 @@
     /* overloads */ OverloadIndex(117),
   },
   {
-    /* [46] */
+    /* [47] */
     /* fn GatherRed[T : fiu32](texture: texture_2d<T>, sampler: sampler, coords: vec2<f32>) -> vec4<T> */
     /* fn GatherRed[T : fiu32](texture: texture_2d<T>, sampler: sampler, coords: vec2<f32>, offset: vec2<i32>) -> vec4<T> */
     /* fn GatherRed[T : fiu32](texture: texture_2d_array<T>, sampler: sampler, coords: vec3<f32>) -> vec4<T> */
@@ -5845,7 +5851,7 @@
     /* overloads */ OverloadIndex(117),
   },
   {
-    /* [47] */
+    /* [48] */
     /* fn GetDimensions[A : access](byte_address_buffer<A>, width: ptr<function, u32, writable>) */
     /* fn GetDimensions[T : fiu32](texture: texture_1d<T>, width: ptr<function, u32, writable>) */
     /* fn GetDimensions[T : fiu32](texture: texture_1d<T>, level: u32, width: ptr<function, u32, writable>, num_levels: ptr<function, u32, writable>) */
@@ -5877,7 +5883,7 @@
     /* overloads */ OverloadIndex(0),
   },
   {
-    /* [48] */
+    /* [49] */
     /* fn Sample(texture: texture_1d<f32>, sampler: sampler, coords: f32) -> vec4<f32> */
     /* fn Sample(texture: texture_2d<f32>, sampler: sampler, coords: vec2<f32>) -> vec4<f32> */
     /* fn Sample(texture: texture_2d<f32>, sampler: sampler, coords: vec2<f32>, offset: vec2<i32>) -> vec4<f32> */
@@ -5897,7 +5903,7 @@
     /* overloads */ OverloadIndex(48),
   },
   {
-    /* [49] */
+    /* [50] */
     /* fn SampleBias(texture: texture_2d<f32>, sampler: sampler, coords: vec2<f32>, bias: f32) -> vec4<f32> */
     /* fn SampleBias(texture: texture_2d<f32>, sampler: sampler, coords: vec2<f32>, bias: f32, offset: vec2<i32>) -> vec4<f32> */
     /* fn SampleBias(texture: texture_2d_array<f32>, sampler: sampler, coords: vec3<f32>, bias: f32) -> vec4<f32> */
@@ -5910,7 +5916,7 @@
     /* overloads */ OverloadIndex(89),
   },
   {
-    /* [50] */
+    /* [51] */
     /* fn SampleCmp(texture: texture_depth_2d, sampler: sampler_comparison, coords: vec2<f32>, value: f32) -> f32 */
     /* fn SampleCmp(texture: texture_depth_2d, sampler: sampler_comparison, coords: vec2<f32>, value: f32, offset: vec2<i32>) -> f32 */
     /* fn SampleCmp(texture: texture_depth_2d_array, sampler: sampler_comparison, coords: vec3<f32>, value: f32) -> f32 */
@@ -5921,7 +5927,7 @@
     /* overloads */ OverloadIndex(123),
   },
   {
-    /* [51] */
+    /* [52] */
     /* fn SampleCmpLevelZero(texture: texture_depth_2d, sampler: sampler_comparison, coords: vec2<f32>, value: f32) -> f32 */
     /* fn SampleCmpLevelZero(texture: texture_depth_2d, sampler: sampler_comparison, coords: vec2<f32>, value: f32, offset: vec2<i32>) -> f32 */
     /* fn SampleCmpLevelZero(texture: texture_depth_2d_array, sampler: sampler_comparison, coords: vec3<f32>, value: f32) -> f32 */
@@ -5932,7 +5938,7 @@
     /* overloads */ OverloadIndex(123),
   },
   {
-    /* [52] */
+    /* [53] */
     /* fn SampleGrad(texture: texture_2d<f32>, sampler: sampler, coords: vec2<f32>, ddx: vec2<f32>, ddy: vec2<f32>) -> vec4<f32> */
     /* fn SampleGrad(texture: texture_2d<f32>, sampler: sampler, coords: vec2<f32>, ddx: vec2<f32>, ddy: vec2<f32>, offset: vec2<i32>) -> vec4<f32> */
     /* fn SampleGrad(texture: texture_2d_array<f32>, sampler: sampler, coords: vec3<f32>, ddx: vec2<f32>, ddy: vec2<f32>) -> vec4<f32> */
@@ -5945,7 +5951,7 @@
     /* overloads */ OverloadIndex(97),
   },
   {
-    /* [53] */
+    /* [54] */
     /* fn SampleLevel(texture: texture_2d<f32>, sampler: sampler, coords: vec2<f32>, level: f32) -> vec4<f32> */
     /* fn SampleLevel(texture: texture_2d<f32>, sampler: sampler, coords: vec2<f32>, level: f32, offset: vec2<i32>) -> vec4<f32> */
     /* fn SampleLevel(texture: texture_2d_array<f32>, sampler: sampler, coords: vec3<f32>, level: f32) -> vec4<f32> */

diff --git a/src/tint/lang/hlsl/writer/raise/shader_io.cc b/src/tint/lang/hlsl/writer/raise/shader_io.cc
index 930ec26..df7e3dc 100644
--- a/src/tint/lang/hlsl/writer/raise/shader_io.cc
+++ b/src/tint/lang/hlsl/writer/raise/shader_io.cc

@@ -35,6 +35,8 @@
 #include "src/tint/lang/core/ir/module.h"
 #include "src/tint/lang/core/ir/transform/shader_io.h"
 #include "src/tint/lang/core/ir/validator.h"
+#include "src/tint/lang/hlsl/builtin_fn.h"
+#include "src/tint/lang/hlsl/ir/builtin_call.h"
 
 using namespace tint::core::fluent_types;     // NOLINT
 using namespace tint::core::number_suffixes;  // NOLINT
@@ -59,6 +61,10 @@
     /// The output values to return from the entry point.
     Vector<core::ir::Value*, 4> output_values;
 
+    // Indices of subgroup invocation id and size, if set
+    std::optional<uint32_t> subgroup_invocation_id_index;
+    std::optional<uint32_t> subgroup_size_index;
+
     /// Constructor
     StateImpl(core::ir::Module& mod, core::ir::Function* f) : ShaderIOBackendState(mod, f) {}
 
@@ -169,23 +175,34 @@
 
         Vector<MemberInfo, 4> input_data;
         for (uint32_t i = 0; i < inputs.Length(); ++i) {
+            // If subgroup invocation id or size, save the index for GetInput
+            if (auto builtin = inputs[i].attributes.builtin) {
+                if (*builtin == core::BuiltinValue::kSubgroupInvocationId) {
+                    subgroup_invocation_id_index = i;
+                    continue;
+                } else if (*builtin == core::BuiltinValue::kSubgroupSize) {
+                    subgroup_size_index = i;
+                    continue;
+                }
+            }
+
             input_data.Push(MemberInfo{inputs[i], i});
         }
 
-        input_indices.Resize(inputs.Length());
+        input_indices.Resize(input_data.Length());
 
         // Sort the struct members to satisfy HLSL interfacing matching rules.
         std::sort(input_data.begin(), input_data.end(),
                   [&](auto& x, auto& y) { return StructMemberComparator(x, y); });
 
-        for (auto input : input_data) {
+        for (auto& input : input_data) {
             input_indices[input.idx] = static_cast<uint32_t>(input_struct_members.Length());
             input_struct_members.Push(input.member);
         }
 
         if (!input_struct_members.IsEmpty()) {
-            auto* input_struct =
-                ty.Struct(ir.symbols.New(ir.NameOf(func).Name() + "_inputs"), input_struct_members);
+            auto* input_struct = ty.Struct(ir.symbols.New(ir.NameOf(func).Name() + "_inputs"),
+                                           std::move(input_struct_members));
             switch (func->Stage()) {
                 case core::ir::Function::PipelineStage::kFragment:
                     input_struct->AddUsage(core::type::PipelineStageUsage::kFragmentInput);
@@ -248,6 +265,17 @@
 
     /// @copydoc ShaderIO::BackendState::GetInput
     core::ir::Value* GetInput(core::ir::Builder& builder, uint32_t idx) override {
+        if (subgroup_invocation_id_index == idx) {
+            return builder
+                .Call<hlsl::ir::BuiltinCall>(ty.u32(), hlsl::BuiltinFn::kWaveGetLaneIndex)
+                ->Result(0);
+        }
+        if (subgroup_size_index == idx) {
+            return builder
+                .Call<hlsl::ir::BuiltinCall>(ty.u32(), hlsl::BuiltinFn::kWaveGetLaneCount)
+                ->Result(0);
+        }
+
         auto index = input_indices[idx];
 
         core::ir::Value* v = builder.Access(inputs[idx].type, input_param, u32(index))->Result(0);

diff --git a/src/tint/lang/hlsl/writer/raise/shader_io_test.cc b/src/tint/lang/hlsl/writer/raise/shader_io_test.cc
index f3d2c82..bd0d91d 100644
--- a/src/tint/lang/hlsl/writer/raise/shader_io_test.cc
+++ b/src/tint/lang/hlsl/writer/raise/shader_io_test.cc

@@ -1069,5 +1069,150 @@
     EXPECT_EQ(expect, str());
 }
 
+TEST_F(HlslWriterTransformTest, ShaderIOParameters_Subgroup_NonStruct) {
+    auto* subgroup_invocation_id = b.FunctionParam("id", ty.u32());
+    subgroup_invocation_id->SetBuiltin(core::BuiltinValue::kSubgroupInvocationId);
+
+    auto* subgroup_size = b.FunctionParam("size", ty.u32());
+    subgroup_size->SetBuiltin(core::BuiltinValue::kSubgroupSize);
+
+    auto* ep = b.Function("foo", ty.u32(), core::ir::Function::PipelineStage::kFragment);
+    ep->SetParams({subgroup_invocation_id, subgroup_size});
+
+    b.Append(ep->Block(), [&] {
+        auto* r = b.Multiply(ty.u32(), subgroup_invocation_id, subgroup_size);
+        b.Return(ep, r);
+    });
+
+    auto* src = R"(
+%foo = @fragment func(%id:u32 [@subgroup_invocation_id], %size:u32 [@subgroup_size]):u32 {
+  $B1: {
+    %4:u32 = mul %id, %size
+    ret %4
+  }
+}
+)";
+    EXPECT_EQ(src, str());
+
+    auto* expect = R"(
+foo_outputs = struct @align(4) {
+  tint_symbol:u32 @offset(0)
+}
+
+%foo_inner = func(%id:u32, %size:u32):u32 {
+  $B1: {
+    %4:u32 = mul %id, %size
+    ret %4
+  }
+}
+%foo = @fragment func():foo_outputs {
+  $B2: {
+    %6:u32 = hlsl.WaveGetLaneIndex
+    %7:u32 = hlsl.WaveGetLaneCount
+    %8:u32 = call %foo_inner, %6, %7
+    %9:foo_outputs = construct %8
+    ret %9
+  }
+}
+)";
+
+    Run(ShaderIO);
+
+    EXPECT_EQ(expect, str());
+}
+
+TEST_F(HlslWriterTransformTest, ShaderIOParameters_Subgroup_Struct) {
+    auto* str_ty = ty.Struct(mod.symbols.New("Inputs"),
+                             {
+                                 {
+                                     mod.symbols.New("id"),
+                                     ty.u32(),
+                                     core::IOAttributes{
+                                         /* location */ std::nullopt,
+                                         /* blend_src */ std::nullopt,
+                                         /* color */ std::nullopt,
+                                         /* builtin */ core::BuiltinValue::kSubgroupInvocationId,
+                                         /* interpolation */ std::nullopt,
+                                         /* invariant */ false,
+                                     },
+                                 },
+                                 {
+                                     mod.symbols.New("size"),
+                                     ty.u32(),
+                                     core::IOAttributes{
+                                         /* location */ std::nullopt,
+                                         /* blend_src */ std::nullopt,
+                                         /* color */ std::nullopt,
+                                         /* builtin */ core::BuiltinValue::kSubgroupSize,
+                                         /* interpolation */ std::nullopt,
+                                         /* invariant */ false,
+                                     },
+                                 },
+                             });
+
+    auto* str_param = b.FunctionParam("inputs", str_ty);
+
+    auto* ep = b.Function("foo", ty.u32(), core::ir::Function::PipelineStage::kFragment);
+    ep->SetParams({str_param});
+
+    b.Append(ep->Block(), [&] {
+        auto* subgroup_invocation_id = b.Access(ty.u32(), str_param, 0_i);
+        auto* subgroup_size = b.Access(ty.u32(), str_param, 1_i);
+        auto* r = b.Multiply(ty.u32(), subgroup_invocation_id, subgroup_size);
+        b.Return(ep, r);
+    });
+
+    auto* src = R"(
+Inputs = struct @align(4) {
+  id:u32 @offset(0), @builtin(subgroup_invocation_id)
+  size:u32 @offset(4), @builtin(subgroup_size)
+}
+
+%foo = @fragment func(%inputs:Inputs):u32 {
+  $B1: {
+    %3:u32 = access %inputs, 0i
+    %4:u32 = access %inputs, 1i
+    %5:u32 = mul %3, %4
+    ret %5
+  }
+}
+)";
+    EXPECT_EQ(src, str());
+
+    auto* expect = R"(
+Inputs = struct @align(4) {
+  id:u32 @offset(0)
+  size:u32 @offset(4)
+}
+
+foo_outputs = struct @align(4) {
+  tint_symbol:u32 @offset(0)
+}
+
+%foo_inner = func(%inputs:Inputs):u32 {
+  $B1: {
+    %3:u32 = access %inputs, 0i
+    %4:u32 = access %inputs, 1i
+    %5:u32 = mul %3, %4
+    ret %5
+  }
+}
+%foo = @fragment func():foo_outputs {
+  $B2: {
+    %7:u32 = hlsl.WaveGetLaneIndex
+    %8:u32 = hlsl.WaveGetLaneCount
+    %9:Inputs = construct %7, %8
+    %10:u32 = call %foo_inner, %9
+    %11:foo_outputs = construct %10
+    ret %11
+  }
+}
+)";
+
+    Run(ShaderIO);
+
+    EXPECT_EQ(expect, str());
+}
+
 }  // namespace
 }  // namespace tint::hlsl::writer::raise

diff --git a/src/tint/lang/msl/writer/raise/shader_io.cc b/src/tint/lang/msl/writer/raise/shader_io.cc
index 0874322..018519c 100644
--- a/src/tint/lang/msl/writer/raise/shader_io.cc
+++ b/src/tint/lang/msl/writer/raise/shader_io.cc

@@ -93,7 +93,7 @@
         core::ir::FunctionParam* input_struct_param = nullptr;
         uint32_t input_struct_param_index = 0xffffffff;
 
-        for (auto input : inputs) {
+        for (auto& input : inputs) {
             if (input.attributes.builtin) {
                 auto* param = b.FunctionParam(input.name.Name(), input.type);
                 param->SetInvariant(input.attributes.invariant);

diff --git a/test/tint/types/functions/shader_io/compute_subgroup_builtins.wgsl.expected.ir.dxc.hlsl b/test/tint/types/functions/shader_io/compute_subgroup_builtins.wgsl.expected.ir.dxc.hlsl
index b0160ea..d067fa8 100644
--- a/test/tint/types/functions/shader_io/compute_subgroup_builtins.wgsl.expected.ir.dxc.hlsl
+++ b/test/tint/types/functions/shader_io/compute_subgroup_builtins.wgsl.expected.ir.dxc.hlsl

@@ -1,11 +1,12 @@
-SKIP: FAILED
 
-..\..\src\tint\lang\hlsl\writer\raise\shader_io.cc:101 internal compiler error: TINT_UNREACHABLE 
-********************************************************************
-*  The tint shader compiler has encountered an unexpected error.   *
-*                                                                  *
-*  Please help us fix this issue by submitting a bug report at     *
-*  crbug.com/tint with the source program that triggered the bug.  *
-********************************************************************
+RWByteAddressBuffer output : register(u0);
+void main_inner(uint subgroup_invocation_id, uint subgroup_size) {
+  output.Store((0u + (uint(subgroup_invocation_id) * 4u)), subgroup_size);
+}
 
-tint executable returned error: exit status 0xc000001d
+[numthreads(1, 1, 1)]
+void main() {
+  uint v = WaveGetLaneIndex();
+  main_inner(v, WaveGetLaneCount());
+}
+

diff --git a/test/tint/types/functions/shader_io/compute_subgroup_builtins.wgsl.expected.ir.fxc.hlsl b/test/tint/types/functions/shader_io/compute_subgroup_builtins.wgsl.expected.ir.fxc.hlsl
index b0160ea..7c1d316 100644
--- a/test/tint/types/functions/shader_io/compute_subgroup_builtins.wgsl.expected.ir.fxc.hlsl
+++ b/test/tint/types/functions/shader_io/compute_subgroup_builtins.wgsl.expected.ir.fxc.hlsl

@@ -1,11 +1,19 @@
-SKIP: FAILED
+SKIP: INVALID
 
-..\..\src\tint\lang\hlsl\writer\raise\shader_io.cc:101 internal compiler error: TINT_UNREACHABLE 
-********************************************************************
-*  The tint shader compiler has encountered an unexpected error.   *
-*                                                                  *
-*  Please help us fix this issue by submitting a bug report at     *
-*  crbug.com/tint with the source program that triggered the bug.  *
-********************************************************************
 
-tint executable returned error: exit status 0xc000001d
+RWByteAddressBuffer output : register(u0);
+void main_inner(uint subgroup_invocation_id, uint subgroup_size) {
+  output.Store((0u + (uint(subgroup_invocation_id) * 4u)), subgroup_size);
+}
+
+[numthreads(1, 1, 1)]
+void main() {
+  uint v = WaveGetLaneIndex();
+  main_inner(v, WaveGetLaneCount());
+}
+
+FXC validation failure:
+<scrubbed_path>(9,12-29): error X3004: undeclared identifier 'WaveGetLaneIndex'
+
+
+tint executable returned error: exit status 1

diff --git a/test/tint/types/functions/shader_io/compute_subgroup_builtins_struct.wgsl.expected.ir.dxc.hlsl b/test/tint/types/functions/shader_io/compute_subgroup_builtins_struct.wgsl.expected.ir.dxc.hlsl
index b0160ea..9434039 100644
--- a/test/tint/types/functions/shader_io/compute_subgroup_builtins_struct.wgsl.expected.ir.dxc.hlsl
+++ b/test/tint/types/functions/shader_io/compute_subgroup_builtins_struct.wgsl.expected.ir.dxc.hlsl

@@ -1,11 +1,18 @@
-SKIP: FAILED
+struct ComputeInputs {
+  uint subgroup_invocation_id;
+  uint subgroup_size;
+};
 
-..\..\src\tint\lang\hlsl\writer\raise\shader_io.cc:101 internal compiler error: TINT_UNREACHABLE 
-********************************************************************
-*  The tint shader compiler has encountered an unexpected error.   *
-*                                                                  *
-*  Please help us fix this issue by submitting a bug report at     *
-*  crbug.com/tint with the source program that triggered the bug.  *
-********************************************************************
 
-tint executable returned error: exit status 0xc000001d
+RWByteAddressBuffer output : register(u0);
+void main_inner(ComputeInputs inputs) {
+  output.Store((0u + (uint(inputs.subgroup_invocation_id) * 4u)), inputs.subgroup_size);
+}
+
+[numthreads(1, 1, 1)]
+void main() {
+  uint v = WaveGetLaneIndex();
+  ComputeInputs v_1 = {v, WaveGetLaneCount()};
+  main_inner(v_1);
+}
+

diff --git a/test/tint/types/functions/shader_io/compute_subgroup_builtins_struct.wgsl.expected.ir.fxc.hlsl b/test/tint/types/functions/shader_io/compute_subgroup_builtins_struct.wgsl.expected.ir.fxc.hlsl
index b0160ea..6a14554 100644
--- a/test/tint/types/functions/shader_io/compute_subgroup_builtins_struct.wgsl.expected.ir.fxc.hlsl
+++ b/test/tint/types/functions/shader_io/compute_subgroup_builtins_struct.wgsl.expected.ir.fxc.hlsl

@@ -1,11 +1,25 @@
-SKIP: FAILED
+SKIP: INVALID
 
-..\..\src\tint\lang\hlsl\writer\raise\shader_io.cc:101 internal compiler error: TINT_UNREACHABLE 
-********************************************************************
-*  The tint shader compiler has encountered an unexpected error.   *
-*                                                                  *
-*  Please help us fix this issue by submitting a bug report at     *
-*  crbug.com/tint with the source program that triggered the bug.  *
-********************************************************************
+struct ComputeInputs {
+  uint subgroup_invocation_id;
+  uint subgroup_size;
+};
 
-tint executable returned error: exit status 0xc000001d
+
+RWByteAddressBuffer output : register(u0);
+void main_inner(ComputeInputs inputs) {
+  output.Store((0u + (uint(inputs.subgroup_invocation_id) * 4u)), inputs.subgroup_size);
+}
+
+[numthreads(1, 1, 1)]
+void main() {
+  uint v = WaveGetLaneIndex();
+  ComputeInputs v_1 = {v, WaveGetLaneCount()};
+  main_inner(v_1);
+}
+
+FXC validation failure:
+<scrubbed_path>(14,12-29): error X3004: undeclared identifier 'WaveGetLaneIndex'
+
+
+tint executable returned error: exit status 1

diff --git a/test/tint/types/functions/shader_io/fragment_subgroup_builtins.wgsl.expected.ir.dxc.hlsl b/test/tint/types/functions/shader_io/fragment_subgroup_builtins.wgsl.expected.ir.dxc.hlsl
index a8c5780..b434635 100644
--- a/test/tint/types/functions/shader_io/fragment_subgroup_builtins.wgsl.expected.ir.dxc.hlsl
+++ b/test/tint/types/functions/shader_io/fragment_subgroup_builtins.wgsl.expected.ir.dxc.hlsl

@@ -1,11 +1,11 @@
-SKIP: FAILED
 
-../../src/tint/lang/hlsl/writer/raise/shader_io.cc:101 internal compiler error: TINT_UNREACHABLE 
-********************************************************************
-*  The tint shader compiler has encountered an unexpected error.   *
-*                                                                  *
-*  Please help us fix this issue by submitting a bug report at     *
-*  crbug.com/tint with the source program that triggered the bug.  *
-********************************************************************
+RWByteAddressBuffer output : register(u0);
+void main_inner(uint subgroup_invocation_id, uint subgroup_size) {
+  output.Store((0u + (uint(subgroup_invocation_id) * 4u)), subgroup_size);
+}
 
-tint executable returned error: signal: illegal instruction
+void main() {
+  uint v = WaveGetLaneIndex();
+  main_inner(v, WaveGetLaneCount());
+}
+

diff --git a/test/tint/types/functions/shader_io/fragment_subgroup_builtins.wgsl.expected.ir.fxc.hlsl b/test/tint/types/functions/shader_io/fragment_subgroup_builtins.wgsl.expected.ir.fxc.hlsl
index a8c5780..9561889 100644
--- a/test/tint/types/functions/shader_io/fragment_subgroup_builtins.wgsl.expected.ir.fxc.hlsl
+++ b/test/tint/types/functions/shader_io/fragment_subgroup_builtins.wgsl.expected.ir.fxc.hlsl

@@ -1,11 +1,18 @@
-SKIP: FAILED
+SKIP: INVALID
 
-../../src/tint/lang/hlsl/writer/raise/shader_io.cc:101 internal compiler error: TINT_UNREACHABLE 
-********************************************************************
-*  The tint shader compiler has encountered an unexpected error.   *
-*                                                                  *
-*  Please help us fix this issue by submitting a bug report at     *
-*  crbug.com/tint with the source program that triggered the bug.  *
-********************************************************************
 
-tint executable returned error: signal: illegal instruction
+RWByteAddressBuffer output : register(u0);
+void main_inner(uint subgroup_invocation_id, uint subgroup_size) {
+  output.Store((0u + (uint(subgroup_invocation_id) * 4u)), subgroup_size);
+}
+
+void main() {
+  uint v = WaveGetLaneIndex();
+  main_inner(v, WaveGetLaneCount());
+}
+
+FXC validation failure:
+<scrubbed_path>(8,12-29): error X3004: undeclared identifier 'WaveGetLaneIndex'
+
+
+tint executable returned error: exit status 1

diff --git a/test/tint/types/functions/shader_io/fragment_subgroup_builtins_struct.wgsl.expected.ir.dxc.hlsl b/test/tint/types/functions/shader_io/fragment_subgroup_builtins_struct.wgsl.expected.ir.dxc.hlsl
index a8c5780..7c314cd 100644
--- a/test/tint/types/functions/shader_io/fragment_subgroup_builtins_struct.wgsl.expected.ir.dxc.hlsl
+++ b/test/tint/types/functions/shader_io/fragment_subgroup_builtins_struct.wgsl.expected.ir.dxc.hlsl

@@ -1,11 +1,17 @@
-SKIP: FAILED
+struct FragmentInputs {
+  uint subgroup_invocation_id;
+  uint subgroup_size;
+};
 
-../../src/tint/lang/hlsl/writer/raise/shader_io.cc:101 internal compiler error: TINT_UNREACHABLE 
-********************************************************************
-*  The tint shader compiler has encountered an unexpected error.   *
-*                                                                  *
-*  Please help us fix this issue by submitting a bug report at     *
-*  crbug.com/tint with the source program that triggered the bug.  *
-********************************************************************
 
-tint executable returned error: signal: illegal instruction
+RWByteAddressBuffer output : register(u0);
+void main_inner(FragmentInputs inputs) {
+  output.Store((0u + (uint(inputs.subgroup_invocation_id) * 4u)), inputs.subgroup_size);
+}
+
+void main() {
+  uint v = WaveGetLaneIndex();
+  FragmentInputs v_1 = {v, WaveGetLaneCount()};
+  main_inner(v_1);
+}
+

diff --git a/test/tint/types/functions/shader_io/fragment_subgroup_builtins_struct.wgsl.expected.ir.fxc.hlsl b/test/tint/types/functions/shader_io/fragment_subgroup_builtins_struct.wgsl.expected.ir.fxc.hlsl
index a8c5780..8d1f051 100644
--- a/test/tint/types/functions/shader_io/fragment_subgroup_builtins_struct.wgsl.expected.ir.fxc.hlsl
+++ b/test/tint/types/functions/shader_io/fragment_subgroup_builtins_struct.wgsl.expected.ir.fxc.hlsl

@@ -1,11 +1,24 @@
-SKIP: FAILED
+SKIP: INVALID
 
-../../src/tint/lang/hlsl/writer/raise/shader_io.cc:101 internal compiler error: TINT_UNREACHABLE 
-********************************************************************
-*  The tint shader compiler has encountered an unexpected error.   *
-*                                                                  *
-*  Please help us fix this issue by submitting a bug report at     *
-*  crbug.com/tint with the source program that triggered the bug.  *
-********************************************************************
+struct FragmentInputs {
+  uint subgroup_invocation_id;
+  uint subgroup_size;
+};
 
-tint executable returned error: signal: illegal instruction
+
+RWByteAddressBuffer output : register(u0);
+void main_inner(FragmentInputs inputs) {
+  output.Store((0u + (uint(inputs.subgroup_invocation_id) * 4u)), inputs.subgroup_size);
+}
+
+void main() {
+  uint v = WaveGetLaneIndex();
+  FragmentInputs v_1 = {v, WaveGetLaneCount()};
+  main_inner(v_1);
+}
+
+FXC validation failure:
+<scrubbed_path>(13,12-29): error X3004: undeclared identifier 'WaveGetLaneIndex'
+
+
+tint executable returned error: exit status 1