[llvm] 8a519b3 - [InstCombine] Ensure constant folding in binop of select fold

[Nikita Popov via llvm-commits]

Author: Nikita Popov
Date: 2022-07-15T11:03:10+02:00
New Revision: 8a519b3c214547dbd7b6a027d83ee5f79b5f09c0

URL: https://github.com/llvm/llvm-project/commit/8a519b3c214547dbd7b6a027d83ee5f79b5f09c0
DIFF: https://github.com/llvm/llvm-project/commit/8a519b3c214547dbd7b6a027d83ee5f79b5f09c0.diff

LOG: [InstCombine] Ensure constant folding in binop of select fold

When folding a binop into a select, we need to ensure that one
of the select arms actually does constant fold, otherwise we'll
create two binop instructions and perform the reverse transform.

Ensure this by performing an explicit constant folding attempt,
and failing the transform if neither side simplifies.

A simple alternative here would have been to limit the fold to
ImmConstants, but given the current representation of scalable
vector splats, this wouldn't be ideal.

Added: 
    

Modified: 
    llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
    llvm/test/Transforms/InstCombine/select.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 75520a0c8d5f6..42c6433270db0 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -994,6 +994,24 @@ Instruction *InstCombinerImpl::foldBinopOfSextBoolToSelect(BinaryOperator &BO) {
   return SelectInst::Create(X, TVal, FVal);
 }
 
+static Constant *constantFoldOperationIntoSelectOperand(
+    Instruction &I, SelectInst *SI, Value *SO) {
+  auto *ConstSO = dyn_cast<Constant>(SO);
+  if (!ConstSO)
+    return nullptr;
+
+  SmallVector<Constant *> ConstOps;
+  for (Value *Op : I.operands()) {
+    if (Op == SI)
+      ConstOps.push_back(ConstSO);
+    else if (auto *C = dyn_cast<Constant>(Op))
+      ConstOps.push_back(C);
+    else
+      llvm_unreachable("Operands should be select or constant");
+  }
+  return ConstantFoldInstOperands(&I, ConstOps, I.getModule()->getDataLayout());
+}
+
 static Value *foldOperationIntoSelectOperand(Instruction &I, Value *SO,
                                              InstCombiner::BuilderTy &Builder) {
   if (auto *Cast = dyn_cast<CastInst>(&I))
@@ -1101,8 +1119,17 @@ Instruction *InstCombinerImpl::FoldOpIntoSelect(Instruction &Op, SelectInst *SI,
     }
   }
 
-  Value *NewTV = foldOperationIntoSelectOperand(Op, TV, Builder);
-  Value *NewFV = foldOperationIntoSelectOperand(Op, FV, Builder);
+  // Make sure that one of the select arms constant folds successfully.
+  Value *NewTV = constantFoldOperationIntoSelectOperand(Op, SI, TV);
+  Value *NewFV = constantFoldOperationIntoSelectOperand(Op, SI, FV);
+  if (!NewTV && !NewFV)
+    return nullptr;
+
+  // Create an instruction for the arm that did not fold.
+  if (!NewTV)
+    NewTV = foldOperationIntoSelectOperand(Op, TV, Builder);
+  if (!NewFV)
+    NewFV = foldOperationIntoSelectOperand(Op, FV, Builder);
   return SelectInst::Create(SI->getCondition(), NewTV, NewFV, "", nullptr, SI);
 }
 

diff  --git a/llvm/test/Transforms/InstCombine/select.ll b/llvm/test/Transforms/InstCombine/select.ll
index 0af3fe523f476..f72e586a302f7 100644
--- a/llvm/test/Transforms/InstCombine/select.ll
+++ b/llvm/test/Transforms/InstCombine/select.ll
@@ -3184,6 +3184,28 @@ define <2 x i8> @ne0_is_all_ones_swap_vec_poison(<2 x i8> %x) {
   ret <2 x i8> %r
 }
 
+define i64 @udiv_of_select_constexpr(i1 %c, i64 %x) {
+; CHECK-LABEL: @udiv_of_select_constexpr(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[C:%.*]], i64 [[X:%.*]], i64 ptrtoint (i32* @glbl to i64)
+; CHECK-NEXT:    [[OP:%.*]] = udiv i64 [[SEL]], 3
+; CHECK-NEXT:    ret i64 [[OP]]
+;
+  %sel = select i1 %c, i64 %x, i64 ptrtoint (i32* @glbl to i64)
+  %op = udiv i64 %sel, 3
+  ret i64 %op
+}
+
+define i64 @udiv_of_select_constexpr_commuted(i1 %c, i64 %x) {
+; CHECK-LABEL: @udiv_of_select_constexpr_commuted(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[C:%.*]], i64 ptrtoint (i32* @glbl to i64), i64 [[X:%.*]]
+; CHECK-NEXT:    [[OP:%.*]] = udiv i64 [[SEL]], 3
+; CHECK-NEXT:    ret i64 [[OP]]
+;
+  %sel = select i1 %c, i64 ptrtoint (i32* @glbl to i64), i64 %x
+  %op = udiv i64 %sel, 3
+  ret i64 %op
+}
+
 declare void @use(i1)
 declare void @use_i8(i8)
 declare void @use_i32(i32)