Fix isRecoverable helper methods to handle precompute relation

Author: weiya711

include/taco/index_notation/provenance_graph.h

@@ -353,9 +353,9 @@ class ProvenanceGraph {
 
   /// Node is recoverable if children appear in defined
   bool isRecoverable(IndexVar indexVar, std::set<IndexVar> defined) const;
-  bool isRecoverableFull(IndexVar indexVar, std::set<IndexVar> defined) const;
-  bool isRecoverableProducer(IndexVar indexVar, std::set<IndexVar> defined) const;
-  bool isRecoverableConsumer(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/provenance_graph.cpp

@@ -1122,62 +1122,35 @@ bool ProvenanceGraph::isRecoverable(taco::IndexVar indexVar, std::set<taco::Inde
   // 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)
-  bool producer = false;
-  bool consumer = false;
+  vector<IndexVar> producers;
+  vector<IndexVar> consumers;
   for (auto& def : defined) {
     if (childRelMap.count(def) && childRelMap.at(def).getRelType() == IndexVarRelType::PRECOMPUTE) {
-      consumer = true;
+      consumers.push_back(def);
     }
     if (parentRelMap.count(def) && parentRelMap.at(def).getRelType() == IndexVarRelType::PRECOMPUTE) {
-      producer = true;
+      producers.push_back(def);
     }
   }
-  if (producer) {
-    return isRecoverableProducer(indexVar, defined);
-  }
-  else if (consumer) {
-    return isRecoverableConsumer(indexVar, defined);
-  } else {
-    return isRecoverableFull(indexVar, defined);
-  }
-}
 
-bool ProvenanceGraph::isRecoverableFull(taco::IndexVar indexVar, std::set<taco::IndexVar> defined) const {
-  // all children are either defined or recoverable from their children
-  // precompute relations are treated as normal relations
-    for (const IndexVar& child : getChildren(indexVar)) {
-      if (!defined.count(child) && (isFullyDerived(child) || !isRecoverableFull(child, defined))) {
-        return false;
-      }
-    }
-  return true;
+  return isRecoverablePrecompute(indexVar, defined, producers, consumers);
 }
 
-bool ProvenanceGraph::isRecoverableConsumer(taco::IndexVar indexVar, std::set<taco::IndexVar> defined) const {
+bool ProvenanceGraph::isRecoverablePrecompute(taco::IndexVar indexVar, std::set<taco::IndexVar> defined, vector<IndexVar> producers, vector<IndexVar> consumers) const {
   vector<IndexVar> childPrecompute;
-  if (childRelMap.count(indexVar) && childRelMap.at(indexVar).getRelType() == IndexVarRelType::PRECOMPUTE) {
-    childPrecompute = childrenMap.at(indexVar);
+  if (std::find(consumers.begin(), consumers.end(), indexVar) != consumers.end()) {
+    return true;
   }
-  for (const IndexVar& child : getChildren(indexVar)) {
-    if (!childPrecompute.empty() && child == childPrecompute[0]) continue;
-    if (!defined.count(child) && (isFullyDerived(child) ||
-    (childRelMap.count(child) && childRelMap.at(child).getRelType() == IndexVarRelType::PRECOMPUTE) || !isRecoverableConsumer(child, defined))) {
-      return false;
+  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);
   }
-  return true;
-}
-
-bool ProvenanceGraph::isRecoverableProducer(taco::IndexVar indexVar, std::set<taco::IndexVar> defined) const {
-  vector<IndexVar> childPrecompute;
   for (const IndexVar& child : getChildren(indexVar)) {
-    if (childRelMap.count(child) && childRelMap.at(child).getRelType() == IndexVarRelType::PRECOMPUTE) {
-      auto precomputeChild = childrenMap.at(child)[0];
-      if (!defined.count(precomputeChild) && (isFullyDerived(precomputeChild) || !isRecoverableProducer(precomputeChild, defined))) {
-        return false;
-      }
-    }
-    else if (!defined.count(child) && (isFullyDerived(child) || !isRecoverableProducer(child, defined))) {
+    if (!defined.count(child) && (isFullyDerived(child)
+                                  || !isRecoverablePrecompute(child, defined, producers, consumers))) {
       return false;
     }
   }

test/tests-workspaces.cpp

@@ -187,46 +187,3 @@ TEST(workspaces, tile_denseMatMul) {
 
 }
 
-TEST(DISABLED_workspaces, multiplePrecomputeIndependentIndexVarsSplit) {
-
-  Tensor<double> A("A", {16}, Format{Dense});
-  Tensor<double> B("B", {16}, Format{Dense});
-  Tensor<double> C("C", {16}, Format{Dense});
-  Tensor<double> D("D", {16}, Format{Dense});
-
-  for (int i = 0; i < 16; i++) {
-    B.insert({i}, (double) i);
-    C.insert({i}, (double) i);
-    D.insert({i}, (double) i);
-  }
-
-  IndexVar i("i");
-  IndexVar iw1("iw1");
-  IndexVar iw2("iw2");
-  IndexVar iw2_outter("iw2_outer");
-  IndexVar iw2_inner("iw2_inner");
-  A(i) = B(i) + C(i) + D(i);
-
-  // Precompute then split iw tensor
-  IndexStmt stmt = A.getAssignment().concretize();
-  TensorVar precomputed1("precomputed1", Type(Float64, {16}), taco::dense);
-  TensorVar precomputed2("precomputed2", Type(Float64, {16}), taco::dense);
-  stmt = stmt.precompute(A.getAssignment().getRhs(), i, iw1, precomputed1);
-  cout << stmt.concretize() << endl;
-  stmt = stmt.precompute(B(iw1)+C(iw1), iw1, iw2, precomputed2);
-             //.split(iw2,iw2_outter, iw2_inner, 8);
-
-  cout << stmt.concretize() << endl;
-  A.compile(stmt.concretize());
-  A.assemble();
-  A.compute();
-
-  Tensor<double> expected("expected", {16}, Format{Dense});
-  expected(i) = B(i) + C(i);
-  expected.compile();
-  expected.assemble();
-  expected.compute();
-
-  ASSERT_TENSOR_EQ(A, expected);
-}
-