[SCEV] Postpone ScalarEvolution::forgetLoop after TripCountSC is expanded when unroll runtime iteration loop

[wmi]

Summary

In llvm::UnrollRuntimeLoopRemainder, if the loop to be unrolled is the inner loop inside a loop nest, the scalar evolution needs to be dropped for its parent loop which is done by ScalarEvolution::forgetLoop. However, we can postpone forgetLoop to the end of UnrollRuntimeLoopRemainder so TripCountSC expansion can still reuse existing value.

[wmi]

Gentle Ping.

[sanjoy]

Hi Wei,

I don't think I'm the best person to review this code, I think you should ask either @mzolotukhin or @evstupac.

[mzolotukhin]

The change looks good to me, thanks!

Michael

test/Analysis/ScalarEvolution/scev-expander-existing-value.ll

1–2

Would it make sense to separate the test into two files maybe? Right now we run opt twice on this file but only check the output of one half each run. Ideally, we could somehow use the same routine for both vectorizer and unroller runs, but I don't know if that's possible.

[wmi]

Sanjoy and Michael, thanks!

Separate the test into two.

lib/Transforms/Utils/LoopUnrollRuntime.cpp

if (!AllowExpensiveTripCount &&
       Expander.isHighCostExpansion(TripCountSC, L, PreHeaderBR))
     return false;
 
   // This constraint lets us deal with an overflowing trip count easily; see the
   // comment on ModVal below.
   if (Log2_32(Count) > BEWidth)
     return false;
 
-  // If this loop is nested, then the loop unroller changes the code in the
-  // parent loop, so the Scalar Evolution pass needs to be run again.
-  if (Loop *ParentLoop = L->getParentLoop())
-    SE->forgetLoop(ParentLoop);
-
   BasicBlock *Latch = L->getLoopLatch();
 
   // Loop structure is the following:
   //
   // PreHeader
   //   Header
   //   ...
   //   Latch
if (UseEpilogRemainder) {
     NewIdx->addIncoming(IdxSub, Latch);
     LatchBR->setCondition(IdxCmp);
   } else {
     // Connect the prolog code to the original loop and update the
     // PHI functions.
     ConnectProlog(L, BECount, Count, PrologExit, PreHeader, NewPreHeader,
                   VMap, DT, LI, PreserveLCSSA);
   }
+
+  // If this loop is nested, then the loop unroller changes the code in the
+  // parent loop, so the Scalar Evolution pass needs to be run again.
+  if (Loop *ParentLoop = L->getParentLoop())
+    SE->forgetLoop(ParentLoop);
+
   NumRuntimeUnrolled++;
   return true;
 }

test/Analysis/ScalarEvolution/scev-expander-existing-value.ll

-; RUN: opt < %s -loop-vectorize -force-vector-width=4 -verify-scev-maps -S |FileCheck %s
+; RUN: opt < %s -loop-vectorize -force-vector-width=4 -verify-scev-maps -S |FileCheck %s --check-prefix=VEC
+; RUN: opt < %s -loop-unroll -unroll-runtime -unroll-count=2 -verify-scev-maps -S | FileCheck %s --check-prefix=UNROLL
 
 ; SCEV expansion uses existing value when the SCEV has no AddRec expr.
-; CHECK: select
-; CHECK-NOT: select
+; VEC-LABEL: @foo_vect(
+; VEC: select
+; VEC-NOT: select
+; VEC: ret
 
 @a = common global [1000 x i16] zeroinitializer, align 16
 
-define i32 @foo(i32 %x, i32 %y) {
+define i32 @foo_vect(i32 %x, i32 %y) {
 entry:
   %cmp = icmp slt i32 %x, %y
   %cond = select i1 %cmp, i32 %x, i32 %y
   %cmp1.10 = icmp sgt i32 %cond, 0
   br i1 %cmp1.10, label %for.body.lr.ph, label %for.end
 
 for.body.lr.ph:                                   ; preds = %entry
   %tmp = sext i32 %cond to i64

 for.end.loopexit:                                 ; preds = %for.body
   %add.lcssa = phi i32 [ %add, %for.body ]
   br label %for.end
 
 for.end:                                          ; preds = %for.end.loopexit, %entry
   %total.0.lcssa = phi i32 [ 0, %entry ], [ %add.lcssa, %for.end.loopexit ]
   ret i32 %total.0.lcssa
 }
+
+; Check SCEV expansion uses existing value when unrolling an inner loop with runtime trip count in a loop nest.
+; UNROLL-LABEL: @foo_unroll(
+; UNROLL: select
+; UNROLL-NOT: select
+; UNROLL: ret
+
+define void @foo_unroll(i32 %xfL, i32 %scaleL) local_unnamed_addr {
+entry:
+  br label %for.body
+
+for.body:                                         ; preds = %for.body5, %for.body, %entry
+  %xfL.addr.033 = phi i32 [ %xfL, %entry ], [ %add, %for.body5 ]
+  %add = add nsw i32 %xfL.addr.033, %scaleL
+  %shr = ashr i32 %add, 16
+  %cmp.i = icmp slt i32 0, %shr
+  %.sroa.speculated = select i1 %cmp.i, i32 0, i32 %shr
+  %cmp425 = icmp slt i32 0, %.sroa.speculated
+  br i1 %cmp425, label %for.body5.preheader, label %for.end
+
+for.body5.preheader:                              ; preds = %for.body
+  %tmp0 = sext i32 %.sroa.speculated to i64
+  br label %for.body5
+
+for.body5:                                        ; preds = %for.body5, %for.body5.preheader
+  %indvars.iv = phi i64 [ 0, %for.body5.preheader ], [ %indvars.iv.next, %for.body5 ]
+  %indvars.iv.next = add nsw i64 %indvars.iv, 1
+  %cmp4 = icmp slt i64 %indvars.iv.next, %tmp0
+  br i1 %cmp4, label %for.body5, label %for.body
+
+for.end:
+  ret void
+}
+