[LoopReroll] Fix reroll root legality checking.

[efriedma]

Summary

The code checked that the first root was an appropriate distance from the base value, but skipped checking the other roots. This could lead to rerolling a loop that can't be legally rerolled (at least, not without rewriting the loop in a non-trivial way).

[efriedma]

Ping

[hfinkel]

LGTM

[fhahn]

LGTM, sorry for the delay!

lib/Transforms/Scalar/LoopRerollPass.cpp

bool LoopReroll::DAGRootTracker::validateRootSet(DAGRootSet &DRS) {
   // Define D = BaseInst@J - BaseInst@J-1, where "@J" means the value at the
   //   loop iteration J.
   //
   // Now, For the loop iterations to be consecutive:
   //   D = d * N
   const auto *ADR = dyn_cast<SCEVAddRecExpr>(SE->getSCEV(DRS.BaseInst));
   if (!ADR)
     return false;
+
+  // Check that the first root is evenly spaced.
   unsigned N = DRS.Roots.size() + 1;
   const SCEV *StepSCEV = SE->getMinusSCEV(SE->getSCEV(DRS.Roots[0]), ADR);
   const SCEV *ScaleSCEV = SE->getConstant(StepSCEV->getType(), N);
   if (ADR->getStepRecurrence(*SE) != SE->getMulExpr(StepSCEV, ScaleSCEV))
     return false;
 
+  // Check that the remainling roots are evenly spaced.
+  for (unsigned i = 1; i < N - 1; ++i) {
+    const SCEV *NewStepSCEV = SE->getMinusSCEV(SE->getSCEV(DRS.Roots[i]),
+                                               SE->getSCEV(DRS.Roots[i-1]));
+    if (NewStepSCEV != StepSCEV)
+      return false;
+  }
+
   return true;
 }
 
 bool LoopReroll::DAGRootTracker::
 findRootsBase(Instruction *IVU, SmallInstructionSet SubsumedInsts) {
   // The base of a RootSet must be an AddRec, so it can be erased.
   const auto *IVU_ADR = dyn_cast<SCEVAddRecExpr>(SE->getSCEV(IVU));
   if (!IVU_ADR || IVU_ADR->getLoop() != L)

test/Transforms/LoopReroll/basic.ll

bb3:                                              ; preds = %bb3, %bb
   %tmp17 = add nuw nsw i64 %tmp, 16
   %tmp18 = icmp eq i64 %tmp17, %arg2
   br i1 %tmp18, label %bb19, label %bb3
 
 bb19:                                             ; preds = %bb3
   ret void
 }
 
+define void @bad_step(i32* nocapture readnone %x) #0 {
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %for.body, %entry
+  %i.08 = phi i32 [ 0, %entry ], [ %add3, %for.body ]
+  %call = tail call i32 @foo(i32 %i.08) #1
+  %add = add nsw i32 %i.08, 2
+  %call1 = tail call i32 @foo(i32 %add) #1
+  %add2 = add nsw i32 %i.08, 3
+  %call3 = tail call i32 @foo(i32 %add2) #1
+  %add3 = add nsw i32 %i.08, 6
+  %exitcond = icmp sge i32 %add3, 500
+  br i1 %exitcond, label %for.end, label %for.body
+
+; CHECK-LABEL: @bad_step
+; CHECK: %add = add nsw i32 %i.08, 2
+; CHECK: %add2 = add nsw i32 %i.08, 3
+; CHECK: %add3 = add nsw i32 %i.08, 6
+
+for.end:                                          ; preds = %for.body
+  ret void
+}
+
 attributes #0 = { nounwind uwtable }
 attributes #1 = { nounwind }