[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

996

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

1006

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.

1025

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

1042

Comment? Where'd WideAR come from?

1079

Why not use auto here?

[sanjoy]

• address review

[sanjoy]

lib/Transforms/Scalar/IndVarSimplify.cpp

1025

Coud not find any such thing.

1042

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

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

1068

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

lib/Transforms/Scalar/IndVarSimplify.cpp

if (OBO->hasNoSignedWrap())
       WideBO->setHasNoSignedWrap();
   }
   return WideBO;
 }
 
 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 with SCEV if any of them are a solution.
+
+  // 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)) {
+    if (!GuessNonIVOperand(!SignExtend))
+      return nullptr;
+    SignExtend = !SignExtend;
+  }
+
+  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();
   }
   return WideBO;

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)