[PIR] Allow to detect parallel loops using the ParallelRegionInfo.

[jdoerfert]

Summary

A parallel loop is a loop that does fork (parts of) its body to a new
task which are joined only after the loop. It is also ensured that
everything that is not forked but part of the loop does not cause
side-effects. These instructions usually compute the exit condition
and the new values for the induction variable(s).

Note: We allow multiple induction variables and complex exit conditions
here. However, for some parallel runtimes these have to be simplified,
e.g. expressed with regards to one canonical induction variable.
TODO: Provide an interface to perform the simplification.

Note: We allow arbitrary code after the loop but before a join
instruction.

include/llvm/Analysis/ParallelRegionInfo.h

 /// could be introduced to maintain this mapping and recompute it if it was not
 /// preserved by a transformation. However, at the moment there are only a small
 /// number of places that require it but a lot of transformations that would
 /// need to be modified to preserve it.
 ///
 /// @TODO The update interface to add/remove parallel regions is missing.
 ///
 class ParallelRegionInfo {
+public:
+  /// The container type used to store top level parallel regions.
+  using ParallelRegionVectorTy = SmallVector<ParallelRegion *, 4>;
+
+  /// Iterator types for top level parallel regions.
+  ///
+  ///{
+  using iterator = ParallelRegionVectorTy::iterator;
+  using const_iterator = ParallelRegionVectorTy::const_iterator;
+  ///}
 
+private:
   /// The parallel regions that are not nested in other parallel regions.
-  SmallVector<ParallelRegion *, 4> TopLevelParallelRegions;
+  ParallelRegionVectorTy TopLevelParallelRegions;
 
 public:
   ParallelRegionInfo() {}
   ParallelRegionInfo(Function &F, const DominatorTree &DT) { recalculate(F, DT); }
   ~ParallelRegionInfo() { releaseMemory(); }
 
   /// Type for the mapping of basic blocks to parallel tasks.
   using ParallelTaskMappingTy = DenseMap<BasicBlock *, const ParallelTask *>;
 
   /// Identify the parallel regions in @p F from scratch.
   ///
   /// Note that only the extend of the parallel regions and there nesting is
   /// stored but not the mapping from basic blocks to the surrounding parallel
   /// task that is returned. Only the former information can be preserved easily
   /// by passes but the latter is generally only valid at creation time. For an
   /// on-demand computation of this mapping use the createMapping() function.
   ///
   /// @returns A mapping from basic blocks to the surrounding parallel task.
   ParallelTaskMappingTy recalculate(Function &F, const DominatorTree &DT);
 
   /// Return the top-level parallel regions in this function.
   ///
   ///{
-  SmallVectorImpl<ParallelRegion *> &getTopLevelParallelRegions() {
+  ParallelRegionVectorTy &getTopLevelParallelRegions() {
     return TopLevelParallelRegions;
   }
-  const SmallVectorImpl<ParallelRegion *> &getTopLevelParallelRegions() const {
+  const ParallelRegionVectorTy &getTopLevelParallelRegions() const {
     return TopLevelParallelRegions;
   }
   ///}
 
+  /// Return true if there is no parallel region in this function.
+  bool empty() const { return TopLevelParallelRegions.empty(); }
+
+  /// Remove all cached parallel regions of this function.
+  void clear() { TopLevelParallelRegions.clear(); }
+
+  /// Return the number of top level parallel regions in this function.
+  ParallelRegionVectorTy::size_type size() const {
+    return TopLevelParallelRegions.size();
+  }
+
+  /// Iterators for parallel regions.
+  ///
+  ///{
+  iterator begin() { return TopLevelParallelRegions.begin(); }
+  iterator end() { return TopLevelParallelRegions.end(); }
+  const_iterator begin() const { return TopLevelParallelRegions.begin(); }
+  const_iterator end() const { return TopLevelParallelRegions.end(); }
+  ///}
+
+  /// Return the parallel region that makes @p L a parallel loop, if any.
+  ///
+  /// A parallel loop is a loop that does fork (parts of) its body to a new
+  /// task which are joined only after the loop. It is also ensured that
+  /// everything that is not forked but part of the loop does not cause
+  /// side-effects. These instructions usually compute the exit condition
+  /// and the new values for the induction variable(s).
+  ///
+  /// Note: We allow multiple induction variables and complex exit conditions
+  /// here. However, for some parallel runtimes these have to be simplified,
+  /// e.g. expressed with regards to one canonical induction variable.
+  /// TODO: Provide an interface to perform the simplification.
+  ///
+  /// Note: We allow arbitrary code after the loop but before a join
+  /// instruction.
+  ParallelRegion *getParallelLoopRegion(const Loop &L,
+                                        const DominatorTree &DT) const;
+
+  /// Return true if @p L is a parallel loop.
+  ///
+  /// See ParallelRegionInfo::getParallelLoopRegion for more information.
+  bool isParallelLoop(const Loop &L, const DominatorTree &DT) const;
+
   /// Check for containment in any parallel region.
   ///{
   bool containedInAny(const BasicBlock *BB, const DominatorTree &DT) const;
   bool containedInAny(const Instruction *I, const DominatorTree &DT) const;
   bool containedInAny(const Loop *L, const DominatorTree &DT) const;
   ///}
 
   /// Check for possible containment in any parallel region.
public:
   ///{
   /// @see ParallelTask::visit(VisitorTy &Visitor, bool Recursive)
   bool visit(ParallelTask::VisitorTy &Visitor, bool Recursive = false) const {
     return ForkedTask.visit(Visitor, Recursive) &&
            ContinuationTask.visit(Visitor, Recursive);
   }
   ///}
 
+  /// Return true if this is a parallel loop region for @p L.
+  ///
+  /// See ParallelRegionInfo::getParallelLoopRegion for more information.
+  bool isParallelLoopRegion(const Loop &L, const DominatorTree &DT) const;
+
   /// @see ParallelTask::contains(const BasicBlock *BB, DominatorTree &DT)
   bool contains(const BasicBlock *BB, const DominatorTree &DT) const;
 
   /// @see ParallelTask::contains(const Instruction *I, DominatorTree &DT)
   bool contains(const Instruction *I, const DominatorTree &DT) const;
 
   /// @see ParallelTask::contains(const Loop *L, DominatorTree &DT)
   bool contains(const Loop *L, const DominatorTree &DT) const;

lib/Analysis/ParallelRegionInfo.cpp

 ParallelRegion::~ParallelRegion() {
   DeleteContainerSeconds(ParallelSubRegions);
 }
 
 void ParallelRegion::addParallelSubRegion(ParallelRegion &ParallelSubRegion) {
   ParallelSubRegions[&ParallelSubRegion.getFork()] = &ParallelSubRegion;
 }
 
+bool ParallelRegion::isParallelLoopRegion(const Loop &L,
+                                          const DominatorTree &DT) const {
+  // Bail if the fork is not part of the loop.
+  if (!L.contains(getFork().getParent()))
+    return false;
+
+  // Bail if a join is part of the loop.
+  const auto &JoinsVector = ContinuationTask.getHaltsOrJoints();
+  if (std::any_of(JoinsVector.begin(), JoinsVector.end(),
+                  [&L](TerminatorInst *TI) { return L.contains(TI); }))
+    return false;
+
+  // Now check for possible side effects outside the forked task. First the
+  // header to the fork and then the part of the continuation which is inside
+  // the loop.
+  SmallVector<BasicBlock *, 8> Worklist;
+  BasicBlock *HeaderBB = L.getHeader();
+  Worklist.push_back(HeaderBB);
+
+  while (!Worklist.empty()) {
+    BasicBlock *BB = Worklist.pop_back_val();
+    assert(L.contains(BB));
+
+    for (Instruction &I : *BB)
+      if (I.mayHaveSideEffects())
+        return false;
+
+    if (&getFork() == BB->getTerminator())
+      continue;
+
+    for (BasicBlock *SuccessorBB : successors(BB))
+      if (L.contains(SuccessorBB))
+        Worklist.push_back(SuccessorBB);
+  }
+
+  BasicBlock &ContinuationEntryBB = ContinuationTask.getEntry();
+  assert(L.contains(&ContinuationEntryBB) &&
+         "Fork instructions should not be used to exit a loop!");
+
+  // We do not support nested loops in the continuation part.
+  if (std::any_of(L.begin(), L.end(), [&](Loop *SubLoop) {
+        return ContinuationTask.contains(SubLoop, DT);
+      }))
+    return false;
+
+  // Traverse all blocks that are in the continuation and in the loop and check
+  // for side-effects.
+  assert(Worklist.empty());
+  Worklist.push_back(&ContinuationEntryBB);
+
+  while (!Worklist.empty()) {
+    BasicBlock *BB = Worklist.pop_back_val();
+    assert(L.contains(BB));
+
+    for (Instruction &I : *BB)
+      if (I.mayHaveSideEffects())
+        return false;
+
+    for (BasicBlock *SuccessorBB : successors(BB))
+      if (L.contains(SuccessorBB) && SuccessorBB != HeaderBB)
+        Worklist.push_back(SuccessorBB);
+  }
+
+  return true;
+}
+
 bool ParallelRegion::contains(const BasicBlock *BB,
                               const DominatorTree &DT) const {
   // All contained blocks are dominated by the fork block.
   if (!DT.properlyDominates(Fork.getParent(), BB))
     return false;
 
   // If the above condition is not met we let the sub-tasks handle it.
   return ForkedTask.contains(BB, DT) || ContinuationTask.contains(BB, DT);
for (ParallelRegion *PR : TopLevelParallelRegions) {
     bool Success = PR->visit(Visitor, true);
     (void)Success;
     assert(Success && "Expeded walk over parallel region to visit all blocks!");
   }
 
   return BB2PTMap;
 }
 
+ParallelRegion *ParallelRegionInfo::getParallelLoopRegion(const Loop &L,
+                                        const DominatorTree &DT) const {
+  if (empty())
+    return nullptr;
+
+  SmallVector<ParallelRegion *, 8> ParallelRegions;
+  ParallelRegions.append(begin(), end());
+
+  // The parallel region we are looking for.
+  ParallelRegion *ParallelLoopRegion = nullptr;
+
+  while (!ParallelRegions.empty()) {
+    ParallelRegion *PR = ParallelRegions.pop_back_val();
+
+    for (const auto &It : PR->getSubRegions())
+      ParallelRegions.push_back(It.getSecond());
+
+    if (!PR->isParallelLoopRegion(L, DT))
+      continue;
+
+    assert(!ParallelLoopRegion &&
+           "Broken nesting resulted in multiple parallel loop regions");
+
+    ParallelLoopRegion = PR;
+
+#ifdef NDEBUG
+    // For debug builds we verify that at most one parallel loop region exists,
+    // otherwise we just assume the nesting was intact.
+    break;
+#endif
+  }
+
+  return ParallelLoopRegion;
+}
+
+bool ParallelRegionInfo::isParallelLoop(const Loop &L,
+                                        const DominatorTree &DT) const {
+  // See ParallelRegionInfo::getParallelLoopRegion(...) for more information.
+  return getParallelLoopRegion(L, DT) != nullptr;
+}
+
 bool ParallelRegionInfo::containedInAny(const BasicBlock *BB,
                                         const DominatorTree &DT) const {
   return std::any_of(
       TopLevelParallelRegions.begin(), TopLevelParallelRegions.end(),
       [BB, &DT](ParallelRegion *PR) { return PR->contains(BB, DT); });
 }
 
 bool ParallelRegionInfo::containedInAny(const Instruction *I,

 bool ParallelRegionInfo::maybeContainedInAny(const Loop *L,
                                              const DominatorTree &DT) const {
   // See ParallelTask::contains(const Loop *, DominatorTree &) for reasoning.
   return maybeContainedInAny(L->getHeader(), DT);
 }
 
 bool ParallelRegionInfo::isSafeToPromote(const AllocaInst &AI,
                                          const DominatorTree &DT) const {
-  if (TopLevelParallelRegions.empty())
+  if (empty())
     return true;
 
   // First check if we know that AI is contained in a parallel region.
   ParallelRegion *AIPR = nullptr;
   if (AI.getParent() != &AI.getFunction()->getEntryBlock()) {
     for (ParallelRegion *PR : TopLevelParallelRegions) {
       // In order to not perform a potentially linear contains check we will
       // skip parallel regions that have multi-exit tasks. This is conservative