[llvm] Pick other one if is/contains (PR #81428)

[via llvm-commits]

https://github.com/AtariDreams created https://github.com/llvm/llvm-project/pull/81428

None

>From df4f19725dd4670eb09444496717c9d372d313cd Mon Sep 17 00:00:00 2001
From: Rose <83477269+AtariDreams at users.noreply.github.com>
Date: Sun, 11 Feb 2024 14:42:54 -0500
Subject: [PATCH] Pick other one if is/contains

---
 .../InstCombine/InstCombineMulDivRem.cpp      | 23 ++++++++++++++-----
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
index f9cee9dfcfadae..bd30ac787032de 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
@@ -1299,18 +1299,29 @@ static Value *takeLog2(IRBuilderBase &Builder, Value *Op, unsigned Depth,
   }
 
   // log2(Cond ? X : Y) -> Cond ? log2(X) : log2(Y)
-  // FIXME: missed optimization: if one of the hands of select is/contains
+  // If one of the hands of select is/contains
   //        undef, just directly pick the other one.
   // FIXME: can both hands contain undef?
   // FIXME: Require one use?
-  if (SelectInst *SI = dyn_cast<SelectInst>(Op))
-    if (Value *LogX = takeLog2(Builder, SI->getOperand(1), Depth,
-                               AssumeNonZero, DoFold))
-      if (Value *LogY = takeLog2(Builder, SI->getOperand(2), Depth,
-                                 AssumeNonZero, DoFold))
+  if (SelectInst *SI = dyn_cast<SelectInst>(Op)) {
+    Value *X = SI->getOperand(1);
+    Value *Y = SI->getOperand(2);
+
+    // If X is undef, directly pick Y
+    if (isa<UndefValue>(X))
+      return takeLog2(Builder, Y, Depth, AssumeNonZero, DoFold);
+
+    // If Y is undef, directly pick X
+    if (isa<UndefValue>(Y))
+      return takeLog2(Builder, X, Depth, AssumeNonZero, DoFold);
+
+    // Otherwise, proceed as before
+    if (Value *LogX = takeLog2(Builder, X, Depth, AssumeNonZero, DoFold))
+      if (Value *LogY = takeLog2(Builder, Y, Depth, AssumeNonZero, DoFold))
         return IfFold([&]() {
           return Builder.CreateSelect(SI->getOperand(0), LogX, LogY);
         });
+  }
 
   // log2(umin(X, Y)) -> umin(log2(X), log2(Y))
   // log2(umax(X, Y)) -> umax(log2(X), log2(Y))