Allow Compilation without Collecting Constants

include/glow/Backends/Backend.h

@@ -43,6 +43,10 @@ class Backend {
   /// Generate code for input function \param F.
   virtual std::unique_ptr<CompiledFunction> compile(Function *F) const = 0;
 
+  /// Generate code for input function \param F but do not collect constants.
+  virtual std::unique_ptr<CompiledFunction>
+  compileWithoutConstants(Function *F) const = 0;
+
   /// Save the bundle for \p F for a later standalone execution
   /// in \p outputDir. Make \p networkName the function name for
   /// the entry point of the network and prepend all generated

include/glow/Backends/BackendUtils.h

@@ -54,6 +54,8 @@ class RuntimeBundle {
   size_t getActivationsSize() const { return activationsMemSize_; }
   /// Get pointer to memory block of constants.
   uint8_t *getConstants() const { return constants_; }
+  /// Set pointer to memory block of constants.
+  void setConstants(uint8_t *constants) { constants_ = constants; }
   /// Helper function, gets offset of \p v.
   size_t getValueOffset(const Named *v) const;
   /// Helper function, gets symbol info for \p v.
@@ -67,7 +69,7 @@ class RuntimeBundle {
   RuntimeBundle() = default;
   RuntimeBundle(std::unordered_map<std::string, RuntimeSymbolInfo> &symbolTable,
                 size_t constWeight, size_t mutableWeight, size_t activations)
-      : symbolTable_(std::move(symbolTable)),
+      : symbolTable_(std::move(symbolTable)), constants_(nullptr),
         constantWeightVarsMemSize_(constWeight),
         mutableWeightVarsMemSize_(mutableWeight),
         activationsMemSize_(activations) {}

lib/Backends/BackendUtils.cpp

@@ -141,8 +141,6 @@ glow::generateRuntimeBundle(const IRFunction &F,
   }
   auto activationsMaxSize = activationsAllocator.getMaxMemoryUsage();
 
-  runtime::RuntimeBundle info(symbolTable, constantMaxSize, placeholderMaxSize,
-                              activationsMaxSize);
-  info.collectConstants(&F);
-  return info;
+  return runtime::RuntimeBundle(symbolTable, constantMaxSize,
+                                placeholderMaxSize, activationsMaxSize);
 }

lib/Backends/CPU/CPUBackend.cpp

@@ -99,14 +99,21 @@ CPUBackend::createIRGen(IRFunction *IR,
 
 std::unique_ptr<CompiledFunction>
 CPUBackend::compileIR(std::unique_ptr<IRFunction> IR) const {
+  auto function = compileIRWithoutConstants(IR.get());
+  static_cast<CPUFunction *>(function.get())->collectConstants(IR.get());
+  return function;
+}
+
+std::unique_ptr<CompiledFunction>
+CPUBackend::compileIRWithoutConstants(IRFunction *IR) const {
   AllocationsInfo allocationsInfo;
-  std::unique_ptr<LLVMIRGen> irgen = createIRGen(IR.get(), allocationsInfo);
+  std::unique_ptr<LLVMIRGen> irgen = createIRGen(IR, allocationsInfo);
   irgen->initTargetMachine(target.empty() ? "" : target.getValue(),
                            llvm::CodeModel::Model::Large);
   irgen->initCodeGen();
   // Perform the address assignment for activations and WeightVars.
 
-  allocateJITMemory(IR.get(), irgen->getAllocationsInfo());
+  allocateJITMemory(IR, irgen->getAllocationsInfo());
   // Create the jitmain function to be invoked by JIT.
   emitJitMain(*irgen);
   // Emit the code for the body of the entry function.
@@ -128,6 +135,12 @@ std::unique_ptr<CompiledFunction> CPUBackend::compile(Function *F) const {
   return compileIR(std::move(IR));
 }
 
+std::unique_ptr<CompiledFunction>
+CPUBackend::compileWithoutConstants(Function *F) const {
+  auto IR = generateAndOptimizeIR(F, shouldShareBuffers());
+  return compileIRWithoutConstants(IR.get());
+}
+
 void CPUBackend::save(Function *F, llvm::StringRef outputDir,
                       llvm::StringRef networkName) const {
   std::string tgt = target.empty() ? "" : target.getValue();

lib/Backends/CPU/CPUBackend.h

@@ -45,8 +45,14 @@ class CPUBackend : public BackendUsingGlowIR {
   std::unique_ptr<CompiledFunction>
   compileIR(std::unique_ptr<IRFunction> IR) const override;
 
+  std::unique_ptr<CompiledFunction>
+  compileIRWithoutConstants(IRFunction *IR) const;
+
   std::unique_ptr<CompiledFunction> compile(Function *F) const override;
 
+  std::unique_ptr<CompiledFunction>
+  compileWithoutConstants(Function *F) const override;
+
   void save(Function *F, llvm::StringRef outputDir,
             llvm::StringRef networkName) const override;
 

lib/Backends/CPU/CPUFunction.cpp

@@ -45,6 +45,10 @@ void CPUFunction::setupRuns() {
   }
 }
 
+void CPUFunction::collectConstants(IRFunction *F) {
+  runtimeBundle_.collectConstants(F);
+}
+
 void CPUFunction::beforeRun(const Context &ctx) {
   // Copy Placeholders into allocated memory.
   for (auto PH : ctx.pairs()) {

lib/Backends/CPU/CPUFunction.h

@@ -38,6 +38,10 @@ class CPUFunction final : public CompiledFunction {
   /// Ctor.
   CPUFunction(std::unique_ptr<llvm::orc::GlowJIT> JIT,
               const runtime::RuntimeBundle &runtimeBundle);
+
+  /// Collects constants for runtime.
+  void collectConstants(IRFunction *F);
+
   /// Allocate Mutable buffers on device this includes Activations and
   /// Placeholders.
   void setupRuns() override;

lib/Backends/Interpreter/Interpreter.cpp

@@ -29,8 +29,22 @@ std::unique_ptr<CompiledFunction> Interpreter::compile(Function *F) const {
   return compileIR(std::move(IR));
 }
 
+std::unique_ptr<CompiledFunction>
+Interpreter::compileWithoutConstants(Function *F) const {
+  auto IR = generateAndOptimizeIR(F, shouldShareBuffers());
+  return compileIRWithoutConstants(std::move(IR));
+}
+
 std::unique_ptr<CompiledFunction>
 Interpreter::compileIR(std::unique_ptr<IRFunction> IR) const {
+  auto function = compileIRWithoutConstants(std::move(IR));
+  auto IFunction = static_cast<InterpreterFunction *>(function.get());
+  IFunction->collectConstants(IFunction->getIR());
+  return function;
+}
+
+std::unique_ptr<CompiledFunction>
+Interpreter::compileIRWithoutConstants(std::unique_ptr<IRFunction> IR) const {
   MemoryAllocator constantWeightsAllocator("ConstantWeights", 0);
   MemoryAllocator placeholderWeightsAllocator("PlaceholderWeights", 0);
   MemoryAllocator activationsAllocator("Activations", 0);

lib/Backends/Interpreter/Interpreter.h

@@ -37,8 +37,14 @@ class Interpreter final : public BackendUsingGlowIR {
   std::unique_ptr<CompiledFunction>
   compileIR(std::unique_ptr<IRFunction> IR) const override;
 
+  std::unique_ptr<CompiledFunction>
+  compileIRWithoutConstants(std::unique_ptr<IRFunction> IR) const;
+
   std::unique_ptr<CompiledFunction> compile(Function *F) const override;
 
+  std::unique_ptr<CompiledFunction>
+  compileWithoutConstants(Function *F) const override;
+
   bool isOpSupported(Kinded::Kind opKind, ElemKind elementTy) const override;
 
   bool shouldLower(const Node *N) const override;

lib/Backends/Interpreter/InterpreterFunction.cpp

@@ -39,6 +39,10 @@ InterpreterFunction::~InterpreterFunction() {
   tearDownRuns();
 }
 
+void InterpreterFunction::collectConstants(IRFunction *F) {
+  runtimeBundle_.collectConstants(F);
+}
+
 void InterpreterFunction::setupRuns() {
   if (!runsSetup_) {
     if (runtimeBundle_.getConstantWeightSize()) {

lib/Backends/Interpreter/InterpreterFunction.h

@@ -62,6 +62,10 @@ class InterpreterFunction final : public CompiledFunction {
 
   /// Does any needed initialization work for the Backend, creates tensors from
   /// constants.
+
+  /// Collects constants for runtime.
+  void collectConstants(IRFunction *F);
+
   void setupRuns() override;
 
   /// Per run setup, adds references for tensors from \p ctx to
@@ -76,6 +80,8 @@ class InterpreterFunction final : public CompiledFunction {
   void tearDownRuns() override;
 
   void execute() override;
+  /// Get reference to IR function.
+  IRFunction *getIR() { return F_.get(); }
   ///@}
 
 private:

lib/Backends/OpenCL/OpenCL.cpp

@@ -1525,15 +1525,31 @@ cl_mem OpenCLFunction::allocDeviceBuffer(uint64_t size) {
 
 void OpenCLFunction::freeDeviceBuffer(cl_mem buf) { clReleaseMemObject(buf); }
 
+void OpenCLFunction::collectConstants(IRFunction *F) {
+  runtimeBundle_.collectConstants(F);
+}
 std::unique_ptr<CompiledFunction>
 OCLBackend::compileIR(std::unique_ptr<IRFunction> IR) const {
+  auto function = compileIRWithoutConstants(std::move(IR));
+  auto OCLFunction = static_cast<OpenCLFunction *>(function.get());
+  OCLFunction->collectConstants(OCLFunction->getIR());
+  return function;
+}
+
+std::unique_ptr<CompiledFunction>
+OCLBackend::compileIRWithoutConstants(std::unique_ptr<IRFunction> IR) 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 {
   auto IR = generateAndOptimizeIR(F, shouldShareBuffers());
   return compileIR(std::move(IR));
 }
+
+std::unique_ptr<CompiledFunction>
+OCLBackend::compileWithoutConstants(Function *F) const {
+  auto IR = generateAndOptimizeIR(F, shouldShareBuffers());
+  return compileIRWithoutConstants(std::move(IR));
+}

lib/Backends/OpenCL/OpenCL.h

@@ -107,6 +107,12 @@ class OpenCLFunction final : public CompiledFunction {
   /// Final cleanup, currently an empty function in OpenCL.
   void tearDownRuns() override;
 
+  /// Returns IR function pointer.
+  IRFunction *getIR() { return F_.get(); }
+
+  /// Collects constants for runtime.
+  void collectConstants(IRFunction *F);
+
 private:
   /// Copy the value from a device to a provided buffer.
   /// \returns number of copied bytes.
@@ -161,8 +167,12 @@ class OCLBackend final : public BackendUsingGlowIR {
 
   std::unique_ptr<CompiledFunction>
   compileIR(std::unique_ptr<IRFunction> IR) const override;
+  std::unique_ptr<CompiledFunction>
+  compileIRWithoutConstants(std::unique_ptr<IRFunction> IR) const;
 
   std::unique_ptr<CompiledFunction> compile(Function *F) const override;
+  std::unique_ptr<CompiledFunction>
+  compileWithoutConstants(Function *F) const override;
 
   bool transformPostLowering(Function *F, CompilationMode mode) const override;
 

tests/unittests/BackendCorrectnessTest.cpp

@@ -239,6 +239,11 @@ class MockCPUBackend : public BackendUsingGlowIR {
   std::unique_ptr<CompiledFunction> compile(Function *F) const override {
     return backend_->compile(F);
   }
+
+  std::unique_ptr<CompiledFunction>
+  compileWithoutConstants(Function *F) const override {
+    return backend_->compile(F);
+  }
   std::unique_ptr<CompiledFunction>
   compileIR(std::unique_ptr<IRFunction> IR) const override {
     return backend_->compileIR(std::move(IR));

tests/unittests/BackendTest.cpp

@@ -168,6 +168,22 @@ TEST_P(BackendTest, debugPrint) {
   function->tearDownRuns();
 }
 
+/// Test the compileWithoutConstants method on the backend completes without
+/// error.
+TEST_P(BackendTest, CompileWithoutConstants) {
+  Module mod;
+  Context ctx;
+  Function *F = mod.createFunction("main");
+  auto *X = mod.createPlaceholder(ElemKind::FloatTy, {3}, "X", false);
+  auto *XTensor = ctx.allocate(X);
+  XTensor->getHandle() = {1., 2., 3.};
+  auto *pow = F->createPow("Pow1", X, 2.0);
+  auto *save = F->createSave("save", pow);
+  ctx.allocate(save->getPlaceholder());
+  std::unique_ptr<Backend> backend(createBackend(GetParam()));
+  auto function = backend->compileWithoutConstants(F);
+}
+
 /// This test checks that we can compile a function without depending on the
 /// graph representation. We compile some function and then delete the function.
 /// Later we execute the code and check that things work.

tests/unittests/BackendTestUtils.h

@@ -32,6 +32,10 @@ class MockBackend : public Backend {
   std::unique_ptr<CompiledFunction> compile(Function *F) const override {
     return llvm::make_unique<MockFunction>();
   }
+  std::unique_ptr<CompiledFunction>
+  compileWithoutConstants(Function *F) const override {
+    return llvm::make_unique<MockFunction>();
+  }
   bool isOpSupported(Kinded::Kind opKind, ElemKind elementTy) const override {
     return false;
   }

tests/unittests/QuantizationTest.cpp

@@ -837,6 +837,11 @@ class MockQuantBackend : public Backend {
   std::unique_ptr<CompiledFunction> compile(Function *F) const override {
     return backend_->compile(F);
   }
+
+  std::unique_ptr<CompiledFunction>
+  compileWithoutConstants(Function *F) const override {
+    return backend_->compile(F);
+  }
   bool isOpSupported(Kinded::Kind opKind, ElemKind elementTy) const override {
     if (opKind == Kinded::Kind::SoftMaxNodeKind ||
         opKind == Kinded::Kind::LocalResponseNormalizationNodeKind) {