Assorted memory-related fixes for HostManager, EE, RecSys

include/glow/ExecutionEngine/ExecutionEngine.h

@@ -47,7 +47,7 @@ class ExecutionEngine final {
   /// this resets the ExecutionEngine.
   std::string backendName_ = "";
 
-  /// Size of device memory, if 0 device default is used.
+  /// Size of device memory in bytes, if 0 device default is used.
   uint64_t deviceMemory_{0};
 
   /// The HostManager for executing the compiled functions.
@@ -61,7 +61,10 @@ class ExecutionEngine final {
   void runInternal(ExecutionContext &context, llvm::StringRef name);
 
 public:
-  ExecutionEngine(llvm::StringRef backend = "Interpreter");
+  /// Constructor for an ExecutionEngine with \p backend and memory \p
+  /// deviceMemory in bytes.
+  ExecutionEngine(llvm::StringRef backend = "Interpreter",
+                  uint64_t deviceMemory = 0);
 
   ~ExecutionEngine();
 

lib/ExecutionEngine/ExecutionEngine.cpp

@@ -27,16 +27,17 @@
 
 using namespace glow;
 
-ExecutionEngine::ExecutionEngine(llvm::StringRef backend) {
+ExecutionEngine::ExecutionEngine(llvm::StringRef backend, uint64_t deviceMemory)
+    : deviceMemory_(deviceMemory) {
   setBackendName(backend);
 }
 
 /// Set the code generator to the given \p backend.
 void ExecutionEngine::setBackendName(llvm::StringRef backend) {
+  clear();
   module_.reset(new Module);
   rawModule_ = module_.get();
   backendName_ = backend;
-  clear();
 
   if (hostManager_) {
     EXIT_ON_ERR(hostManager_->clearHost());
@@ -61,6 +62,7 @@ void ExecutionEngine::clear() {
     EXIT_ON_ERR(hostManager_->clearHost());
   }
   compiledFunctions_.clear();
+  module_.reset(nullptr);
 }
 
 void glow::updateInputPlaceholders(PlaceholderBindings &bindings,

lib/Runtime/HostManager/HostManager.cpp

@@ -226,20 +226,18 @@ llvm::Error HostManager::clearHost() {
   DCHECK_EQ(activeRequestCount_, 0)
       << "All requests should be finished when shutting down HostManager.";
 
+  // Remove all networks from the host and device(s).
+  while (networks_.size() != 0) {
+    RETURN_IF_ERR(removeNetwork(networks_.begin()->first));
+  }
+
+  // Now it's safe to stop the DeviceManagers.
   std::lock_guard<std::mutex> networkLock(networkLock_);
   OneErrOnly errContainer;
   for (auto &it : devices_) {
     errContainer.set(it.second->stop());
   }
 
-  for (auto &network : networks_) {
-    for (auto &node : network.second.dag.nodes) {
-      for (auto device : node->deviceIDs) {
-        devices_[device]->evictNetwork(node->name);
-      }
-    }
-  }
-  networks_.clear();
   return errContainer.get();
 }
 

tests/unittests/BackendTestUtils.h

@@ -74,7 +74,10 @@ DECLARE_STATELESS_BACKEND_TEST(BackendStatelessTest, std::tuple<std::string>);
 
 class BackendTest : public BackendStatelessTest {
 public:
-  BackendTest() : mod_(EE_.getModule()) { F_ = mod_.createFunction("main"); }
+  BackendTest(uint64_t deviceMemory = 0)
+      : EE_(getBackendName(), deviceMemory), mod_(EE_.getModule()) {
+    F_ = mod_.createFunction("main");
+  }
 
 protected:
   ExecutionEngine EE_{getBackendName()};