From 93f3d7f7b3a902ad9e90ef0d9bf147582a2cf620 Mon Sep 17 00:00:00 2001
From: Roman Lebedev <lebedev.ri@gmail.com>
Date: Wed, 4 Nov 2020 14:33:11 +0300
Subject: [PATCH] [Reassociate] Guard `add`-like `or` conversion into an `add`
 with profitability check

This is slightly better compile-time wise,
since we avoid potentially-costly knownbits analysis that will
ultimately not allow us to actually do anything with said `add`.
---
 llvm/lib/Transforms/Scalar/Reassociate.cpp      | 23 +++++++++++++++++++++++
 llvm/test/Transforms/Reassociate/add-like-or.ll | 14 +++++++++-----
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/Reassociate.cpp b/llvm/lib/Transforms/Scalar/Reassociate.cpp
index 164971c..73fdd6e 100644
--- a/llvm/lib/Transforms/Scalar/Reassociate.cpp
+++ b/llvm/lib/Transforms/Scalar/Reassociate.cpp
@@ -920,6 +920,28 @@ static Value *NegateValue(Value *V, Instruction *BI,
   return NewNeg;
 }
 
+/// Return true if it may be profitable to convert this (X|Y) into (X+Y).
+static bool ShouldConvertOrWithNoCommonBitsToAdd(Instruction *Or) {
+  // Don't bother to convert this up unless either the LHS is an associable add
+  // or subtract or mul or if this is only used by one of the above.
+  // This is only a compile-time improvement, it is not needed for correctness!
+  auto isInteresting = [](Value *V) {
+    for (auto Op : {Instruction::Add, Instruction::Sub, Instruction::Mul})
+      if (isReassociableOp(V, Op))
+        return true;
+    return false;
+  };
+
+  if (any_of(Or->operands(), isInteresting))
+    return true;
+
+  Value *VB = Or->user_back();
+  if (Or->hasOneUse() && isInteresting(VB))
+    return true;
+
+  return false;
+}
+
 /// If we have (X|Y), and iff X and Y have no common bits set,
 /// transform this into (X+Y) to allow arithmetics reassociation.
 static BinaryOperator *ConvertOrWithNoCommonBitsToAdd(Instruction *Or) {
@@ -2137,6 +2159,7 @@ void ReassociatePass::OptimizeInst(Instruction *I) {
   // If this is a bitwise or instruction of operands
   // with no common bits set, convert it to X+Y.
   if (I->getOpcode() == Instruction::Or &&
+      ShouldConvertOrWithNoCommonBitsToAdd(I) &&
       haveNoCommonBitsSet(I->getOperand(0), I->getOperand(1),
                           I->getModule()->getDataLayout(), /*AC=*/nullptr, I,
                           /*DT=*/nullptr)) {
diff --git a/llvm/test/Transforms/Reassociate/add-like-or.ll b/llvm/test/Transforms/Reassociate/add-like-or.ll
index 13477ad..dec1b3f 100644
--- a/llvm/test/Transforms/Reassociate/add-like-or.ll
+++ b/llvm/test/Transforms/Reassociate/add-like-or.ll
@@ -6,23 +6,27 @@
 define i32 @test1(i32 %a, i32 %b) {
 ; CHECK-LABEL: @test1(
 ; CHECK-NEXT:    [[C:%.*]] = or i32 [[B:%.*]], [[A:%.*]]
-; CHECK-NEXT:    ret i32 [[C]]
+; CHECK-NEXT:    [[C_PLUS_ONE:%.*]] = add i32 [[C]], 1
+; CHECK-NEXT:    ret i32 [[C_PLUS_ONE]]
 ;
   %c = or i32 %a, %b
-  ret i32 %c
+  %c.plus.one = add i32 %c, 1
+  ret i32 %c.plus.one
 }
 
 ; But if we *do* know  that operands have no common bits set,
 ; we *can* convert the `or` into an `add`.
-define i32 @test2(i32 %x) {
+define i32 @test2(i32 %x, i32 %y) {
 ; CHECK-LABEL: @test2(
 ; CHECK-NEXT:    [[X_NUMLZ:%.*]] = tail call i32 @llvm.ctlz.i32(i32 [[X:%.*]], i1 true), [[RNG0:!range !.*]]
 ; CHECK-NEXT:    [[RES:%.*]] = add nuw nsw i32 [[X_NUMLZ]], -32
-; CHECK-NEXT:    ret i32 [[RES]]
+; CHECK-NEXT:    [[RES_PLUS_ONE:%.*]] = add i32 [[RES]], [[Y:%.*]]
+; CHECK-NEXT:    ret i32 [[RES_PLUS_ONE]]
 ;
   %x.numlz = tail call i32 @llvm.ctlz.i32(i32 %x, i1 true), !range !0
   %res = or i32 %x.numlz, -32
-  ret i32 %res
+  %res.plus.one = add i32 %res, %y
+  ret i32 %res.plus.one
 }
 
 ; And that allows reassociation in general.
-- 
2.7.4