[Reassociate] Skip reassociation for IV if its def instruction contains NSW/NUW flags to enable better IndVar and LSR

[wmi]

Summary

Reassociation will remove nsw/nuw flags after change the order of computation. It may block IndVar from doing IV use widening and also block the optimization of LSR.

For %sub1 = %i.0 + %size - 1 in the following testcase, reassociation will convert it to %sub1 = %size - 1 + %i.0 and remove all the nsw flags. Then IndVar cannot do widening for %idxprom and cannot remove the sext. In addition, LSR cannot regard %arrayidx as an induction variable and it may produce suboptimal result.

for.body:

%add = add nsw i32 %i.0, %size
%sub1 = add nsw i32 %add, -1
%idxprom = sext i32 %sub1 to i64
%arrayidx = getelementptr inbounds [1000 x i32], [1000 x i32]* @maxarray, i64 0, i64 %idxprom
%tmp0 = load i32, i32* %arrayidx, align 4
%tmp1 = load i32, i32* @total, align 4
%add2 = add nsw i32 %tmp1, %tmp0
store i32 %add2, i32* @total, align 4
%inc = add nsw i32 %i.0, 1
br label %for.cond

The patch skips reassociation for IV use when it contains NSW/NUW flags so as to let IndVar and LSR do better jobs. Because LSR will also do reassociation work in its Formula canonicalization, we won't lose performance opportunity mostly.

[wmi]

Switch the order of reassociate and looprotate so reassociate will know more precisely whether the loop will be in simplified form and suitable for LSR. I expect the pass order change will not have noticeable impact.

[sanjoy]

This can potentially be a compile time issue. I think the SCEV dependency itself is fine, since SCEV will do things lazily, but computing the LoopInfo can be a problem.

Can this be solved by some changes to our pass order? I.e. first run indvars and then re-associate?

[mehdi_amini]

What about having two modes for reassociate: one where we don't do any transformation that drop nsw/nuw flag, and another where it is permitted. Depending on the position in the pipeline we allow or not to drop the flags.

include/llvm/Transforms/Scalar/Reassociate.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_TRANSFORMS_SCALAR_REASSOCIATE_H
 #define LLVM_TRANSFORMS_SCALAR_REASSOCIATE_H
 
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/SetVector.h"
+#include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 
 /// A private "module" namespace for types and utilities used by Reassociate.
 /// These are implementation details and should not be used by clients.
