Change precompute(...) algorithm to use workspaces paper algorithm

Author: weiya711

include/taco/index_notation/index_notation.h

@@ -551,7 +551,7 @@ class IndexStmt : public util::IntrusivePtr<const IndexStmtNode> {
 
   /// Takes any index notation and concretizes unknowns to make it concrete notation
   /// given a Provenance Graph of indexVars
-  IndexStmt concretizeScheduled(ProvenanceGraph provGraph) const;
+  IndexStmt concretizeScheduled(ProvenanceGraph provGraph, std::vector<IndexVar> forallIndexVarList) const;
 
   /// The \code{split} transformation splits (strip-mines) an index
   /// variable into two nested index variables, where the size of the
@@ -1162,7 +1162,7 @@ IndexStmt makeReductionNotationScheduled(IndexStmt, ProvenanceGraph);
 /// Convert reduction notation to concrete notation, by inserting forall nodes,
 /// replacing reduction nodes by compound assignments, and inserting temporaries
 /// as needed while taking into account a schedule given by the Provenance Graph.
-IndexStmt makeConcreteNotationScheduled(IndexStmt, ProvenanceGraph);
+IndexStmt makeConcreteNotationScheduled(IndexStmt, ProvenanceGraph, std::vector<IndexVar> forallIndexVars);
 
 /// Returns the results of the index statement, in the order they appear.
 std::vector<TensorVar> getResults(IndexStmt stmt);

src/error/error_checks.cpp

@@ -41,7 +41,9 @@ std::pair<bool, string> dimensionsTypecheck(const std::vector<IndexVar>& resultV
   for (size_t mode = 0; mode < resultVars.size(); mode++) {
     IndexVar var = resultVars[mode];
     auto dimension = shape.getDimension(mode);
-    if (util::contains(indexVarDims,var) && indexVarDims.at(var) != dimension) {
+    if (util::contains(indexVarDims,var) && indexVarDims.at(var) != dimension &&
+        !(indexVarDims.at(var).isIndexVarSized() && indexVarDims.at(var).getIndexVarSize() == var) &&
+        !(dimension.isIndexVarSized() && dimension.getIndexVarSize() == var)) {
       errors.push_back(addDimensionError(var, indexVarDims.at(var), dimension));
     } else {
       indexVarDims.insert({var, dimension});
@@ -63,7 +65,9 @@ std::pair<bool, string> dimensionsTypecheck(const std::vector<IndexVar>& resultV
         dimension = Dimension(a.getIndexSet(mode).size());
       }
 
-      if (util::contains(indexVarDims,var) && indexVarDims.at(var) != dimension) {
+      if (util::contains(indexVarDims,var) && indexVarDims.at(var) != dimension &&
+        !(indexVarDims.at(var).isIndexVarSized() && indexVarDims.at(var).getIndexVarSize() == var) &&
+        !(dimension.isIndexVarSized() && dimension.getIndexVarSize() == var)) {
         errors.push_back(addDimensionError(var, indexVarDims.at(var), dimension));
       } else {
         indexVarDims.insert({var, dimension});

src/index_notation/index_notation.cpp

@@ -1463,14 +1463,18 @@ map<IndexVar,Dimension> IndexStmt::getIndexVarDomains() const {
 
 
 
-IndexStmt IndexStmt::concretizeScheduled(ProvenanceGraph provGraph) const {
+IndexStmt IndexStmt::concretizeScheduled(ProvenanceGraph provGraph, vector<IndexVar> forallIndexVarList) const {
   IndexStmt stmt = *this;
   string r;
-  if (isEinsumNotation(stmt)) {
+  cout << "Pre concretized stmt: " << stmt << endl;
+  if (isEinsumNotation(stmt, &r)) {
     stmt = makeReductionNotationScheduled(stmt, provGraph);
+    cout << "Post Reduction Stmt: " << stmt << endl;
   }
+  cout << r << endl;
   if (isReductionNotationScheduled(stmt, provGraph, &r)) {
-    stmt = makeConcreteNotationScheduled(stmt, provGraph);
+    stmt = makeConcreteNotationScheduled(stmt, provGraph, forallIndexVarList);
+    cout << "Post Concretize Stmt: " << stmt << endl;
   }
   return stmt;
 }
@@ -2775,17 +2779,22 @@ IndexStmt makeReductionNotationScheduled(IndexStmt stmt, ProvenanceGraph provGra
   return makeReductionNotationScheduled(to<Assignment>(stmt), provGraph);
 }
 
-IndexStmt makeConcreteNotationScheduled(IndexStmt stmt, ProvenanceGraph provGraph) {
+IndexStmt makeConcreteNotationScheduled(IndexStmt stmt, ProvenanceGraph provGraph, vector<IndexVar> forallIndexVars) {
   std::string reason;
   taco_iassert(isReductionNotationScheduled(stmt, provGraph, &reason))
     << "Not reduction notation: " << stmt << std::endl << reason;
   taco_iassert(isa<Assignment>(stmt));
 
   // Free variables and reductions covering the whole rhs become top level loops
   vector<IndexVar> freeVars = to<Assignment>(stmt).getFreeVars();
+  vector<IndexVar> reductionAndFreeVars;
 
   struct RemoveTopLevelReductions : IndexNotationRewriter {
     using IndexNotationRewriter::visit;
+    vector<IndexVar> forallIndexVars;
+    vector<IndexVar> reductionAndFreeVars;
+
+    RemoveTopLevelReductions(vector<IndexVar> forallIndexVars) : forallIndexVars(forallIndexVars) {}
 
     void visit(const AssignmentNode* node) {
       // Easiest to just walk down the reduction node until we find something
@@ -2800,43 +2809,86 @@ IndexStmt makeConcreteNotationScheduled(IndexStmt stmt, ProvenanceGraph provGrap
 
       if (rhs != node->rhs) {
         stmt = Assignment(node->lhs, rhs, Add());
-        for (auto& i : util::reverse(topLevelReductions)) {
-          stmt = forall(i, stmt);
+        if (forallIndexVars.empty()) {
+          for (auto &i : util::reverse(topLevelReductions)) {
+            stmt = forall(i, stmt);
+          }
+        } else {
+          reductionAndFreeVars.insert(reductionAndFreeVars.end(), topLevelReductions.begin(), topLevelReductions.end());
         }
       }
       else {
         stmt = node;
       }
     }
   };
-  stmt = RemoveTopLevelReductions().rewrite(stmt);
-
+  auto rewriter = RemoveTopLevelReductions(forallIndexVars);
+  stmt = rewriter.rewrite(stmt);
+  reductionAndFreeVars = rewriter.reductionAndFreeVars;
   // This gets the list of indexVars on the rhs of an assignment
   // TODO: check to make sure that we want to get ALL rhs indexVars (not just the upper level)
   vector<IndexVar> rhsVars;
   match(stmt,
         function<void(const AccessNode*, Matcher*)>([&](const AccessNode* op, Matcher* ctx) {
-          rhsVars.insert(rhsVars.end(), op->indexVars.begin(), op->indexVars.end());
+          for (auto &i : op->indexVars) {
+            if (std::find(rhsVars.begin(), rhsVars.end(), i) == rhsVars.end()) {
+              rhsVars.push_back(i);
+            }
+          }
         }),
         function<void(const AssignmentNode*, Matcher*)>([&](const AssignmentNode* op, Matcher* ctx) {
           ctx->match(op->rhs);
         })
   );
 
+  cout << "freeVars: ";
+  for (auto &i : freeVars) {
+    cout << i << ", ";
+  }
+  cout << endl;
+
+  cout << "rhsVars: ";
+  for (auto &i : rhsVars) {
+    cout << i << ", ";
+  }
+  cout << endl;
+
+  cout << "forallIndexVars: ";
+  for (auto &i : forallIndexVars) {
+    cout << i << ", ";
+  }
+  cout << endl;
   // Emit the freeVars as foralls if the freeVars are fully derived
   // else emit the fully derived descendant of the freeVar found in rhsVars
-  for (auto& i : util::reverse(freeVars)) {
-    if (provGraph.isFullyDerived(i))
-      stmt = forall(i, stmt);
-    else {
-      auto derivedVars = provGraph.getFullyDerivedDescendants(i);
-      IndexVar derivedI = *rhsVars.begin();
-      for (auto& derivedVar : derivedVars) {
-        if (std::find(rhsVars.begin(), rhsVars.end(), derivedVar) != rhsVars.end()) {
-          derivedI = derivedVar;
+  if (forallIndexVars.empty()) {
+    for (auto &i : util::reverse(freeVars)) {
+      if (provGraph.isFullyDerived(i))
+        stmt = forall(i, stmt);
+      else {
+        auto derivedVars = provGraph.getFullyDerivedDescendants(i);
+        IndexVar derivedI = *rhsVars.begin();
+        for (auto &derivedVar : derivedVars) {
+          if (std::find(rhsVars.begin(), rhsVars.end(), derivedVar) != rhsVars.end()) {
+            derivedI = derivedVar;
+          }
+        }
+        stmt = forall(derivedI, stmt);
+      }
+    }
+  } else {
+    reductionAndFreeVars.insert(reductionAndFreeVars.end(), freeVars.begin(), freeVars.end());
+    for (auto &i : util::reverse(forallIndexVars)) {
+      if (std::find(reductionAndFreeVars.begin(), reductionAndFreeVars.end(), i) != reductionAndFreeVars.end())
+        stmt = forall(i, stmt);
+      else {
+        auto ancestorVars = provGraph.getUnderivedAncestors(i);
+        IndexVar ancestorI = *reductionAndFreeVars.begin();
+        for (auto &ancestorVar : ancestorVars) {
+          if (std::find(reductionAndFreeVars.begin(), reductionAndFreeVars.end(), ancestorVar) != reductionAndFreeVars.end()) {
+            stmt = forall(i, stmt);
+          }
         }
       }
-      stmt = forall(derivedI, stmt);
     }
   }