naga: fix panic parsing SPIR-V subgroup reduce/scan/ballot ops

Author: nazar-pc

CHANGELOG.md

@@ -84,6 +84,7 @@ Bottom level categories:
 #### naga
 
 - Fix panics when shader `var<immediate>` size is larger than 256 bytes. By @beicause in [#9725](https://github.com/gfx-rs/wgpu/pull/9725).
+- Fix a panic in the SPIR-V frontend when a subgroup collective operation (e.g. `OpGroupNonUniformUMin`) or `OpGroupNonUniformBallot` used an argument whose value needed to be spilled to a temporary variable, such as when the argument was computed inside a loop. By @nazar-pc in [#9957](https://github.com/gfx-rs/wgpu/issues/9957).
 
 #### Validation
 

naga/src/front/spv/next_block.rs

@@ -2533,7 +2533,6 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                 }
                 Op::GroupNonUniformBallot => {
                     inst.expect(5)?;
-                    block.extend(emitter.finish(ctx.expressions));
                     let result_type_id = self.next()?;
                     let result_id = self.next()?;
                     let exec_scope_id = self.next()?;
@@ -2564,6 +2563,7 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                         Some(predicate_handle)
                     };
 
+                    block.extend(emitter.finish(ctx.expressions));
                     let result_handle = ctx
                         .expressions
                         .append(crate::Expression::SubgroupBallotResult, span);
@@ -2603,7 +2603,6 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                 | Op::GroupNonUniformLogicalAnd
                 | Op::GroupNonUniformLogicalOr
                 | Op::GroupNonUniformLogicalXor => {
-                    block.extend(emitter.finish(ctx.expressions));
                     inst.expect(
                         if matches!(inst.op, Op::GroupNonUniformAll | Op::GroupNonUniformAny) {
                             5
@@ -2639,6 +2638,8 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                     let argument_lookup = self.lookup_expression.lookup(argument_id)?;
                     let argument_handle = get_expr_handle!(argument_id, argument_lookup);
 
+                    block.extend(emitter.finish(ctx.expressions));
+
                     let exec_scope_const = self.lookup_constant.lookup(exec_scope_id)?;
                     let _exec_scope = resolve_constant(ctx.gctx(), &exec_scope_const.inner)
                         .filter(|exec_scope| *exec_scope == spirv::Scope::Subgroup as u32)
@@ -2711,7 +2712,6 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                     } else {
                         6
                     })?;
-                    block.extend(emitter.finish(ctx.expressions));
                     let result_type_id = self.next()?;
                     let result_id = self.next()?;
                     let exec_scope_id = self.next()?;
@@ -2752,6 +2752,7 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                         }
                     };
 
+                    block.extend(emitter.finish(ctx.expressions));
                     let result_type = self.lookup_type.lookup(result_type_id)?;
 
                     let result_handle = ctx.expressions.append(
@@ -2781,7 +2782,6 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                 }
                 Op::GroupNonUniformQuadSwap => {
                     inst.expect(6)?;
-                    block.extend(emitter.finish(ctx.expressions));
                     let result_type_id = self.next()?;
                     let result_id = self.next()?;
                     let exec_scope_id = self.next()?;
@@ -2791,6 +2791,8 @@ impl<I: Iterator<Item = u32>> Frontend<I> {
                     let argument_lookup = self.lookup_expression.lookup(argument_id)?;
                     let argument_handle = get_expr_handle!(argument_id, argument_lookup);
 
+                    block.extend(emitter.finish(ctx.expressions));
+
                     let exec_scope_const = self.lookup_constant.lookup(exec_scope_id)?;
                     let _exec_scope = resolve_constant(ctx.gctx(), &exec_scope_const.inner)
                         .filter(|exec_scope| *exec_scope == spirv::Scope::Subgroup as u32)

naga/tests/in/spv/subgroup-reduce-issue-8389.spvasm

@@ -0,0 +1,54 @@
+; SPIR-V
+; Version: 1.3
+;
+; Regression test for https://github.com/gfx-rs/wgpu/issues/8389
+;
+; The `naga` SPIR-V frontend panicked with
+; `called Option::unwrap() on a None value` at `naga::proc::emitter::Emitter::finish`
+; when a `OpGroupNonUniform*` reduction/scan instruction (e.g. `OpGroupNonUniformUMin`)
+; used an argument whose defining expression came from an already-closed control-flow
+; body (here, a value computed inside a loop and used after the loop merges), which
+; requires spilling the value through a temporary local variable.
+               OpCapability Shader
+               OpCapability GroupNonUniform
+               OpCapability GroupNonUniformArithmetic
+       %ext = OpExtInstImport "GLSL.std.450"
+               OpMemoryModel Logical GLSL450
+               OpEntryPoint GLCompute %main "main" %local_id
+               OpExecutionMode %main LocalSize 1 1 1
+               OpDecorate %local_id BuiltIn LocalInvocationIndex
+
+       %void = OpTypeVoid
+     %voidfn = OpTypeFunction %void
+       %uint = OpTypeInt 32 0
+       %bool = OpTypeBool
+%_ptr_Input_uint = OpTypePointer Input %uint
+   %local_id = OpVariable %_ptr_Input_uint Input
+     %uint_0 = OpConstant %uint 0
+     %uint_1 = OpConstant %uint 1
+     %uint_4 = OpConstant %uint 4
+%subgroup_scope = OpConstant %uint 3
+
+       %main = OpFunction %void None %voidfn
+      %entry = OpLabel
+         %id = OpLoad %uint %local_id
+               OpBranch %loop_header
+
+%loop_header = OpLabel
+          %i = OpPhi %uint %uint_0 %entry %i_next %continue
+               OpLoopMerge %merge %continue None
+               OpBranch %body
+
+       %body = OpLabel
+        %acc = OpIAdd %uint %id %i
+               OpBranch %continue
+
+   %continue = OpLabel
+     %i_next = OpIAdd %uint %i %uint_1
+       %cond = OpULessThan %bool %i_next %uint_4
+               OpBranchConditional %cond %loop_header %merge
+
+      %merge = OpLabel
+        %min = OpGroupNonUniformUMin %uint %subgroup_scope Reduce %acc
+               OpReturn
+               OpFunctionEnd

naga/tests/in/spv/subgroup-reduce-issue-8389.toml

@@ -0,0 +1,5 @@
+capabilities = "SUBGROUP"
+targets = "WGSL"
+
+[spv]
+version = [1, 3]

naga/tests/out/wgsl/spv-subgroup-reduce-issue-8389.wgsl

@@ -0,0 +1,28 @@
+var<private> global: u32;
+
+fn function_() {
+    var phi_16_: u32;
+    var local: u32;
+
+    let _e4 = global;
+    phi_16_ = 0u;
+    loop {
+        let _e6 = phi_16_;
+        local = (_e4 + _e6);
+        continue;
+        continuing {
+            let _e8 = (_e6 + 1u);
+            phi_16_ = _e8;
+            break if !((_e8 < 4u));
+        }
+    }
+    let _e12 = local;
+    let _e13 = subgroupMin(_e12);
+    return;
+}
+
+@compute @workgroup_size(1, 1, 1) 
+fn main(@builtin(local_invocation_index) param: u32) {
+    global = param;
+    function_();
+}