[IndVarSimplify] Rotate exit comparisons to make them more invariant

[reames]

Summary

(If anyone has better terminology for this than "rotation", please throw it out.)

If we have a loop exit test which involves a loop varying binary op compared to a loop invariant value, see if we can invert that binary op on both sides of the comparison. This doesn't directly remove any instructions, but does make more of the computation loop invariant. This both helps directly for long running loops, and also simplifies the loop structure such that other transforms may run.

[reames]

Hm, it looks like there may be an overlap w/LFTR. With a slightly less reduced test case, we end up being able to prove nsw/nuw on several of the adds if we allow CVP to do so. We may be better off fixing https://bugs.llvm.org/show_bug.cgi?id=31181, and then extending LFTR slightly to catch the remaining cases.

I'd appreciate review in case I missed something, but I think I'm going to delay landing this until I've dug into that option a bit.

[reames]

In D61922#1502199, @reames wrote:

Hm, it looks like there may be an overlap w/LFTR. With a slightly less reduced test case, we end up being able to prove nsw/nuw on several of the adds if we allow CVP to do so. We may be better off fixing https://bugs.llvm.org/show_bug.cgi?id=31181, and then extending LFTR slightly to catch the remaining cases.

I'd appreciate review in case I missed something, but I think I'm going to delay landing this until I've dug into that option a bit.

There's a lot less overlap then there first looked to be. My test cases turned out to be over-reduced in another manner. :) This approach has the advantage of a) not requiring a statically computable loop exit count and b) working on multiple exit loops. LFTR could probably be extended to handle (b), but (a) is quite useful on it's own. I think having both is quite useful.

[reames]

No longer actively pursuing. The idea is reasonable, and only partly overlaps with LFTR, but all of the non-reduced motivating cases are covered by multiple exit LFTR.

If anyone wants to pick this up, feel free. There's probably some commonality with the existing simplifyIVUsers code, particularly the truncate path. Might be worth reframing in those terms. Or maybe I'll get back to this someday.

lib/Transforms/Scalar/IndVarSimplify.cpp

   bool handleFloatingPointIV(Loop *L, PHINode *PH);
   bool rewriteNonIntegerIVs(Loop *L);
 
   bool simplifyAndExtend(Loop *L, SCEVExpander &Rewriter, LoopInfo *LI);
 
   bool canLoopBeDeleted(Loop *L, SmallVector<RewritePhi, 8> &RewritePhiSet);
   bool rewriteLoopExitValues(Loop *L, SCEVExpander &Rewriter);
   bool rewriteFirstIterationLoopExitValues(Loop *L);
+  bool rotateLoopExitTests(Loop *L);
   bool hasHardUserWithinLoop(const Loop *L, const Instruction *I) const;
 
   bool linearFunctionTestReplace(Loop *L, const SCEV *BackedgeTakenCount,
                                  PHINode *IndVar, SCEVExpander &Rewriter);
 
   bool sinkUnusedInvariants(Loop *L);
 
 public:

   }
 
   // The insertion point instruction may have been deleted; clear it out
   // so that the rewriter doesn't trip over it later.
   Rewriter.clearInsertPoint();
   return Changed;
 }
 
+bool IndVarSimplify::rotateLoopExitTests(Loop *L) {
+  bool Changed = false;
+  SmallVector<BasicBlock*, 16> ExitingBlocks;
+  L->getExitingBlocks(ExitingBlocks);
+  for (BasicBlock *ExitingBB : ExitingBlocks) {
+    auto *BI = dyn_cast<BranchInst>(ExitingBB->getTerminator());
+    if (!BI || !BI->isConditional())
+      continue;
+    
+    auto *ICmp =  dyn_cast<ICmpInst>(BI->getCondition());
+    if (!ICmp || !L->contains(ICmp->getParent()) ||
+        !ICmp->hasOneUse())
+      continue;
+
+    // If not canonical form, then ignore
+    auto *LHS = dyn_cast<Instruction>(ICmp->getOperand(0));
+    auto *RHS = ICmp->getOperand(1);
+    if (!LHS || !isa<BinaryOperator>(LHS) || !LHS->hasOneUse() ||
+        !L->isLoopInvariant(RHS))
+      continue;
+
+    // We're looking to remove the shl (multiply) on the LHS, by right shifting
+    // (dividing) each both sides such that the shl cancels.  To do so, we need
+    // to a) ensure there's no overflow in either computation, and b) ensure we
+    // use a correctly signed shift (divide).  We don't to worry about
+    // out-of-bounds or zero shift amounts.  Zero shifts are noops (to both
+    // sides), and out-of-bounds shifts simply move poison from one cmp operand
+    // to the other.
+    
+    // TODO: generalize for other binary operators (e.g. add, sub, div)
+    Value *ShOp0;
+    Value *ShOp1;
+    using namespace PatternMatch;
+    if ((ICmp->isSigned() &&
+         match(LHS, m_NSWShl(m_Value(ShOp0), m_Value(ShOp1)))) ||
+        (!ICmp->isSigned() &&
+         match(LHS, m_NUWShl(m_Value(ShOp0), m_Value(ShOp1)))))
+      if (L->isLoopInvariant(ShOp1)) {
+        auto NewOp = ICmp->isSigned() ? Instruction::AShr : Instruction::LShr;
+        auto *NewRHS = BinaryOperator::Create(NewOp, RHS, LHS->getOperand(1),
+                                              "", ICmp);
+        ICmp->setOperand(0, ShOp0);
+        ICmp->setOperand(1, NewRHS);
+        LHS->eraseFromParent();
+        Changed = true;
+      }    
+  }
+      
+  return Changed;
+}
+
 //===---------------------------------------------------------------------===//
 // rewriteFirstIterationLoopExitValues: Rewrite loop exit values if we know
 // they will exit at the first iteration.
 //===---------------------------------------------------------------------===//
 
 /// Check to see if this loop has loop invariant conditions which lead to loop
 /// exits. If so, we know that if the exit path is taken, it is at the first
 /// loop iteration. This lets us predict exit values of PHI nodes that live in

   // loop may be sunk below the loop to reduce register pressure.
   Changed |= sinkUnusedInvariants(L);
 
   // rewriteFirstIterationLoopExitValues does not rely on the computation of
   // trip count and therefore can further simplify exit values in addition to
   // rewriteLoopExitValues.
   Changed |= rewriteFirstIterationLoopExitValues(L);
 
