[LICM] Reassociate & hoist add expressions

[mkazantsev]

Summary

This patch allows LICM to reassociate and hoist following expressions:

loop:
  %sum = add nsw %iv, %C1
  %cmp = icmp <signed pred> %sum, C2

where C1 and C2 are loop invariants. The reassociated version looks like

preheader:
  %inv_sum = C2 - C1
...
loop:
  %cmp = icmp <signed pred> %iv, %inv_sum

In order to prove legality, we need both initial addition and the newly created subtraction
to happen without overflow.

[skatkov]

llvm/lib/Transforms/Scalar/LICM.cpp

2627

Any comments?
Why LHS invariant is not interested to you? Covered in other place? Assert instead of if?

Is hasOneUse a cost model? We want to remove original add?

2638

Comment?

2646

Some words about, why it is legal if both add and sub are nsw?

[mkazantsev]

llvm/lib/Transforms/Scalar/LICM.cpp

2627

Yes, we want to remove the original add.

[mkazantsev]

Added even more comments.

[anna]

llvm/lib/Transforms/Scalar/LICM.cpp

2639

Nit: LHS itself is an invariant.

2653

Don't we need to check for isa<SCEVCouldNotCompute>(RHSS) here and below for InvariantOp before getting NegativeSCEV?

[nikic]

Proof: https://alive2.llvm.org/ce/z/imgdNt

Somewhat skeptical about introducing SCEV use in LICM. As you are actually reasoning about loop invariants here, do we need SCEV? Would using computeOverflowForSignedSub() from ValueTracking work, possibly even provide better results?

llvm/lib/Transforms/Scalar/LICM.cpp

2654–2657

This seems more obvious and matches the generated sub?

[skatkov]

llvm/lib/Transforms/Scalar/LICM.cpp

2651

Will SCEV more powerful to check that LHS is NSW then just check of flag on instruction?

[mkazantsev]

llvm/lib/Transforms/Scalar/LICM.cpp

2654–2657

Yes, and it provides worse results on the tests that are here.

[mkazantsev]

llvm/lib/Transforms/Scalar/LICM.cpp

2653

No, value's SCEV is never SCEVCouldNotCompute. You can only get SCEVCouldNotCompute when computing some loop exit count, for example.

[mkazantsev]

llvm/lib/Transforms/Scalar/LICM.cpp

2651

In theory, yes.

[mkazantsev]

In D149132#4296348, @nikic wrote:

Proof: https://alive2.llvm.org/ce/z/imgdNt

Somewhat skeptical about introducing SCEV use in LICM. As you are actually reasoning about loop invariants here, do we need SCEV? Would using computeOverflowForSignedSub() from ValueTracking work, possibly even provide better results?

Let me check if it achieves the same result.

[mkazantsev]

Reworked w/o SCEV, using ValueTracking's API as Nikita has proposed.

llvm/lib/Transforms/Scalar/LICM.cpp

2639

I must have been crazy when I wrote this. It's not just a bunch of typos, but it's also "variant" instead of "invariant" :)

2651

Added one more overflow check, leaving nsw as shortcut.

2654–2657

No longer actual after API change.

[nikic]

llvm/lib/Transforms/Scalar/LICM.cpp

2561

If it is

2574–2575

InstCombine will do exactly this inference already, so I think we should only be checking the nsw flag here.

[mkazantsev]

Typo fixed, check reduced to nsw check.

[nikic]

llvm/test/Transforms/LICM/hoist-add-sub.ll

394

TODO incorrectly dropped here. But I think maybe the test isn't right and %x is supposed to use range !1 not !0? Same for the pair of tests above.

444

Drop TODO.

[mkazantsev]

Err, indeed I messed up with test TODOs.

llvm/test/Transforms/LICM/hoist-add-sub.ll

394

No, range !0 is good enough here. The original add doesn't overflow because of nsw, and the new sub doesn't overflow because it's a subtraction of two non-negatives.

[mkazantsev]

This comment was removed by mkazantsev.

[mkazantsev]

Fixed bug (typo) that caused messup with checks in tests. Bug was

if (ProvedNoOverflowAfterReassociate)
  return false;

instead of

if (!ProvedNoOverflowAfterReassociate)
  return false;

Friday evening refactoring is evil.

[mkazantsev]

Some NFC reafactoring due to request to untwine sub and add logic.

[mkazantsev]

NFC. Used m_NSWAdd.

[mkazantsev]

I'm planning to roll back to SCEV version. Value Tracking doesn't infer things from dominating conditions.

[mkazantsev]

Ok, getting SCEV here brings potential correctness issues... Let's keep Value tracking.

[mkazantsev]

D149648 shows that SCEV is more powerful than value tracking (the latter cannot figure out the range from dominating checks).

[mkazantsev]

Ping

[mkazantsev]

Ping

[skatkov]

LGTM, @nikic, any remaining concerns?

llvm/lib/Transforms/Scalar/LICM.cpp

 STATISTIC(NumMovedCalls, "Number of call insts hoisted or sunk");
 STATISTIC(NumPromotionCandidates, "Number of promotion candidates");
 STATISTIC(NumLoadPromoted, "Number of load-only promotions");
 STATISTIC(NumLoadStorePromoted, "Number of load and store promotions");
 STATISTIC(NumMinMaxHoisted,
           "Number of min/max expressions hoisted out of the loop");
 STATISTIC(NumGEPsHoisted,
           "Number of geps reassociated and hoisted out of the loop");
+STATISTIC(NumAddSubHoisted, "Number of add/subtract expressions reassociated "
+                            "and hoisted out of the loop");
 
 /// Memory promotion is enabled by default.
 static cl::opt<bool>
     DisablePromotion("disable-licm-promotion", cl::Hidden, cl::init(false),
                      cl::desc("Disable memory promotion in LICM pass"));
 
 static cl::opt<bool> ControlFlowHoisting(
     "licm-control-flow-hoisting", cl::Hidden, cl::init(false),
Value *NewGEP = Builder.CreateGEP(Src->getSourceElementType(), NewSrc,
                                     SmallVector<Value *>(Src->indices()), "gep",
                                     IsInBounds);
   GEP->replaceAllUsesWith(NewGEP);
   eraseInstruction(*GEP, SafetyInfo, MSSAU);
   eraseInstruction(*Src, SafetyInfo, MSSAU);
   return true;
 }
 
