Merge pull request #474 from tensor-compiler/test-pr470

Modified PR 470: update isRecoverable check and iterators to fix precompute bugs

Author: weiya711

include/taco/index_notation/provenance_graph.h

@@ -354,6 +354,9 @@ class ProvenanceGraph {
   /// Node is recoverable if children appear in defined
   bool isRecoverable(IndexVar indexVar, std::set<IndexVar> defined) const;
 
+  /// isRecoverable helper method to handle precompute relations and where statements in the provenance graph
+  bool isRecoverablePrecompute(IndexVar indexVar, std::set<IndexVar> defined, std::vector<IndexVar> producers, std::vector<IndexVar> consumers) const;
+
   /// Node is recoverable if at most 1 unknown variable in relationship (parents + siblings)
   bool isChildRecoverable(taco::IndexVar indexVar, std::set<taco::IndexVar> defined) const;
 

src/index_notation/index_notation.cpp

@@ -1514,6 +1514,14 @@ IndexStmt IndexStmt::precompute(IndexExpr expr, IndexVar i, IndexVar iw, TensorV
   IndexStmt transformed = *this;
   string reason;
 
+  if (i != iw) {
+    IndexVarRel rel = IndexVarRel(new PrecomputeRelNode(i, iw));
+    transformed = Transformation(AddSuchThatPredicates({rel})).apply(transformed, &reason);
+    if (!transformed.defined()) {
+      taco_uerror << reason;
+    }
+  }
+
   transformed = Transformation(Precompute(expr, i, iw, workspace)).apply(transformed, &reason);
   if (!transformed.defined()) {
     taco_uerror << reason;

src/index_notation/provenance_graph.cpp

@@ -1120,8 +1120,39 @@ bool ProvenanceGraph::isAvailable(IndexVar indexVar, std::set<IndexVar> defined)
 
 bool ProvenanceGraph::isRecoverable(taco::IndexVar indexVar, std::set<taco::IndexVar> defined) const {
   // all children are either defined or recoverable from their children
+  // This checks the definedVars list to determine where in the statement the variables are trying to be
+  // recovered from ( either on the producer or consumer side of a where stmt or not in a where stmt)
+  vector<IndexVar> producers;
+  vector<IndexVar> consumers;
+  for (auto& def : defined) {
+    if (childRelMap.count(def) && childRelMap.at(def).getRelType() == IndexVarRelType::PRECOMPUTE) {
+      consumers.push_back(def);
+    }
+    if (parentRelMap.count(def) && parentRelMap.at(def).getRelType() == IndexVarRelType::PRECOMPUTE) {
+      producers.push_back(def);
+    }
+  }
+
+  return isRecoverablePrecompute(indexVar, defined, producers, consumers);
+}
+
+bool ProvenanceGraph::isRecoverablePrecompute(taco::IndexVar indexVar, std::set<taco::IndexVar> defined,
+                                              vector<IndexVar> producers, vector<IndexVar> consumers) const {
+  vector<IndexVar> childPrecompute;
+  if (std::find(consumers.begin(), consumers.end(), indexVar) != consumers.end()) {
+    return true;
+  }
+  if (!producers.empty() && (childRelMap.count(indexVar) &&
+                             childRelMap.at(indexVar).getRelType() == IndexVarRelType::PRECOMPUTE)) {
+    auto precomputeChild = getChildren(indexVar)[0];
+    if (std::find(producers.begin(), producers.end(), precomputeChild) != producers.end()) {
+      return true;
+    }
+    return isRecoverablePrecompute(precomputeChild, defined, producers, consumers);
+  }
   for (const IndexVar& child : getChildren(indexVar)) {
-    if (!defined.count(child) && (isFullyDerived(child) || !isRecoverable(child, defined))) {
+    if (!defined.count(child) && (isFullyDerived(child) ||
+                                  !isRecoverablePrecompute(child, defined, producers, consumers))) {
       return false;
     }
   }

src/lower/iterator.cpp

@@ -531,6 +531,15 @@ Iterators::Iterators(IndexStmt stmt, const map<TensorVar, Expr>& tensorVars)
           underivedAdded.insert(underived);
         }
       }
+
+      // Insert all children of current index variable into iterators as well
+      for (const IndexVar& child : provGraph.getChildren(n->indexVar)) {
+        if (!underivedAdded.count(child)) {
+          content->modeIterators.insert({child, child});
+          underivedAdded.insert(child);
+        }
+      }
+
       m->match(n->stmt);
     })
   );