[LICM] Reassociate & hoist sub expressions

[mkazantsev]

Summary

LICM could reassociate mixed variant/invariant comparison/arithmetic operations
and hoist invariant parts out of loop if it can prove that they can be computed
without overflow. Motivating example here:

INV1 - VAR1 < INV2

can be turned into

VAR > INV1 - INV2

if we can prove no-signed-overflow here. Then INV1 - INV2 can be computed
out of loop, so we save one arithmetic operation in-loop.

I tried to make it as general as possible in this patch. For some reason it only
works when no-overflow facts are known from non-contextual facts, but this
is maybe a limitation of SCEV. It can be a subject of further improvement in
SCEV.

[nikic]

I'll check, but I'm pretty sure LICM shouldn't use SCEV because it is part of LPM1, not LPM2.

[nikic]

Could you please rebase this patch to current main?

[mkazantsev]

planned rebase

[mkazantsev]

Wow, I'm so lucky to make this change in this exact place in this exact moment :)

[mkazantsev]

llvm/lib/Transforms/Scalar/LICM.cpp

2759

Should have used eraseInstruction

[mkazantsev]

Rebased

[mkazantsev]

Ping

[mkazantsev]

Ping

[anna]

llvm/lib/Transforms/Scalar/LICM.cpp

2727

This gets confusing enough to read for me .. will take a while to parse this :)

Can I suggest to start with the minimal pattern you wanted to support? I see you have couple of swaps earlier as well to handle all general forms.

[mkazantsev]

@anna I have split off the minimul piece of functionality as D149132. Will rebase this one now.

[mkazantsev]

Rebased on top of split off piece. I also tried to simplify the explanation and math formulae.

[mkazantsev]

llvm/lib/Transforms/Scalar/LICM.cpp

2727

I tried to do it more reader-friendly. :)

[skatkov]

General question: don't you want to separate sub and add handling in completely different functions.
To be honest, there will be small code duplication but in this case no need to mixup add and sub logic.
It will be easier to review and support in the future.

If duplication is a problem you can create functions for sub and add to keep first and last lines are in common.

[anna]

In D148001#4295213, @skatkov wrote:

General question: don't you want to separate sub and add handling in completely different functions.
To be honest, there will be small code duplication but in this case no need to mixup add and sub logic.
It will be easier to review and support in the future.

If duplication is a problem you can create functions for sub and add to keep first and last lines are in common.

Strongly agree with this :)

[mkazantsev]

In D148001#4295213, @skatkov wrote:

General question: don't you want to separate sub and add handling in completely different functions.
To be honest, there will be small code duplication but in this case no need to mixup add and sub logic.
It will be easier to review and support in the future.

If duplication is a problem you can create functions for sub and add to keep first and last lines are in common.

But code for sub is exactly same as the code that handles both add and sub, with 2 lines of difference.

[mkazantsev]

This will require rebase.

[skatkov]

In D148001#4297996, @mkazantsev wrote:

In D148001#4295213, @skatkov wrote:

General question: don't you want to separate sub and add handling in completely different functions.
To be honest, there will be small code duplication but in this case no need to mixup add and sub logic.
It will be easier to review and support in the future.

If duplication is a problem you can create functions for sub and add to keep first and last lines are in common.

But code for sub is exactly same as the code that handles both add and sub, with 2 lines of difference.

ok, will review after rebase.

[mkazantsev]

Rebased & split off sub code from add as Serguei and Anna have proposed.

[mkazantsev]

Rolling back to SCEV version.

[mkazantsev]

Bringing SCEV to LICM brings potential correctness issues. I'll rather try to support this in value tracking.

[mkazantsev]

Ping

[mkazantsev]

Ping

[skatkov]

LGTM with comment for consideration - I do not insist.

llvm/lib/Transforms/Scalar/LICM.cpp

2610

Not sure it is more readable (may be only for me) but consider

// C1 - LV < C2 case, so we need C1 - C2 does not overflow.
if (VariantSubtracted && computeOverflowForSignedSub(InvariantOp, InvariantRHS, DL, AC,
                                         &ICmp, DT) != llvm::OverflowResult::NeverOverflows)
  return false;
// LV - C1 < C2 case, so we need C1 + C2 does not overflow.
if (!ValueSubrtacted && computeOverflowForSignedAdd(InvariantOp, InvariantRHS, DL, AC,
                                         &ICmp, DT) != llvm::OverflowResult::NeverOverflows)
  return false;

