[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

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

1001

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

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

1037

Comment? Where'd WideAR come from?

1056

Why not use auto here?

[sanjoy]

• address review

[sanjoy]

lib/Transforms/Scalar/IndVarSimplify.cpp

1020

Coud not find any such thing.

1037

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

1055

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

1062

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

lib/Transforms/Scalar/IndVarSimplify.cpp

if (const auto *OBO = dyn_cast<OverflowingBinaryOperator>(NarrowBO)) {
     if (OBO->hasNoSignedWrap())
       WideBO->setHasNoSignedWrap();
   }
   return WideBO;
 }
 
 Instruction *WidenIV::cloneArithmeticIVUser(NarrowIVDefUse DU,
                                             const SCEVAddRecExpr *WideAR) {
-  DEBUG(dbgs() << "Cloning arithmetic IVUser: " << *DU.NarrowUse << "\n");
+  DEBUG(dbgs() << "Cloning Arithmetic IVUser: " << *DU.NarrowUse << "\n");
+  Instruction *NarrowUse = DU.NarrowUse;
+  Instruction *NarrowDef = DU.NarrowDef;
+
+  unsigned IVOpIdx = (NarrowUse->getOperand(0) == NarrowDef) ? 0 : 1;
+
+  auto GuessNonIVOperand = [&](bool SignExt) {
+    const SCEV *WideLHS;
+    const SCEV *WideRHS;
+
+    auto GetExtend = SignExt ? &ScalarEvolution::getSignExtendExpr
+                             : &ScalarEvolution::getZeroExtendExpr;
+
+    if (IVOpIdx == 0) {
+      WideLHS = SE->getSCEV(DU.WideDef);
+      const SCEV *NarrowRHS = SE->getSCEV(NarrowUse->getOperand(1));
+      WideRHS = (SE->*GetExtend)(NarrowRHS, WideType);
+    } else {
+      const SCEV *NarrowLHS = SE->getSCEV(NarrowUse->getOperand(0));
+      WideLHS = (SE->*GetExtend)(NarrowLHS, WideType);
+      WideRHS = SE->getSCEV(DU.WideDef);
+    }
 
-  // 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 = (DU.NarrowUse->getOperand(0) == DU.NarrowDef)
+    const SCEV *WideOp = nullptr;
+
+    switch (NarrowUse->getOpcode()) {
+    default:
+      llvm_unreachable("No other possibility!");
+
+    case Instruction::Add:
+      WideOp = SE->getAddExpr(WideLHS, WideRHS);
+      break;
+
+    case Instruction::Mul:
+      WideOp = SE->getMulExpr(WideLHS, WideRHS);
+      break;
+
+    case Instruction::UDiv:
+      WideOp = SE->getUDivExpr(WideLHS, WideRHS);
+      break;
+
+    case Instruction::Sub:
+      WideOp = SE->getMinusSCEV(WideLHS, WideRHS);
+      break;
+    }
+
+    return WideOp == WideAR;
+  };
+
+  bool SignExtend = IsSigned;
+  if (!GuessNonIVOperand(SignExtend)) {
+    if (!GuessNonIVOperand(!SignExtend))
+      return nullptr;
+    SignExtend = !SignExtend;
+  }
+
+  Value *LHS = (NarrowUse->getOperand(0) == NarrowDef)
                    ? DU.WideDef
-                   : getExtend(DU.NarrowUse->getOperand(0), WideType, IsSigned,
-                               DU.NarrowUse);
-  Value *RHS = (DU.NarrowUse->getOperand(1) == DU.NarrowDef)
+                   : getExtend(NarrowUse->getOperand(0), WideType, SignExtend,
+                               NarrowUse);
+  Value *RHS = (NarrowUse->getOperand(1) == NarrowDef)
                    ? DU.WideDef
-                   : getExtend(DU.NarrowUse->getOperand(1), WideType, IsSigned,
-                               DU.NarrowUse);
+                   : getExtend(NarrowUse->getOperand(1), WideType, SignExtend,
+                               NarrowUse);
 
-  auto *NarrowBO = cast<BinaryOperator>(DU.NarrowUse);
-  auto *WideBO = BinaryOperator::Create(NarrowBO->getOpcode(), LHS, RHS,
-                                        NarrowBO->getName());
-  IRBuilder<> Builder(DU.NarrowUse);
+  BinaryOperator *NarrowBO = cast<BinaryOperator>(DU.NarrowUse);
+  BinaryOperator *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)