[llvm] [Reg2Mem] Handle CallBr instructions (PR #90953)

[via llvm-commits]

https://github.com/XChy created https://github.com/llvm/llvm-project/pull/90953

Fixes #90900

>From f88217b66139ba83d6592258a84ddc02e8896ef5 Mon Sep 17 00:00:00 2001
From: XChy <xxs_chy at outlook.com>
Date: Fri, 3 May 2024 16:37:48 +0800
Subject: [PATCH] [Reg2Mem] Handle CallBr instructions

---
 .../lib/Transforms/Utils/DemoteRegToStack.cpp |  7 +++--
 llvm/test/Transforms/Reg2Mem/callbr-crash.ll  | 31 +++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/Reg2Mem/callbr-crash.ll

diff --git a/llvm/lib/Transforms/Utils/DemoteRegToStack.cpp b/llvm/lib/Transforms/Utils/DemoteRegToStack.cpp
index b2a88eadd3dee8..c25c878282ba60 100644
--- a/llvm/lib/Transforms/Utils/DemoteRegToStack.cpp
+++ b/llvm/lib/Transforms/Utils/DemoteRegToStack.cpp
@@ -102,9 +102,12 @@ AllocaInst *llvm::DemoteRegToStack(Instruction &I, bool VolatileLoads,
         new StoreInst(&I, Slot, Handler->getFirstInsertionPt());
       return Slot;
     }
+  } else if (InvokeInst *II = dyn_cast<InvokeInst>(&I)) {
+    InsertPt = II->getNormalDest()->getFirstInsertionPt();
+  } else if (CallBrInst *CBI = dyn_cast<CallBrInst>(&I)) {
+    InsertPt = CBI->getDefaultDest()->getFirstInsertionPt();
   } else {
-    InvokeInst &II = cast<InvokeInst>(I);
-    InsertPt = II.getNormalDest()->getFirstInsertionPt();
+    llvm_unreachable("Unsupported terminator for Reg2Mem");
   }
 
   new StoreInst(&I, Slot, InsertPt);
diff --git a/llvm/test/Transforms/Reg2Mem/callbr-crash.ll b/llvm/test/Transforms/Reg2Mem/callbr-crash.ll
new file mode 100644
index 00000000000000..cdf52f4fa8821b
--- /dev/null
+++ b/llvm/test/Transforms/Reg2Mem/callbr-crash.ll
@@ -0,0 +1,31 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=reg2mem -S < %s | FileCheck %s
+
+define void @crash() {
+; CHECK-LABEL: @crash(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[A_REG2MEM:%.*]] = alloca i64, align 8
+; CHECK-NEXT:    %"reg2mem alloca point" = bitcast i32 0 to i32
+; CHECK-NEXT:    [[A:%.*]] = callbr i64 asm "", "=r,r,!i"(i64 0)
+; CHECK-NEXT:            to label [[THEN:%.*]] [label %entry.else_crit_edge]
+; CHECK:       entry.else_crit_edge:
+; CHECK-NEXT:    br label [[ELSE:%.*]]
+; CHECK:       then:
+; CHECK-NEXT:    store i64 [[A]], ptr [[A_REG2MEM]], align 4
+; CHECK-NEXT:    [[A_RELOAD:%.*]] = load i64, ptr [[A_REG2MEM]], align 4
+; CHECK-NEXT:    [[B:%.*]] = inttoptr i64 [[A_RELOAD]] to ptr
+; CHECK-NEXT:    br label [[ELSE]]
+; CHECK:       else:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %a = callbr i64 asm "", "=r,r,!i"(i64 0)
+  to label %then [label %else]
+
+then:
+  %b = inttoptr i64 %a to ptr
+  br label %else
+
+else:
+  ret void
+}