llvm/lib/Transforms/Scalar/LICM.cpp

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;
 }
 
+/// Try to reassociate and hoist the following two patterns:
+/// LV - C1 < C2 --> LV < C1 + C2,
+/// C1 - LV < C2 --> LV > C1 - C2.
+static bool hoistSub(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_NSWSub(m_Value(VariantOp), m_Value(InvariantOp))))
+    return false;
+
+  bool VariantSubtracted = 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
+  // the variant operand goes with minus, we use a slightly different scheme.
+  if (L.isLoopInvariant(VariantOp)) {
+    std::swap(VariantOp, InvariantOp);
+    VariantSubtracted = true;
+    Pred = ICmpInst::getSwappedPredicate(Pred);
+  }
+  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. Likewise, for "C1 - LV < C2" we need to prove that
+  // "C1 - C2" does not overflow.
+  auto &DL = L.getHeader()->getModule()->getDataLayout();
+  if (VariantSubtracted) {
+    // C1 - LV < C2 --> LV > C1 - C2
+    if (computeOverflowForSignedSub(InvariantOp, InvariantRHS, DL, AC, &ICmp,
+                                    DT) != llvm::OverflowResult::NeverOverflows)
+      return false;
+  } else {
+    // LV - C1 < C2 --> LV < C1 + C2
+    if (computeOverflowForSignedAdd(InvariantOp, InvariantRHS, DL, AC, &ICmp,
+                                    DT) != llvm::OverflowResult::NeverOverflows)
+      return false;
+  }
+  auto *Preheader = L.getLoopPreheader();
+  assert(Preheader && "Loop is not in simplify form?");
+  IRBuilder<> Builder(Preheader->getTerminator());
+  Value *NewCmpOp =
+      VariantSubtracted
+          ? Builder.CreateSub(InvariantOp, InvariantRHS, "invariant.op",
+                              /*HasNUW*/ false, /*HasNSW*/ true)
+          : Builder.CreateAdd(InvariantOp, InvariantRHS, "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))))
static bool hoistAddSub(Instruction &I, Loop &L, ICFLoopSafetyInfo &SafetyInfo,
   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.
+  if (hoistSub(Pred, LHS, RHS, cast<ICmpInst>(I), L, SafetyInfo, MSSAU, AC, DT))
+    return true;
 
   return false;
 }
 
 static bool hoistArithmetics(Instruction &I, Loop &L,
                              ICFLoopSafetyInfo &SafetyInfo,
                              MemorySSAUpdater &MSSAU, AssumptionCache *AC,
                              DominatorTree *DT) {

 }
 
 /// 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

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 2
 ; RUN: opt -S -passes=licm -verify-memoryssa < %s | FileCheck %s
 
-; TODO: x - iv < 4 ==> iv > x - 4
+; x - iv < 4 ==> iv > x - 4
 define i32 @test_01(ptr %p, ptr %x_p, ptr %length_p) {
 ; CHECK-LABEL: define i32 @test_01
 ; 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:![0-9]+]]
 ; CHECK-NEXT:    [[LENGTH:%.*]] = load i32, ptr [[LENGTH_P]], align 4, !range [[RNG0]]
+; CHECK-NEXT:    [[INVARIANT_OP:%.*]] = sub nsw i32 [[X]], 4
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ]
-; CHECK-NEXT:    [[ARITH:%.*]] = sub nsw i32 [[X]], [[IV]]
-; CHECK-NEXT:    [[X_CHECK:%.*]] = icmp slt i32 [[ARITH]], 4
+; CHECK-NEXT:    [[X_CHECK:%.*]] = icmp sgt 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_03(ptr %p, ptr %x_p, ptr %length_p) {
 ; CHECK-LABEL: define i32 @test_03
 ; CHECK-SAME: (ptr [[P:%.*]], ptr [[X_P:%.*]], ptr [[LENGTH_P:%.*]]) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[X:%.*]] = load i32, ptr [[X_P]], align 4, !range [[RNG1:![0-9]+]]
 ; CHECK-NEXT:    [[LENGTH:%.*]] = load i32, ptr [[LENGTH_P]], align 4, !range [[RNG0]]
+; CHECK-NEXT:    [[INVARIANT_OP:%.*]] = add nsw i32 [[X]], 4
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[BACKEDGE:%.*]] ]
-; CHECK-NEXT:    [[ARITH:%.*]] = sub 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: