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

(On Diff #301209)

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 ↗

(On Diff #301209)

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

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

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

assert((NarrowUse->getOperand(0) == NarrowDef ||
           NarrowUse->getOperand(1) == 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));
+    if (!SE->isKnownNegative(RHS))
+      return false;
+    bool ProvedSubNUW = SE->isKnownPredicateAt(
+        ICmpInst::ICMP_UGE, 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;
 
   for (Use &U : NarrowUse->uses()) {
     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,
-                                      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;

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:    [[FOO:%.*]] = add i32 [[TMP1]], -1
-; CHECK-NEXT:    [[INDEX:%.*]] = zext i32 [[FOO]] 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
 ;