Add in fixes for isRecoverable() based on if there is a where stmt

Author: weiya711

include/taco/index_notation/provenance_graph.h

@@ -353,6 +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;
 
   /// 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

@@ -2366,8 +2366,12 @@ bool isConcreteNotation(IndexStmt stmt, std::string* reason) {
       for (auto& var : op->indexVars) {
         // non underived variables may appear in temporaries, but we don't check these
         if (!boundVars.contains(var) && provGraph.isUnderived(var) && (provGraph.isFullyDerived(var) || !provGraph.isRecoverable(var, definedVars))) {
+          string string2 = "definedvars: ";
+          for (auto& d : definedVars)
+            string2.append(d.getName() + ", ");
+
           *reason = "all variables in concrete notation must be bound by a "
-                    "forall statement";
+                    "forall statement" + var.getName() + string2;
           isConcrete = false;
         }
       }

src/index_notation/provenance_graph.cpp

@@ -1120,13 +1120,66 @@ 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
-  // precompute relations are always recoverable since their children never appear in the same loop
-  if (!(childRelMap.count(indexVar) && childRelMap.at(indexVar).getRelType() == IndexVarRelType::PRECOMPUTE)) {
+  // 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;
+  for (auto& def : defined) {
+    if (childRelMap.count(def) && childRelMap.at(def).getRelType() == IndexVarRelType::PRECOMPUTE) {
+      consumer = true;
+    }
+    if (parentRelMap.count(def) && parentRelMap.at(def).getRelType() == IndexVarRelType::PRECOMPUTE) {
+      producer = true;
+    }
+  }
+  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) || !isRecoverable(child, defined))) {
+      if (!defined.count(child) && (isFullyDerived(child) || !isRecoverableFull(child, defined))) {
+        return false;
+      }
+    }
+  return true;
+}
+
+bool ProvenanceGraph::isRecoverableConsumer(taco::IndexVar indexVar, std::set<taco::IndexVar> defined) const {
+  vector<IndexVar> childPrecompute;
+  if (childRelMap.count(indexVar) && childRelMap.at(indexVar).getRelType() == IndexVarRelType::PRECOMPUTE) {
+    childPrecompute = childrenMap.at(indexVar);
+  }
+  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;
+    }
+  }
+  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))) {
+      return false;
+    }
   }
   return true;
 }
@@ -1289,6 +1342,12 @@ bool ProvenanceGraph::hasExactBound(IndexVar indexVar) const {
   {
     return rel.getNode<BoundRelNode>()->getBoundType() == BoundType::MaxExact;
   }
+//  else if (rel.getRelType() == SPLIT)
+//  {
+//    return rel.getNode<SplitRelNode>()->getInnerVar() == indexVar;
+//  } else if (rel.getRelType() == PRECOMPUTE) {
+//    return hasExactBound(rel.getNode<PrecomputeRelNode>()->getParentVar());
+//  }
   // TODO: include non-irregular variables
   return false;
 }

src/lower/lowerer_impl_imperative.cpp

@@ -840,7 +840,6 @@ Stmt LowererImplImperative::lowerForall(Forall forall)
 
 Stmt LowererImplImperative::lowerForallCloned(Forall forall) {
   // want to emit guards outside of loop to prevent unstructured loop exits
-
   // construct guard
   // underived or pos variables that have a descendant that has not been defined yet
   vector<IndexVar> varsWithGuard;
@@ -858,7 +857,6 @@ Stmt LowererImplImperative::lowerForallCloned(Forall forall) {
       }
     }
   }
-
   // determine min and max values for vars given already defined variables.
   // we do a recovery where we fill in undefined variables with either 0's or the max of their iteration
   std::map<IndexVar, Expr> minVarValues;
@@ -903,7 +901,6 @@ Stmt LowererImplImperative::lowerForallCloned(Forall forall) {
     minVarValues[var] = provGraph.recoverVariable(var, currentDefinedVarOrder, underivedBounds, minChildValues, iterators);
     maxVarValues[var] = provGraph.recoverVariable(var, currentDefinedVarOrder, underivedBounds, maxChildValues, iterators);
   }
-
   // Build guards
   Expr guardCondition;
   for (auto var : varsWithGuard) {

test/tests-workspaces.cpp

@@ -187,3 +187,46 @@ 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);
+}
+