[IndVars] Recognize 'sub nuw' expressed as 'add' for widening

[mkazantsev]

Summary

InstCombine canonicalizes 'sub nuw' instructions to 'add' without the
nuw flag. The typical case where we see it is decrementing induction
variables. For them, IndVars fails to prove that it's legal to widen them,
and inserts unprofitable zext's.

This patch adds recognition of such pattern using SCEV.

[mkazantsev]

Ping.

[fhahn]

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1158

This seems independent of the add x, -y -> sub nuw? Can this be split off into a separate patch? IIUC the test case for the detection without this change would be something like

loop:
  %iv = phi i32 [%start, %entry], [%iv.next.1, %backedge]
  %cond = icmp eq i32 %iv, 0
  br i1 %cond, label %exit, label %backedge

backedge:
  %foo = add i32 %iv, -1
  %index = zext i32 %foo to i64
  %store.addr = getelementptr i32, i32* %p, i64 %index
  store i32 1, i32* %store.addr
  %load.addr = getelementptr i32, i32* %q, i64 %index
  %stop = load i32, i32* %q
  %loop.cond = icmp eq i32 %stop, 0
  %iv.next.1 = add i32 %iv, -1
  br i1 %loop.cond, label %loop, label %failure

[mkazantsev]

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

1158

Good point, thanks for catching! I'll split it out.

[mkazantsev]

Split out widen increment logic into separate change.

[skatkov]

lgtm

[fhahn]

Thanks for the update! LGTM as well.

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

1552 ↗

(On Diff #303343)

tiny nit: perhaps move just before use? Or inline.

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

bool WidenIV::widenWithVariantUse(NarrowIVDefUse DU) {
   assert(DU.NarrowUse->getOperand(1 - ExtendOperIdx) == DU.NarrowDef &&
          "bad DU");
 
   const OverflowingBinaryOperator *OBO =
     cast<OverflowingBinaryOperator>(NarrowUse);
   ExtendKind ExtKind = getExtendKind(NarrowDef);
   bool CanSignExtend = ExtKind == SignExtended && OBO->hasNoSignedWrap();
   bool CanZeroExtend = ExtKind == ZeroExtended && OBO->hasNoUnsignedWrap();
-  if (!CanSignExtend && !CanZeroExtend)
-    return false;
+  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);
+    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).
   if (ExtKind == SignExtended) {
     for (Use &U : NarrowUse->uses()) {
+      if (U.getUser() == NarrowDef)
+        continue;
       SExtInst *User = dyn_cast<SExtInst>(U.getUser());
       if (!User || User->getType() != WideType)
         return false;
     }
   } else { // ExtKind == ZeroExtended
     for (Use &U : NarrowUse->uses()) {
+      if (U.getUser() == NarrowDef)
+        continue;
       ZExtInst *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,
-                                      ExtKind, NarrowUse);
+                                      AnotherOpExtKind, NarrowUse);
   Value *RHS = (NarrowUse->getOperand(1) == NarrowDef)
                    ? WideDef
                    : createExtendInst(NarrowUse->getOperand(1), WideType,
-                                      ExtKind, NarrowUse);
+                                      AnotherOpExtKind, NarrowUse);
 
   auto *NarrowBO = cast<BinaryOperator>(NarrowUse);
   auto *WideBO = BinaryOperator::Create(NarrowBO->getOpcode(), LHS, RHS,
                                         NarrowBO->getName());
   IRBuilder<> Builder(NarrowUse);
   Builder.Insert(WideBO);
   WideBO->copyIRFlags(NarrowBO);
   ExtendKindMap[NarrowUse] = ExtKind;
 
   for (Use &U : NarrowUse->uses()) {
+    // Ignore narrow def: it will be removed after the transform.
+    if (U.getUser() == NarrowDef)
+      continue;
     Instruction *User = nullptr;
     if (ExtKind == SignExtended)
       User = cast<SExtInst>(U.getUser());
     else
       User = cast<ZExtInst>(U.getUser());
     assert(User->getType() == WideType && "Checked before!");
     LLVM_DEBUG(dbgs() << "INDVARS: eliminating " << *User << " replaced by "
                       << *WideBO << "\n");

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

 ; 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:    br i1 [[COND]], label [[EXIT:%.*]], label [[BACKEDGE]]
 ; CHECK:       backedge:
-; CHECK-NEXT:    [[TMP1:%.*]] = trunc i64 [[INDVARS_IV]] to i32
-; CHECK-NEXT:    [[IV_NEXT:%.*]] = add i32 [[TMP1]], -1
-; CHECK-NEXT:    [[INDEX:%.*]] = zext i32 [[IV_NEXT]] to i64
-; CHECK-NEXT:    [[STORE_ADDR:%.*]] = getelementptr i32, i32* [[P:%.*]], i64 [[INDEX]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add nsw i64 [[INDVARS_IV]], -1
+; 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
 ;