Make our data-locality schedule tree transforms externally accessible

Other passes which perform different optimizations might be interested in
also applying data-locality transformations as part of their overall
transformation.

llvm-svn: 245824
diff --git a/polly/lib/Transform/ScheduleOptimizer.cpp b/polly/lib/Transform/ScheduleOptimizer.cpp
index 0296af4..3193e26 100644
--- a/polly/lib/Transform/ScheduleOptimizer.cpp
+++ b/polly/lib/Transform/ScheduleOptimizer.cpp
@@ -160,135 +160,10 @@
                       cl::Hidden, cl::ZeroOrMore, cl::CommaSeparated,
                       cl::cat(PollyCategory));
 
-namespace {
-
-class IslScheduleOptimizer : public ScopPass {
-public:
-  static char ID;
-  explicit IslScheduleOptimizer() : ScopPass(ID) { LastSchedule = nullptr; }
-
-  ~IslScheduleOptimizer() { isl_schedule_free(LastSchedule); }
-
-  bool runOnScop(Scop &S) override;
-  void printScop(raw_ostream &OS, Scop &S) const override;
-  void getAnalysisUsage(AnalysisUsage &AU) const override;
-
-private:
-  isl_schedule *LastSchedule;
-
-  /// @brief Decide if the @p NewSchedule is profitable for @p S.
-  ///
-  /// @param S           The SCoP we optimize.
-  /// @param NewSchedule The new schedule we computed.
-  ///
-  /// @return True, if we believe @p NewSchedule is an improvement for @p S.
-  bool isProfitableSchedule(Scop &S, __isl_keep isl_union_map *NewSchedule);
-
-  /// @brief Tile a schedule node.
-  ///
-  /// @param Node            The node to tile.
-  /// @param Identifier      An name that identifies this kind of tiling and
-  ///                        that is used to mark the tiled loops in the
-  ///                        generated AST.
-  /// @param TileSizes       A vector of tile sizes that should be used for
-  ///                        tiling.
-  /// @param DefaultTileSize A default tile size that is used for dimensions
-  ///                        that are not covered by the TileSizes vector.
-  static __isl_give isl_schedule_node *
-  tileNode(__isl_take isl_schedule_node *Node, const char *Identifier,
-           ArrayRef<int> TileSizes, int DefaultTileSize);
-
-  /// @brief Check if this node is a band node we want to tile.
-  ///
-  /// We look for innermost band nodes where individual dimensions are marked as
-  /// permutable.
-  ///
-  /// @param Node The node to check.
-  static bool isTileableBandNode(__isl_keep isl_schedule_node *Node);
-
-  /// @brief Pre-vectorizes one scheduling dimension of a schedule band.
-  ///
-  /// prevectSchedBand splits out the dimension DimToVectorize, tiles it and
-  /// sinks the resulting point loop.
-  ///
-  /// Example (DimToVectorize=0, VectorWidth=4):
-  ///
-  /// | Before transformation:
-  /// |
-  /// | A[i,j] -> [i,j]
-  /// |
-  /// | for (i = 0; i < 128; i++)
-  /// |    for (j = 0; j < 128; j++)
-  /// |      A(i,j);
-  ///
-  /// | After transformation:
-  /// |
-  /// | for (it = 0; it < 32; it+=1)
-  /// |    for (j = 0; j < 128; j++)
-  /// |      for (ip = 0; ip <= 3; ip++)
-  /// |        A(4 * it + ip,j);
-  ///
-  /// The goal of this transformation is to create a trivially vectorizable
-  /// loop.  This means a parallel loop at the innermost level that has a
-  /// constant number of iterations corresponding to the target vector width.
-  ///
-  /// This transformation creates a loop at the innermost level. The loop has
-  /// a constant number of iterations, if the number of loop iterations at
-  /// DimToVectorize can be divided by VectorWidth. The default VectorWidth is
-  /// currently constant and not yet target specific. This function does not
-  /// reason about parallelism.
-  static __isl_give isl_schedule_node *
-  prevectSchedBand(__isl_take isl_schedule_node *Node, unsigned DimToVectorize,
-                   int VectorWidth);
-
-  /// @brief Apply additional optimizations on the bands in the schedule tree.
-  ///
-  /// We are looking for an innermost band node and apply the following
-  /// transformations:
-  ///
-  ///  - Tile the band
-  ///      - if the band is tileable
-  ///      - if the band has more than one loop dimension
-  ///
-  ///  - Prevectorize the schedule of the band (or the point loop in case of
-  ///    tiling).
-  ///      - if vectorization is enabled
-  ///
-  /// @param Node The schedule node to (possibly) optimize.
-  /// @param User A pointer to forward some use information (currently unused).
-  static isl_schedule_node *optimizeBand(isl_schedule_node *Node, void *User);
-
-  /// @brief Apply post-scheduling transformations.
-  ///
-  /// This function applies a set of additional local transformations on the
-  /// schedule tree as it computed by the isl scheduler. Local transformations
-  /// applied include:
-  ///
-  ///   - Tiling
-  ///   - Prevectorization
-  ///
-  /// @param Schedule The schedule object post-transformations will be applied
-  ///                 on.
-  /// @returns        The transformed schedule.
-  static __isl_give isl_schedule *
-  addPostTransforms(__isl_take isl_schedule *Schedule);
-
-  using llvm::Pass::doFinalization;
-
-  virtual bool doFinalization() override {
-    isl_schedule_free(LastSchedule);
-    LastSchedule = nullptr;
-    return true;
-  }
-};
-}
-
-char IslScheduleOptimizer::ID = 0;
-
 __isl_give isl_schedule_node *
