lower: fix a bug causing undefined variables when applying fuse

src/lower/lowerer_impl.cpp

@@ -2250,35 +2250,50 @@ Stmt LowererImpl::declLocatePosVars(vector<Iterator> locators) {
   for (Iterator& locator : locators) {
     accessibleIterators.insert(locator);
 
-    bool doLocate = true;
+    // Pull out some logic for constructing the locators for a given iterator.
+    auto addLocator = [&](const Iterator& iter) {
+      ModeFunction locate = iter.locate(coordinates(iter));
+      taco_iassert(isValue(locate.getResults()[1], true));
+      Stmt declarePosVar = VarDecl::make(iter.getPosVar(), locate.getResults()[0]);
+      result.push_back(declarePosVar);
+    };
+
+    // Look through all of the parent iterators. If any of these iterators
+    // are not accessible, we need to construct their accessors before emitting
+    // locator's accessors. This is because locator may use the ancestor's
+    // variables in its accessors. We add these ancestors into a vector and reverse
+    // it so that the highest parent in the tree's accessors get declared first.
+    std::vector<Iterator> ancestors;
     for (Iterator ancestorIterator = locator.getParent();
          !ancestorIterator.isRoot() && ancestorIterator.hasLocate();
          ancestorIterator = ancestorIterator.getParent()) {
       if (!accessibleIterators.contains(ancestorIterator)) {
-        doLocate = false;
+        // Since we're going to emit the locators for this iterator, add it to
+        // accessibleIterators so that other locators with this as an ancestor
+        // don't do the same.
+        accessibleIterators.insert(ancestorIterator);
+        ancestors.push_back(ancestorIterator);
       }
     }
+    for (auto it = ancestors.rbegin(); it != ancestors.rend(); it++) addLocator(*it);
 
-    if (doLocate) {
-      Iterator locateIterator = locator;
-      if (locateIterator.hasPosIter()) {
-        taco_iassert(!provGraph.isUnderived(locateIterator.getIndexVar()));
-        continue; // these will be recovered with separate procedure
-      }
-      do {
-        ModeFunction locate = locateIterator.locate(coordinates(locateIterator));
-        taco_iassert(isValue(locate.getResults()[1], true));
-        Stmt declarePosVar = VarDecl::make(locateIterator.getPosVar(),
-                                           locate.getResults()[0]);
-        result.push_back(declarePosVar);
-
-        if (locateIterator.isLeaf()) {
-          break;
-        }
-
-        locateIterator = locateIterator.getChild();
-      } while (accessibleIterators.contains(locateIterator));
+    Iterator locateIterator = locator;
+    // Position iterators will be recovered with a separate procedure, so
+    // don't emit anything if locator is one.
+    if (locateIterator.hasPosIter()) {
+      taco_iassert(!provGraph.isUnderived(locateIterator.getIndexVar()));
+      continue;
     }
+
+    // Once all parent locators have been declared, add the target and all
+    // children locators.
+    do {
+      addLocator(locateIterator);
+      if (locateIterator.isLeaf()) {
+        break;
+      }
+      locateIterator = locateIterator.getChild();
+    } while (accessibleIterators.contains(locateIterator));
   }
   return result.empty() ? Stmt() : Block::make(result);
 }

test/tests-scheduling.cpp

@@ -72,6 +72,56 @@ TEST(scheduling, splitIndexStmt) {
   ASSERT_TRUE(equals(a(i) = b(i), i2Forall.getStmt()));
 }
 
+TEST(scheduling, fuseDenseLoops) {
+    auto dim = 4;
+    Tensor<int> A("A", {dim, dim, dim}, {Dense, Dense, Dense});
+    Tensor<int> B("B", {dim, dim, dim}, {Dense, Dense, Dense});
+    Tensor<int> expected("expected", {dim, dim, dim}, {Dense, Dense, Dense});
+    IndexVar f("f"), g("g");
+    for (int i = 0; i < dim; i++) {
+        for (int j = 0; j < dim; j++) {
+            for (int k = 0; k < dim; k++) {
+                A.insert({i, j, k}, i + j + k);
+                B.insert({i, j, k}, i + j + k);
+                expected.insert({i, j, k}, 2 * (i + j + k));
+            }
+        }
+    }
+    A.pack();
+    B.pack();
+    expected.pack();
+
+    // Helper function to evaluate the target statement and verify the results.
+    // It takes in a function that applies some scheduling transforms to the
+    // input IndexStmt, and applies to the point-wise tensor addition below.
+    // The test is structured this way as TACO does its best to avoid re-compilation
+    // whenever possible. I.e. changing the stmt that a tensor is compiled with
+    // doesn't cause compilation to occur again.
+    auto testFn = [&](IndexStmt modifier (IndexStmt)) {
+        Tensor<int> C("C", {dim, dim, dim}, {Dense, Dense, Dense});
+        C(i, j, k) = A(i, j, k) + B(i, j, k);
+        auto stmt = C.getAssignment().concretize();
+        C.compile(modifier(stmt));
+        C.evaluate();
+        ASSERT_TRUE(equals(C, expected)) << endl << C << endl << expected << endl;
+    };
+
+    // First, a sanity check with no transformations.
+    testFn([](IndexStmt stmt) { return stmt; });
+    // Next, fuse the outer two loops. This tests the original bug in #355.
+    testFn([](IndexStmt stmt) {
+        IndexVar f("f");
+        return stmt.fuse(i, j, f);
+    });
+    // Lastly, fuse all of the loops into a single loop. This ensures that
+    // locators with a chain of ancestors have all of their dependencies
+    // generated in a valid ordering.
+    testFn([](IndexStmt stmt) {
+        IndexVar f("f"), g("g");
+        return stmt.fuse(i, j, f).fuse(f, k, g);
+    });
+}
+
 TEST(scheduling, lowerDenseMatrixMul) {
   Tensor<double> A("A", {4, 4}, {Dense, Dense});
   Tensor<double> B("B", {4, 4}, {Dense, Dense});