Merge remote-tracking branch 'origin' into multidim-workspace

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;
 

include/taco/lower/lowerer_impl.h

@@ -1534,15 +1534,18 @@ IndexStmt IndexStmt::precompute(IndexExpr expr, std::vector<IndexVar> i_vars,
   // TODO: need to assert they are same length
   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;
-  //   }
-  // }
+
+ // FIXME: need to re-enable this later
+ // 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_vars, iw_vars, workspace)).apply(transformed, &reason);
+
   if (!transformed.defined()) {
     taco_uerror << reason;
   }

include/taco/lower/lowerer_impl_imperative.h

@@ -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/index_notation/index_notation.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);
     })
   );

src/index_notation/provenance_graph.cpp

@@ -15,6 +15,7 @@
 #include "taco/ir/ir_printer.h"
 
 #include "taco/lower/lowerer_impl.h"
+#include "taco/lower/lowerer_impl_imperative.h"
 #include "taco/lower/iterator.h"
 #include "mode_access.h"
 
@@ -33,7 +34,7 @@ namespace taco {
 
 
 // class Lowerer
-Lowerer::Lowerer() : impl(new LowererImpl()) {
+Lowerer::Lowerer() : impl(new LowererImplImperative()) {
 }
 
 Lowerer::Lowerer(LowererImpl* impl) : impl(impl) {