[SimplifyCFG] Improve FoldTwoEntryPHINode when one of phi values is undef

[dtcxzyw]

When converting select-like branches into a select, we can just use one phi value instead of creating a select if it is guaranteed not to be poison and another phi value is undef. After converting to select, we will lose information from the dominator tree.

Alive2: https://alive2.llvm.org/ce/z/rheXK4

Fixes #67342.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Yingwei Zheng (dtcxzyw)

Changes

When converting select-like branches into a select, we can just use one phi value instead of creating a select if it is guaranteed not to be poison and another phi value is undef. After converting to select, we will lose information from the dominator tree.

Alive2: https://alive2.llvm.org/ce/z/rheXK4

Fixes #67342.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Utils/SimplifyCFG.cpp (+16-3)
• (added) llvm/test/Transforms/SimplifyCFG/pr67342.ll (+50)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 35fead111aa9666..bab5363e63a85eb 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -3543,9 +3543,22 @@ static bool FoldTwoEntryPHINode(PHINode *PN, const TargetTransformInfo &TTI,
     Value *TrueVal = PN->getIncomingValueForBlock(IfTrue);
     Value *FalseVal = PN->getIncomingValueForBlock(IfFalse);
 
-    Value *Sel = Builder.CreateSelect(IfCond, TrueVal, FalseVal, "", DomBI);
-    PN->replaceAllUsesWith(Sel);
-    Sel->takeName(PN);
+    Value *ReplaceVal = nullptr;
+    if (isa<UndefValue>(TrueVal) &&
+        isGuaranteedNotToBePoison(FalseVal, /*AC*/ nullptr,
+                                  IfFalse->getTerminator(),
+                                  DTU ? &DTU->getDomTree() : nullptr, 0)) {
+      ReplaceVal = FalseVal;
+    } else if (isa<UndefValue>(FalseVal) &&
+               isGuaranteedNotToBePoison(
+                   TrueVal, /*AC*/ nullptr, IfFalse->getTerminator(),
+                   DTU ? &DTU->getDomTree() : nullptr, 0)) {
+      ReplaceVal = TrueVal;
+    } else {
+      ReplaceVal = Builder.CreateSelect(IfCond, TrueVal, FalseVal, "", DomBI);
+      ReplaceVal->takeName(PN);
+    }
+    PN->replaceAllUsesWith(ReplaceVal);
     PN->eraseFromParent();
   }
 
diff --git a/llvm/test/Transforms/SimplifyCFG/pr67342.ll b/llvm/test/Transforms/SimplifyCFG/pr67342.ll
new file mode 100644
index 000000000000000..c2c9f460d8e5f9e
--- /dev/null
+++ b/llvm/test/Transforms/SimplifyCFG/pr67342.ll
@@ -0,0 +1,50 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
+; RUN: opt -S -passes=simplifycfg < %s | FileCheck %s
+
+; Tests from PR67342
+define i16 @test1(i32 %err) {
+; CHECK-LABEL: define i16 @test1(
+; CHECK-SAME: i32 [[ERR:%.*]]) {
+; CHECK-NEXT:  bb3:
+; CHECK-NEXT:    [[_3:%.*]] = icmp slt i32 [[ERR]], 0
+; CHECK-NEXT:    [[OK:%.*]] = trunc i32 [[ERR]] to i16
+; CHECK-NEXT:    ret i16 [[OK]]
+;
+  %_3 = icmp slt i32 %err, 0
+  br i1 %_3, label %bb1, label %bb2
+
+bb2:
+  %ok = trunc i32 %err to i16
+  br label %bb3
+
+bb1:
+  br label %bb3
+
+bb3:
+  %r.sroa.3.0 = phi i16 [ undef, %bb1 ], [ %ok, %bb2 ]
+  ret i16 %r.sroa.3.0
+}
+
+; commuted test
+define i16 @test2(i32 %err) {
+; CHECK-LABEL: define i16 @test2(
+; CHECK-SAME: i32 [[ERR:%.*]]) {
+; CHECK-NEXT:  bb3:
+; CHECK-NEXT:    [[_3:%.*]] = icmp slt i32 [[ERR]], 0
+; CHECK-NEXT:    [[OK:%.*]] = trunc i32 [[ERR]] to i16
+; CHECK-NEXT:    ret i16 [[OK]]
+;
+  %_3 = icmp slt i32 %err, 0
+  br i1 %_3, label %bb1, label %bb2
+
+bb2:
+  br label %bb3
+
+bb1:
+  %ok = trunc i32 %err to i16
+  br label %bb3
+
+bb3:
+  %r.sroa.3.0 = phi i16 [ %ok, %bb1 ], [ undef, %bb2 ]
+  ret i16 %r.sroa.3.0
+}

[nikic]

You are not allowed to use the DominatorTree inside SimplifyCFG, only to preserve it.

[dtcxzyw]

Ping.

[nikic]

So this now works via programUndefinedIfPoison, which is going to make this quite fragile. It's going to work if this is the first branch in the function, but not in a later one.

Also, while #67342 reported this without the noundef attribute, the original in rust-lang/rust#116150 does have it, so I don't think you're even fixing the right issue here.