Fix invalid SPIRV from continue inside switch inside for loop (#9585, #10198)

For a for-loop with a switch containing a `continue`, the IR has this shape:

  loop(target=%body, break=%after, continue=%incr)
  %body:
    switch(x, break=%post_switch, ...)
      case 0: unconditionalBranch(%incr)        <- 'continue'
      case 1: unconditionalBranch(%post_switch) <- 'break'
  %post_switch:
    unconditionalBranch(%incr)   <- normal post-switch flow
  %incr: i++; unconditionalBranch(%body)  <- back-edge
  %after: ...                              <- loop break

The `continue` inside the switch branches to %incr (the loop's continueBlock).
Because continueBlock (%incr) != targetBlock (%body) in a for-loop, this is a
branch that exits the switch region to reach an exit block of the enclosing loop
region -- a multi-level branch that must be transformed for valid SPIRV.

The bug: populateExitBlocks() stored continueBlock but did not add it to
region->exitBlocks, so mapExitBlockToRegion never mapped %incr to the loop
region. gatherInfo() therefore never flagged the branch from case 0 to %incr as
a multi-level branch, needsContinueElimination stayed false, and the raw IR
branch reached the SPIRV emitter -- producing unstructured control flow.

The fix: add continueBlock to region->exitBlocks whenever continueBlock !=
targetBlock (i.e. for for-loops). This makes %incr visible in
mapExitBlockToRegion, gatherInfo() detects the multi-level branch, and
eliminateContinueBlocksInFunc() is called to wrap the loop body in an inner
breakable region. After that transformation the `continue` becomes a break from
the inner region -- a valid SPIRV structured exit to its merge block -- and the
outer loop handles the actual iteration.

Also fix a stack leak: after processing a loop region, pop continueBlock from
the global exitBlocks stack unconditionally. For for-loops it is already in
info.exitBlocks and removed by the existing loop; for while-loops (where
continueBlock == targetBlock and is not in info.exitBlocks) it was pushed to
the global stack but never popped, which could affect sibling constructs.

Fixes #9585, #10198.

Author: jkwak-work

source/slang/slang-ir-eliminate-multilevel-break.cpp

@@ -71,8 +71,33 @@ struct EliminateMultiLevelBreakContext
             // If this is a loop, store the continue block.
             // We add it to the exitBlocks stack separately in collectBreakableRegionBlocks
             // so that nested constructs treat it as an exit point.
-            if (as<IRLoop>(headerInst))
+            if (auto loop = as<IRLoop>(headerInst))
+            {
                 continueBlock = getContinueBlock();
+                SLANG_ASSERT(continueBlock);
+
+                // For a for-loop with a switch statement inside, the IR looks like:
+                //
+                //   loop(target=%body, break=%after, continue=%incr)
+                //   %body: ...
+                //     switch(x, break=%post_switch, ...)
+                //       case 0: unconditionalBranch(%incr)   <- 'continue'
+                //       case 1: unconditionalBranch(%post_switch)  <- 'break'
+                //   %post_switch: ...
+                //     unconditionalBranch(%incr)             <- normal loop flow
+                //   %incr: i++; unconditionalBranch(%body)  <- back-edge
+                //   %after: ...                              <- loop break
+                //
+                // The 'continue' inside the switch branches to %incr (continueBlock).
+                // This is a multi-level branch: it exits the switch region to reach
+                // an exit block of the enclosing loop region.
+                //
+                // In order to handle the multi-level branch properly, the continueBlock
+                // needs to be added to exitBlocks.
+                //
+                if (continueBlock != loop->getTargetBlock())
+                    exitBlocks.add(continueBlock);
+            }
         }
 
         void replaceBreakBlock(IRBuilder* builder, IRBlock* block)
@@ -177,9 +202,11 @@ struct EliminateMultiLevelBreakContext
                 }
             }
 
-            // Pop the exit blocks.
+            // Pop the exit blocks that were pushed at the top of this function.
             for (auto exitBlock : info.exitBlocks)
                 exitBlocks.remove(exitBlock);
+            if (info.continueBlock)
+                exitBlocks.remove(info.continueBlock);
         }
 
         void gatherInfo(IRGlobalValueWithCode* func)

tests/bugs/nested-switch-continue-in-loop.slang

@@ -0,0 +1,66 @@
+//TEST:SIMPLE(filecheck=SPIRV): -target spirv-asm -stage compute -entry computeMain
+
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-dx12 -output-using-type
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type
+
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-dx12 -output-using-type -Xslang -DWHILE_LOOP
+//TEST(compute):COMPARE_COMPUTE(filecheck-buffer=CHECK):-vk -output-using-type -Xslang -DWHILE_LOOP
+
+// Disable dx11/FXC: "error X3708: 'continue' cannot be used in a switch"
+//DISABLE_TEST(compute):COMPARE_COMPUTE:-dx11 -output-using-type
+//DISABLE_TEST(compute):COMPARE_COMPUTE:-dx11 -output-using-type -Xslang -DWHILE_LOOP
+
+// Test `continue` statment inside of switch inside of a loop.
+
+// SPIRV: OpEntryPoint
+
+//TEST_INPUT:ubuffer(data=[0 0 0], stride=4):out,name=outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+[[noinline]]
+int test(int value)
+{
+    int result = 0;
+#if defined(WHILE_LOOP)
+    uint i = 0;
+    while(i < 3)
+    {
+        i++;
+#else
+    for (int i = 0; i < 3; i++)
+    {
+#endif
+        switch (value)
+        {
+        case 0:
+            result += 1;
+            continue; // triggers the bug: continue from switch inside for-loop
+        case 1:
+            result += 2;
+            break;
+        default:
+            result += 3;
+            break;
+        }
+        result += 10;
+    }
+    return result;
+}
+
+[shader("compute")]
+[numthreads(1, 1, 1)]
+void computeMain()
+{
+    // CHECK: 3
+    // test(0): case 0 hits, result += 1, continue (skips +10) -> 3 iterations -> result = 3
+    outputBuffer[0] = test(0);
+
+    // CHECK: 36
+    // test(1): case 1 hits, result += 2, break, then +10 -> 3 iterations -> result = 36
+    outputBuffer[1] = test(1);
+
+    // CHECK: 39
+    // test(2): default hits, result += 3, break, then +10 -> 3 iterations -> result = 39
+    outputBuffer[2] = test(2);
+}
+