[llvm] [SPIRV] Support switch terminators when merging region exits (PR #209859)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-spir-v

Author: Farzon Lotfi (farzonl)

<details>
<summary>Changes</summary>

fixes #<!-- -->209310

This change fixes the `llvm_unreachable("Unhandled terminator type.")` crash in `createExitVariable` when a convergence region's exit block ends in a switch instruction.

To handle the SwitchInst properly we needed to:

1. Build a select chain keyed on the switch condition, using the default destination as the fallback value and skipping successors that are internal to the region. This is an almost exact copy of what already exists for `CondBrInst`.
2. Modify the SPIRVStructurizer.cpp `splitSwitchCases`to removes a case whose target is the default destination instead of splitting it into a new successor block.

Assisted by Claude Opus 4.8

---
Full diff: https://github.com/llvm/llvm-project/pull/209859.diff


3 Files Affected:

- (modified) llvm/lib/Target/SPIRV/SPIRVStructurizer.cpp (+10) 
- (modified) llvm/lib/Target/SPIRV/SPIRVUtils.cpp (+20-1) 
- (added) llvm/test/CodeGen/SPIRV/structurizer/merge-exit-switch.ll (+60) 


``````````diff
diff --git a/llvm/lib/Target/SPIRV/SPIRVStructurizer.cpp b/llvm/lib/Target/SPIRV/SPIRVStructurizer.cpp
index f02bd261244aa..df18997fd7a51 100644
--- a/llvm/lib/Target/SPIRV/SPIRVStructurizer.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVStructurizer.cpp
@@ -935,6 +935,16 @@ class SPIRVStructurizer : public FunctionPass {
       auto It = SI->case_begin();
       while (It != SI->case_end()) {
         BasicBlock *Target = It->getCaseSuccessor();
+
+        // Don't Split. Just remove cases branching to the default destination
+        // to prevent spurious extra successors thus preserving single-exit
+        // convergence regions (i.e. if a merged exit is default & a case).
+        if (Target == SI->getDefaultDest()) {
+          Modified = true;
+          It = SI->removeCase(It);
+          continue;
+        }
+
         if (Seen.count(Target) == 0) {
           Seen.insert(Target);
           ++It;
diff --git a/llvm/lib/Target/SPIRV/SPIRVUtils.cpp b/llvm/lib/Target/SPIRV/SPIRVUtils.cpp
index 1a1247fbe3e18..236da61192f1a 100644
--- a/llvm/lib/Target/SPIRV/SPIRVUtils.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVUtils.cpp
@@ -891,7 +891,26 @@ createExitVariable(BasicBlock *BB,
     return Builder.CreateSelect(BI->getCondition(), LHS, RHS);
   }
 
-  // TODO: add support for switch cases.
+  if (auto *SI = dyn_cast<SwitchInst>(T)) {
+    Value *Condition = SI->getCondition();
+    // The default destination acts as the fallback value of the select chain.
+    Value *Result = TargetToValue.lookup(SI->getDefaultDest());
+    for (const auto &Case : SI->cases()) {
+      Value *CaseValue = TargetToValue.lookup(Case.getCaseSuccessor());
+      // Successors that are internal to the region have no exit value.
+      if (CaseValue == nullptr)
+        continue;
+      // The first known exit value becomes the base of the select chain.
+      if (Result == nullptr) {
+        Result = CaseValue;
+        continue;
+      }
+      Value *Cmp = Builder.CreateICmpEQ(Condition, Case.getCaseValue());
+      Result = Builder.CreateSelect(Cmp, CaseValue, Result);
+    }
+    return Result;
+  }
+
   llvm_unreachable("Unhandled terminator type.");
 }
 
diff --git a/llvm/test/CodeGen/SPIRV/structurizer/merge-exit-switch.ll b/llvm/test/CodeGen/SPIRV/structurizer/merge-exit-switch.ll
new file mode 100644
index 0000000000000..ccb809ea1183e
--- /dev/null
+++ b/llvm/test/CodeGen/SPIRV/structurizer/merge-exit-switch.ll
@@ -0,0 +1,60 @@
+; RUN: llc -mtriple=spirv-unknown-vulkan-compute -O0 %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv-unknown-vulkan-compute %s -o - -filetype=obj | spirv-val %}
+
+; The loop header switch has a single merged exit and a single latch.
+; CHECK: OpLoopMerge %[[#exit:]] %[[#latch:]] None
+; CHECK: OpSelectionMerge %[[#merge:]] None
+
+; The switch keeps exactly three cases. This confirms the redundant case that
+; duplicated the default exit is removed not split into a fourth case / extra
+; exit block. The end-of-line anchor ensures no additional case is appended.
+; CHECK: OpSwitch %[[#]] %[[#merge]] 2 %[[#]] 0 %[[#]] 1 %[[#]]{{$}}
+
+; The selection merge routes to the single loop exit or back to the latch.
+; CHECK: OpBranchConditional %[[#]] %[[#exit]] %[[#latch]]
+
+define void @main() #0 {
+entry:
+  %t0 = tail call token @llvm.experimental.convergence.entry()
+  br label %for.cond
+
+for.cond:
+  %I = phi i32 [ 0, %entry ], [ %inc, %latch ]
+  %tl = tail call token @llvm.experimental.convergence.loop() [ "convergencectrl"(token %t0) ]
+  switch i32 %I, label %exit [
+    i32 4, label %exit
+    i32 0, label %c0
+    i32 1, label %c1
+    i32 2, label %c2
+  ]
+
+c0:
+  %v0 = call i32 @get(i32 0) [ "convergencectrl"(token %tl) ]
+  br label %latch
+
+c1:
+  %v1 = call i32 @get(i32 1) [ "convergencectrl"(token %tl) ]
+  br label %latch
+
+c2:
+  %v2 = call i32 @get(i32 2) [ "convergencectrl"(token %tl) ]
+  br label %latch
+
+latch:
+  %acc = phi i32 [ %v0, %c0 ], [ %v1, %c1 ], [ %v2, %c2 ]
+  call void @put(i32 %I, i32 %acc) [ "convergencectrl"(token %tl) ]
+  %inc = add nuw nsw i32 %I, 1
+  br label %for.cond
+
+exit:
+  ret void
+}
+
+declare token @llvm.experimental.convergence.entry() #1
+declare token @llvm.experimental.convergence.loop() #1
+declare i32 @get(i32) #2
+declare void @put(i32, i32) #2
+
+attributes #0 = { convergent noinline norecurse nounwind "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }
+attributes #1 = { convergent nocallback nofree nosync nounwind willreturn memory(none) }
+attributes #2 = { convergent }

``````````

</details>


https://github.com/llvm/llvm-project/pull/209859