add reordering for index stmt with branches

Author: adhithadias

include/taco/index_notation/index_notation.h

@@ -668,6 +668,9 @@ 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;
 
+  /// reorders the index variables in a nested structure with where clauses
+  IndexStmt reorder(std::vector<int> path, 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.

include/taco/index_notation/transformations.h

@@ -67,10 +67,12 @@ class Reorder : public TransformationInterface {
 public:
   Reorder(IndexVar i, IndexVar j);
   Reorder(std::vector<IndexVar> replacePattern);
+  Reorder(std::vector<int> path, std::vector<IndexVar> replacePattern);
 
   IndexVar geti() const;
   IndexVar getj() const;
   const std::vector<IndexVar>& getreplacepattern() const;
+  const std::vector<int>& getpath() const;
 
   /// Apply the reorder optimization to a concrete index statement.  Returns
   /// an undefined statement and a reason if the statement cannot be lowered.

src/index_notation/index_notation.cpp

@@ -1927,6 +1927,17 @@ IndexStmt IndexStmt::reorder(std::vector<IndexVar> reorderedvars) const {
   return transformed;
 }
 
+IndexStmt IndexStmt::reorder(std::vector<int> path, std::vector<IndexVar> reorderedvars) const {
+  string reason;
+  cout << "Index statement path: " << util::join(path) << endl;
+  cout << "Index statement reorderedvars: " << reorderedvars << endl;
+  IndexStmt transformed = Reorder(path, reorderedvars).apply(*this, &reason);
+  if (!transformed.defined()) {
+    taco_uerror << reason;
+  }
+  return transformed;
+}
+
 IndexStmt IndexStmt::mergeby(IndexVar i, MergeStrategy strategy) const {
   string reason;
   IndexStmt transformed = SetMergeStrategy(i, strategy).apply(*this, &reason);

src/index_notation/transformations.cpp

@@ -64,6 +64,7 @@ std::ostream& operator<<(std::ostream& os, const Transformation& t) {
 
 // class Reorder
 struct Reorder::Content {
+  std::vector<int> path;
   std::vector<IndexVar> replacePattern;
   bool pattern_ordered; // In case of Reorder(i, j) need to change replacePattern ordering to actually reorder
 };
@@ -78,6 +79,12 @@ Reorder::Reorder(std::vector<taco::IndexVar> replacePattern) : content(new Conte
   content->pattern_ordered = true;
 }
 
+Reorder::Reorder(std::vector<int> path, std::vector<taco::IndexVar> replacePattern) : content(new Content) {
+  content->path = path;
+  content->replacePattern = replacePattern;
+  content->pattern_ordered = true;
+}
+
 IndexVar Reorder::geti() const {
   return content->replacePattern[0];
 }
@@ -93,13 +100,66 @@ const std::vector<IndexVar>& Reorder::getreplacepattern() const {
   return content->replacePattern;
 }
 
+const std::vector<int>& Reorder::getpath() const {
+  return content->path;
+}
+
 IndexStmt Reorder::apply(IndexStmt stmt, string* reason) const {
   INIT_REASON(reason);
 
   string r;
-  if (!isConcreteNotation(stmt, &r)) {
-    *reason = "The index statement is not valid concrete index notation: " + r;
-    return IndexStmt();
+
+  // TODO - Add a different check for concrete index notation with branching
+  // if (!isConcreteNotation(stmt, &r)) {
+  //   *reason = "The index statement is not valid concrete index notation: " + r;
+  //   return IndexStmt();
+  // }
+
+  IndexStmt originalStmt = stmt;
+  struct ReorderVisitor : public IndexNotationVisitor {
+    using IndexNotationVisitor::visit;
+    vector<int>& path;
+    unsigned int pathIdx = 0;
+    IndexStmt innerStmt;
+
+    ReorderVisitor(vector<int>& path) : path(path) {}
+
+    void visit(const ForallNode* node) {
+      if (pathIdx == path.size()) {
+        innerStmt = IndexStmt(node);
+        return;
+      }
+      IndexNotationVisitor::visit(node);
+    }
+
+    void visit(const WhereNode* node) {
+
+      Where where(node);
+
+      if (pathIdx == path.size()) {
+        innerStmt = IndexStmt(node);
+        return;
+      }
+
+      if (!path[pathIdx]) {
+        pathIdx++;
+        IndexNotationVisitor::visit(node->producer);
+      } else {
+        pathIdx++;
+        IndexNotationVisitor::visit(node->consumer);
+      }
+    }
+  };
+
+  cout << "original statement: " << originalStmt << endl;
+  ReorderVisitor reorderVisitor(content->path);
+
+  auto p = getpath();
+  cout << "path: " << util::join(p) << endl;
+  if (p.size() > 0) {
+    originalStmt.accept(&reorderVisitor);
+    cout << "reordering statment: " << reorderVisitor.innerStmt << endl;
+    stmt = reorderVisitor.innerStmt;
   }
 
   // collect current ordering of IndexVars
@@ -130,7 +190,52 @@ IndexStmt Reorder::apply(IndexStmt stmt, string* reason) const {
     *reason = "The foralls of reorder pattern: " + util::join(getreplacepattern()) + " were not directly nested.";
     return IndexStmt();
   }
-  return ForAllReplace(currentOrdering, getreplacepattern()).apply(stmt, reason);
+
+  cout << "replacePattern: " << util::join(getreplacepattern()) << endl;
+  auto reorderedStmt = ForAllReplace(currentOrdering, getreplacepattern()).apply(stmt, reason);
+
+
+  struct ReorderedRewriter : public IndexNotationRewriter {
+    using IndexNotationRewriter::visit;
+
+    IndexStmt reorderedStmt;
+    vector<int>& path;
+    vector<int> visited;
+
+    ReorderedRewriter(IndexStmt reorderedStmt, vector<int>& path) : reorderedStmt(reorderedStmt), path(path) {}
+
+    void visit(const ForallNode* node) {
+      // at the end of the path, rewrite should happen using the producer and consumer
+      if (visited == path) {
+        stmt = reorderedStmt;
+        return;
+      }
+      IndexNotationRewriter::visit(node);
+    }
+
+    void visit(const WhereNode* node) {
+      Where where(node);
+
+      // add 0 to visited if the producer is visited and 1 if the consumer is visited
+      visited.push_back(0);
+      IndexStmt producer = rewrite(node->producer);
+      visited.pop_back();
+      visited.push_back(1);
+      IndexStmt consumer = rewrite(node->consumer);
+      visited.pop_back();
+      if (producer == node->producer && consumer == node->consumer) {
+        stmt = node;
+      }
+      else {
+        stmt = new WhereNode(consumer, producer);
+      }
+
+    }
+  };
+  ReorderedRewriter reorderedRewriter(reorderedStmt, content->path);
+  stmt = reorderedRewriter.rewrite(originalStmt);  
+
+  return stmt;
 }
 
 void Reorder::print(std::ostream& os) const {
@@ -1068,10 +1173,12 @@ IndexStmt ForAllReplace::apply(IndexStmt stmt, string* reason) const {
   INIT_REASON(reason);
 
   string r;
-  if (!isConcreteNotation(stmt, &r)) {
-    *reason = "The index statement is not valid concrete index notation: " + r;
-    return IndexStmt();
-  }
+
+  // TODO - Add a different check for concrete index notation with branching
+  // if (!isConcreteNotation(stmt, &r)) {
+  //   *reason = "The index statement is not valid concrete index notation: " + r;
+  //   return IndexStmt();
+  // }
 
   /// Since all IndexVars can only appear once, assume replacement will work and error if it doesn't
   struct ForAllReplaceRewriter : public IndexNotationRewriter {

test/tests-workspaces.cpp

@@ -757,6 +757,60 @@ TEST(workspaces, loopcontractfuse) {
   ASSERT_TENSOR_EQ(expected, A);
 }
 
+TEST(workspaces, loopreordercontractfuse) {
+  int N = 16;
+  Tensor<double> A("A", {N, N, N}, Format{Dense, Dense, Dense});
+  Tensor<double> B("B", {N, N, N}, Format{Dense, Sparse, Sparse});
+  Tensor<double> C("C", {N, N}, Format{Dense, Dense});
+  Tensor<double> D("D", {N, N}, Format{Dense, Dense});
+  Tensor<double> E("E", {N, N}, Format{Dense, Dense});
+
+  for (int i = 0; i < N; i++) {
+    for (int j = 0; j < N; j++) {
+      for (int k = 0; k < N; k++) {
+        B.insert({i, j, k}, (double) i);
+      }
+      C.insert({i, j}, (double) j);
+      E.insert({i, j}, (double) i*j);
+      D.insert({i, j}, (double) i*j);
+    }
+  }
+
+  IndexVar i("i"), j("j"), k("k"), l("l"), m("m"), n("n");
+  A(l,m,n) = B(i,j,k) * C(i,l) * D(j,m) * E(k,n);
+
+  IndexStmt stmt = A.getAssignment().concretize();
+
+  std::cout << stmt << endl;
+  vector<int> path1;
+  vector<int> path2 = {1};
+  stmt = stmt
+    .reorder({l,i,m, j, k, n})
+    .loopfuse(2, true, path1)
+    .reorder(path2, {m,k,j,n})
+    .loopfuse(2, true, path2)
+    ;
+  stmt = stmt
+    .parallelize(l, ParallelUnit::CPUThread, OutputRaceStrategy::NoRaces)
+    ;
+
+
+  stmt = stmt.concretize();
+  cout << "final stmt: " << stmt << endl;
+  printCodeToFile("loopreordercontractfuse", stmt);
+
+  A.compile(stmt.concretize());
+  A.assemble();
+  A.compute();
+
+  Tensor<double> expected("expected", {N, N, N}, Format{Dense, Dense, Dense});
+  expected(l,m,n) = B(i,j,k) * C(i,l) * D(j,m) * E(k,n);
+  expected.compile();
+  expected.assemble();
+  expected.compute();
+  ASSERT_TENSOR_EQ(expected, A);
+}
+
 TEST(workspaces, precompute2D_mul) {
   int N = 16;
   Tensor<double> A("A", {N, N}, Format{Dense, Dense});