Add mergeby scheduling directive.

This commit addes a new scheduling directive called mergeby. This directive specifies if the iterators of a given variable is merged by Two Finger merge or Galloping. The default strategy is Two Finger merge which is the same as the old behavior. Galloping merges the iterators with exponential search, which can be more efficient if the iterator sizes are skewed.

Author: remysucre

include/taco/index_notation/index_notation.h

@@ -634,6 +634,23 @@ class IndexStmt : public util::IntrusivePtr<const IndexStmtNode> {
   /// reorder takes a new ordering for a set of index variables that are directly nested in the iteration order
   IndexStmt reorder(std::vector<IndexVar> reorderedvars) const;
 
+  /// The mergeby transformation specifies how to merge iterators on
+  /// the given index variable. By default, if an iterator is used for windowing
+  /// it will be merged with the "gallop" strategy.
+  /// All other iterators are merged with the "two finger" strategy.
+  /// The two finger strategy merges by advancing each iterator one at a time, 
+  /// while the gallop strategy implements the exponential search algorithm.
+  /// 
+  /// Preconditions:
+  /// This command applies to variables involving sparse iterators only;
+  /// it is a no-op if the variable invovles any dense iterators.
+  /// Any variable can be merged with the two finger strategy, whereas gallop
+  /// only applies to a variable if its merge lattice has a single point 
+  /// (i.e. an intersection). For example, if a variable involves multiplications
+  /// only, it can be merged with gallop.
+  /// Furthermore, all iterators must be ordered for gallop to apply.
+  IndexStmt mergeby(IndexVar i, MergeStrategy strategy) const;
+
   /// The parallelize
   /// transformation tags an index variable for parallel execution.  The
   /// transformation takes as an argument the type of parallel hardware
@@ -829,13 +846,14 @@ class Forall : public IndexStmt {
   Forall() = default;
   Forall(const ForallNode*);
   Forall(IndexVar indexVar, IndexStmt stmt);
-  Forall(IndexVar indexVar, IndexStmt stmt, ParallelUnit parallel_unit, OutputRaceStrategy output_race_strategy, size_t unrollFactor = 0);
+  Forall(IndexVar indexVar, IndexStmt stmt, MergeStrategy merge_strategy, ParallelUnit parallel_unit, OutputRaceStrategy output_race_strategy, size_t unrollFactor = 0);
 
   IndexVar getIndexVar() const;
   IndexStmt getStmt() const;
 
   ParallelUnit getParallelUnit() const;
   OutputRaceStrategy getOutputRaceStrategy() const;
+  MergeStrategy getMergeStrategy() const;
 
   size_t getUnrollFactor() const;
 
@@ -844,7 +862,7 @@ class Forall : public IndexStmt {
 
 /// Create a forall index statement.
 Forall forall(IndexVar i, IndexStmt stmt);
-Forall forall(IndexVar i, IndexStmt stmt, ParallelUnit parallel_unit, OutputRaceStrategy output_race_strategy, size_t unrollFactor = 0);
+Forall forall(IndexVar i, IndexStmt stmt, MergeStrategy merge_strategy, ParallelUnit parallel_unit, OutputRaceStrategy output_race_strategy, size_t unrollFactor = 0);
 
 
 /// A where statment has a producer statement that binds a tensor variable in

include/taco/index_notation/index_notation_nodes.h

@@ -398,15 +398,16 @@ struct YieldNode : public IndexStmtNode {
 };
 
 struct ForallNode : public IndexStmtNode {
-  ForallNode(IndexVar indexVar, IndexStmt stmt, ParallelUnit parallel_unit, OutputRaceStrategy  output_race_strategy, size_t unrollFactor = 0)
-      : indexVar(indexVar), stmt(stmt), parallel_unit(parallel_unit), output_race_strategy(output_race_strategy), unrollFactor(unrollFactor) {}
+  ForallNode(IndexVar indexVar, IndexStmt stmt, MergeStrategy merge_strategy, ParallelUnit parallel_unit, OutputRaceStrategy  output_race_strategy, size_t unrollFactor = 0)
+      : indexVar(indexVar), stmt(stmt), merge_strategy(merge_strategy), parallel_unit(parallel_unit), output_race_strategy(output_race_strategy), unrollFactor(unrollFactor) {}
 
   void accept(IndexStmtVisitorStrict* v) const {
     v->visit(this);
   }
 
   IndexVar indexVar;
   IndexStmt stmt;
+  MergeStrategy merge_strategy;
   ParallelUnit parallel_unit;
   OutputRaceStrategy  output_race_strategy;
   size_t unrollFactor = 0;

include/taco/index_notation/transformations.h

@@ -22,6 +22,7 @@ class AddSuchThatPredicates;
 class Parallelize;
 class TopoReorder;
 class SetAssembleStrategy;
+class SetMergeStrategy;
 
 /// A transformation is an optimization that transforms a statement in the
 /// concrete index notation into a new statement that computes the same result
@@ -36,6 +37,7 @@ class Transformation {
   Transformation(TopoReorder);
   Transformation(AddSuchThatPredicates);
   Transformation(SetAssembleStrategy);
+  Transformation(SetMergeStrategy);
 
   IndexStmt apply(IndexStmt stmt, std::string *reason = nullptr) const;
 
@@ -206,6 +208,25 @@ class SetAssembleStrategy : public TransformationInterface {
 /// Print a SetAssembleStrategy command.
 std::ostream &operator<<(std::ostream &, const SetAssembleStrategy&);
 
+class SetMergeStrategy : public TransformationInterface {
+public:
+  SetMergeStrategy(IndexVar i, MergeStrategy strategy);
+
+  IndexVar geti() const;
+  MergeStrategy getMergeStrategy() const;
+
+  IndexStmt apply(IndexStmt stmt, std::string *reason = nullptr) const;
+
+  void print(std::ostream &os) const;
+
+private:
+  struct Content;
+  std::shared_ptr<Content> content;
+};
+
+/// Print a SetMergeStrategy command.
+std::ostream &operator<<(std::ostream &, const SetMergeStrategy&);
+
 // Autoscheduling functions
 
 /**

include/taco/ir_tags.h

@@ -33,6 +33,13 @@ enum class AssembleStrategy {
 };
 extern const char *AssembleStrategy_NAMES[];
 
+/// MergeStrategy::TwoFinger merges iterators by incrementing one at a time
+/// MergeStrategy::Galloping merges iterators by exponential search (galloping)
+enum class MergeStrategy {
+  TwoFinger, Gallop
+};
+extern const char *MergeStrategy_NAMES[];
+
 }
 
 #endif //TACO_IR_TAGS_H

include/taco/lower/lowerer_impl_imperative.h

@@ -146,15 +146,18 @@ class LowererImplImperative : public LowererImpl {
      * \param statement
      *      A concrete index notation statement to compute at the points in the
      *      sparse iteration space described by the merge lattice.
+     * \param mergeStrategy
+     *      A strategy for merging iterators. One of TwoFinger or Gallop.
      *
      * \return
      *       IR code to compute the forall loop.
      */
   virtual ir::Stmt lowerMergeLattice(MergeLattice lattice, IndexVar coordinateVar,
                                      IndexStmt statement, 
-                                     const std::set<Access>& reducedAccesses);
+                                     const std::set<Access>& reducedAccesses, 
+                                     MergeStrategy mergeStrategy);
 
-  virtual ir::Stmt resolveCoordinate(std::vector<Iterator> mergers, ir::Expr coordinate, bool emitVarDecl);
+  virtual ir::Stmt resolveCoordinate(std::vector<Iterator> mergers, ir::Expr coordinate, bool emitVarDecl, bool mergeWithMax);
 
     /**
      * Lower the merge point at the top of the given lattice to code that iterates
@@ -169,23 +172,29 @@ class LowererImplImperative : public LowererImpl {
      *      coordinate the merge point is at.
      *      A concrete index notation statement to compute at the points in the
      *      sparse iteration space region described by the merge point.
+     * \param mergeWithMax
+     *      A boolean indicating whether coordinates should be combined with MAX instead of MIN.
+     *      MAX is needed when the iterators are merged with the Gallop strategy.
      */
   virtual ir::Stmt lowerMergePoint(MergeLattice pointLattice,
                                    ir::Expr coordinate, IndexVar coordinateVar, IndexStmt statement,
-                                   const std::set<Access>& reducedAccesses, bool resolvedCoordDeclared);
+                                   const std::set<Access>& reducedAccesses, bool resolvedCoordDeclared, 
+                                   MergeStrategy mergestrategy);
 
   /// Lower a merge lattice to cases.
   virtual ir::Stmt lowerMergeCases(ir::Expr coordinate, IndexVar coordinateVar, IndexStmt stmt,
                                    MergeLattice lattice,
-                                   const std::set<Access>& reducedAccesses);
+                                   const std::set<Access>& reducedAccesses, 
+                                   MergeStrategy mergeStrategy);
 
   /// Lower a forall loop body.
   virtual ir::Stmt lowerForallBody(ir::Expr coordinate, IndexStmt stmt,
                                    std::vector<Iterator> locaters,
                                    std::vector<Iterator> inserters,
                                    std::vector<Iterator> appenders,
                                    MergeLattice caseLattice,
-                                   const std::set<Access>& reducedAccesses);
+                                   const std::set<Access>& reducedAccesses, 
+                                   MergeStrategy mergeStrategy);
 
 
   /// Lower a where statement.
@@ -375,7 +384,7 @@ class LowererImplImperative : public LowererImpl {
 
   /// Conditionally increment iterator position variables.
   ir::Stmt codeToIncIteratorVars(ir::Expr coordinate, IndexVar coordinateVar,
-          std::vector<Iterator> iterators, std::vector<Iterator> mergers);
+          std::vector<Iterator> iterators, std::vector<Iterator> mergers, MergeStrategy strategy);
 
   ir::Stmt codeToLoadCoordinatesFromPosIterators(std::vector<Iterator> iterators, bool declVars);
 
@@ -410,7 +419,8 @@ class LowererImplImperative : public LowererImpl {
   /// Lowers a merge lattice to cases assuming there are no more loops to be emitted in stmt.
   /// Will emit checks for explicit zeros for each mode iterator and each locator in the lattice.
   ir::Stmt lowerMergeCasesWithExplicitZeroChecks(ir::Expr coordinate, IndexVar coordinateVar, IndexStmt stmt,
-                                                 MergeLattice lattice, const std::set<Access>& reducedAccesses);
+                                                 MergeLattice lattice, const std::set<Access>& reducedAccesses, 
+                                                 MergeStrategy mergeStrategy);
 
   /// Constructs cases comparing the coordVar for each iterator to the resolved coordinate.
   /// Returns a vector where coordComparisons[i] corresponds to a case for iters[i]
@@ -444,7 +454,7 @@ class LowererImplImperative : public LowererImpl {
   /// The map must be of iterators to exprs of boolean types
   std::vector<ir::Stmt> lowerCasesFromMap(std::map<Iterator, ir::Expr> iteratorToCondition,
                                           ir::Expr coordinate, IndexStmt stmt, const MergeLattice& lattice,
-                                          const std::set<Access>& reducedAccesses);
+                                          const std::set<Access>& reducedAccesses, MergeStrategy mergeStrategy);
 
   /// Constructs an expression which checks if this access is "zero"
   ir::Expr constructCheckForAccessZero(Access);

src/codegen/codegen_c.cpp

@@ -62,6 +62,28 @@ const string cHeaders =
   "int cmp(const void *a, const void *b) {\n"
   "  return *((const int*)a) - *((const int*)b);\n"
   "}\n"
+  // Increment arrayStart until array[arrayStart] >= target or arrayStart >= arrayEnd
+  // using an exponential search algorithm: https://en.wikipedia.org/wiki/Exponential_search.
+  "int taco_gallop(int *array, int arrayStart, int arrayEnd, int target) {\n"
+  "  if (array[arrayStart] >= target || arrayStart >= arrayEnd) {\n"
+  "    return arrayStart;\n"
+  "  }\n"
+  "  int step = 1;\n"
+  "  int curr = arrayStart;\n"
+  "  while (curr + step < arrayEnd && array[curr + step] < target) {\n"
+  "    curr += step;\n"
+  "    step = step * 2;\n"
+  "  }\n"
+  "\n"
+  "  step = step / 2;\n"
+  "  while (step > 0) {\n"
+  "    if (curr + step < arrayEnd && array[curr + step] < target) {\n"
+  "      curr += step;\n"
+  "    }\n"
+  "    step = step / 2;\n"
+  "  }\n"
+  "  return curr+1;\n"
+  "}\n"
   "int taco_binarySearchAfter(int *array, int arrayStart, int arrayEnd, int target) {\n"
   "  if (array[arrayStart] >= target) {\n"
   "    return arrayStart;\n"

src/index_notation/index_notation.cpp

@@ -1907,6 +1907,15 @@ IndexStmt IndexStmt::reorder(std::vector<IndexVar> reorderedvars) const {
   return transformed;
 }
 
+IndexStmt IndexStmt::mergeby(IndexVar i, MergeStrategy strategy) const {
+  string reason;
+  IndexStmt transformed = SetMergeStrategy(i, strategy).apply(*this, &reason);
+  if (!transformed.defined()) {
+    taco_uerror << reason;
+  }
+  return transformed;
+}
+
 IndexStmt IndexStmt::parallelize(IndexVar i, ParallelUnit parallel_unit, OutputRaceStrategy output_race_strategy) const {
   string reason;
   IndexStmt transformed = Parallelize(i, parallel_unit, output_race_strategy).apply(*this, &reason);
@@ -2017,7 +2026,7 @@ IndexStmt IndexStmt::unroll(IndexVar i, size_t unrollFactor) const {
 
     void visit(const ForallNode* node) {
       if (node->indexVar == i) {
-        stmt = Forall(i, rewrite(node->stmt), node->parallel_unit, node->output_race_strategy, unrollFactor);
+        stmt = Forall(i, rewrite(node->stmt), node->merge_strategy, node->parallel_unit, node->output_race_strategy, unrollFactor);
       }
       else {
         IndexNotationRewriter::visit(node);
@@ -2125,11 +2134,11 @@ Forall::Forall(const ForallNode* n) : IndexStmt(n) {
 }
 
 Forall::Forall(IndexVar indexVar, IndexStmt stmt)
-    : Forall(indexVar, stmt, ParallelUnit::NotParallel, OutputRaceStrategy::IgnoreRaces) {
+    : Forall(indexVar, stmt, MergeStrategy::TwoFinger, ParallelUnit::NotParallel, OutputRaceStrategy::IgnoreRaces) {
 }
 
-Forall::Forall(IndexVar indexVar, IndexStmt stmt, ParallelUnit parallel_unit, OutputRaceStrategy output_race_strategy, size_t unrollFactor)
-        : Forall(new ForallNode(indexVar, stmt, parallel_unit, output_race_strategy, unrollFactor)) {
+Forall::Forall(IndexVar indexVar, IndexStmt stmt, MergeStrategy merge_strategy, ParallelUnit parallel_unit, OutputRaceStrategy output_race_strategy, size_t unrollFactor)
+        : Forall(new ForallNode(indexVar, stmt, merge_strategy, parallel_unit, output_race_strategy, unrollFactor)) {
 }
 
 IndexVar Forall::getIndexVar() const {
@@ -2148,6 +2157,10 @@ OutputRaceStrategy Forall::getOutputRaceStrategy() const {
   return getNode(*this)->output_race_strategy;
 }
 
+MergeStrategy Forall::getMergeStrategy() const {
+  return getNode(*this)->merge_strategy;
+}
+
 size_t Forall::getUnrollFactor() const {
   return getNode(*this)->unrollFactor;
 }
@@ -2156,8 +2169,8 @@ Forall forall(IndexVar i, IndexStmt stmt) {
   return Forall(i, stmt);
 }
 
-Forall forall(IndexVar i, IndexStmt stmt, ParallelUnit parallel_unit, OutputRaceStrategy output_race_strategy, size_t unrollFactor) {
-  return Forall(i, stmt, parallel_unit, output_race_strategy, unrollFactor);
+Forall forall(IndexVar i, IndexStmt stmt, MergeStrategy merge_strategy, ParallelUnit parallel_unit, OutputRaceStrategy output_race_strategy, size_t unrollFactor) {
+  return Forall(i, stmt, merge_strategy, parallel_unit, output_race_strategy, unrollFactor);
 }
 
 template <> bool isa<Forall>(IndexStmt s) {
@@ -3938,7 +3951,7 @@ struct Zero : public IndexNotationRewriterStrict {
       stmt = op;
     }
     else {
-      stmt = new ForallNode(op->indexVar, body, op->parallel_unit, op->output_race_strategy, op->unrollFactor);
+      stmt = new ForallNode(op->indexVar, body, op->merge_strategy, op->parallel_unit, op->output_race_strategy, op->unrollFactor);
     }
   }
 

src/index_notation/index_notation_rewriter.cpp

@@ -185,7 +185,7 @@ void IndexNotationRewriter::visit(const ForallNode* op) {
     stmt = op;
   }
   else {
-    stmt = new ForallNode(op->indexVar, s, op->parallel_unit, op->output_race_strategy, op->unrollFactor);
+    stmt = new ForallNode(op->indexVar, s, op->merge_strategy, op->parallel_unit, op->output_race_strategy, op->unrollFactor);
   }
 }
 
@@ -406,7 +406,7 @@ struct ReplaceIndexVars : public IndexNotationRewriter {
       stmt = op;
     }
     else {
-      stmt = new ForallNode(iv, s, op->parallel_unit, op->output_race_strategy, 
+      stmt = new ForallNode(iv, s, op->merge_strategy, op->parallel_unit, op->output_race_strategy, 
                             op->unrollFactor);
     }
   }