[IndVars] LCSSA Phi users should not prevent widening

[mkazantsev]

Summary

When widening an IndVar that has LCSSA Phi users outside
the loop, we can safely widen it as usual and then truncate
the result outside the loop without hurting the performance.

[skatkov]

lgtm

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

1646

I doubt that someone can change the content of this set after creation :)

1649

".wide"?

[mkazantsev]

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

1646

Just in case if this code changes later.

[mkazantsev]

Rebased, fixed '.wide name.

[yrouban]

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

1654

must be

Builder.SetInsertPoint(&*User->getParent()->getFirstInsertionPt());

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

bool WidenIV::widenWithVariantUse(WidenIV::NarrowIVDefUse DU) {
   const OverflowingBinaryOperator *OBO =
     cast<OverflowingBinaryOperator>(NarrowUse);
   ExtendKind ExtKind = getExtendKind(NarrowDef);
   bool CanSignExtend = ExtKind == SignExtended && OBO->hasNoSignedWrap();
   bool CanZeroExtend = ExtKind == ZeroExtended && OBO->hasNoUnsignedWrap();
   auto AnotherOpExtKind = ExtKind;
 
   // The facts we are going to prove need to be true at the point that
-  // dominates all users. Find this point while checking the users.
+  // dominates all ext users. Find this point while checking the users.
   const Instruction *Context = nullptr;
-  // Check that all uses are either s/zext, or narrow def (in case of we are
-  // widening the IV increment).
+  // Check that all uses are either:
+  // - s/zext that we plan to get rid of;
+  // - narrow def (in case of we are widening the IV increment);
+  // - LCSSA Phis outside the loop.
   SmallVector<Instruction *, 4> ExtUsers;
+  SmallVector<PHINode *, 4> LCSSAPhiUsers;
   for (Use &U : NarrowUse->uses()) {
-    if (U.getUser() == NarrowDef)
+    Instruction *User = dyn_cast<Instruction>(U.getUser());
+    if (User == NarrowDef)
       continue;
-    Instruction *User = nullptr;
+    if (!L->contains(User)) {
+      auto *LCSSAPhi = cast<PHINode>(User);
+      // Make sure there is only 1 input, so that we don't have to split
+      // critical edges.
+      if (LCSSAPhi->getNumOperands() != 1)
+        return false;
+      LCSSAPhiUsers.push_back(LCSSAPhi);
+      continue;
+    }
     if (ExtKind == SignExtended)
-      User = dyn_cast<SExtInst>(U.getUser());
+      User = dyn_cast<SExtInst>(User);
     else
-      User = dyn_cast<ZExtInst>(U.getUser());
+      User = dyn_cast<ZExtInst>(User);
     if (!User || User->getType() != WideType)
       return false;
     ExtUsers.push_back(User);
     if (!Context)
        // Pick first user as context.
        Context = User;
      else if (DT->dominates(User, Context))
        Context = User;
bool WidenIV::widenWithVariantUse(WidenIV::NarrowIVDefUse DU) {
   for (Instruction *User : ExtUsers) {
     assert(User->getType() == WideType && "Checked before!");
     LLVM_DEBUG(dbgs() << "INDVARS: eliminating " << *User << " replaced by "
                       << *WideBO << "\n");
     ++NumElimExt;
     User->replaceAllUsesWith(WideBO);
     DeadInsts.emplace_back(User);
   }
+
+  for (PHINode *User : LCSSAPhiUsers) {
+    assert(User->getNumOperands() == 1 && "Checked before!");
+    Builder.SetInsertPoint(User);
+    auto *WidePN =
+        Builder.CreatePHI(WideBO->getType(), 1, User->getName() + "wide");
+    BasicBlock *LoopExitingBlock = User->getParent()->getSinglePredecessor();
+    assert(LoopExitingBlock && L->contains(LoopExitingBlock) &&
+           "Not a LCSSA Phi?");
+    WidePN->addIncoming(WideBO, LoopExitingBlock);
+    Builder.SetInsertPoint(User->getParent()->getFirstNonPHI());
+    auto *TruncPN = Builder.CreateTrunc(WidePN, User->getType());
+    User->replaceAllUsesWith(TruncPN);
+    DeadInsts.emplace_back(User);
+  }
   return true;
 }
 
 /// Determine whether an individual user of the narrow IV can be widened. If so,
 /// return the wide clone of the user.
 Instruction *WidenIV::widenIVUse(WidenIV::NarrowIVDefUse DU, SCEVExpander &Rewriter) {
   assert(ExtendKindMap.count(DU.NarrowDef) &&
          "Should already know the kind of extension used to widen NarrowDef");

llvm/test/Transforms/IndVarSimplify/widen-loop-comp.ll

 define i32 @test14(i32 %start, i32* %p, i32* %q) {
 ; CHECK-LABEL: @test14(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = zext i32 [[START:%.*]] to i64
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[INDVARS_IV_NEXT:%.*]], [[BACKEDGE:%.*]] ], [ [[TMP0]], [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[COND:%.*]] = icmp eq i64 [[INDVARS_IV]], 0
-; CHECK-NEXT:    [[TMP1:%.*]] = trunc i64 [[INDVARS_IV]] to i32
-; CHECK-NEXT:    [[FOO:%.*]] = add i32 [[TMP1]], -1
+; CHECK-NEXT:    [[TMP1:%.*]] = add nsw i64 [[INDVARS_IV]], -1
 ; CHECK-NEXT:    br i1 [[COND]], label [[EXIT:%.*]], label [[BACKEDGE]]
 ; CHECK:       backedge:
-; CHECK-NEXT:    [[INDEX:%.*]] = zext i32 [[FOO]] to i64
-; CHECK-NEXT:    [[STORE_ADDR:%.*]] = getelementptr i32, i32* [[P:%.*]], i64 [[INDEX]]
+; CHECK-NEXT:    [[STORE_ADDR:%.*]] = getelementptr i32, i32* [[P:%.*]], i64 [[TMP1]]
 ; CHECK-NEXT:    store i32 1, i32* [[STORE_ADDR]], align 4
-; CHECK-NEXT:    [[LOAD_ADDR:%.*]] = getelementptr i32, i32* [[Q:%.*]], i64 [[INDEX]]
-; CHECK-NEXT:    [[STOP:%.*]] = load i32, i32* [[Q]], align 4
+; CHECK-NEXT:    [[STOP:%.*]] = load i32, i32* [[Q:%.*]], align 4
 ; CHECK-NEXT:    [[LOOP_COND:%.*]] = icmp eq i32 [[STOP]], 0
 ; CHECK-NEXT:    [[INDVARS_IV_NEXT]] = add nsw i64 [[INDVARS_IV]], -1
 ; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[LOOP]], label [[FAILURE:%.*]]
 ; CHECK:       exit:
-; CHECK-NEXT:    ret i32 -1
+; CHECK-NEXT:    [[TMP2:%.*]] = trunc i64 -1 to i32
+; CHECK-NEXT:    ret i32 [[TMP2]]
 ; CHECK:       failure:
 ; CHECK-NEXT:    unreachable
 ;
 entry:
   br label %loop
 
 loop:
   %iv = phi i32 [%start, %entry], [%iv.next.1, %backedge]

 define i32 @test15(i32 %start, i32* %p, i32* %q) {
 ; CHECK-LABEL: @test15(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = zext i32 [[START:%.*]] to i64
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[INDVARS_IV_NEXT:%.*]], [[BACKEDGE:%.*]] ], [ [[TMP0]], [[ENTRY:%.*]] ]
 ; CHECK-NEXT:    [[COND:%.*]] = icmp eq i64 [[INDVARS_IV]], 0
-; CHECK-NEXT:    [[TMP1:%.*]] = trunc i64 [[INDVARS_IV]] to i32
-; CHECK-NEXT:    [[FOO:%.*]] = add i32 [[TMP1]], -1
+; CHECK-NEXT:    [[TMP1:%.*]] = add nsw i64 [[INDVARS_IV]], -1
 ; CHECK-NEXT:    br i1 [[COND]], label [[EXIT:%.*]], label [[BACKEDGE]]
 ; CHECK:       backedge:
-; CHECK-NEXT:    [[INDEX:%.*]] = zext i32 [[FOO]] to i64
-; CHECK-NEXT:    [[STORE_ADDR:%.*]] = getelementptr i32, i32* [[P:%.*]], i64 [[INDEX]]
+; CHECK-NEXT:    [[STORE_ADDR:%.*]] = getelementptr i32, i32* [[P:%.*]], i64 [[TMP1]]
 ; CHECK-NEXT:    store i32 1, i32* [[STORE_ADDR]], align 4
-; CHECK-NEXT:    [[LOAD_ADDR:%.*]] = getelementptr i32, i32* [[Q:%.*]], i64 [[INDEX]]
-; CHECK-NEXT:    [[STOP:%.*]] = load i32, i32* [[Q]], align 4
+; CHECK-NEXT:    [[STOP:%.*]] = load i32, i32* [[Q:%.*]], align 4
 ; CHECK-NEXT:    [[LOOP_COND:%.*]] = icmp eq i32 [[STOP]], 0
 ; CHECK-NEXT:    [[INDVARS_IV_NEXT]] = add nsw i64 [[INDVARS_IV]], -1
 ; CHECK-NEXT:    br i1 [[LOOP_COND]], label [[LOOP]], label [[FAILURE:%.*]]
 ; CHECK:       exit:
-; CHECK-NEXT:    call void @use(i32 -1)
-; CHECK-NEXT:    ret i32 -1
+; CHECK-NEXT:    [[TMP2:%.*]] = trunc i64 -1 to i32
+; CHECK-NEXT:    call void @use(i32 [[TMP2]])
+; CHECK-NEXT:    ret i32 [[TMP2]]
 ; CHECK:       failure:
-; CHECK-NEXT:    [[FOO_LCSSA1:%.*]] = phi i32 [ [[FOO]], [[BACKEDGE]] ]
-; CHECK-NEXT:    call void @use(i32 [[FOO_LCSSA1]])
+; CHECK-NEXT:    [[FOO_LCSSA1WIDE:%.*]] = phi i64 [ [[TMP1]], [[BACKEDGE]] ]
+; CHECK-NEXT:    [[TMP3:%.*]] = trunc i64 [[FOO_LCSSA1WIDE]] to i32
+; CHECK-NEXT:    call void @use(i32 [[TMP3]])
 ; CHECK-NEXT:    unreachable
 ;
 entry:
   br label %loop
 
 loop:
   %iv = phi i32 [%start, %entry], [%iv.next.1, %backedge]
   %cond = icmp eq i32 %iv, 0