[NaryReassociate] Transform expression (X << C1) + C2 to (X + C3) << C1,

[AlexM]

Summary

where C1, C2 and C3 are integer constants, such that C2 = C3 << C1.

Some transformations in Instruction Combining may block GVN later.
For example,

(X + const1) << 22 || (X + const1) >> 10

may become

(X << 22 + const2) || (X + const1) >> 10

which will disable the GVN on x + const1.

This patch implements reverse transformation of such expressions thus, unfolding a common code to be eliminated by either GVN or CSE. In the default compilation pipeline GVN/CSE are run several times: after Instruction Combining, but before NaryReassociate and after NaryReassociate. Therefore, with this patch we try both patterns for GVN/CSE.

[mkazantsev]

llvm/include/llvm/Transforms/Scalar/NaryReassociate.h

152

// Tries to match (X << C1) + C2...

?

llvm/lib/Transforms/Scalar/NaryReassociate.cpp

510

If the initial operations had nuw/nsw flags, will they now be lost? Can we preserve them somehow?

llvm/test/Transforms/NaryReassociate/NVPTX/const-shl.ll

1

Could you please precommit the test to show what your patch is changing?

[mkazantsev]

Do you know why at all InstCombine would want to transform (X + const1) << 22 into (X << 22 + const2)? By itself it doesn't seem profitable. Does it enable any useful further transforms?

[ebrevnov]

As I see it, you are trying to solve pass ordering issue (InstCombine prevents GVN/CSE). But your solution assumes there is GVN/CSE follows NaryReassociation what creates another potential pass ordering dependence.
Another point is the way the transformation is implemented right now doesn't fit intent of NaryReassociation pass. In the heart of NaryReassociation pass is findClosestMatchingDominator call to find and delete common expressions after reassociation. Since you are not checking for existence of common expressions you can easily do the transformation in some other place (ordinary Reassociate pass, or InstSimplify/InstCombine)
In order to solve both mentioned problems and cover your case you will need to use findClosestMatchingDominator and possibly extend it to look for post dominating expressions as well.

[nikic]

As the first step, please add a phase ordering test demonstrating the issue. Shouldn't EarlyCSE remove these already?

[AlexM]

In D146928#4232813, @mkazantsev wrote:

Do you know why at all InstCombine would want to transform (X + const1) << 22 into (X << 22 + const2)?

This functionality was added in 44bd392cbff238e75d84ee757189d5f62cf91679 which says simply that it helps by "opening opportunities for future improvements."

In D146928#4232900, @ebrevnov wrote:

As I see it, you are trying to solve pass ordering issue (InstCombine prevents GVN/CSE). But your solution assumes there is GVN/CSE follows NaryReassociation what creates another potential pass ordering dependence.
Another point is the way the transformation is implemented right now doesn't fit intent of NaryReassociation pass. In the heart of NaryReassociation pass is findClosestMatchingDominator call to find and delete common expressions after reassociation. Since you are not checking for existence of common expressions you can easily do the transformation in some other place (ordinary Reassociate pass, or InstSimplify/InstCombine)
In order to solve both mentioned problems and cover your case you will need to use findClosestMatchingDominator and possibly extend it to look for post dominating expressions as well.

Hmm, this is a pretty fundamental issue, let me take some time to figure out how I can integrate this more deeply into the pass or move it somewhere more appropriate.

llvm/lib/Transforms/Scalar/NaryReassociate.cpp

510

This dose seem like an issue, might require a bit more code but certainly seems feasible. I did check InstCombine, which does the opposite transformation and these flags are completely discarded there. I'm very unfamiliar with the semantics of these flags but if dropping them is a problem here is it also a problem in InstCombine? Is dropping them a correctness issue or does it potentially prevent other optimizations?

llvm/test/Transforms/NaryReassociate/NVPTX/const-shl.ll

1

This is an embarrassingly ignorant question but what do you mean by this?

llvm/include/llvm/Transforms/Scalar/NaryReassociate.h

Instruction *tryReassociateBinaryOp(Value *LHS, Value *RHS,
                                       BinaryOperator *I);
   // Rewrites I to (LHS op RHS) if LHS is computed already.
   Instruction *tryReassociatedBinaryOp(const SCEV *LHS, Value *RHS,
                                        BinaryOperator *I);
 
   // Tries to match Op1 and Op2 by using V.
   bool matchTernaryOp(BinaryOperator *I, Value *V, Value *&Op1, Value *&Op2);
 
+  // Tries to match X, C1 and C2 by using LHS and RHS.
+  bool matchConstLShiftOp(BinaryOperator *I, Value *LHS, Value *RHS, Value *&X,
+                          const APInt *&C1, const APInt *&C2);
+
   // Gets SCEV for (LHS op RHS).
   const SCEV *getBinarySCEV(BinaryOperator *I, const SCEV *LHS,
                             const SCEV *RHS);
 
   // Returns the closest dominator of \c Dominatee that computes
   // \c CandidateExpr. Returns null if not found.
   Instruction *findClosestMatchingDominator(const SCEV *CandidateExpr,
                                             Instruction *Dominatee);

llvm/lib/Transforms/Scalar/NaryReassociate.cpp

if (auto *NewI = tryReassociateBinaryOp(LHS, RHS, I))
     return NewI;
   if (auto *NewI = tryReassociateBinaryOp(RHS, LHS, I))
     return NewI;
   return nullptr;
 }
 
 Instruction *NaryReassociatePass::tryReassociateBinaryOp(Value *LHS, Value *RHS,
                                                          BinaryOperator *I) {
+  // To be conservative, we reassociate I only when it is the only user of LHS.
+  if (!LHS->hasOneUse())
+    return nullptr;
+
   Value *A = nullptr, *B = nullptr;
-  // To be conservative, we reassociate I only when it is the only user of (A op
-  // B).
-  if (LHS->hasOneUse() && matchTernaryOp(I, LHS, A, B)) {
+  if (matchTernaryOp(I, LHS, A, B)) {
     // I = (A op B) op RHS
     //   = (A op RHS) op B or (B op RHS) op A
     const SCEV *AExpr = SE->getSCEV(A), *BExpr = SE->getSCEV(B);
     const SCEV *RHSExpr = SE->getSCEV(RHS);
     if (BExpr != RHSExpr) {
       if (auto *NewI =
               tryReassociatedBinaryOp(getBinarySCEV(I, AExpr, RHSExpr), B, I))
         return NewI;
     }
     if (AExpr != RHSExpr) {
       if (auto *NewI =
               tryReassociatedBinaryOp(getBinarySCEV(I, BExpr, RHSExpr), A, I))
         return NewI;
     }
+    return nullptr;
+  }
+
+  // Handle special case: I = (X << C1) + C2, where C1 and C2 are constant
+  // values such that I can be represented as (X + C3) << C1, C2 = C3 << C1.
+  // Then (X + C3) expression can be CSE'd later.
+  Value *X = nullptr;
+  const APInt *C1 = nullptr, *C2 = nullptr;
+  if (matchConstLShiftOp(I, LHS, RHS, X, C1, C2)) {
+    auto *C3Val = ConstantInt::get(I->getType(), C2->lshr(*C1));
+    auto *Add = BinaryOperator::CreateAdd(X, C3Val, "add.nary", I);
+    auto *C1Val = ConstantInt::get(I->getType(), *C1);
+    return BinaryOperator::CreateShl(Add, C1Val, "shl.nary", I);
   }
   return nullptr;
 }
 
 Instruction *NaryReassociatePass::tryReassociatedBinaryOp(const SCEV *LHSExpr,
                                                           Value *RHS,
                                                           BinaryOperator *I) {
   // Look for the closest dominator LHS of I that computes LHSExpr, and replace
case Instruction::Mul:
     break;
   default:
     llvm_unreachable("Unexpected instruction.");
   }
   NewI->takeName(I);
   return NewI;
 }
 