+/// Try to turn things like "LV + C1 < C2" into "LV < C2 - C1". Here
+/// C1 and C2 are loop invariants and LV is a loop-variant.
+static bool hoistAdd(ICmpInst::Predicate Pred, Value *VariantLHS,
+                     Value *InvariantRHS, ICmpInst &ICmp, Loop &L,
+                     ICFLoopSafetyInfo &SafetyInfo, MemorySSAUpdater &MSSAU,
+                     AssumptionCache *AC, DominatorTree *DT) {
+  assert(ICmpInst::isSigned(Pred) && "Not supported yet!");
+  assert(!L.isLoopInvariant(VariantLHS) && "Precondition.");
+  assert(L.isLoopInvariant(InvariantRHS) && "Precondition.");
+
+  // Try to represent VariantLHS as sum of invariant and variant operands.
+  using namespace PatternMatch;
+  Value *VariantOp, *InvariantOp;
+  if (!match(VariantLHS, m_NSWAdd(m_Value(VariantOp), m_Value(InvariantOp))))
+    return false;
+
+  // LHS itself is a loop-variant, try to represent it in the form:
+  // "VariantOp + InvariantOp". If it is possible, then we can reassociate.
+  if (L.isLoopInvariant(VariantOp))
+    std::swap(VariantOp, InvariantOp);
+  if (L.isLoopInvariant(VariantOp) || !L.isLoopInvariant(InvariantOp))
+    return false;
+
+  // In order to turn "LV + C1 < C2" into "LV < C2 - C1", we need to be able to
+  // freely move values from left side of inequality to right side (just as in
+  // normal linear arithmetics). Overflows make things much more complicated, so
+  // we want to avoid this.
+  auto &DL = L.getHeader()->getModule()->getDataLayout();
+  bool ProvedNoOverflowAfterReassociate =
+      computeOverflowForSignedSub(InvariantRHS, InvariantOp, DL, AC, &ICmp,
+                                  DT) == llvm::OverflowResult::NeverOverflows;
+  if (!ProvedNoOverflowAfterReassociate)
+    return false;
+  auto *Preheader = L.getLoopPreheader();
+  assert(Preheader && "Loop is not in simplify form?");
+  IRBuilder<> Builder(Preheader->getTerminator());
+  Value *NewCmpOp = Builder.CreateSub(InvariantRHS, InvariantOp, "invariant.op",
+                                      /*HasNUW*/ false, /*HasNSW*/ true);
+  ICmp.setPredicate(Pred);
+  ICmp.setOperand(0, VariantOp);
+  ICmp.setOperand(1, NewCmpOp);
+  eraseInstruction(cast<Instruction>(*VariantLHS), SafetyInfo, MSSAU);
+  return true;
+}
+
+/// Reassociate and hoist add/sub expressions.
+static bool hoistAddSub(Instruction &I, Loop &L, ICFLoopSafetyInfo &SafetyInfo,
+                        MemorySSAUpdater &MSSAU, AssumptionCache *AC,
+                        DominatorTree *DT) {
+  using namespace PatternMatch;
+  ICmpInst::Predicate Pred;
+  Value *LHS, *RHS;
+  if (!match(&I, m_ICmp(Pred, m_Value(LHS), m_Value(RHS))))
+    return false;
+
+  // TODO: Support unsigned predicates?
+  if (!ICmpInst::isSigned(Pred))
+    return false;
+
+  // Put variant operand to LHS position.
+  if (L.isLoopInvariant(LHS)) {
+    std::swap(LHS, RHS);
+    Pred = ICmpInst::getSwappedPredicate(Pred);
+  }
+  // We want to delete the initial operation after reassociation, so only do it
+  // if it has no other uses.
+  if (L.isLoopInvariant(LHS) || !L.isLoopInvariant(RHS) || !LHS->hasOneUse())
+    return false;
+
+  // TODO: We could go with smarter context, taking common dominator of all I's
+  // users instead of I itself.
+  if (hoistAdd(Pred, LHS, RHS, cast<ICmpInst>(I), L, SafetyInfo, MSSAU, AC, DT))
+    return true;
+
+  // TODO: Support Sub.
+
+  return false;
+}
+
 static bool hoistArithmetics(Instruction &I, Loop &L,
                              ICFLoopSafetyInfo &SafetyInfo,
-                             MemorySSAUpdater &MSSAU,
-                             AssumptionCache *AC, DominatorTree *DT) {
+                             MemorySSAUpdater &MSSAU, AssumptionCache *AC,
+                             DominatorTree *DT) {
   // Optimize complex patterns, such as (x < INV1 && x < INV2), turning them
   // into (x < min(INV1, INV2)), and hoisting the invariant part of this
   // expression out of the loop.
   if (hoistMinMax(I, L, SafetyInfo, MSSAU)) {
     ++NumHoisted;
     ++NumMinMaxHoisted;
     return true;
   }
 
   // Try to hoist GEPs by reassociation.
   if (hoistGEP(I, L, SafetyInfo, MSSAU, AC, DT)) {
     ++NumHoisted;
     ++NumGEPsHoisted;
     return true;
   }
 
+  // Try to hoist add/sub's by reassociation.
+  if (hoistAddSub(I, L, SafetyInfo, MSSAU, AC, DT)) {
+    ++NumHoisted;
+    ++NumAddSubHoisted;
+    return true;
+  }
+
   return false;
 }
 
 /// Little predicate that returns true if the specified basic block is in
 /// a subloop of the current one, not the current one itself.
 ///
 static bool inSubLoop(BasicBlock *BB, Loop *CurLoop, LoopInfo *LI) {
   assert(CurLoop->contains(BB) && "Only valid if BB is IN the loop");
   return LI->getLoopFor(BB) != CurLoop;
 }

llvm/test/Transforms/LICM/hoist-add-sub.ll

 exit:
   ret i32 %iv.next
 
 out_of_bounds:
   ret i32 -1
 }
 
 
-; TODO: x + iv < 4 ==> iv < 4 - x
+; x + iv < 4 ==> iv < 4 - x
 define i32 @test_02(ptr %p, ptr %x_p, ptr %length_p) {
 ; CHECK-LABEL: define i32 @test_02
 ; CHECK-SAME: (ptr [[P:%.*]], ptr [[X_P:%.*]], ptr [[LENGTH_P:%.*]]) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[X:%.*]] = load i32, ptr [[X_P]], align 4, !range [[RNG0]]
 ; CHECK-NEXT:    [[LENGTH:%.*]] = load i32, ptr [[LENGTH_P]], align 4, !range [[RNG0]]
