Merge pull request #524 from remysucre/master

Intersect skewed sparse iterators by galloping

Author: rohany

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);
     }
   }