[WIP] Add Provisioner

include/glow/Backends/Backend.h

@@ -41,7 +41,8 @@ class Backend {
   virtual ~Backend() = default;
 
   /// Generate code for input function \param F.
-  virtual std::unique_ptr<CompiledFunction> compile(Function *F) const = 0;
+  virtual std::unique_ptr<CompiledFunction>
+  compile(Function *F, bool collectConstants = true) const = 0;
 
   /// Save the bundle for \p F for a later standalone execution
   /// in \p outputDir. Make \p networkName the function name for
@@ -91,7 +92,8 @@ class BackendUsingGlowIR : public Backend {
   /// maps the graph to the concrete execution environment for a specific
   /// function. This is used only for unit testing.
   virtual std::unique_ptr<CompiledFunction>
-  compileIR(std::unique_ptr<IRFunction> IR) const = 0;
+  compileIR(std::unique_ptr<IRFunction> IR,
+            bool collectConstants = true) const = 0;
 };
 
 } // namespace glow

include/glow/Backends/BackendUtils.h

@@ -16,7 +16,6 @@
 #ifndef GLOW_BACKENDS_BACKENDUTILS_H
 #define GLOW_BACKENDS_BACKENDUTILS_H
 
-#include "glow/Backends/CompiledFunction.h"
 #include "glow/CodeGen/MemoryAllocator.h"
 #include "glow/IR/IR.h"
 
@@ -79,7 +78,8 @@ class RuntimeBundle {
 runtime::RuntimeBundle
 generateRuntimeBundle(const IRFunction &F, MemoryAllocator &constantAllocator,
                       MemoryAllocator &placeholderAllocator,
-                      MemoryAllocator &activationsAllocator);
+                      MemoryAllocator &activationsAllocator,
+                      bool collectConstants = true);
 
 } // end namespace glow
 #endif // GLOW_BACKENDS_BACKENDUTILS_H

include/glow/Backends/CompiledFunction.h

@@ -16,12 +16,15 @@
 #ifndef GLOW_BACKENDS_COMPILEDFUNCTION_H
 #define GLOW_BACKENDS_COMPILEDFUNCTION_H
 
+#include "glow/Backends/BackendUtils.h"
 #include "glow/Graph/Nodes.h"