struct Factor {
   Factor(Value *Base, unsigned Power) : Base(Base), Power(Power) {}
 };
 
 class XorOpnd;
 }
 
 /// Reassociate commutative expressions.
 class ReassociatePass : public PassInfoMixin<ReassociatePass> {
+  LoopInfo *LI;
+  ScalarEvolution *SE;
   DenseMap<BasicBlock *, unsigned> RankMap;
   DenseMap<AssertingVH<Value>, unsigned> ValueRankMap;
   SetVector<AssertingVH<Instruction>> RedoInsts;
   bool MadeChange;
 
 public:
-  PreservedAnalyses run(Function &F, FunctionAnalysisManager &);
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+  PreservedAnalyses runImpl(Function &F, LoopInfo *LI_, ScalarEvolution *SE_);
 
 private:
   void BuildRankMap(Function &F);
   unsigned getRank(Value *V);
   void canonicalizeOperands(Instruction *I);
   void ReassociateExpression(BinaryOperator *I);
   void RewriteExprTree(BinaryOperator *I,
                        SmallVectorImpl<reassociate::ValueEntry> &Ops);

lib/Transforms/IPO/PassManagerBuilder.cpp

void PassManagerBuilder::addFunctionSimplificationPasses(
   MPM.add(createCorrelatedValuePropagationPass()); // Propagate conditionals
   MPM.add(createCFGSimplificationPass());     // Merge & remove BBs
   // Combine silly seq's
   addInstructionCombiningPass(MPM);
   addExtensionsToPM(EP_Peephole, MPM);
 
   MPM.add(createTailCallEliminationPass()); // Eliminate tail calls
   MPM.add(createCFGSimplificationPass());     // Merge & remove BBs
-  MPM.add(createReassociatePass());           // Reassociate expressions
   // Rotate Loop - disable header duplication at -Oz
   MPM.add(createLoopRotatePass(SizeLevel == 2 ? 0 : -1));
+  MPM.add(createReassociatePass());           // Reassociate expressions
   MPM.add(createLICMPass());                  // Hoist loop invariants
   MPM.add(createLoopUnswitchPass(SizeLevel || OptLevel < 3));
   MPM.add(createCFGSimplificationPass());
   addInstructionCombiningPass(MPM);
   MPM.add(createIndVarSimplifyPass());        // Canonicalize indvars
   MPM.add(createLoopIdiomPass());             // Recognize idioms like memset.
   MPM.add(createLoopDeletionPass());          // Delete dead loops
   if (EnableLoopInterchange) {

lib/Transforms/Scalar/Reassociate.cpp

 
 #include "llvm/Transforms/Scalar/Reassociate.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/GlobalsModRef.h"
+#include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"

 
 /// Inspect and optimize the given instruction. Note that erasing
 /// instructions is not allowed.
 void ReassociatePass::OptimizeInst(Instruction *I) {
   // Only consider operations that we understand.
   if (!isa<BinaryOperator>(I))
     return;
 
+  // If I contains NSW/NUW flag and it is potentially an IV use, don't do
+  // reassociation for it so that the NSW/NUW will be kept and IndVar may do
+  // better IV use Widening. LSR can also do reassociation for IV use later
+  // so we will not lose the benefit of reassociation.
+  Loop *Loop = LI->getLoopFor(I->getParent());
+  if (Loop && Loop->isLoopSimplifyForm() && isa<OverflowingBinaryOperator>(I) &&
+      (I->hasNoSignedWrap() || I->hasNoUnsignedWrap()) &&
+      SE->isSCEVable(I->getType()) && isa<SCEVAddRecExpr>(SE->getSCEV(I)))
+    return;
+
   if (I->getOpcode() == Instruction::Shl && isa<ConstantInt>(I->getOperand(1)))
     // If an operand of this shift is a reassociable multiply, or if the shift
     // is used by a reassociable multiply or add, turn into a multiply.
     if (isReassociableOp(I->getOperand(0), Instruction::Mul) ||
         (I->hasOneUse() &&
          (isReassociableOp(I->user_back(), Instruction::Mul) ||
           isReassociableOp(I->user_back(), Instruction::Add)))) {
       Instruction *NI = ConvertShiftToMul(I);
if (Ops.size() == 1) {
     return;
   }
 
   // Now that we ordered and optimized the expressions, splat them back into
   // the expression tree, removing any unneeded nodes.
   RewriteExprTree(I, Ops);
 }
 
-PreservedAnalyses ReassociatePass::run(Function &F, FunctionAnalysisManager &) {
+PreservedAnalyses ReassociatePass::run(Function &F,
+                                       FunctionAnalysisManager &AM) {
+  return runImpl(F, &AM.getResult<LoopAnalysis>(F),
+                 &AM.getResult<ScalarEvolutionAnalysis>(F));
+}
+
+PreservedAnalyses ReassociatePass::runImpl(Function &F, LoopInfo *LI_,
+                                           ScalarEvolution *SE_) {
+  LI = LI_;
+  SE = SE_;
+
   // Calculate the rank map for F.
   BuildRankMap(F);
 
   MadeChange = false;
   for (Function::iterator BI = F.begin(), BE = F.end(); BI != BE; ++BI) {
     // Optimize every instruction in the basic block.
     for (BasicBlock::iterator II = BI->begin(), IE = BI->end(); II != IE;)
       if (isInstructionTriviallyDead(&*II)) {
public:
     ReassociateLegacyPass() : FunctionPass(ID) {
       initializeReassociateLegacyPassPass(*PassRegistry::getPassRegistry());
     }
 
     bool runOnFunction(Function &F) override {
       if (skipFunction(F))
         return false;
 
-      FunctionAnalysisManager DummyFAM;
-      auto PA = Impl.run(F, DummyFAM);
+      auto PA =
+          Impl.runImpl(F, &getAnalysis<LoopInfoWrapperPass>().getLoopInfo(),
+                       &getAnalysis<ScalarEvolutionWrapperPass>().getSE());
       return !PA.areAllPreserved();
     }
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
+      AU.addRequired<LoopInfoWrapperPass>();
+      AU.addRequired<ScalarEvolutionWrapperPass>();
+
       AU.setPreservesCFG();
       AU.addPreserved<GlobalsAAWrapperPass>();
     }
   };
 }
 
 char ReassociateLegacyPass::ID = 0;
-INITIALIZE_PASS(ReassociateLegacyPass, "reassociate",
+INITIALIZE_PASS_BEGIN(ReassociateLegacyPass, "reassociate",
+                      "Reassociate expressions", false, false)
+INITIALIZE_PASS_DEPENDENCY(LoopPass)
+INITIALIZE_PASS_DEPENDENCY(ScalarEvolutionWrapperPass)
+INITIALIZE_PASS_END(ReassociateLegacyPass, "reassociate",
                     "Reassociate expressions", false, false)
 
 // Public interface to the Reassociate pass
 FunctionPass *llvm::createReassociatePass() {
   return new ReassociateLegacyPass();
 }

test/Transforms/Reassociate/ivuse.ll

+; RUN: opt < %s -S -reassociate | FileCheck %s
+; RUN: opt < %s -passes='reassociate' -S | FileCheck %s
+; Check %add and %sub1 are not reassociated so their nsw flags are kept. It is better for IndVar to do iv use widening.
+
+define void @foo(i32 %size) {
+entry:
+  br label %for.cond
+
+for.cond:
+  %i.0 = phi i32 [ 1, %entry ], [ %inc, %for.body ]
+  %sub = add nsw i32 %size, -1
+  %cmp = icmp slt i32 %i.0, %sub
+  br i1 %cmp, label %for.body, label %for.end
+
+; CHECK-LABEL: @foo
+; CHECK: for.body:
+; CHECK-NEXT: %add = add nsw i32 %i.0, %size
+; CHECK-NEXT: %sub1 = add nsw i32 %add, -1
+
+for.body:
+  %add = add nsw i32 %i.0, %size
+  %sub1 = add nsw i32 %add, -1
+  %idxprom = sext i32 %sub1 to i64
+  %arrayidx = getelementptr inbounds [1000 x i32], [1000 x i32]* @maxarray, i64 0, i64 %idxprom
+  %tmp0 = load i32, i32* %arrayidx, align 4
+  %tmp1 = load i32, i32* @total, align 4
+  %add2 = add nsw i32 %tmp1, %tmp0
+  store i32 %add2, i32* @total, align 4
+  %inc = add nsw i32 %i.0, 1
+  br label %for.cond
+
+for.end:
+  ret void
+}
+
+@maxarray = common local_unnamed_addr global [1000 x i32] zeroinitializer, align 16
+@total = common local_unnamed_addr global i32 0, align 4
+