+  // Try to rotate loop tests so as to make more of the computation feeding the
+  // test loop invariant.
+  Changed |= rotateLoopExitTests(L);
+
   // Clean up dead instructions.
   Changed |= DeleteDeadPHIs(L->getHeader(), TLI);
 
   // Check a post-condition.
   assert(L->isRecursivelyLCSSAForm(*DT, *LI) &&
          "Indvars did not preserve LCSSA!");
 
   // Verify that LFTR, and any other change have not interfered with SCEV's

test/Transforms/IndVarSimplify/rotate-exit-test.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -indvars -S < %s | FileCheck %s
+
+target datalayout = "e-m:e-p:32:32-i64:64-v128:64:128-a:0:32-n32-S64"
+
+@Sink = external global i32
+
+;; TODO: Should be able to infer overflow fact here
+define i32 @test_signed() {
+; CHECK-LABEL: @test_signed(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add nuw nsw i32 [[IV]], 1
+; CHECK-NEXT:    store volatile i32 [[IV]], i32* @Sink
+; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp ne i32 [[IV_NEXT]], 751
+; CHECK-NEXT:    br i1 [[EXITCOND]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret i32 750
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i32 [0, %entry], [%iv.next, %loop]
+  %iv.next = add i32 %iv, 1
+  store volatile i32 %iv, i32* @Sink
+  %shl = shl i32 %iv, 2
+  %cmp1 = icmp slt i32 %shl, 2999
+  br i1 %cmp1, label %loop, label %exit
+
+exit:
+  ret i32 %iv
+}
+
+;; TODO: Should be able to infer overflow fact here
+define i32 @test_unsigned() {
+; CHECK-LABEL: @test_unsigned(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add nuw nsw i32 [[IV]], 1
+; CHECK-NEXT:    store volatile i32 [[IV]], i32* @Sink
+; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp ne i32 [[IV_NEXT]], 751
+; CHECK-NEXT:    br i1 [[EXITCOND]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret i32 750
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i32 [0, %entry], [%iv.next, %loop]
+  %iv.next = add i32 %iv, 1
+  store volatile i32 %iv, i32* @Sink
+  %shl = shl i32 %iv, 2
+  %cmp1 = icmp ult i32 %shl, 2999
+  br i1 %cmp1, label %loop, label %exit
+
+exit:
+  ret i32 %iv
+}
+
+;; TODO: Should be able to infer overflow fact here
+define i32 @test_ne() {
+; CHECK-LABEL: @test_ne(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add nuw nsw i32 [[IV]], 1
+; CHECK-NEXT:    store volatile i32 [[IV]], i32* @Sink
+; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp ne i32 [[IV_NEXT]], 501
+; CHECK-NEXT:    br i1 [[EXITCOND]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret i32 500
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i32 [0, %entry], [%iv.next, %loop]
+  %iv.next = add i32 %iv, 1
+  store volatile i32 %iv, i32* @Sink
+  %shl = shl i32 %iv, 1
+  %cmp1 = icmp ne i32 %shl, 1000
+  br i1 %cmp1, label %loop, label %exit
+
+exit:
+  ret i32 %iv
+}
+
+; Negative test, unknown bounds
+define i32 @neg_unknown_bounds(i32 %N, i32 %shiftamt) {
+; CHECK-LABEL: @neg_unknown_bounds(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add i32 [[IV]], 1
+; CHECK-NEXT:    store volatile i32 [[IV]], i32* @Sink
+; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[IV]], [[SHIFTAMT:%.*]]
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp slt i32 [[SHL]], [[N:%.*]]
+; CHECK-NEXT:    br i1 [[CMP1]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[IV_LCSSA:%.*]] = phi i32 [ [[IV]], [[LOOP]] ]
+; CHECK-NEXT:    ret i32 [[IV_LCSSA]]
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i32 [0, %entry], [%iv.next, %loop]
+  %iv.next = add i32 %iv, 1
+  store volatile i32 %iv, i32* @Sink
+  %shl = shl i32 %iv, %shiftamt
+  %cmp1 = icmp slt i32 %shl, %N
+  br i1 %cmp1, label %loop, label %exit
+
+exit:
+  ret i32 %iv
+}
+
+define i32 @test_nsw(i32 %N, i32 %shiftamt) {
+; CHECK-LABEL: @test_nsw(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add i32 [[IV]], 1
+; CHECK-NEXT:    store volatile i32 [[IV]], i32* @Sink
+; CHECK-NEXT:    [[TMP0:%.*]] = ashr i32 [[N:%.*]], [[SHIFTAMT:%.*]]
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp slt i32 [[IV]], [[TMP0]]
+; CHECK-NEXT:    br i1 [[CMP1]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[IV_LCSSA:%.*]] = phi i32 [ [[IV]], [[LOOP]] ]
+; CHECK-NEXT:    ret i32 [[IV_LCSSA]]
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i32 [0, %entry], [%iv.next, %loop]
+  %iv.next = add i32 %iv, 1
+  store volatile i32 %iv, i32* @Sink
+  %shl = shl nsw i32 %iv, %shiftamt
+  %cmp1 = icmp slt i32 %shl, %N
+  br i1 %cmp1, label %loop, label %exit
+
+exit:
+  ret i32 %iv
+}
+
+define i32 @test_nuw(i32 %N, i32 %shiftamt) {
+; CHECK-LABEL: @test_nuw(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add i32 [[IV]], 1
+; CHECK-NEXT:    store volatile i32 [[IV]], i32* @Sink
+; CHECK-NEXT:    [[TMP0:%.*]] = lshr i32 [[N:%.*]], [[SHIFTAMT:%.*]]
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp ult i32 [[IV]], [[TMP0]]
+; CHECK-NEXT:    br i1 [[CMP1]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[IV_LCSSA:%.*]] = phi i32 [ [[IV]], [[LOOP]] ]
+; CHECK-NEXT:    ret i32 [[IV_LCSSA]]
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i32 [0, %entry], [%iv.next, %loop]
+  %iv.next = add i32 %iv, 1
+  store volatile i32 %iv, i32* @Sink
+  %shl = shl nuw i32 %iv, %shiftamt
+  %cmp1 = icmp ult i32 %shl, %N
+  br i1 %cmp1, label %loop, label %exit
+
+exit:
+  ret i32 %iv
+}
+
+define i32 @test_ne_nuw(i32 %N, i32 %shiftamt) {
+; CHECK-LABEL: @test_ne_nuw(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add i32 [[IV]], 1
+; CHECK-NEXT:    store volatile i32 [[IV]], i32* @Sink
+; CHECK-NEXT:    [[TMP0:%.*]] = lshr i32 [[N:%.*]], [[SHIFTAMT:%.*]]
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp ne i32 [[IV]], [[TMP0]]
+; CHECK-NEXT:    br i1 [[CMP1]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[IV_LCSSA:%.*]] = phi i32 [ [[IV]], [[LOOP]] ]
+; CHECK-NEXT:    ret i32 [[IV_LCSSA]]
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i32 [0, %entry], [%iv.next, %loop]
+  %iv.next = add i32 %iv, 1
+  store volatile i32 %iv, i32* @Sink
+  %shl = shl nuw i32 %iv, %shiftamt
+  %cmp1 = icmp ne i32 %shl, %N
+  br i1 %cmp1, label %loop, label %exit
+
+exit:
+  ret i32 %iv
+}
+
+define i32 @neg_ne_mayoverflow(i32 %N, i32 %shiftamt) {
+; CHECK-LABEL: @neg_ne_mayoverflow(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[IV:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[IV_NEXT:%.*]], [[LOOP]] ]
+; CHECK-NEXT:    [[IV_NEXT]] = add i32 [[IV]], 1
+; CHECK-NEXT:    store volatile i32 [[IV]], i32* @Sink
+; CHECK-NEXT:    [[SHL:%.*]] = shl i32 [[IV]], [[SHIFTAMT:%.*]]
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp ne i32 [[SHL]], [[N:%.*]]
+; CHECK-NEXT:    br i1 [[CMP1]], label [[LOOP]], label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    [[IV_LCSSA:%.*]] = phi i32 [ [[IV]], [[LOOP]] ]
+; CHECK-NEXT:    ret i32 [[IV_LCSSA]]
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i32 [0, %entry], [%iv.next, %loop]
+  %iv.next = add i32 %iv, 1
+  store volatile i32 %iv, i32* @Sink
+  %shl = shl i32 %iv, %shiftamt
+  %cmp1 = icmp ne i32 %shl, %N
+  br i1 %cmp1, label %loop, label %exit
+
+exit:
+  ret i32 %iv
+}