[Transforms][Reassociate] "Reassociate expressions" pass optimizations not always profitable

[pawosm01]

Summary

Consider the following piece of code:

void innermost_loop(int i, double d1, double d2, double delta, int n, double cells[n])
{
  int j;
  const double d1d = d1 * delta;
  const double d2d = d2 * delta;

  for (j = 0; j <= i; j++)
    cells[j] = d1d * cells[j + 1] + d2d * cells[j];
}

When compiling at -Ofast level, after the "Reassociate expressions"
pass, this code is transformed into an equivalent of:

int j;

for (j = 0; j <= i; j++)
  cells[j] = (d1 * cells[j + 1] + d2 * cells[j]) * delta;

Effectively, the computation of those loop invariants isn't done
before the loop anymore, we have one extra multiplication on each
loop iteration instead. Sadly, this results in a significant
performance hit.

This patch makes the OptimizeAdd() function of the "Reassociate
expressions" pass aware of modifying operations in a loop and
bails out when some of the operands are pulled from the outside
of the loop.

See also: https://github.com/llvm/llvm-project/issues/62736

[pawosm01]

After the precommit comes the real thing.

[tschuett]

llvm/lib/Transforms/Scalar/Reassociate.cpp

1707

for (unsigned i = 0, size_t e = Ops.size(); i != e; ++i) {
}

[nikic]

Just to get the obvious issues out of the way...

llvm/lib/Transforms/Scalar/Reassociate.cpp

1699

These need to be fetched from the pass manager instead.

1701

Should be just LoopInfo.

1732

This looks like a very complicated and inefficient re-implementation of LI->getLoopFor().

llvm/test/Transforms/Reassociate/reassociate-not-from-the-outside-of-the-loop.ll

1

Should use update_test_checks.py.

[pawosm01]

llvm/test/Transforms/Reassociate/reassociate-not-from-the-outside-of-the-loop.ll

1

It will make it look like this:

--- a/llvm/test/Transforms/Reassociate/reassociate-not-from-the-outside-of-the-loop.ll
+++ b/llvm/test/Transforms/Reassociate/reassociate-not-from-the-outside-of-the-loop.ll
@@ -1,3 +1,4 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
 ; RUN: opt -passes=reassociate -S < %s | FileCheck %s

 ; This test is to ensure that no computations are pulled into a loop
