[llvm] 2fc0d17 - [FuncSpec] Avoid crashing when SwitchInst doesn't see ConstantInt

[Vincent Lee via llvm-commits]

Author: Vincent Lee
Date: 2023-06-30T21:02:50-07:00
New Revision: 2fc0d17e8b3d62ea20d11a509a9005c564a29755

URL: https://github.com/llvm/llvm-project/commit/2fc0d17e8b3d62ea20d11a509a9005c564a29755
DIFF: https://github.com/llvm/llvm-project/commit/2fc0d17e8b3d62ea20d11a509a9005c564a29755.diff

LOG: [FuncSpec] Avoid crashing when SwitchInst doesn't see ConstantInt

D150464 updated the cost model for function specialization. Unfortunately, this
also crashes when trying to build stage2 LLD with thinLTO and assertions. It looks
like the issue is caused by a mishandling of the Constant in a SwitchInst since the
Constant cannot always be assumed to safely casted to a ConstantInt. In the case
of the crash, Constant was a ConstantExpr which triggered the assertion.

Reviewed By: ChuanqiXu

Differential Revision: https://reviews.llvm.org/D154159

Added: 
    llvm/test/Transforms/FunctionSpecialization/function-specialization-constant-expression6.ll

Modified: 
    llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp b/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
index 93d2eaaddf310e..5180e3ffa09777 100644
--- a/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
+++ b/llvm/lib/Transforms/IPO/FunctionSpecialization.cpp
@@ -190,7 +190,10 @@ Cost InstCostVisitor::estimateSwitchInst(SwitchInst &I) {
   if (I.getCondition() != LastVisited->first)
     return 0;
 
-  auto *C = cast<ConstantInt>(LastVisited->second);
+  auto *C = dyn_cast<ConstantInt>(LastVisited->second);
+  if (!C)
+    return 0;
+
   BasicBlock *Succ = I.findCaseValue(C)->getCaseSuccessor();
   // Initialize the worklist with the dead basic blocks. These are the
   // destination labels which are 
diff erent from the one corresponding

diff  --git a/llvm/test/Transforms/FunctionSpecialization/function-specialization-constant-expression6.ll b/llvm/test/Transforms/FunctionSpecialization/function-specialization-constant-expression6.ll
new file mode 100644
index 00000000000000..9c9611364b7665
--- /dev/null
+++ b/llvm/test/Transforms/FunctionSpecialization/function-specialization-constant-expression6.ll
@@ -0,0 +1,22 @@
+; RUN: opt -passes="ipsccp<func-spec>" -force-specialization -S < %s
+; Check that we don't crash when SwitchInst Constant is not ConstantInt.
+
+ at S = external constant [1 x i8]
+
+define i1 @foo() {
+entry:
+  %tmp = call i32 @bar(ptr @S)
+  ret i1 0
+}
+
+define i32 @bar(ptr %arg) {
+entry:
+  %magicptr = ptrtoint ptr %arg to i64
+  switch i64 %magicptr, label %bb2 [
+    i64 0, label %bb1
+  ]
+bb1:
+  ret i32 0
+bb2:
+  ret i32 1
+}