[IndVars] Have `cloneArithmeticIVUser` guess better

[sanjoy]

Summary

cloneArithmeticIVUser currently trips over expression like `add %iv,
-1` when %iv is being zero extended -- it tries to construct the
widened use as add %iv.zext, zext(-1) and (correctly) fails to prove
equivalence to zext(add %iv, -1) (here the SCEV for %iv is
{1,+,1}).

This change teaches IndVars to try sign extending the non-IV operand
if that makes the newly constructed IV use equivalent to the widened
narrow IV use.

[reames]

lib/Transforms/Scalar/IndVarSimplify.cpp

992 ↗

(On Diff #37315)

Pulling out these variables is an NFC change. Can you separate and commit? It'll make the diff easier.

1001 ↗

(On Diff #37315)

This would be more clear as a lambda which captures SE

Actually, doesn't SCEV have a GetExtend w/a signed param? I know IR builder does. Adding one wouldn't be unreasonable.

1020 ↗

(On Diff #37315)

Isn't there something along the lines of SE->getBinOp?

1037 ↗

(On Diff #37315)

Comment? Where'd WideAR come from?

1056 ↗

(On Diff #37315)

Why not use auto here?

[sanjoy]

• address review

[sanjoy]

lib/Transforms/Scalar/IndVarSimplify.cpp

1020 ↗

(On Diff #37315)

Coud not find any such thing.

1037 ↗

(On Diff #37315)

Added a comment explaining what's going on.

[reames]

I have only minor comments, but I'm not familiar enough with the surrounding code to spot any subtle mistakes. I'll say LGTM, but take that for what it's worth.

lib/Transforms/Scalar/IndVarSimplify.cpp

1061 ↗

(On Diff #37381)

I'd restructure this to put the assignment before the second guess. It would make it more obvious what's going on.

1068 ↗

(On Diff #37381)

The fact this function is also called getExtend is confusing. getExtend != GetExtend?

llvm/trunk/lib/Transforms/Scalar/IndVarSimplify.cpp

 Instruction *WidenIV::cloneArithmeticIVUser(NarrowIVDefUse DU,
                                             const SCEVAddRecExpr *WideAR) {
   Instruction *NarrowUse = DU.NarrowUse;
   Instruction *NarrowDef = DU.NarrowDef;
   Instruction *WideDef = DU.WideDef;
 
   DEBUG(dbgs() << "Cloning arithmetic IVUser: " << *NarrowUse << "\n");
 
-  // Replace NarrowDef operands with WideDef. Otherwise, we don't know anything
-  // about the narrow operand yet so must insert a [sz]ext. It is probably loop
-  // invariant and will be folded or hoisted. If it actually comes from a
-  // widened IV, it should be removed during a future call to widenIVUse.
-  Value *LHS =
-      (NarrowUse->getOperand(0) == NarrowDef)
+  unsigned IVOpIdx = (NarrowUse->getOperand(0) == NarrowDef) ? 0 : 1;
+
+  // We're trying to find X such that
+  //
+  //  Widen(NarrowDef `op` NonIVNarrowDef) == WideAR == WideDef `op.wide` X
+  //
+  // We guess two solutions to X, sext(NonIVNarrowDef) and zext(NonIVNarrowDef),
+  // and check using SCEV if any of them are correct.
+
+  // Returns true if extending NonIVNarrowDef according to `SignExt` is a
+  // correct solution to X.
+  auto GuessNonIVOperand = [&](bool SignExt) {
+    const SCEV *WideLHS;
+    const SCEV *WideRHS;
+
+    auto GetExtend = [this, SignExt](const SCEV *S, Type *Ty) {
+      if (SignExt)
+        return SE->getSignExtendExpr(S, Ty);
+      return SE->getZeroExtendExpr(S, Ty);
+    };
+
+    if (IVOpIdx == 0) {
+      WideLHS = SE->getSCEV(WideDef);
+      const SCEV *NarrowRHS = SE->getSCEV(NarrowUse->getOperand(1));
+      WideRHS = GetExtend(NarrowRHS, WideType);
+    } else {
+      const SCEV *NarrowLHS = SE->getSCEV(NarrowUse->getOperand(0));
+      WideLHS = GetExtend(NarrowLHS, WideType);
+      WideRHS = SE->getSCEV(WideDef);
+    }
+
+    // WideUse is "WideDef `op.wide` X" as described in the comment.
+    const SCEV *WideUse = nullptr;
+
+    switch (NarrowUse->getOpcode()) {
+    default:
+      llvm_unreachable("No other possibility!");
+
+    case Instruction::Add:
+      WideUse = SE->getAddExpr(WideLHS, WideRHS);
+      break;
+
+    case Instruction::Mul:
+      WideUse = SE->getMulExpr(WideLHS, WideRHS);
+      break;
+
+    case Instruction::UDiv:
+      WideUse = SE->getUDivExpr(WideLHS, WideRHS);
+      break;
+
+    case Instruction::Sub:
+      WideUse = SE->getMinusSCEV(WideLHS, WideRHS);
+      break;
+    }
+
+    return WideUse == WideAR;
+  };
+
+  bool SignExtend = IsSigned;
+  if (!GuessNonIVOperand(SignExtend)) {
+    SignExtend = !SignExtend;
+    if (!GuessNonIVOperand(SignExtend))
+      return nullptr;
+  }
+
+  Value *LHS = (NarrowUse->getOperand(0) == NarrowDef)
                    ? WideDef
-          : getExtend(NarrowUse->getOperand(0), WideType, IsSigned, NarrowUse);
-  Value *RHS =
-      (NarrowUse->getOperand(1) == NarrowDef)
+                   : getExtend(NarrowUse->getOperand(0), WideType, SignExtend,
+                               NarrowUse);
+  Value *RHS = (NarrowUse->getOperand(1) == NarrowDef)
                    ? WideDef
-          : getExtend(NarrowUse->getOperand(1), WideType, IsSigned, NarrowUse);
+                   : getExtend(NarrowUse->getOperand(1), WideType, SignExtend,
+                               NarrowUse);
 
   auto *NarrowBO = cast<BinaryOperator>(NarrowUse);
   auto *WideBO = BinaryOperator::Create(NarrowBO->getOpcode(), LHS, RHS,
                                         NarrowBO->getName());
+
   IRBuilder<> Builder(NarrowUse);
   Builder.Insert(WideBO);
   if (const auto *OBO = dyn_cast<OverflowingBinaryOperator>(NarrowBO)) {
     if (OBO->hasNoUnsignedWrap())
       WideBO->setHasNoUnsignedWrap();
     if (OBO->hasNoSignedWrap())
       WideBO->setHasNoSignedWrap();
   }

llvm/trunk/test/Transforms/IndVarSimplify/iv-widen.ll

 ; RUN: opt < %s -indvars -S | FileCheck %s
 
 ; Provide legal integer types.
 target datalayout = "n8:16:32:64"
 
 
 target triple = "x86_64-apple-darwin"
 
-; CHECK-LABEL: @sloop
+; CHECK-LABEL: @loop_0
 ; CHECK-LABEL: B18:
 ; Only one phi now.
 ; CHECK: phi
 ; CHECK-NOT: phi
 ; One trunc for the gep.
 ; CHECK: trunc i64 %indvars.iv to i32
 ; One trunc for the dummy() call.
 ; CHECK-LABEL: exit24:
 ; CHECK: trunc i64 {{.*}}lcssa.wide to i32
-define void @sloop(i32* %a) {
+define void @loop_0(i32* %a) {
 Prologue:
   br i1 undef, label %B18, label %B6
 
 B18:                                        ; preds = %B24, %Prologue
   %.02 = phi i32 [ 0, %Prologue ], [ %tmp33, %B24 ]
   %tmp23 = zext i32 %.02 to i64
   %tmp33 = add i32 %.02, 1
   %o = getelementptr i32, i32* %a, i32 %.02
   %v = load i32, i32* %o
   %t = icmp eq i32 %v, 0
   br i1 %t, label %exit24, label %B24
 
 B24:                                        ; preds = %B18
   %t2 = icmp eq i32 %tmp33, 20
   br i1 %t2, label %B6, label %B18
 
 B6:                                       ; preds = %Prologue
   ret void
 
 exit24:                      ; preds = %B18
   call void @dummy(i32 %.02)
   unreachable
 }
 
+define void @loop_1(i32 %lim) {
+; CHECK-LABEL: @loop_1(
+ entry:
+  %entry.cond = icmp ne i32 %lim, 0
+  br i1 %entry.cond, label %loop, label %leave
+
+ loop:
+; CHECK: loop:
+; CHECK:  %indvars.iv = phi i64 [ 1, %loop.preheader ], [ %indvars.iv.next, %loop ]
+; CHECK:  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+; CHECK:  [[IV_INC:%[^ ]+]] = add nsw i64 %indvars.iv, -1
+; CHECK:  call void @dummy.i64(i64 [[IV_INC]])
+
+  %iv = phi i32 [ 1, %entry ], [ %iv.inc, %loop ]
+  %iv.inc = add i32 %iv, 1
+  %iv.inc.sub = add i32 %iv, -1
+  %iv.inc.sub.zext = zext i32 %iv.inc.sub to i64
+  call void @dummy.i64(i64 %iv.inc.sub.zext)
+  %be.cond = icmp ult i32 %iv.inc, %lim
+  br i1 %be.cond, label %loop, label %leave
+
+ leave:
+  ret void
+}
+
 declare void @dummy(i32)
+declare void @dummy.i64(i64)