Intial commit for new fix

Author: nirvikbaruah

.vscode/settings.json

@@ -0,0 +1,69 @@
+{
+    "files.associations": {
+        "array": "cpp",
+        "atomic": "cpp",
+        "bit": "cpp",
+        "*.tcc": "cpp",
+        "cctype": "cpp",
+        "chrono": "cpp",
+        "clocale": "cpp",
+        "cmath": "cpp",
+        "compare": "cpp",
+        "complex": "cpp",
+        "concepts": "cpp",
+        "condition_variable": "cpp",
+        "cstdarg": "cpp",
+        "cstddef": "cpp",
+        "cstdint": "cpp",
+        "cstdio": "cpp",
+        "cstdlib": "cpp",
+        "cstring": "cpp",
+        "ctime": "cpp",
+        "cwchar": "cpp",
+        "cwctype": "cpp",
+        "deque": "cpp",
+        "forward_list": "cpp",
+        "list": "cpp",
+        "map": "cpp",
+        "set": "cpp",
+        "unordered_map": "cpp",
+        "unordered_set": "cpp",
+        "vector": "cpp",
+        "exception": "cpp",
+        "algorithm": "cpp",
+        "functional": "cpp",
+        "iterator": "cpp",
+        "memory": "cpp",
+        "memory_resource": "cpp",
+        "numeric": "cpp",
+        "optional": "cpp",
+        "random": "cpp",
+        "ratio": "cpp",
+        "string": "cpp",
+        "string_view": "cpp",
+        "system_error": "cpp",
+        "tuple": "cpp",
+        "type_traits": "cpp",
+        "utility": "cpp",
+        "fstream": "cpp",
+        "initializer_list": "cpp",
+        "iomanip": "cpp",
+        "iosfwd": "cpp",
+        "iostream": "cpp",
+        "istream": "cpp",
+        "limits": "cpp",
+        "mutex": "cpp",
+        "new": "cpp",
+        "ostream": "cpp",
+        "ranges": "cpp",
+        "sstream": "cpp",
+        "stdexcept": "cpp",
+        "stop_token": "cpp",
+        "streambuf": "cpp",
+        "thread": "cpp",
+        "typeindex": "cpp",
+        "typeinfo": "cpp",
+        "valarray": "cpp",
+        "variant": "cpp"
+    }
+}

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;
@@ -2358,6 +2366,7 @@ 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))) {
+          // cout << "Variable: " << var << " Statement: " << stmt << endl;
           *reason = "all variables in concrete notation must be bound by a "
                     "forall statement";
           isConcrete = false;

src/index_notation/provenance_graph.cpp

@@ -1120,9 +1120,12 @@ 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
-  for (const IndexVar& child : getChildren(indexVar)) {
-    if (!defined.count(child) && (isFullyDerived(child) || !isRecoverable(child, defined))) {
-      return false;
+  // precompute relations are always recoverable since their children never appear in the same loop
+  if (!(childRelMap.at(indexVar).getRelType() == IndexVarRelType::PRECOMPUTE)) {
+    for (const IndexVar& child : getChildren(indexVar)) {
+      if (!defined.count(child) && (isFullyDerived(child) || !isRecoverable(child, defined))) {
+        return false;
+      }
     }
   }
   return true;

src/lower/iterator.cpp

@@ -521,16 +521,28 @@ Iterators::Iterators(IndexStmt stmt, const map<TensorVar, Expr>& tensorVars)
 {
   ProvenanceGraph provGraph = ProvenanceGraph(stmt);
   set<IndexVar> underivedAdded;
+
   // Create dimension iterators
   match(stmt,
     function<void(const ForallNode*, Matcher*)>([&](auto n, auto m) {
       content->modeIterators.insert({n->indexVar, Iterator(n->indexVar, !provGraph.hasCoordBounds(n->indexVar) && provGraph.isCoordVariable(n->indexVar))});
+      cout << "Adding following index var to iterators: " << n->indexVar << " for statement (" << n->stmt << ")" << endl;
       for (const IndexVar& underived : provGraph.getUnderivedAncestors(n->indexVar)) {
         if (!underivedAdded.count(underived)) {
+          cout << "Adding following underived ancestor to iterators: " << underived << endl;
           content->modeIterators.insert({underived, underived});
           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);
     })
   );
@@ -553,6 +565,8 @@ Iterators::Iterators(IndexStmt stmt, const map<TensorVar, Expr>& tensorVars)
   for (auto& iterator : content->levelIterators) {
     content->modeAccesses.insert({iterator.second, iterator.first});
   }
+
+  // cout << "FINISHED ITERATORS BUILDING" << endl;
 }
 
 
@@ -662,6 +676,8 @@ ModeAccess Iterators::modeAccess(Iterator iterator) const
 Iterator Iterators::modeIterator(IndexVar indexVar) const
 {
   taco_iassert(content != nullptr);
+  cout << "Searching for " << indexVar << " in "
+      << util::join(content->modeIterators) << endl;
   taco_iassert(util::contains(content->modeIterators, indexVar));
   return content->modeIterators.at(indexVar);
 }

src/lower/merge_lattice.cpp

@@ -613,7 +613,9 @@ MergeLattice MergeLattice::make(Forall forall, Iterators iterators, ProvenanceGr
 
   vector<IndexVar> underivedAncestors = provGraph.getUnderivedAncestors(indexVar);
   for (auto ancestor : underivedAncestors) {
+    // cout << "Is recoverable from merge lattice: " << ancestor << endl;
     if(!provGraph.isRecoverable(ancestor, definedIndexVars)) {
+      // cout << "Is not recoverable for ancestor " << ancestor << endl;
       return MergeLattice({MergePoint({iterators.modeIterator(indexVar)}, {}, {})});
     }
   }

test/tests-scheduling.cpp

@@ -276,7 +276,7 @@ TEST(scheduling, lowerSparseMulSparse) {
   //  codegen->compile(compute, true);
 }
 
-TEST(scheduling, precomputeIndependentIndexVars) {
+TEST(scheduling, NIRVIK_TEST) {
   Tensor<double> A("A", {16}, Format{Dense});
   Tensor<double> B("B", {16}, Format{Dense});
   Tensor<double> C("C", {16}, Format{Dense});
@@ -312,7 +312,7 @@ TEST(scheduling, precomputeIndependentIndexVars) {
   ASSERT_TENSOR_EQ(A, expected);
 }
 
-TEST(scheduling, precomputeIndependentIndexVarsSplit) {
+TEST(scheduling, FAILED_TEST) {
   Tensor<double> A("A", {16}, Format{Dense});
   Tensor<double> B("B", {16}, Format{Dense});
   Tensor<double> C("C", {16}, Format{Dense});
@@ -347,6 +347,11 @@ TEST(scheduling, precomputeIndependentIndexVarsSplit) {
   expected.assemble();
   expected.compute();
 
+  ir::IRPrinter irp = ir::IRPrinter(cout);
+  ir::Stmt compute = lower(stmt, "compute", false, true);
+  cout << "Imperative IR" << endl;
+  irp.print(compute);
+
   ASSERT_TENSOR_EQ(A, expected);
 }