[IndVars] Use more precise context when eliminating narrowing

[mkazantsev]

Summary

When deciding to widen narrow use, we may need to prove some facts
about it. For proof, the context is used. Currently we take the instruction
being widened as the context.

However, we may be more precise here if we take as context the point that
dominates all users of instruction being widened.

[mkazantsev]

Rebased.

[anna]

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

1564

dominates is a more expensive check than comesBefore. We can special case for uses in the same basic block, i.e. :

if (User->getParent() == Context->getParent() && User->comesBefore(Context))
    Context = User

[skatkov]

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

1564

comesBefore is used in DT->dominates...

1569

What do you think about introducing the method in DT which finds nearestCommonDominator for an array of values?
Here, just to collect all users and trigger the utility method?

Alternatively, introduce findNearestCommonDominator which accepts Instructions or values in DomTree?
And use just it instead of this dispatching.

1573

If all users are NarrowDef, should we mark it as dead?

[mkazantsev]

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

1569

It's a good idea.

1573

NarrowDef will be deleted in the end in any case.

[anna]

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

1564

yes I know :), but there are a bunch of other checks in DT-dominates() which are first checked that are not needed for the same basic block case. Maybe it doesn't make a difference...

[mkazantsev]

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

1564

But why do you think that the users will be in the same block? In either case, DT checks cost is nothing compared to SCEV creation cost.

[mkazantsev]

Rebased, with dom tree logic factored out into separate patch (added dependency).

[skatkov]

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

1546

Can you land this code movement as NFC in a separate patch and update this one?

[mkazantsev]

Done

[mkazantsev]

Roll back to older code concept (without generatization into DomTree). It seems unreasonably difficult to check in code there.

[skatkov]

It is pity that we did not introduce an useful utility function.

Still lgtm.

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

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;
 
   // Check that all uses are either s/zext, or narrow def (in case of we are
   // widening the IV increment).
   SmallVector<Instruction *, 4> ExtUsers;
   for (Use &U : NarrowUse->uses()) {
     if (U.getUser() == NarrowDef)
       continue;
     Instruction *User = nullptr;
     if (ExtKind == SignExtended)
       User = dyn_cast<SExtInst>(U.getUser());
     else
       User = dyn_cast<ZExtInst>(U.getUser());
     if (!User || User->getType() != WideType)
       return false;
     ExtUsers.push_back(User);
   }
   if (ExtUsers.empty()) {
     DeadInsts.emplace_back(NarrowUse);
     return true;
   }
 
+  // We'll prove some facts that should be true in the context of ext users. If
+  // there is no users, we are done now. If there are some, pick their common
+  // dominator as context.
+  Instruction *Context = nullptr;
+  for (auto *Ext : ExtUsers) {
+    if (!Context || DT->dominates(Ext, Context))
+      Context = Ext;
+    else if (!DT->dominates(Context, Ext))
+      // For users that don't have dominance relation, use common dominator.
+      Context =
+          DT->findNearestCommonDominator(Context->getParent(), Ext->getParent())
+              ->getTerminator();
+  }
+  assert(Context && "Context not found?");
+
   if (!CanSignExtend && !CanZeroExtend) {
     // Because InstCombine turns 'sub nuw' to 'add' losing the no-wrap flag, we
     // will most likely not see it. Let's try to prove it.
     if (OpCode != Instruction::Add)
       return false;
     if (ExtKind != ZeroExtended)
       return false;
     const SCEV *LHS = SE->getSCEV(OBO->getOperand(0));
     const SCEV *RHS = SE->getSCEV(OBO->getOperand(1));
     if (!SE->isKnownNegative(RHS))
       return false;
     bool ProvedSubNUW = SE->isKnownPredicateAt(
-        ICmpInst::ICMP_UGE, LHS, SE->getNegativeSCEV(RHS), NarrowUse);
+        ICmpInst::ICMP_UGE, LHS, SE->getNegativeSCEV(RHS), Context);
     if (!ProvedSubNUW)
       return false;
     // In fact, our 'add' is 'sub nuw'. We will need to widen the 2nd operand as
     // neg(zext(neg(op))), which is basically sext(op).
     AnotherOpExtKind = SignExtended;
   }
 
   // Verifying that Defining operand is an AddRec

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

 ; its value is not used.
 define i32 @test11(i32 %start, i32* %p, i32* %q) {
 ; CHECK-LABEL: @test11(
 ; 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:    [[TMP1:%.*]] = trunc i64 [[INDVARS_IV]] to i32
-; CHECK-NEXT:    [[IV_NEXT:%.*]] = add i32 [[TMP1]], -1
+; CHECK-NEXT:    [[TMP1:%.*]] = add nsw i64 [[INDVARS_IV]], -1
 ; CHECK-NEXT:    [[COND:%.*]] = icmp eq i64 [[INDVARS_IV]], 0
 ; CHECK-NEXT:    br i1 [[COND]], label [[EXIT:%.*]], label [[BACKEDGE]]
 ; CHECK:       backedge:
-; CHECK-NEXT:    [[INDEX:%.*]] = zext i32 [[IV_NEXT]] 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 0
 ; CHECK:       failure:
 ; CHECK-NEXT:    unreachable
 ;