From 3ba7ebde5598a1bc03e5062b6149f4b7d9ce2daa Mon Sep 17 00:00:00 2001
From: YongKang Zhu <yongzhu@fb.com>
Date: Fri, 2 Aug 2024 09:07:48 -0700
Subject: [PATCH] [InstCombine] Remove transformation on call instruction where
 return value need void to non-void conversion (#98536)

Skip simplification on call instruction where a non-void return value is
expected but the callee returns void, which is undefined behavior and
could lead to non-determinism or crashes.
---
 .../InstCombine/InstCombineCalls.cpp          | 23 +++++--------
 .../skip-opt-void-to-non-void-conversion.ll   | 33 +++++++++++++++++++
 2 files changed, 42 insertions(+), 14 deletions(-)
 create mode 100644 llvm/test/Transforms/InstCombine/skip-opt-void-to-non-void-conversion.ll

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index cc68fd4cf1c1b6..bf7c91bf363062 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -4145,9 +4145,7 @@ bool InstCombinerImpl::transformConstExprCastCall(CallBase &Call) {
       if (Callee->isDeclaration())
         return false;   // Cannot transform this return value.
 
-      if (!Caller->use_empty() &&
-          // void -> non-void is handled specially
-          !NewRetTy->isVoidTy())
+      if (!Caller->use_empty())
         return false;   // Cannot transform this return value.
     }
 
@@ -4338,17 +4336,14 @@ bool InstCombinerImpl::transformConstExprCastCall(CallBase &Call) {
   Instruction *NC = NewCall;
   Value *NV = NC;
   if (OldRetTy != NV->getType() && !Caller->use_empty()) {
-    if (!NV->getType()->isVoidTy()) {
-      NV = NC = CastInst::CreateBitOrPointerCast(NC, OldRetTy);
-      NC->setDebugLoc(Caller->getDebugLoc());
-
-      auto OptInsertPt = NewCall->getInsertionPointAfterDef();
-      assert(OptInsertPt && "No place to insert cast");
-      InsertNewInstBefore(NC, *OptInsertPt);
-      Worklist.pushUsersToWorkList(*Caller);
-    } else {
-      NV = PoisonValue::get(Caller->getType());
-    }
+    assert(!NV->getType()->isVoidTy());
+    NV = NC = CastInst::CreateBitOrPointerCast(NC, OldRetTy);
+    NC->setDebugLoc(Caller->getDebugLoc());
+
+    auto OptInsertPt = NewCall->getInsertionPointAfterDef();
+    assert(OptInsertPt && "No place to insert cast");
+    InsertNewInstBefore(NC, *OptInsertPt);
+    Worklist.pushUsersToWorkList(*Caller);
   }
 
   if (!Caller->use_empty())
diff --git a/llvm/test/Transforms/InstCombine/skip-opt-void-to-non-void-conversion.ll b/llvm/test/Transforms/InstCombine/skip-opt-void-to-non-void-conversion.ll
new file mode 100644
index 00000000000000..8108201b4db774
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/skip-opt-void-to-non-void-conversion.ll
@@ -0,0 +1,33 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt --passes=instcombine -S < %s | FileCheck %s
+
+define void @foo() {
+; CHECK-LABEL: define void @foo() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  ret void
+}
+
+define i32 @bar() {
+; CHECK-LABEL: define i32 @bar() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = tail call i32 @foo()
+; CHECK-NEXT:    ret i32 [[TMP0]]
+;
+entry:
+  %1 = tail call i32 @foo()
+  ret i32 %1
+}
+
+define void @goo() {
+; CHECK-LABEL: define void @goo() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    call void @foo()
+; CHECK-NEXT:    ret void
+;
+entry:
+  %res = call i32 @foo()
+  ret void
+}