Use ISL to Determine Loop Trip Count

[mssimpso]

Summary

Use ISL to compute the loop trip count when scalar evolution is
unable to do so.

Patch originally by Johannes Doerfert!

[jdoerfert]

Some comments are inlined. Except those and the test coverage we could improve the patch LGTM.

We should probably add a test case for the nsw problem, e.g., that we do not check for nsw and therefore can assume an unconstrainted loop. The example should probably look similar to the one below, however we would need to force ScalarEvolution to not recognize the loop trip count.

// Infinite loop for N < 0 in a non-wrapping world.
for (int i = 0; i < N; i++)
  A[i] = i;

lib/Analysis/ScopDetection.cpp

708 ↗

(On Diff #24829)

move this in the if of the next line

969 ↗

(On Diff #24829)

I don't think we need this as the loops we now allow are "affine" hence the hasAffineLoops should be set to true if we encountered one. If not we want to bail even if we would allow such loops.

lib/Analysis/ScopInfo.cpp

77 ↗

(On Diff #24829)

I know there was no documentation for this function before, but maybe you could write a short one that explains that this computes not the real loop-depth (in the sence of Loop/LoopInfo) but a relative one in the SCoP.

742 ↗

(On Diff #24829)

I don't get this change. If the input space (setDomain) was 0-dimensional we return a map like: [ ] -> [ ] ?

Could you give me some intuition why?

922 ↗

(On Diff #24829)

Our virtual loops (i0,i1,..) always start at 0. We should keep it that way otherwise we will soon end up with an offset by 1 mess. This should also simplify this code block a bit.

933 ↗

(On Diff #24829)

We should not enforce unsigned conditions, they are/will be supported soon (at least under a flag).

test/ScopInfo/isl_trip_count.ll

3 ↗

(On Diff #24829)

see offset-by-one comment above.

[grosser]

Is this patch waiting for some kind of feedback from my side? Johannes already gave a good review so feel free to follow up with him.

One comment:

1. We should enable this by default to ensure we get enough test coverage.

[mssimpso]

Rebased and addressed comments. Thanks for the reviews!

@grosser:

Flag is now enabled by default.

@jdoerfert:

My apologies for waiting so long to update this patch. I've addressed your
comments, added an additional test case, and performed some additional
refactoring. Regarding the actual computation of the trip count, I made a small
adjustment to the expression given in the PET paper. I now simplify the
condition on the domain to remove unnecessary constraints. This corrects the
"offset-by-one" issue you pointed out in the original test case and also fixes
some cases we saw in our internal testing.

[mssimpso]

lib/Analysis/ScopDetection.cpp

804 ↗

(On Diff #32987)

I moved this to canUseISLTripCount.

1039 ↗

(On Diff #32987)

Agreed. I now set Context.hasAffineLoops = true if canUseISLTripCount succeeds.

lib/Analysis/ScopInfo.cpp

99 ↗

(On Diff #32987)

I made this utility a member of Scop since SCEVAffinator was factored out since the initial revision.

602 ↗

(On Diff #32987)

This was a mistake.

750 ↗

(On Diff #32987)

See comment above. I simplify the condition on the domain to remove unnecessary constraints.

[jdoerfert]

LGTM, if it passes LNT (and maybe even your internal test). I will work on the domain modeling again soon but this is a really good starting point.

lib/Analysis/ScopDetection.cpp

784 ↗

(On Diff #32987)

For now we can be restrictive but could you explain to me why do we need this part? I want to refactor the Domain creation soon and any information you have would be usefull.

[mssimpso]

lib/Analysis/ScopDetection.cpp

784 ↗

(On Diff #32987)

Without the invariant check, we were hitting cases like for (int i=0; i < A[i]; i++) (e.g., in test/ScopDetect/non-affine-loop.ll). This was previously rejected because the trip count could not be computed by SCEV.

798 ↗

(On Diff #32987)

I think the refactoring was a bit too aggressive here. If the trip count is computable by SCEV, but it is not affine, we still need to reject. I will post an update shortly.

[mssimpso]

Removed some refactoring to restore correct behavior.

@jdoerfert:

Please feel free to commit this change when you're satisfied with it, as I do
not yet have commit access. Thanks!

polly/trunk/include/polly/ScopDetection.h

private:
   /// @return True if the loop is valid in the region.
   bool isValidLoop(Loop *L, DetectionContext &Context) const;
 
   /// @brief Check if the function @p F is marked as invalid.
   ///
   /// @note An OpenMP subfunction will be marked as invalid.
   bool isValidFunction(llvm::Function &F);
 
+  /// @brief Can ISL compute the trip count of a loop.
+  ///
+  /// @param L The loop to check.
+  /// @param Context The context of scop detection.
+  ///
+  /// @return True if ISL can compute the trip count of the loop.
+  bool canUseISLTripCount(Loop *L, DetectionContext &Context) const;
+
   /// @brief Print the locations of all detected scops.
   void printLocations(llvm::Function &F);
 
   /// @brief Track diagnostics for invalid scops.
   ///
   /// @param Context The context of scop detection.
   /// @param Assert Throw an assert in verify mode or not.
   /// @param Args Argument list that gets passed to the constructor of RR.

polly/trunk/include/polly/ScopInfo.h

private:
 
   std::string BaseName;
 
   /// Build the statement.
   //@{
   __isl_give isl_set *buildConditionSet(const Comparison &Cmp);
   void addConditionsToDomain(TempScop &tempScop, const Region &CurRegion);
   void addLoopBoundsToDomain(TempScop &tempScop);
+  void addLoopTripCountToDomain(const Loop *L);
   void buildDomain(TempScop &tempScop, const Region &CurRegion);
 
   /// @brief Create the accesses for instructions in @p Block.
   ///
   /// @param tempScop       The template SCoP.
   /// @param Block          The basic block for which accesses should be
   ///                       created.
   /// @param isApproximated Flag to indicate blocks that might not be executed,
public:
   ///
   /// @brief NewSchedule The new schedule (given as schedule tree).
   void setScheduleTree(__isl_take isl_schedule *NewSchedule);
 
   /// @brief Intersects the domains of all statements in the SCoP.
   ///
   /// @return true if a change was made
   bool restrictDomains(__isl_take isl_union_set *Domain);
+
+  /// @brief Get the depth of a loop relative to the outermost loop in the Scop.
+  unsigned getRelativeLoopDepth(const Loop *L) const;
 };
 
 /// @brief Print Scop scop to raw_ostream O.
 static inline raw_ostream &operator<<(raw_ostream &O, const Scop &scop) {
   scop.print(O);
   return O;
 }
 

polly/trunk/include/polly/Support/SCEVAffinator.h

private:
 
   Scop *S;
   isl_ctx *Ctx;
   unsigned NumIterators;
   const llvm::Region &R;
   llvm::ScalarEvolution &SE;
   const ScopStmt *Stmt;
 
-  int getLoopDepth(const llvm::Loop *L);
-
   __isl_give isl_pw_aff *visit(const llvm::SCEV *E);
   __isl_give isl_pw_aff *visitConstant(const llvm::SCEVConstant *E);
   __isl_give isl_pw_aff *visitTruncateExpr(const llvm::SCEVTruncateExpr *E);
   __isl_give isl_pw_aff *visitZeroExtendExpr(const llvm::SCEVZeroExtendExpr *E);
   __isl_give isl_pw_aff *visitSignExtendExpr(const llvm::SCEVSignExtendExpr *E);
   __isl_give isl_pw_aff *visitAddExpr(const llvm::SCEVAddExpr *E);
   __isl_give isl_pw_aff *visitMulExpr(const llvm::SCEVMulExpr *E);
   __isl_give isl_pw_aff *visitUDivExpr(const llvm::SCEVUDivExpr *E);

polly/trunk/lib/Analysis/ScopDetection.cpp

PollyDelinearizeX("polly-delinearize",
                       cl::ZeroOrMore, cl::init(true), cl::cat(PollyCategory));
 
 static cl::opt<bool>
     VerifyScops("polly-detect-verify",
                 cl::desc("Verify the detected SCoPs after each transformation"),
                 cl::Hidden, cl::init(false), cl::ZeroOrMore,
                 cl::cat(PollyCategory));
 
+static cl::opt<bool> AllowNonSCEVBackedgeTakenCount(
+    "polly-allow-non-scev-backedge-taken-count",
+    cl::desc("Allow loops even if SCEV cannot provide a trip count"),
+    cl::Hidden, cl::init(true), cl::ZeroOrMore, cl::cat(PollyCategory));
+
 bool polly::PollyTrackFailures = false;
 bool polly::PollyDelinearize = false;
 StringRef polly::PollySkipFnAttr = "polly.skip.fn";
 
 //===----------------------------------------------------------------------===//
 // Statistics.
 
 STATISTIC(ValidRegion, "Number of regions that a valid part of Scop");
if (isa<LoadInst>(Inst) || isa<StoreInst>(Inst)) {
     Context.hasLoads |= isa<LoadInst>(Inst);
     return isValidMemoryAccess(Inst, Context);
   }
 
   // We do not know this instruction, therefore we assume it is invalid.
   return invalid<ReportUnknownInst>(Context, /*Assert=*/true, &Inst);
 }
 
+bool ScopDetection::canUseISLTripCount(Loop *L,
+                                       DetectionContext &Context) const {
+
+  Region &CurRegion = Context.CurRegion;
+
+  // Ensure the loop has a single back edge.
+  if (L->getNumBackEdges() != 1)
+    return false;
+
+  // Ensure the loop has a single exiting block.
+  BasicBlock *ExitingBB = L->getExitingBlock();
+  if (!ExitingBB)
+    return false;
+
+  // Ensure the exiting block is terminated by a conditional branch.
+  BranchInst *Term = dyn_cast<BranchInst>(ExitingBB->getTerminator());
+  if (!Term || !Term->isConditional())
+    return false;
+
+  Value *Cond = Term->getCondition();
+
+  // If the terminating condition is an integer comparison, ensure that it is a
+  // comparison between a recurrence and an invariant value.
+  if (ICmpInst *I = dyn_cast<ICmpInst>(Cond)) {
+    const Value *Op0 = I->getOperand(0);
+    const Value *Op1 = I->getOperand(1);
+    const SCEV *LHS = SE->getSCEVAtScope(const_cast<Value *>(Op0), L);
+    const SCEV *RHS = SE->getSCEVAtScope(const_cast<Value *>(Op1), L);
+    if ((isa<SCEVAddRecExpr>(LHS) && !isInvariant(*Op1, CurRegion)) ||
+        (isa<SCEVAddRecExpr>(RHS) && !isInvariant(*Op0, CurRegion)))
+      return false;
+  }
+
+  // If the terminating condition is not an integer comparison, ensure that it
+  // is a constant.
+  else if (!isa<ConstantInt>(Cond))
+    return false;
+
+  // We can use ISL to compute the trip count of L.
+  return true;
+}
+
 bool ScopDetection::isValidLoop(Loop *L, DetectionContext &Context) const {
   // Is the loop count affine?
+  bool IsLoopCountAffine = false;
   const SCEV *LoopCount = SE->getBackedgeTakenCount(L);
-  if (isAffineExpr(&Context.CurRegion, LoopCount, *SE)) {
+  if (!isa<SCEVCouldNotCompute>(LoopCount))
+    IsLoopCountAffine = isAffineExpr(&Context.CurRegion, LoopCount, *SE);
+  else
+    IsLoopCountAffine = canUseISLTripCount(L, Context);
+  if (IsLoopCountAffine) {
     Context.hasAffineLoops = true;
     return true;
   }
 
   if (AllowNonAffineSubRegions) {
     Region *R = RI->getRegionFor(L->getHeader());
     if (R->contains(L))
       if (addOverApproximatedRegion(R, Context))

polly/trunk/lib/Analysis/ScopInfo.cpp

 
 void ScopStmt::realignParams() {
   for (MemoryAccess *MA : *this)
     MA->realignParams();
 
   Domain = isl_set_align_params(Domain, Parent.getParamSpace());
 }
 
+void ScopStmt::addLoopTripCountToDomain(const Loop *L) {
+
+  unsigned loopDimension = getParent()->getRelativeLoopDepth(L);
+  ScalarEvolution *SE = getParent()->getSE();
+  isl_space *DomSpace = isl_set_get_space(Domain);
+
+  isl_space *MapSpace = isl_space_map_from_set(isl_space_copy(DomSpace));
+  isl_multi_aff *LoopMAff = isl_multi_aff_identity(MapSpace);
+  isl_aff *LoopAff = isl_multi_aff_get_aff(LoopMAff, loopDimension);
+  LoopAff = isl_aff_add_constant_si(LoopAff, 1);
+  LoopMAff = isl_multi_aff_set_aff(LoopMAff, loopDimension, LoopAff);
+  isl_map *TranslationMap = isl_map_from_multi_aff(LoopMAff);
+
+  BasicBlock *ExitingBB = L->getExitingBlock();
+  assert(ExitingBB && "Loop has more than one exiting block");
+
+  BranchInst *Term = dyn_cast<BranchInst>(ExitingBB->getTerminator());
+  assert(Term && Term->isConditional() && "Terminator is not conditional");
+
+  const SCEV *LHS = nullptr;
+  const SCEV *RHS = nullptr;
+  Value *Cond = Term->getCondition();
+  CmpInst::Predicate Pred = CmpInst::Predicate::BAD_ICMP_PREDICATE;
+
+  ICmpInst *CondICmpInst = dyn_cast<ICmpInst>(Cond);
+  ConstantInt *CondConstant = dyn_cast<ConstantInt>(Cond);
+  if (CondICmpInst) {
+    LHS = SE->getSCEVAtScope(CondICmpInst->getOperand(0), L);
+    RHS = SE->getSCEVAtScope(CondICmpInst->getOperand(1), L);
+    Pred = CondICmpInst->getPredicate();
+  } else if (CondConstant) {
+    LHS = SE->getConstant(CondConstant);
+    RHS = SE->getConstant(ConstantInt::getTrue(SE->getContext()));
+    Pred = CmpInst::Predicate::ICMP_EQ;
+  } else {
+    llvm_unreachable("Condition is neither a ConstantInt nor a ICmpInst");
+  }
+
+  if (!L->contains(Term->getSuccessor(0)))
+    Pred = ICmpInst::getInversePredicate(Pred);
+  Comparison Comp(LHS, RHS, Pred);
+
+  isl_set *CondSet = buildConditionSet(Comp);
+  isl_map *ForwardMap = isl_map_lex_le(isl_space_copy(DomSpace));
+  for (unsigned i = 0; i < isl_set_n_dim(Domain); i++)
+    if (i != loopDimension)
+      ForwardMap = isl_map_equate(ForwardMap, isl_dim_in, i, isl_dim_out, i);
+
+  ForwardMap = isl_map_apply_range(ForwardMap, isl_map_copy(TranslationMap));
+  isl_set *CondDom = isl_set_subtract(isl_set_copy(Domain), CondSet);
+  isl_set *ForwardCond = isl_set_apply(CondDom, isl_map_copy(ForwardMap));
+  isl_set *ForwardDomain = isl_set_apply(isl_set_copy(Domain), ForwardMap);
+  ForwardCond = isl_set_gist(ForwardCond, ForwardDomain);
+  Domain = isl_set_subtract(Domain, ForwardCond);
+
+  isl_map_free(TranslationMap);
+  isl_space_free(DomSpace);
+}
+
 __isl_give isl_set *ScopStmt::buildConditionSet(const Comparison &Comp) {
   isl_pw_aff *L = getPwAff(Comp.getLHS());
   isl_pw_aff *R = getPwAff(Comp.getRHS());
 
   switch (Comp.getPred()) {
   case ICmpInst::ICMP_EQ:
     return isl_pw_aff_eq_set(L, R);
   case ICmpInst::ICMP_NE:
for (int i = 0, e = getNumIterators(); i != e; ++i) {
 
     // 0 <= IV.
     isl_set *LowerBound = isl_pw_aff_nonneg_set(isl_pw_aff_copy(IV));
     Domain = isl_set_intersect(Domain, LowerBound);
 
     // IV <= LatchExecutions.
     const Loop *L = getLoopForDimension(i);
     const SCEV *LatchExecutions = SE->getBackedgeTakenCount(L);
+    if (!isa<SCEVCouldNotCompute>(LatchExecutions)) {
       isl_pw_aff *UpperBound = getPwAff(LatchExecutions);
       isl_set *UpperBoundSet = isl_pw_aff_le_set(IV, UpperBound);
       Domain = isl_set_intersect(Domain, UpperBoundSet);
+    } else {
+      // If SCEV cannot provide a loop trip count we compute it with ISL.
+      addLoopTripCountToDomain(L);
+      isl_pw_aff_free(IV);
+    }
   }
 
   isl_local_space_free(LocalSpace);
 }
 
 void ScopStmt::addConditionsToDomain(TempScop &tempScop,
                                      const Region &CurRegion) {
   const Region *TopRegion = tempScop.getMaxRegion().getParent(),

 
 ScopStmt *Scop::getStmtForBasicBlock(BasicBlock *BB) const {
   auto StmtMapIt = StmtMap.find(BB);
   if (StmtMapIt == StmtMap.end())
     return nullptr;
   return StmtMapIt->second;
 }
 
+unsigned Scop::getRelativeLoopDepth(const Loop *L) const {
+  Loop *OuterLoop = R.outermostLoopInRegion(const_cast<Loop *>(L));
+  assert(OuterLoop && "Scop does not contain this loop");
+  return L->getLoopDepth() - OuterLoop->getLoopDepth();
+}
+
 //===----------------------------------------------------------------------===//
 ScopInfo::ScopInfo() : RegionPass(ID), scop(0) {
   ctx = isl_ctx_alloc();
   isl_options_set_on_error(ctx, ISL_ON_ERROR_ABORT);
 }
 
 ScopInfo::~ScopInfo() {
   clear();

polly/trunk/lib/Support/SCEVAffinator.cpp

SCEVAffinator::visitAddRecExpr(const SCEVAddRecExpr *Expr) {
   if (Expr->getStart()->isZero()) {
     assert(S->getRegion().contains(Expr->getLoop()) &&
            "Scop does not contain the loop referenced in this AddRec");
 
     isl_pw_aff *Step = visit(Expr->getOperand(1));
     isl_space *Space = isl_space_set_alloc(Ctx, 0, NumIterators);
     isl_local_space *LocalSpace = isl_local_space_from_space(Space);
 
-    int loopDimension = getLoopDepth(Expr->getLoop());
+    unsigned loopDimension = S->getRelativeLoopDepth(Expr->getLoop());
 
     isl_aff *LAff = isl_aff_set_coefficient_si(
         isl_aff_zero_on_domain(LocalSpace), isl_dim_in, loopDimension, 1);
     isl_pw_aff *LPwAff = isl_pw_aff_from_aff(LAff);
 
     // TODO: Do we need to check for NSW and NUW?
     return isl_pw_aff_mul(Step, LPwAff);
   }
if (Instruction *I = dyn_cast<Instruction>(Expr->getValue())) {
     default:
       break; // Fall through.
     }
   }
 
   llvm_unreachable(
       "Unknowns SCEV was neither parameter nor a valid instruction.");
 }
-
-int SCEVAffinator::getLoopDepth(const Loop *L) {
-  Loop *outerLoop = S->getRegion().outermostLoopInRegion(const_cast<Loop *>(L));
-  assert(outerLoop && "Scop does not contain this loop");
-  return L->getLoopDepth() - outerLoop->getLoopDepth();
-}

polly/trunk/test/ScopInfo/isl_trip_count_01.ll

+; RUN: opt %loadPolly -polly-detect-unprofitable -polly-allow-non-scev-backedge-taken-count -polly-scops -analyze < %s | FileCheck %s
+;
+; CHECK: [M, N] -> { Stmt_while_body[i0] : i0 >= 0 and 4i0 <= -M + N }
+;
+;   void f(int *A, int N, int M) {
+;     int i = 0;
+;     while (M <= N) {
+;       A[i++] = 1;
+;       M += 4;
+;     }
+;   }
+;
+target datalayout = "e-m:e-p:32:32-i64:64-v128:64:128-a:0:32-n32-S64"
+
+define void @f(i32* nocapture %A, i32 %N, i32 %M) {
+entry:
+  %cmp3 = icmp sgt i32 %M, %N
+  br i1 %cmp3, label %while.end, label %while.body.preheader
+
+while.body.preheader:
+  br label %while.body
+
+while.body:
+  %i.05 = phi i32 [ %inc, %while.body ], [ 0, %while.body.preheader ]
+  %M.addr.04 = phi i32 [ %add, %while.body ], [ %M, %while.body.preheader ]
+  %inc = add nuw nsw i32 %i.05, 1
+  %arrayidx = getelementptr inbounds i32, i32* %A, i32 %i.05
+  store i32 1, i32* %arrayidx, align 4
+  %add = add nsw i32 %M.addr.04, 4
+  %cmp = icmp sgt i32 %add, %N
+  br i1 %cmp, label %while.end.loopexit, label %while.body
+
+while.end.loopexit:
+  br label %while.end
+
+while.end:
+  ret void
+}

polly/trunk/test/ScopInfo/isl_trip_count_02.ll

+; RUN: opt %loadPolly -polly-detect-unprofitable -polly-allow-non-scev-backedge-taken-count -polly-scops -analyze < %s | FileCheck %s
+;
+; CHECK: [M, N] -> { Stmt_for_body[i0] : i0 >= 0 and N <= -1 + M };
+;
+;   void f(int *A, int N, int M) {
+;     for (int i = M; i > N; i++)
+;       A[i] = i;
+;   }
+;
+target datalayout = "e-m:e-p:32:32-i64:64-v128:64:128-a:0:32-n32-S64"
+
+define void @f(i32* %A, i32 %N, i32 %M) {
+entry:
+  br label %entry.split
+
+entry.split:
+  %cmp.1 = icmp sgt i32 %M, %N
+  br i1 %cmp.1, label %for.body, label %for.end
+
+for.body:
+  %indvars.iv = phi i32 [ %indvars.iv.next, %for.body ], [ %M, %entry.split ]
+  %arrayidx = getelementptr inbounds i32, i32* %A, i32 %indvars.iv
+  store i32 %indvars.iv, i32* %arrayidx, align 4
+  %cmp = icmp slt i32 %M, %N
+  %indvars.iv.next = add i32 %indvars.iv, 1
+  br i1 %cmp, label %for.cond.for.end_crit_edge, label %for.body
+
+for.cond.for.end_crit_edge:
+  br label %for.end
+
+for.end:
+  ret void
+}