+; CHECK-NEXT:    [[INVARIANT_OP:%.*]] = sub nsw i32 4, [[X]]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ]
-; CHECK-NEXT:    [[ARITH:%.*]] = add nsw i32 [[X]], [[IV]]
-; CHECK-NEXT:    [[X_CHECK:%.*]] = icmp slt i32 [[ARITH]], 4
+; CHECK-NEXT:    [[X_CHECK:%.*]] = icmp slt i32 [[IV]], [[INVARIANT_OP]]
 ; CHECK-NEXT:    br i1 [[X_CHECK]], label [[OUT_OF_BOUNDS:%.*]], label [[BACKEDGE]]
 ; CHECK:       backedge:
 ; CHECK-NEXT:    [[EL_PTR:%.*]] = getelementptr i32, ptr [[P]], i32 [[IV]]
 ; CHECK-NEXT:    store i32 1, ptr [[EL_PTR]], align 4
 ; CHECK-NEXT:    [[IV_NEXT]] = add nuw nsw i32 [[IV]], 4
 ; CHECK-NEXT:    [[LOOP_COND:%.*]] = icmp slt i32 [[IV_NEXT]], [[LENGTH]]
 ; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[LOOP]], label [[EXIT:%.*]]
 ; CHECK:       exit:

 
 out_of_bounds:
   ret i32 -1
 
 failed:
   ret i32 -2
 }
 
-; TODO: iv + x < 4 ==> iv < 4 - x
+; iv + x < 4 ==> iv < 4 - x
 define i32 @test_04(ptr %p, ptr %x_p, ptr %length_p) {
 ; CHECK-LABEL: define i32 @test_04
 ; CHECK-SAME: (ptr [[P:%.*]], ptr [[X_P:%.*]], ptr [[LENGTH_P:%.*]]) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[X:%.*]] = load i32, ptr [[X_P]], align 4, !range [[RNG0]]
 ; CHECK-NEXT:    [[LENGTH:%.*]] = load i32, ptr [[LENGTH_P]], align 4, !range [[RNG0]]
+; CHECK-NEXT:    [[INVARIANT_OP:%.*]] = sub nsw i32 4, [[X]]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ]
-; CHECK-NEXT:    [[ARITH:%.*]] = add nsw i32 [[IV]], [[X]]
-; CHECK-NEXT:    [[X_CHECK:%.*]] = icmp slt i32 [[ARITH]], 4
+; CHECK-NEXT:    [[X_CHECK:%.*]] = icmp slt i32 [[IV]], [[INVARIANT_OP]]
 ; CHECK-NEXT:    br i1 [[X_CHECK]], label [[OUT_OF_BOUNDS:%.*]], label [[BACKEDGE]]
 ; CHECK:       backedge:
 ; CHECK-NEXT:    [[EL_PTR:%.*]] = getelementptr i32, ptr [[P]], i32 [[IV]]
 ; CHECK-NEXT:    store i32 1, ptr [[EL_PTR]], align 4
 ; CHECK-NEXT:    [[IV_NEXT]] = add nuw nsw i32 [[IV]], 4
 ; CHECK-NEXT:    [[LOOP_COND:%.*]] = icmp slt i32 [[IV_NEXT]], [[LENGTH]]
 ; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[LOOP]], label [[EXIT:%.*]]
 ; CHECK:       exit:

 
 exit:
   ret i32 %iv.next
 
 out_of_bounds:
   ret i32 -1
 }
 
 ; TODO: iv + x < 4 ==> iv < 4 - x
 define i32 @test_04a(ptr %p, ptr %x_p, ptr %length_p) {
 ; CHECK-LABEL: define i32 @test_04a
 ; CHECK-SAME: (ptr [[P:%.*]], ptr [[X_P:%.*]], ptr [[LENGTH_P:%.*]]) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[X:%.*]] = load i32, ptr [[X_P]], align 4
 ; CHECK-NEXT:    [[LENGTH:%.*]] = load i32, ptr [[LENGTH_P]], align 4
 ; CHECK-NEXT:    [[PRECOND_1:%.*]] = icmp sge i32 [[X]], 0
 ; CHECK-NEXT:    [[PRECOND_2:%.*]] = icmp sge i32 [[LENGTH]], 0