@@ -9,12 +10,35 @@
 ; Reassociate pass.

 define void @innermost_loop(i32 %i, double %d1, double %d2, double %delta, ptr %cells) {
-; CHECK-LABEL: @innermost_loop(
+; CHECK-LABEL: define void @innermost_loop
+; CHECK-SAME: (i32 [[I:%.*]], double [[D1:%.*]], double [[D2:%.*]], double [[DELTA:%.*]], ptr [[CELLS:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[MUL:%.*]] = fmul fast double [[DELTA]], [[D1]]
+; CHECK-NEXT:    [[MUL1:%.*]] = fmul fast double [[DELTA]], [[D2]]
+; CHECK-NEXT:    br label [[FOR_COND:%.*]]
+; CHECK:       for.cond:
+; CHECK-NEXT:    [[J_0:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[ADD:%.*]], [[FOR_BODY:%.*]] ]
+; CHECK-NEXT:    [[CMP_NOT:%.*]] = icmp sgt i32 [[J_0]], [[I]]
+; CHECK-NEXT:    br i1 [[CMP_NOT]], label [[FOR_END:%.*]], label [[FOR_BODY]]
+; CHECK:       for.body:
+; CHECK-NEXT:    [[ADD]] = add nuw nsw i32 [[J_0]], 1
+; CHECK-NEXT:    [[IDXPROM:%.*]] = zext i32 [[ADD]] to i64
+; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds double, ptr [[CELLS]], i64 [[IDXPROM]]
+; CHECK-NEXT:    [[TMP0:%.*]] = load double, ptr [[ARRAYIDX]], align 8
+; CHECK-NEXT:    [[MUL2:%.*]] = fmul fast double [[MUL]], [[TMP0]]
+; CHECK-NEXT:    [[IDXPROM3:%.*]] = zext i32 [[J_0]] to i64
+; CHECK-NEXT:    [[ARRAYIDX4:%.*]] = getelementptr inbounds double, ptr [[CELLS]], i64 [[IDXPROM3]]
+; CHECK-NEXT:    [[TMP1:%.*]] = load double, ptr [[ARRAYIDX4]], align 8
+; CHECK-NEXT:    [[MUL5:%.*]] = fmul fast double [[MUL1]], [[TMP1]]
+; CHECK-NEXT:    [[ADD6:%.*]] = fadd fast double [[MUL5]], [[MUL2]]
+; CHECK-NEXT:    store double [[ADD6]], ptr [[ARRAYIDX4]], align 8
+; CHECK-NEXT:    br label [[FOR_COND]]
+; CHECK:       for.end:
+; CHECK-NEXT:    ret void
+;

Is that what is expected here? It doesn't seem focusing on the matter of this patch, and it doesn't seem futureproof: many things not related to this patch can change that can affect this test case in this form...

[pawosm01]

llvm/test/Transforms/Reassociate/reassociate-not-from-the-outside-of-the-loop.ll

1

...and where those CHECK-NOT's has gone?

[pawosm01]

Review comments addressed.

[pawosm01]

llvm/lib/Transforms/Scalar/Reassociate.cpp

1699

OK, used pass manager instead.

1701

Although this place has gone, I've considered your suggestion in the other place below.

1707

Thanks for spotting that!

1732

using LI->getLoopFor() instead now.

[fhahn]

There's been some work to reassoicate in LICM to enable hoisting: D148001. It would probably be more general to extend the logic in LICM to handle those cases, which would also trigger if the user wrote the expression like reassociate creates at the moment.

[pawosm01]

In D151616#4381099, @fhahn wrote:

There's been some work to reassoicate in LICM to enable hoisting: D148001. It would probably be more general to extend the logic in LICM to handle those cases, which would also trigger if the user wrote the expression like reassociate creates at the moment.

So effectively in certain circumstances LICM should undo what Reassociate have done?

[pawosm01]

In D151616#4381099, @fhahn wrote:

There's been some work to reassoicate in LICM to enable hoisting: D148001. It would probably be more general to extend the logic in LICM to handle those cases, which would also trigger if the user wrote the expression like reassociate creates at the moment.

You were right about the user written code like that: we're doing much worse than GCC here! Yet still, addressing it in LICM and not blocking it in Reassociate will create one more place for competing optimization, and LLVM is already haunted by those.

[qcolombet]

llvm/lib/Transforms/Scalar/Reassociate.cpp

1705

Although it is less performance sensitive, pulling adds in a loop could also be problematic.

1709

Early exit to reduce indentation

1719

Shouldn't PHIs with the same loop as L return true here?

1742

TL;DR I think the model should be more general than that but as it is, it is probably a good proxy for not shooting ourselves in the foot.

The problem is not pulling something into a loop, the problem is replacing X something with 1 * frequency-of-the-target-basic-block something.

I.e., let's imagine we have:
v1 = s1 * delta
v2 = s2 * delta
...
vN = sN *delta
for () {

v1 + ... + vN // with loop dependent stuff

}

Transforming this into:
for () {

(s1 + ... + sN) * delta

}

Could actually be beneficial as long as the loop count is smaller than N.

Similarly, something like:
for () {

if (unlikely)
 v1 + ... + vN

}

Is even more likely to be beneficial since the expected frequency of the basic block may be smaller than the source basic block.

[pawosm01]

llvm/lib/Transforms/Scalar/Reassociate.cpp

1742

Transforming this into: Could actually be beneficial as long as the loop count is smaller than N.

Yes, I was afraid that with this change I'm cutting off some of the beneficial transformations, and making the logic of this pass more complicated than necessary.

Therefore I tend to agree with Florian, that this problem could be addressed in LICM instead. In the end, it's not a wrong thing for a transformation pass to undo whatever previous pass(es) did when there is a new knowledge available (and here, the new knowledge is that we are in a loop). I'm very close to giving up this commit and work on extending hoistArithmetics() (of LICM) with the ability to spot this code pattern. As I mentioned in the other comment, currently LICM can't figure it out even if it wasn't caused by Reassociate pass, namely a simple C code that I specifically crafted to replicate this problem was enough to show that GCC produces faster executable than LLVM. But there's a problem with LICM too. This hoistArithmetics() function and its surrounding is getting more and more complicated already. And I assume any attempt to extend it with ability to deal with the situation discussed here, will create a lot of code bloat. But maybe that's something inevitable.

[pawosm01]

Overriden by D152281

llvm/include/llvm/Transforms/Scalar/Reassociate.h

 class Function;
 class Instruction;
 class IRBuilderBase;
+class LoopInfo;
 class Value;
 
 /// A private "module" namespace for types and utilities used by Reassociate.

 public:
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &);
 
+  // Glue for old PM.
+  bool runImpl(Function &F, LoopInfo &LI);
+
 private:
   void BuildRankMap(Function &F, ReversePostOrderTraversal<Function *> &RPOT);
   unsigned getRank(Value *V);
   void canonicalizeOperands(Instruction *I);
-  void ReassociateExpression(BinaryOperator *I);
+  void ReassociateExpression(BinaryOperator *I, LoopInfo &LI);
   void RewriteExprTree(BinaryOperator *I,
                        SmallVectorImpl<reassociate::ValueEntry> &Ops);
   Value *OptimizeExpression(BinaryOperator *I,
-                            SmallVectorImpl<reassociate::ValueEntry> &Ops);
+                            SmallVectorImpl<reassociate::ValueEntry> &Ops,
+                            LoopInfo &LI);
   Value *OptimizeAdd(Instruction *I,
-                     SmallVectorImpl<reassociate::ValueEntry> &Ops);
+                     SmallVectorImpl<reassociate::ValueEntry> &Ops,
+                     LoopInfo &LI);
   Value *OptimizeXor(Instruction *I,
                      SmallVectorImpl<reassociate::ValueEntry> &Ops);
   bool CombineXorOpnd(Instruction *I, reassociate::XorOpnd *Opnd1,

   Value *RemoveFactorFromExpression(Value *V, Value *Factor);
   void EraseInst(Instruction *I);
   void RecursivelyEraseDeadInsts(Instruction *I, OrderedSet &Insts);
-  void OptimizeInst(Instruction *I);
+  void OptimizeInst(Instruction *I, LoopInfo &LI);
   Instruction *canonicalizeNegFPConstantsForOp(Instruction *I, Instruction *Op,
                                                Value *OtherOp);
   Instruction *canonicalizeNegFPConstants(Instruction *I);

llvm/lib/Transforms/Scalar/Reassociate.cpp

 #include "llvm/Analysis/BasicAliasAnalysis.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/GlobalsModRef.h"
+#include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/Argument.h"
 #include "llvm/IR/BasicBlock.h"

 #include "llvm/Transforms/Utils/Local.h"
 #include <algorithm>
 #include <cassert>
+#include <functional>
 #include <utility>
 
 using namespace llvm;

 /// optimizes based on identities.  If it can be reduced to a single Value, it
 /// is returned, otherwise the Ops list is mutated as necessary.
 Value *ReassociatePass::OptimizeAdd(Instruction *I,
-                                    SmallVectorImpl<ValueEntry> &Ops) {
+                                    SmallVectorImpl<ValueEntry> &Ops,
+                                    LoopInfo &LI) {
   // Scan the operand lists looking for X and -X pairs.  If we find any, we
   // can simplify expressions like X+-X == 0 and X+~X ==-1.  While we're at it,
   // scan for any

                       << '\n');
     ++NumFactor;
 
+    // This transformation can potentially move a loop invariant computation
+    // into the loop. This should be avoided as it increases the number of
+    // computations and has potential to prevent other optimizations that would
+    // need to know how to revert this transformation in order to succeed.
+
+    // Go through the operations and see if all the operands belong to the same
+    // loop.
+    for (unsigned i = 0, e = Ops.size(); i != e; ++i) {
+      // Only consider expressions we're allowed to transform.
+      BinaryOperator *BOp =
+          isReassociableOp(Ops[i].Op, Instruction::Mul, Instruction::FMul);
+      if (!BOp)
+        continue;
+
+      if (Loop *L = LI.getLoopFor(BOp->getParent())) {
+        // Make sure we're not pulling operands from the outside of the loop.
+        for (Value *Operand : BOp->operands()) {
+          std::function<bool(LoopInfo &, Value *, Loop *)> isInOneLoop =
+              [&isInOneLoop](LoopInfo &LI, Value *V, Loop *L) -> bool {
+            Instruction *I = dyn_cast<Instruction>(V);
+
+            if (!I)
+              return true;
+            if ((LI.getLoopFor(I->getParent()) != L) || (dyn_cast<PHINode>(I)))
+              return false;
+            for (Value *Operand : I->operands()) {
+              if (!isInOneLoop(LI, Operand, L))
+                return false;
+            }
+            return true;
+          };
+          std::function<bool(Value *, Value *)> hasFactor =
+              [&hasFactor](Value *V, Value *Factor) -> bool {
+            BinaryOperator *BO =
+                isReassociableOp(V, Instruction::Mul, Instruction::FMul);
+
+            if (!BO)
+              return V == Factor;
+            for (Value *Operand : BO->operands()) {
+              if (hasFactor(Operand, Factor))
+                return true;
+            }
+            return false;
+          };
+
+          // Only the operands using the factor are being transformed.
+          if ((hasFactor(Operand, MaxOccVal)) && (!isInOneLoop(LI, Operand, L)))
+            return nullptr; // Transformation not profitable, giving up.
+        }
+      }
+    }
+
     // Create a new instruction that uses the MaxOccVal twice.  If we don't do
     // this, we could otherwise run into situations where removing a factor
     // from an expression will drop a use of maxocc, and this can cause

 }
 
 Value *ReassociatePass::OptimizeExpression(BinaryOperator *I,
-                                           SmallVectorImpl<ValueEntry> &Ops) {
+                                           SmallVectorImpl<ValueEntry> &Ops,
+                                           LoopInfo &LI) {
   // Now that we have the linearized expression tree, try to optimize it.
   // Start by folding any constants that we found.
   const DataLayout &DL = I->getModule()->getDataLayout();

 
   case Instruction::Add:
   case Instruction::FAdd:
-    if (Value *Result = OptimizeAdd(I, Ops))
+    if (Value *Result = OptimizeAdd(I, Ops, LI))
       return Result;
     break;
 

   }
 
   if (Ops.size() != NumOps)
-    return OptimizeExpression(I, Ops);
+    return OptimizeExpression(I, Ops, LI);
   return nullptr;
 }
 

 
 /// Inspect and optimize the given instruction. Note that erasing
 /// instructions is not allowed.