+// Match if I = (X << C1) + C2, such that C2 = C3 << C1,
+// where C1, C2 and C3 are constant values.
+bool NaryReassociatePass::matchConstLShiftOp(BinaryOperator *I, Value *LHS,
+                                             Value *RHS, Value *&X,
+                                             const APInt *&C1,
+                                             const APInt *&C2) {
+  if (I->getOpcode() == Instruction::Add) {
+    if (match(LHS, m_Shl(m_Value(X), m_APInt(C1))) && match(RHS, m_APInt(C2)) &&
+        (C2->countTrailingZeros() >= C1->getZExtValue()))
+      return true;
+  }
+  return false;
+}
+
 bool NaryReassociatePass::matchTernaryOp(BinaryOperator *I, Value *V,
                                          Value *&Op1, Value *&Op2) {
   switch (I->getOpcode()) {
   case Instruction::Add:
     return match(V, m_Add(m_Value(Op1), m_Value(Op2)));
   case Instruction::Mul:
     return match(V, m_Mul(m_Value(Op1), m_Value(Op2)));
   default:

llvm/test/Transforms/NaryReassociate/NVPTX/const-shl.ll

+; RUN: opt < %s -passes=instcombine,nary-reassociate -S | FileCheck %s
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s --check-prefix=INSTCOMB
+
+define i32 @foo(i32 %x) {
+; CHECK-LABEL: @foo
+; CHECK-NEXT: %add.nary = add i32 %x, 787
+; CHECK-NEXT: %shl.nary = shl i32 %add.nary, 22
+; CHECK-NEXT: %add.nary1 = add i32 %x, 787
+; CHECK-NEXT: %shl.nary2 = shl i32 %add.nary1, 11
+; CHECK-NEXT: %add.nary3 = add i32 %x, 787
+; CHECK-NEXT: %shl.nary4 = shl i32 %add.nary3, 3
+; CHECK-NEXT: %add.nary5 = add i32 %x, 787
+; CHECK-NEXT: %shl.nary6 = shl i32 %add.nary5, 13
+; CHECK-NEXT: %or1 = or i32 %shl.nary, %shl.nary2
+; CHECK-NEXT: %or2 = or i32 %shl.nary4, %shl.nary6
+; CHECK-NEXT: %or3 = or i32 %or1, %or2
+; CHECK-NEXT: ret i32 %or3
+
+  %add1 = shl i32 %x, 22
+  %shl1 = add i32 %add1, -994050048
+  %add2 = shl i32 %x, 11
+  %shl2 = add i32 %add2, 1611776
+  %add3 = shl i32 %x, 3
+  %shl3 = add i32 %add3, 6296
+  %add4 = shl i32 %x, 13
+  %shl4 = add i32 %add4, 6447104
+  %or1 = or i32 %shl1, %shl2
+  %or2 = or i32 %shl3, %shl4
+  %or3 = or i32 %or1, %or2
+  ret i32 %or3
+}
+
+; An example of how instcombine performs transformations
+; that block GVN/CSE from happening.
+
+define i32 @bar(i32 %x) {
+; INSTCOMB-LABEL: @bar
+; INSTCOMB-NEXT: %add.nary = shl i32 %x, 22
+; INSTCOMB-NEXT: %shl.nary = add i32 %add.nary, -994050048
+; INSTCOMB-NEXT: %add.nary1 = shl i32 %x, 11
+; INSTCOMB-NEXT: %shl.nary2 = add i32 %add.nary1, 1611776
+; INSTCOMB-NEXT: %add.nary3 = shl i32 %x, 3
+; INSTCOMB-NEXT: %shl.nary4 = add i32 %add.nary3, 6296
+; INSTCOMB-NEXT: %add.nary5 = shl i32 %x, 13
+; INSTCOMB-NEXT: %shl.nary6 = add i32 %add.nary5, 6447104
+; INSTCOMB-NEXT: %or1 = or i32 %shl.nary, %shl.nary2
+; INSTCOMB-NEXT: %or2 = or i32 %shl.nary4, %shl.nary6
+; INSTCOMB-NEXT: %or3 = or i32 %or1, %or2
+; INSTCOMB-NEXT: ret i32 %or3
+
+  %add.nary = add i32 %x, 787
+  %shl.nary = shl i32 %add.nary, 22
+  %add.nary1 = add i32 %x, 787
+  %shl.nary2 = shl i32 %add.nary1, 11
+  %add.nary3 = add i32 %x, 787
+  %shl.nary4 = shl i32 %add.nary3, 3
+  %add.nary5 = add i32 %x, 787
+  %shl.nary6 = shl i32 %add.nary5, 13
+  %or1 = or i32 %shl.nary, %shl.nary2
+  %or2 = or i32 %shl.nary4, %shl.nary6
+  %or3 = or i32 %or1, %or2
+  ret i32 %or3
+}