Merge pull request #361 from tensor-compiler/revert-360-fuse-bug

Revert "lower: fix a bug causing undefined variables when applying fuse"

Author: stephenchouca

src/lower/lowerer_impl.cpp

@@ -2250,50 +2250,35 @@ Stmt LowererImpl::declLocatePosVars(vector<Iterator> locators) {
   for (Iterator& locator : locators) {
     accessibleIterators.insert(locator);
 
-    // 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;
+    bool doLocate = true;
     for (Iterator ancestorIterator = locator.getParent();
          !ancestorIterator.isRoot() && ancestorIterator.hasLocate();
          ancestorIterator = ancestorIterator.getParent()) {
       if (!accessibleIterators.contains(ancestorIterator)) {
-        // 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);
+        doLocate = false;
       }
     }
-    for (auto it = ancestors.rbegin(); it != ancestors.rend(); it++) addLocator(*it);
-
-    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;
+    if (doLocate) {
+      Iterator locateIterator = locator;
+      if (locateIterator.hasPosIter()) {
+        taco_iassert(!provGraph.isUnderived(locateIterator.getIndexVar()));
+        continue; // these will be recovered with separate procedure
       }
-      locateIterator = locateIterator.getChild();
-    } while (accessibleIterators.contains(locateIterator));
+      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));
+    }
   }
   return result.empty() ? Stmt() : Block::make(result);
 }

test/tests-scheduling.cpp

@@ -72,56 +72,6 @@ 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});