-void ReassociatePass::OptimizeInst(Instruction *I) {
+void ReassociatePass::OptimizeInst(Instruction *I, LoopInfo &LI) {
   // Only consider operations that we understand.
   if (!isa<UnaryOperator>(I) && !isa<BinaryOperator>(I))
     return;

       cast<Instruction>(BO->user_back())->getOpcode() == Instruction::FSub)
     return;
 
-  ReassociateExpression(BO);
+  ReassociateExpression(BO, LI);
 }
 
-void ReassociatePass::ReassociateExpression(BinaryOperator *I) {
+void ReassociatePass::ReassociateExpression(BinaryOperator *I, LoopInfo &LI) {
   // First, walk the expression tree, linearizing the tree, collecting the
   // operand information.
   SmallVector<RepeatedValue, 8> Tree;

 
   // Now that we have the expression tree in a convenient
   // sorted form, optimize it globally if possible.
-  if (Value *V = OptimizeExpression(I, Ops)) {
+  if (Value *V = OptimizeExpression(I, Ops, LI)) {
     if (V == I)
       // Self-referential expression in unreachable code.
       return;

   }
 }
 
-PreservedAnalyses ReassociatePass::run(Function &F, FunctionAnalysisManager &) {
+bool ReassociatePass::runImpl(Function &F, LoopInfo &LI) {
   // Get the functions basic blocks in Reverse Post Order. This order is used by
   // BuildRankMap to pre calculate ranks correctly. It also excludes dead basic
   // blocks (it has been seen that the analysis in this pass could hang when

       if (isInstructionTriviallyDead(&*II)) {
         EraseInst(&*II++);
       } else {
-        OptimizeInst(&*II);
+        OptimizeInst(&*II, LI);
         assert(II->getParent() == &*BI && "Moved to a different block!");
         ++II;
       }

       if (isInstructionTriviallyDead(I))
         EraseInst(I);
       else
-        OptimizeInst(I);
+        OptimizeInst(I, LI);
     }
   }
 

   for (auto &Entry : PairMap)
     Entry.clear();
 
-  if (MadeChange) {
-    PreservedAnalyses PA;
-    PA.preserveSet<CFGAnalyses>();
-    return PA;
-  }
+  return MadeChange;
+}
+
+PreservedAnalyses ReassociatePass::run(Function &F,
+                                       FunctionAnalysisManager &AM) {
+  auto &LI = AM.getResult<LoopAnalysis>(F);
 
-  return PreservedAnalyses::all();
+  if (!runImpl(F, LI))
+    return PreservedAnalyses::all();
+
+  PreservedAnalyses PA;
+  PA.preserveSet<CFGAnalyses>();
+  return PA;
 }
 
 namespace {

     bool runOnFunction(Function &F) override {
       if (skipFunction(F))
         return false;
-
-      FunctionAnalysisManager DummyFAM;
-      auto PA = Impl.run(F, DummyFAM);
-      return !PA.areAllPreserved();
+      return Impl.runImpl(F, getAnalysis<LoopInfoWrapperPass>().getLoopInfo());
     }
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {

       AU.addPreserved<AAResultsWrapperPass>();
       AU.addPreserved<BasicAAWrapperPass>();
       AU.addPreserved<GlobalsAAWrapperPass>();
+      AU.addPreserved<LoopInfoWrapperPass>();
     }
   };
 

llvm/test/Other/new-pm-defaults.ll

 ; CHECK-O23SZ-NEXT: Running pass: TailCallElimPass
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass
 ; CHECK-O-NEXT: Running pass: ReassociatePass
-; CHECK-O-NEXT: Running pass: LoopSimplifyPass
 ; CHECK-O-NEXT: Running analysis: LoopAnalysis
+; CHECK-O-NEXT: Running pass: LoopSimplifyPass
 ; CHECK-O-NEXT: Running pass: LCSSAPass
 ; CHECK-O-NEXT: Running analysis: ScalarEvolutionAnalysis
 ; CHECK-O-NEXT: Running analysis: InnerAnalysisManagerProxy

llvm/test/Other/new-pm-thinlto-postlink-defaults.ll

 ; CHECK-O23SZ-NEXT: Running pass: TailCallElimPass
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass
 ; CHECK-O-NEXT: Running pass: ReassociatePass
-; CHECK-O-NEXT: Running pass: LoopSimplifyPass
 ; CHECK-O-NEXT: Running analysis: LoopAnalysis
+; CHECK-O-NEXT: Running pass: LoopSimplifyPass
 ; CHECK-O-NEXT: Running pass: LCSSAPass
 ; CHECK-O-NEXT: Running analysis: ScalarEvolutionAnalysis
 ; CHECK-O-NEXT: Running analysis: InnerAnalysisManagerProxy

llvm/test/Other/new-pm-thinlto-prelink-defaults.ll

 ; CHECK-O23SZ-NEXT: Running pass: TailCallElimPass
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass
 ; CHECK-O-NEXT: Running pass: ReassociatePass
-; CHECK-O-NEXT: Running pass: LoopSimplifyPass
 ; CHECK-O-NEXT: Running analysis: LoopAnalysis
+; CHECK-O-NEXT: Running pass: LoopSimplifyPass
 ; CHECK-O-NEXT: Running pass: LCSSAPass
 ; CHECK-O-NEXT: Running analysis: ScalarEvolutionAnalysis
 ; CHECK-O-NEXT: Running analysis: InnerAnalysisManagerProxy

llvm/test/Transforms/Reassociate/reassociate-not-from-the-outside-of-the-loop.ll

 ; computed before the loop anymore. In case of this test, it would add an extra
 ; multiplication into the loop.
 
-; FIXME: the checks below need to be inverted to confirm the change to the
-; Reassociate pass.
-
 define void @innermost_loop(i32 %i, double %d1, double %d2, double %delta, ptr %cells) {
 ; CHECK-LABEL: @innermost_loop(
 entry:
 ; CHECK-LABEL: entry:
   %mul = fmul fast double %d1, %delta
+; CHECK: %{{.*}} = fmul {{.*}} %delta
   %mul1 = fmul fast double %d2, %delta
-; CHECK-NOT: %{{.*}} = fmul {{.*}} %delta
+; CHECK: %{{.*}} = fmul {{.*}} %delta
   br label %for.cond
 
 for.cond:

   %1 = load double, ptr %arrayidx4
   %mul5 = fmul fast double %mul1, %1
   %add6 = fadd fast double %mul2, %mul5
-; CHECK: %reass{{.*}} = fadd
-; CHECK-NEXT: %reass{{.*}} = fmul {{.*}} %delta
+; CHECK-NOT: %reass{{.*}} = fadd
+; CHECK-NOT: %reass{{.*}} = fmul
   store double %add6, ptr %arrayidx4
   br label %for.cond