+
 #include <unordered_map>
 
 namespace glow {
 
 class Context;
+
 /// Interface for executing a compiled function.
 class CompiledFunction {
 public:
@@ -41,6 +44,8 @@ class CompiledFunction {
   virtual void afterRun(const Context &ctx) = 0;
   /// Final cleanup. Release memory, reset device.
   virtual void tearDownRuns() = 0;
+  /// Return RuntimeBundle containing symbol information.
+  virtual runtime::RuntimeBundle getRuntimeBundle();
 };
 } // end namespace glow
 

include/glow/Runtime/Provisioner/Provisioner.h

@@ -0,0 +1,102 @@
+/**
+ * Copyright (c) 2017-present, Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#ifndef GLOW_RUNTIME_PROVISIONER_H
+#define GLOW_RUNTIME_PROVISIONER_H
+
+// #include "glow/Runtime/RuntimeTypes.h"
+
+//////////////////////////////////////////////////////////
+#include "glow/Backends/Backend.h"
+#include "glow/Backends/BackendUtils.h"
+#include "glow/Backends/CompiledFunction.h"
+#include "glow/Graph/Graph.h"
+
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+namespace glow {
+namespace runtime {
+
+using NetworkIDty = size_t;
+using DeviceIDty = size_t;
+
+/// Enum to communicate results when communicating with device at initialization
+/// and runtime.
+enum ResultCode { READY, EXECUTED, FAILED, CANCELLED };
+
+/// Data structure that contains device constraint information for each device.
+/// Used to communicate memory constraints and later costs to the Partitioner.
+struct DeviceInfo {
+  /// Available memory on device in bytes.
+  uint64_t availableMemory;
+};
+
+/// Individual Node in the DAG for a given network. This contains all the
+/// information needed to run the sub-network at inference time.
+struct DAGNode {
+  /// The children of this node, these are nodes that depend on the current
+  /// node.
+  std::vector<DAGNode> children;
+  /// Pointers to the parents of this node. This is used by the executor for
+  /// determining if a given node has all dependencies met.
+  std::vector<DAGNode *> parents;
+  /// ID of the deviceManager that this network is assigned to.
+  DeviceIDty deviceID;
+  /// The logicalDevice is an output of the Partitioner to indicate that two
+  /// networks should be assigned to the same device.
+  DeviceIDty logicalDevice;
+  /// Name assigned to the sub-network, this is the id that will be passed to
+  /// the DeviceManager when requesting a run of the network.
+  std::string name;
+  /// Runtime bundle containing all the symbol information for this network at
+  /// runtime.
+  RuntimeBundle runtimeBundle;
+};
+
+} // namespace runtime
+} // namespace glow
+///////////////////////////////////////////////
+
+#include <unordered_map>
+#include <vector>
+
+namespace glow {
+namespace runtime {
+
+class DeviceManager;
+/// The Provisioner is responsible for assigning networks to an actual device.
+/// It is a stateless class, relying on information being passed in by the
+/// caller.
+class Provisioner final {
+public:
+  Provisioner();
+  ~Provisioner();
+  /// Walks \p networks and assigns each module to a DeviceManager in \p
+  /// devices. The Provisioner calls the addNetwork method for each
+  /// DeviceManager and uses the returned networkID to populate \p runDAG.
+  /// Returns a ResultCode indicating if the operation was a success.
+  ResultCode provision(std::vector<DAGNode> &networks,
+                       std::unordered_map<DeviceIDty, DeviceManager> &devices,
+                       Module &module);
+
+private:
+  std::unique_ptr<Backend> backend_;
+};
+} // namespace runtime
+} // namespace glow
+
+#endif // GLOW_RUNTIME_PROVISIONER_H

lib/Backends/BackendUtils.cpp

@@ -55,11 +55,10 @@ runtime::RuntimeBundle::getSymbolInfo(const Named *v) const {
   return it->second;
 }
 
-runtime::RuntimeBundle
-glow::generateRuntimeBundle(const IRFunction &F,
-                            MemoryAllocator &constantAllocator,
-                            MemoryAllocator &placeholderAllocator,
-                            MemoryAllocator &activationsAllocator) {
+runtime::RuntimeBundle glow::generateRuntimeBundle(
+    const IRFunction &F, MemoryAllocator &constantAllocator,
+    MemoryAllocator &placeholderAllocator,
+    MemoryAllocator &activationsAllocator, bool collectConstants) {
   // Handle Constants, Placeholders, and Activations, in that order.
   // Symbol table mapping symbol name to offset for runtime.
   std::unordered_map<std::string, runtime::RuntimeSymbolInfo> symbolTable;
@@ -143,6 +142,8 @@ glow::generateRuntimeBundle(const IRFunction &F,
 
   runtime::RuntimeBundle info(symbolTable, constantMaxSize, placeholderMaxSize,
                               activationsMaxSize);
-  info.collectConstants(&F);
+  if (collectConstants) {
+    info.collectConstants(&F);
+  }
   return info;
 }

lib/Backends/CPU/CPUBackend.cpp

@@ -98,7 +98,8 @@ CPUBackend::createIRGen(IRFunction *IR,
 }
 
 std::unique_ptr<CompiledFunction>
-CPUBackend::compileIR(std::unique_ptr<IRFunction> IR) const {
+CPUBackend::compileIR(std::unique_ptr<IRFunction> IR,
+                      bool collectConstants) const {
   AllocationsInfo allocationsInfo;
   std::unique_ptr<LLVMIRGen> irgen = createIRGen(IR.get(), allocationsInfo);
   irgen->initTargetMachine(target.empty() ? "" : target.getValue(),
@@ -118,14 +119,16 @@ CPUBackend::compileIR(std::unique_ptr<IRFunction> IR) const {
   MemoryAllocator constantAllocator("ConstantWeights", 0);
   MemoryAllocator placeholderAllocator("Placeholders", 0);
   MemoryAllocator activationsAllocator("Activations", 0);
-  runtime::RuntimeBundle runtimeInfo = generateRuntimeBundle(
-      *IR, constantAllocator, placeholderAllocator, activationsAllocator);
+  runtime::RuntimeBundle runtimeInfo =
+      generateRuntimeBundle(*IR, constantAllocator, placeholderAllocator,
+                            activationsAllocator, collectConstants);
   return llvm::make_unique<CPUFunction>(std::move(JIT), runtimeInfo);
 }
 
-std::unique_ptr<CompiledFunction> CPUBackend::compile(Function *F) const {
+std::unique_ptr<CompiledFunction>
+CPUBackend::compile(Function *F, bool collectConstants) const {
   auto IR = generateAndOptimizeIR(F, shouldShareBuffers());
-  return compileIR(std::move(IR));
+  return compileIR(std::move(IR), collectConstants);
 }
 
 void CPUBackend::save(Function *F, llvm::StringRef outputDir,

lib/Backends/CPU/CPUBackend.h

@@ -43,9 +43,11 @@ class CPUBackend : public BackendUsingGlowIR {
   ~CPUBackend() override = default;
 
   std::unique_ptr<CompiledFunction>
-  compileIR(std::unique_ptr<IRFunction> IR) const override;
+  compileIR(std::unique_ptr<IRFunction> IR,
+            bool collectConstants = true) const override;
 
-  std::unique_ptr<CompiledFunction> compile(Function *F) const override;
+  std::unique_ptr<CompiledFunction>
+  compile(Function *F, bool collectConstants = true) const override;
 
   void save(Function *F, llvm::StringRef outputDir,
             llvm::StringRef networkName) const override;

lib/Backends/Interpreter/Interpreter.cpp

@@ -24,13 +24,15 @@
 #include "glow/IR/IR.h"
 
 using namespace glow;
-std::unique_ptr<CompiledFunction> Interpreter::compile(Function *F) const {
+std::unique_ptr<CompiledFunction>
+Interpreter::compile(Function *F, bool collectConstants) const {
   auto IR = generateAndOptimizeIR(F, shouldShareBuffers());
   return compileIR(std::move(IR));
 }
 
 std::unique_ptr<CompiledFunction>
-Interpreter::compileIR(std::unique_ptr<IRFunction> IR) const {
+Interpreter::compileIR(std::unique_ptr<IRFunction> IR,
+                       bool collectConstants) const {
   MemoryAllocator constantWeightsAllocator("ConstantWeights", 0);
   MemoryAllocator placeholderWeightsAllocator("PlaceholderWeights", 0);
   MemoryAllocator activationsAllocator("Activations", 0);

lib/Backends/Interpreter/Interpreter.h

@@ -35,9 +35,11 @@ class Interpreter final : public BackendUsingGlowIR {
   ~Interpreter() override = default;
 
   std::unique_ptr<CompiledFunction>
-  compileIR(std::unique_ptr<IRFunction> IR) const override;
+  compileIR(std::unique_ptr<IRFunction> IR,
+            bool collectConstants = true) const override;
 
-  std::unique_ptr<CompiledFunction> compile(Function *F) const override;
+  std::unique_ptr<CompiledFunction>
+  compile(Function *F, bool collectConstants = true) const override;
 
   bool isOpSupported(Kinded::Kind opKind, ElemKind elementTy) const override;
 

lib/Backends/OpenCL/OpenCL.cpp

@@ -1514,14 +1514,16 @@ cl_mem OpenCLFunction::allocDeviceBuffer(uint64_t size) {
 void OpenCLFunction::freeDeviceBuffer(cl_mem buf) { clReleaseMemObject(buf); }
 
 std::unique_ptr<CompiledFunction>
-OCLBackend::compileIR(std::unique_ptr<IRFunction> IR) const {
+OCLBackend::compileIR(std::unique_ptr<IRFunction> IR,
+                      bool collectConstants) const {
   MemoryAllocator allocator("GPU", 0xFFFFFFFF);
   runtime::RuntimeBundle bundle =
       generateRuntimeBundle(*IR, allocator, allocator, allocator);
   return llvm::make_unique<OpenCLFunction>(std::move(IR), bundle);
 }
 
-std::unique_ptr<CompiledFunction> OCLBackend::compile(Function *F) const {
+std::unique_ptr<CompiledFunction>
+OCLBackend::compile(Function *F, bool collectConstants) const {
   auto IR = generateAndOptimizeIR(F, shouldShareBuffers());
   return compileIR(std::move(IR));
 }

lib/Backends/OpenCL/OpenCL.h

@@ -162,9 +162,11 @@ class OCLBackend final : public BackendUsingGlowIR {
   ~OCLBackend() override = default;
 
   std::unique_ptr<CompiledFunction>
-  compileIR(std::unique_ptr<IRFunction> IR) const override;
+  compileIR(std::unique_ptr<IRFunction> IR,
+            bool collectConstants = true) const override;
 
-  std::unique_ptr<CompiledFunction> compile(Function *F) const override;
+  std::unique_ptr<CompiledFunction>
+  compile(Function *F, bool collectConstants = true) const override;
 
   bool transformPostLowering(Function *F, CompilationMode mode) const override;