[llvm] [InstCombine] Use the select condition to try to constant fold binops into select (PR #84696)

via llvm-commits llvm-commits at lists.llvm.org
Mon Mar 11 08:13:07 PDT 2024


https://github.com/goldsteinn updated https://github.com/llvm/llvm-project/pull/84696

>From ecda11e84eadaa8cb00ad0bad439a2c0a247c596 Mon Sep 17 00:00:00 2001
From: Noah Goldstein <goldstein.w.n at gmail.com>
Date: Sun, 10 Mar 2024 16:30:24 -0500
Subject: [PATCH 1/2] [InstCombine] Add more tests for folding rem/div/mul with
 select; NFC

---
 .../Transforms/InstCombine/binop-select.ll    | 98 +++++++++++++++++++
 1 file changed, 98 insertions(+)

diff --git a/llvm/test/Transforms/InstCombine/binop-select.ll b/llvm/test/Transforms/InstCombine/binop-select.ll
index 6cd4132eadd77b..22796405a67f5f 100644
--- a/llvm/test/Transforms/InstCombine/binop-select.ll
+++ b/llvm/test/Transforms/InstCombine/binop-select.ll
@@ -403,3 +403,101 @@ define i32 @ashr_sel_op1_use(i1 %b) {
   %r = ashr i32 -2, %s
   ret i32 %r
 }
+
+
+define i32 @test_mul_to_const_Cmul(i32 %x) {
+; CHECK-LABEL: @test_mul_to_const_Cmul(
+; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[X:%.*]], 61
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C]], i32 9, i32 14
+; CHECK-NEXT:    [[R:%.*]] = mul i32 [[COND]], [[X]]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %c = icmp eq i32 %x, 61
+  %cond = select i1 %c, i32 9, i32 14
+  %r = mul i32 %x, %cond
+  ret i32 %r
+}
+
+define i32 @test_mul_to_const_mul(i32 %x, i32 %y) {
+; CHECK-LABEL: @test_mul_to_const_mul(
+; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[X:%.*]], 61
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C]], i32 9, i32 [[Y:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = mul i32 [[COND]], [[X]]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %c = icmp eq i32 %x, 61
+  %cond = select i1 %c, i32 9, i32 %y
+  %r = mul i32 %x, %cond
+  ret i32 %r
+}
+
+
+define i32 @test_mul_to_const_Cmul_fail_multiuse(i32 %x, i32 %y) {
+; CHECK-LABEL: @test_mul_to_const_Cmul_fail_multiuse(
+; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[X:%.*]], 61
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C]], i32 9, i32 14
+; CHECK-NEXT:    [[R:%.*]] = mul i32 [[COND]], [[X]]
+; CHECK-NEXT:    call void @use(i32 [[COND]])
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %c = icmp eq i32 %x, 61
+  %cond = select i1 %c, i32 9, i32 14
+  %r = mul i32 %x, %cond
+  call void @use(i32 %cond)
+  ret i32 %r
+}
+
+
+define i32 @test_div_to_const_div_fail_non_speculatable(i32 %x, i32 %y) {
+; CHECK-LABEL: @test_div_to_const_div_fail_non_speculatable(
+; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[X:%.*]], 61
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C]], i32 9, i32 [[Y:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = udiv i32 [[X]], [[COND]]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %c = icmp eq i32 %x, 61
+  %cond = select i1 %c, i32 9, i32 %y
+  %r = udiv i32 %x, %cond
+  ret i32 %r
+}
+
+
+define i32 @test_div_to_const_Cdiv_todo(i32 %x, i32 %y) {
+; CHECK-LABEL: @test_div_to_const_Cdiv_todo(
+; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[X:%.*]], 61
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C]], i32 9, i32 14
+; CHECK-NEXT:    [[R:%.*]] = udiv i32 [[X]], [[COND]]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %c = icmp eq i32 %x, 61
+  %cond = select i1 %c, i32 9, i32 14
+  %r = udiv i32 %x, %cond
+  ret i32 %r
+}
+
+
+define i32 @test_or_with_multiuse_fail(i1 %cond) {
+; CHECK-LABEL: @test_or_with_multiuse_fail(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 32, i32 0
+; CHECK-NEXT:    call void @use(i32 [[SEL]])
+; CHECK-NEXT:    [[RET:%.*]] = or disjoint i32 [[SEL]], 22
+; CHECK-NEXT:    ret i32 [[RET]]
+;
+  %sel = select i1 %cond, i32 32, i32 0
+  call void @use(i32 %sel)
+  %ret = or disjoint i32 %sel, 22
+  ret i32 %ret
+}
+
+define i32 @test_div_with_multiuse(i1 %cond) {
+; CHECK-LABEL: @test_div_with_multiuse(
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 132, i32 66
+; CHECK-NEXT:    call void @use(i32 [[SEL]])
+; CHECK-NEXT:    [[RET:%.*]] = udiv i32 [[SEL]], 22
+; CHECK-NEXT:    ret i32 [[RET]]
+;
+  %sel = select i1 %cond, i32 132, i32 66
+  call void @use(i32 %sel)
+  %ret = sdiv i32 %sel, 22
+  ret i32 %ret
+}

>From 008d412b944001a2ca53bb42e61d6d9239ba1cf6 Mon Sep 17 00:00:00 2001
From: Noah Goldstein <goldstein.w.n at gmail.com>
Date: Sun, 10 Mar 2024 15:14:25 -0500
Subject: [PATCH 2/2] [InstCombine] Use the select condition to try to constant
 fold binops into select

The select condition may allow us to constant fold binops on
non-constant arms if the condition implies one of the binop operand is constant.
For example if we have:
```
%c = icmp eq i8 %y, 10
%s = select i1 %c, i8 123, i8 %x
%r = mul i8 %s, %y
```

We can replace substitate `10` in for `%y` on the true arm and do:
```
%c = icmp eq i8 %y, 10
%mul = mul i8 %x, %y
%r = select i1 %c, i8 1230, i8 %mul
```
---
 .../InstCombine/InstCombineInternal.h         |  9 ++-
 .../InstCombine/InstCombineMulDivRem.cpp      | 24 ++++----
 .../InstCombine/InstructionCombining.cpp      | 56 +++++++++++++++++--
 .../Transforms/InstCombine/binop-select.ll    | 10 ++--
 llvm/test/Transforms/InstCombine/pr72433.ll   |  3 +-
 5 files changed, 74 insertions(+), 28 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
index 6a1ef6edeb4077..63e8d4eca90f43 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
+++ b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
@@ -610,8 +610,13 @@ class LLVM_LIBRARY_VISIBILITY InstCombinerImpl final
   Instruction *FoldOpIntoSelect(Instruction &Op, SelectInst *SI,
                                 bool FoldWithMultiUse = false);
 
-  /// This is a convenience wrapper function for the above two functions.
-  Instruction *foldBinOpIntoSelectOrPhi(BinaryOperator &I);
+  /// This is a convenience wrapper function for the above function.
+  Instruction *foldBinOpIntoSelect(BinaryOperator &I,
+                                   bool AllowMultiUse = false);
+
+  /// This is a convenience wrapper function for the above three functions.
+  Instruction *foldBinOpIntoSelectOrPhi(BinaryOperator &I,
+                                        bool AllowMultiUseSelect = false);
 
   Instruction *foldAddWithConstant(BinaryOperator &Add);
 
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
index 278be6233f4b8a..624183becdcbd1 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
@@ -287,7 +287,8 @@ Instruction *InstCombinerImpl::visitMul(BinaryOperator &I) {
     }
   }
 
-  if (Instruction *FoldedMul = foldBinOpIntoSelectOrPhi(I))
+  if (Instruction *FoldedMul =
+          foldBinOpIntoSelectOrPhi(I, /*AllowMultiUseSelect=*/true))
     return FoldedMul;
 
   if (Value *FoldedMul = foldMulSelectToNegate(I, Builder))
@@ -1048,12 +1049,9 @@ Instruction *InstCombinerImpl::commonIDivTransforms(BinaryOperator &I) {
   // If the divisor is a select-of-constants, try to constant fold all div ops:
   // C / (select Cond, TrueC, FalseC) --> select Cond, (C / TrueC), (C / FalseC)
   // TODO: Adapt simplifyDivRemOfSelectWithZeroOp to allow this and other folds.
-  if (match(Op0, m_ImmConstant()) &&
-      match(Op1, m_Select(m_Value(), m_ImmConstant(), m_ImmConstant()))) {
-    if (Instruction *R = FoldOpIntoSelect(I, cast<SelectInst>(Op1),
-                                          /*FoldWithMultiUse*/ true))
-      return R;
-  }
+  if (Instruction *R = foldBinOpIntoSelect(I,
+                                           /*AllowMultiUse=*/true))
+    return R;
 
   const APInt *C2;
   if (match(Op1, m_APInt(C2))) {
@@ -1136,7 +1134,8 @@ Instruction *InstCombinerImpl::commonIDivTransforms(BinaryOperator &I) {
     }
 
     if (!C2->isZero()) // avoid X udiv 0
-      if (Instruction *FoldedDiv = foldBinOpIntoSelectOrPhi(I))
+      if (Instruction *FoldedDiv =
+              foldBinOpIntoSelectOrPhi(I, /*AllowMultiUseSelect=*/true))
         return FoldedDiv;
   }
 
@@ -2001,12 +2000,9 @@ Instruction *InstCombinerImpl::commonIRemTransforms(BinaryOperator &I) {
   // If the divisor is a select-of-constants, try to constant fold all rem ops:
   // C % (select Cond, TrueC, FalseC) --> select Cond, (C % TrueC), (C % FalseC)
   // TODO: Adapt simplifyDivRemOfSelectWithZeroOp to allow this and other folds.
-  if (match(Op0, m_ImmConstant()) &&
-      match(Op1, m_Select(m_Value(), m_ImmConstant(), m_ImmConstant()))) {
-    if (Instruction *R = FoldOpIntoSelect(I, cast<SelectInst>(Op1),
-                                          /*FoldWithMultiUse*/ true))
-      return R;
-  }
+  if (Instruction *R = foldBinOpIntoSelect(I,
+                                           /*AllowMultiUse=*/true))
+    return R;
 
   if (isa<Constant>(Op1)) {
     if (Instruction *Op0I = dyn_cast<Instruction>(Op0)) {
diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 1a831805dc72a0..aab4fe8ce87f79 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -1951,14 +1951,60 @@ Instruction *InstCombinerImpl::foldBinopWithPhiOperands(BinaryOperator &BO) {
   return NewPhi;
 }
 
-Instruction *InstCombinerImpl::foldBinOpIntoSelectOrPhi(BinaryOperator &I) {
+// Return std::nullopt if we should not fold. Return true if we should fold
+// multi-use select and false for single-use select.
+static std::optional<bool> shouldFoldOpIntoSelect(BinaryOperator &I, Value *Op,
+                                                  Value *OpOther,
+                                                  bool AllowMultiUse) {
+  if (!AllowMultiUse && !Op->hasOneUse())
+    return std::nullopt;
+  if (isa<SelectInst>(Op))
+    // If we will be able to constant fold the incorperated binop, then
+    // multi-use. Otherwise single-use.
+    return match(OpOther, m_ImmConstant()) &&
+           match(Op, m_Select(m_Value(), m_ImmConstant(), m_ImmConstant()));
+
+  return std::nullopt;
+}
+
+Instruction *InstCombinerImpl::foldBinOpIntoSelect(BinaryOperator &I,
+                                                   bool AllowMultiUse) {
+  std::optional<bool> CanSpeculativelyExecuteRes;
+  for (unsigned OpIdx = 0; OpIdx < 2; ++OpIdx) {
+    // Slightly more involved logic for select. For select we use the condition
+    // to to infer information about the arm. This allows us to constant-fold
+    // even when the select arm(s) are not constant. For example if we have: `(X
+    // == 10 ? 19 : Y) * X`, we can entirely contant fold the true arm as `X ==
+    // 10` dominates it. So we end up with `X == 10 ? 190 : (X * Y))`.
+    if (auto MultiUse = shouldFoldOpIntoSelect(
+            I, I.getOperand(OpIdx), I.getOperand(1 - OpIdx), AllowMultiUse)) {
+      if (!*MultiUse) {
+        if (!CanSpeculativelyExecuteRes.has_value()) {
+          const SimplifyQuery Q = SQ.getWithInstruction(&I);
+          CanSpeculativelyExecuteRes =
+              isSafeToSpeculativelyExecute(&I, Q.CxtI, Q.AC, Q.DT, &TLI);
+        }
+        if (!*CanSpeculativelyExecuteRes)
+          return nullptr;
+      }
+      if (Instruction *NewSel = FoldOpIntoSelect(
+              I, cast<SelectInst>(I.getOperand(OpIdx)), *MultiUse))
+        return NewSel;
+    }
+  }
+  return nullptr;
+}
+
+Instruction *
+InstCombinerImpl::foldBinOpIntoSelectOrPhi(BinaryOperator &I,
+                                           bool AllowMultiUseSelect) {
+  if (auto *SI = foldBinOpIntoSelect(I, AllowMultiUseSelect))
+    return SI;
+
   if (!isa<Constant>(I.getOperand(1)))
     return nullptr;
 
-  if (auto *Sel = dyn_cast<SelectInst>(I.getOperand(0))) {
-    if (Instruction *NewSel = FoldOpIntoSelect(I, Sel))
-      return NewSel;
-  } else if (auto *PN = dyn_cast<PHINode>(I.getOperand(0))) {
+  if (auto *PN = dyn_cast<PHINode>(I.getOperand(0))) {
     if (Instruction *NewPhi = foldOpIntoPhi(I, PN))
       return NewPhi;
   }
diff --git a/llvm/test/Transforms/InstCombine/binop-select.ll b/llvm/test/Transforms/InstCombine/binop-select.ll
index 22796405a67f5f..35ae545c458405 100644
--- a/llvm/test/Transforms/InstCombine/binop-select.ll
+++ b/llvm/test/Transforms/InstCombine/binop-select.ll
@@ -408,8 +408,8 @@ define i32 @ashr_sel_op1_use(i1 %b) {
 define i32 @test_mul_to_const_Cmul(i32 %x) {
 ; CHECK-LABEL: @test_mul_to_const_Cmul(
 ; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[X:%.*]], 61
-; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C]], i32 9, i32 14
-; CHECK-NEXT:    [[R:%.*]] = mul i32 [[COND]], [[X]]
+; CHECK-NEXT:    [[TMP1:%.*]] = mul i32 [[X]], 14
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[C]], i32 549, i32 [[TMP1]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %c = icmp eq i32 %x, 61
@@ -421,8 +421,8 @@ define i32 @test_mul_to_const_Cmul(i32 %x) {
 define i32 @test_mul_to_const_mul(i32 %x, i32 %y) {
 ; CHECK-LABEL: @test_mul_to_const_mul(
 ; CHECK-NEXT:    [[C:%.*]] = icmp eq i32 [[X:%.*]], 61
-; CHECK-NEXT:    [[COND:%.*]] = select i1 [[C]], i32 9, i32 [[Y:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = mul i32 [[COND]], [[X]]
+; CHECK-NEXT:    [[TMP1:%.*]] = mul i32 [[Y:%.*]], [[X]]
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[C]], i32 549, i32 [[TMP1]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %c = icmp eq i32 %x, 61
@@ -493,7 +493,7 @@ define i32 @test_div_with_multiuse(i1 %cond) {
 ; CHECK-LABEL: @test_div_with_multiuse(
 ; CHECK-NEXT:    [[SEL:%.*]] = select i1 [[COND:%.*]], i32 132, i32 66
 ; CHECK-NEXT:    call void @use(i32 [[SEL]])
-; CHECK-NEXT:    [[RET:%.*]] = udiv i32 [[SEL]], 22
+; CHECK-NEXT:    [[RET:%.*]] = select i1 [[COND]], i32 6, i32 3
 ; CHECK-NEXT:    ret i32 [[RET]]
 ;
   %sel = select i1 %cond, i32 132, i32 66
diff --git a/llvm/test/Transforms/InstCombine/pr72433.ll b/llvm/test/Transforms/InstCombine/pr72433.ll
index c6e74582a13d30..1633885075e872 100644
--- a/llvm/test/Transforms/InstCombine/pr72433.ll
+++ b/llvm/test/Transforms/InstCombine/pr72433.ll
@@ -6,8 +6,7 @@ define i32 @widget(i32 %arg, i32 %arg1) {
 ; CHECK-SAME: i32 [[ARG:%.*]], i32 [[ARG1:%.*]]) {
 ; CHECK-NEXT:  bb:
 ; CHECK-NEXT:    [[ICMP:%.*]] = icmp ne i32 [[ARG]], 0
-; CHECK-NEXT:    [[TMP0:%.*]] = zext i1 [[ICMP]] to i32
-; CHECK-NEXT:    [[MUL:%.*]] = shl nuw nsw i32 20, [[TMP0]]
+; CHECK-NEXT:    [[MUL:%.*]] = select i1 [[ICMP]], i32 40, i32 20
 ; CHECK-NEXT:    [[XOR:%.*]] = zext i1 [[ICMP]] to i32
 ; CHECK-NEXT:    [[ADD9:%.*]] = or disjoint i32 [[MUL]], [[XOR]]
 ; CHECK-NEXT:    [[TMP1:%.*]] = zext i1 [[ICMP]] to i32



More information about the llvm-commits mailing list