[GraphScheduler] Honor memory dependencies implied by the SaveNode

When a variable is used as output in a save node, this is expected
to be its last use. Make sure the GraphSchedule looks at the implicit
dependencies between all the uses of a variable and its save node
when walking the dependencies.

This fixes issue #1283.

Author: Quentin Colombet

lib/IR/GraphScheduler.cpp

@@ -140,6 +140,27 @@ class ChildMemSizeBasedScheduler : public Scheduler {
       orderedChildren.push_back(N->getPredicate());
     }
 
+    // SaveNode hack:
+    // We don't model memory dependencies, but we still need to honor them.
+    // Make sure the SaveNode happens after the last use of the output variable.
+    if (auto *save = dyn_cast<SaveNode>(N)) {
+      Variable *output = save->getVariable();
+      for (NodeUse &use : output->getUsers()) {
+        Node *user = use.getUser();
+        if (user == save) {
+          continue;
+        }
+        // Variables may have users scattered across different functions.
+        // Only accounts for the ones in that function.
+        if (&G_ != user->getParent()) {
+          continue;
+        }
+        assert(!isa<SaveNode>(user) &&
+               "Variables must be saved at most once in each function");
+        orderedChildren.push_back(user);
+      }
+    }
+
     // Order children by (maxSize - resultSize). It gives more
     // priority to the nodes that free more memory after
     // their computation.

tests/unittests/graphTest.cpp

@@ -546,8 +546,7 @@ TEST(Graph, schedulingOfSaves) {
     }
   }
   EXPECT_TRUE(A->getPayload().isEqual(zero->getPayload(), 0.0));
-  // FIXME: This should be true, but right now it is not!
-  EXPECT_FALSE /*TRUE*/ (allEqual);
+  EXPECT_TRUE(allEqual);
 }
 
 /// Check that the parent link is properly updated while tweaking