-IslScheduleOptimizer::prevectSchedBand(__isl_take isl_schedule_node *Node,
-                                       unsigned DimToVectorize,
-                                       int VectorWidth) {
+ScheduleTreeOptimizer::prevectSchedBand(__isl_take isl_schedule_node *Node,
+                                        unsigned DimToVectorize,
+                                        int VectorWidth) {
   assert(isl_schedule_node_get_type(Node) == isl_schedule_node_band);
 
   auto Space = isl_schedule_node_band_get_space(Node);
@@ -319,9 +194,9 @@
 }
 
 __isl_give isl_schedule_node *
-IslScheduleOptimizer::tileNode(__isl_take isl_schedule_node *Node,
-                               const char *Identifier, ArrayRef<int> TileSizes,
-                               int DefaultTileSize) {
+ScheduleTreeOptimizer::tileNode(__isl_take isl_schedule_node *Node,
+                                const char *Identifier, ArrayRef<int> TileSizes,
+                                int DefaultTileSize) {
   auto Ctx = isl_schedule_node_get_ctx(Node);
   auto Space = isl_schedule_node_band_get_space(Node);
   auto Dims = isl_space_dim(Space, isl_dim_set);
@@ -346,7 +221,7 @@
   return Node;
 }
 
-bool IslScheduleOptimizer::isTileableBandNode(
+bool ScheduleTreeOptimizer::isTileableBandNode(
     __isl_keep isl_schedule_node *Node) {
   if (isl_schedule_node_get_type(Node) != isl_schedule_node_band)
     return false;
@@ -375,8 +250,8 @@
 }
 
 __isl_give isl_schedule_node *
-IslScheduleOptimizer::optimizeBand(__isl_take isl_schedule_node *Node,
-                                   void *User) {
+ScheduleTreeOptimizer::optimizeBand(__isl_take isl_schedule_node *Node,
+                                    void *User) {
   if (!isTileableBandNode(Node))
     return Node;
 
@@ -405,7 +280,7 @@
 
   for (int i = Dims - 1; i >= 0; i--)
     if (isl_schedule_node_band_member_get_coincident(Node, i)) {
-      Node = IslScheduleOptimizer::prevectSchedBand(Node, i, PrevectorWidth);
+      Node = prevectSchedBand(Node, i, PrevectorWidth);
       break;
     }
 
@@ -413,17 +288,22 @@
 }
 
 __isl_give isl_schedule *
-IslScheduleOptimizer::addPostTransforms(__isl_take isl_schedule *Schedule) {
+ScheduleTreeOptimizer::optimizeSchedule(__isl_take isl_schedule *Schedule) {
   isl_schedule_node *Root = isl_schedule_get_root(Schedule);
+  Root = optimizeScheduleNode(Root);
   isl_schedule_free(Schedule);
-  Root = isl_schedule_node_map_descendant_bottom_up(
-      Root, IslScheduleOptimizer::optimizeBand, NULL);
   auto S = isl_schedule_node_get_schedule(Root);
   isl_schedule_node_free(Root);
   return S;
 }
 
-bool IslScheduleOptimizer::isProfitableSchedule(
+__isl_give isl_schedule_node *ScheduleTreeOptimizer::optimizeScheduleNode(
+    __isl_take isl_schedule_node *Node) {
+  Node = isl_schedule_node_map_descendant_bottom_up(Node, optimizeBand, NULL);
+  return Node;
+}
+
+bool ScheduleTreeOptimizer::isProfitableSchedule(
     Scop &S, __isl_keep isl_union_map *NewSchedule) {
   // To understand if the schedule has been optimized we check if the schedule
   // has changed at all.
@@ -440,6 +320,33 @@
   return changed;
 }
 
+namespace {
+class IslScheduleOptimizer : public ScopPass {
+public:
+  static char ID;
+  explicit IslScheduleOptimizer() : ScopPass(ID) { LastSchedule = nullptr; }
+
+  ~IslScheduleOptimizer() { isl_schedule_free(LastSchedule); }
+
+  bool runOnScop(Scop &S) override;
+  void printScop(raw_ostream &OS, Scop &S) const override;
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
+
+private:
+  isl_schedule *LastSchedule;
+
+  using llvm::Pass::doFinalization;
+
+  virtual bool doFinalization() override {
+    isl_schedule_free(LastSchedule);
+    LastSchedule = nullptr;
+    return true;
+  }
+};
+}
+
+char IslScheduleOptimizer::ID = 0;
+
 bool IslScheduleOptimizer::runOnScop(Scop &S) {
 
   // Skip empty SCoPs but still allow code generation as it will delete the
@@ -562,10 +469,10 @@
     isl_printer_free(P);
   });
 
-  isl_schedule *NewSchedule = addPostTransforms(Schedule);
+  isl_schedule *NewSchedule = ScheduleTreeOptimizer::optimizeSchedule(Schedule);
   isl_union_map *NewScheduleMap = isl_schedule_get_map(NewSchedule);
 
-  if (!isProfitableSchedule(S, NewScheduleMap)) {
+  if (!ScheduleTreeOptimizer::isProfitableSchedule(S, NewScheduleMap)) {
     isl_union_map_free(NewScheduleMap);
     isl_schedule_free(NewSchedule);
     return false;