[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 ↗

(On Diff #305667)

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 ↗

(On Diff #305667)

comesBefore is used in DT->dominates...

1569 ↗

(On Diff #305667)

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 ↗

(On Diff #305667)

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

[mkazantsev]

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

1569 ↗

(On Diff #305667)

It's a good idea.

1573 ↗

(On Diff #305667)

NarrowDef will be deleted in the end in any case.

[anna]

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

1564 ↗

(On Diff #305667)

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 ↗

(On Diff #305667)

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 ↗

(On Diff #306317)

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/Scalar/IndVarSimplify.cpp

bool WidenIV::widenWithVariantUse(NarrowIVDefUse DU) {
   // other operand.
   assert((NarrowUse->getOperand(0) == NarrowDef ||
           NarrowUse->getOperand(1) == NarrowDef) &&
          "bad DU");
 
   const OverflowingBinaryOperator *OBO =
     cast<OverflowingBinaryOperator>(NarrowUse);
   ExtendKind ExtKind = getExtendKind(NarrowDef);
+  const Instruction *Context = nullptr;
+  // Check that all uses are either s/zext, or narrow def (in case of we are
+  // widening the IV increment).
+  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;
+    if (!Context)
+      // Pick first user as context.
+      Context = User;
+    else if (DT->dominates(User, Context))
+      Context = User;
+    else if (!DT->dominates(Context, User))
+      // For users that don't have dominance relation, use common dominator.
+      Context =
+          DT->findNearestCommonDominator(Context->getParent(),
+                                         User->getParent())->getTerminator();
+  }
+
   bool CanSignExtend = ExtKind == SignExtended && OBO->hasNoSignedWrap();
   bool CanZeroExtend = ExtKind == ZeroExtended && OBO->hasNoUnsignedWrap();
   auto AnotherOpExtKind = ExtKind;
   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));
     auto GEPred = ICmpInst::ICMP_UGE;
     if (!SE->isKnownNegative(RHS))
       return false;
-    bool ProvedSubNUW = SE->isKnownPredicateAt(
-        GEPred, LHS, SE->getNegativeSCEV(RHS), NarrowUse);
+    bool ProvedSubNUW =
+        SE->isKnownPredicateAt(GEPred, 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
   const SCEV *Op1 = SE->getSCEV(WideDef);
   const SCEVAddRecExpr *AddRecOp1 = dyn_cast<SCEVAddRecExpr>(Op1);
   if (!AddRecOp1 || AddRecOp1->getLoop() != L)
     return false;
 
-  // Check that all uses are either s/zext, or narrow def (in case of we are
-  // widening the IV increment).
-  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;
-  }
-
   LLVM_DEBUG(dbgs() << "Cloning arithmetic IVUser: " << *NarrowUse << "\n");
 
   // Generating a widening use instruction.
   Value *LHS = (NarrowUse->getOperand(0) == NarrowDef)
                    ? WideDef
                    : createExtendInst(NarrowUse->getOperand(0), WideType,
                                       AnotherOpExtKind, NarrowUse);
   Value *RHS = (NarrowUse->getOperand(1) == NarrowDef)

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

;
   %ext = sext i32 %iv to i64
   call void @consume.i64(i64 %ext)
   br i1 %be.cond, label %loop, label %leave
 
   leave:
   ret i32 22
 }
 
-; TODO: We don't really need trunc/zext here because when iv.next overflows,
-; 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
 ;