diff --git a/include/glow/Backends/Backend.h b/include/glow/Backends/Backend.h index 8750aced84..d7d68987a3 100644 --- a/include/glow/Backends/Backend.h +++ b/include/glow/Backends/Backend.h @@ -43,6 +43,10 @@ class Backend { /// Generate code for input function \param F. virtual std::unique_ptr compile(Function *F) const = 0; + /// Generate code for input function \param F but do not collect constants. + virtual std::unique_ptr + 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 diff --git a/include/glow/Backends/BackendUtils.h b/include/glow/Backends/BackendUtils.h index e3d08e5bf4..d7ac1919a5 100644 --- a/include/glow/Backends/BackendUtils.h +++ b/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 &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) {} diff --git a/lib/Backends/BackendUtils.cpp b/lib/Backends/BackendUtils.cpp index ec1d2eb23c..077c21df69 100644 --- a/lib/Backends/BackendUtils.cpp +++ b/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); } diff --git a/lib/Backends/CPU/CPUBackend.cpp b/lib/Backends/CPU/CPUBackend.cpp index e4cbde9c79..db8bee6b80 100644 --- a/lib/Backends/CPU/CPUBackend.cpp +++ b/lib/Backends/CPU/CPUBackend.cpp @@ -99,14 +99,21 @@ CPUBackend::createIRGen(IRFunction *IR, std::unique_ptr CPUBackend::compileIR(std::unique_ptr IR) const { + auto function = compileIRWithoutConstants(IR.get()); + static_cast(function.get())->collectConstants(IR.get()); + return function; +} + +std::unique_ptr +CPUBackend::compileIRWithoutConstants(IRFunction *IR) const { AllocationsInfo allocationsInfo; - std::unique_ptr irgen = createIRGen(IR.get(), allocationsInfo); + std::unique_ptr 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 CPUBackend::compile(Function *F) const { return compileIR(std::move(IR)); } +std::unique_ptr +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(); diff --git a/lib/Backends/CPU/CPUBackend.h b/lib/Backends/CPU/CPUBackend.h index ebcc765826..efda99d4c3 100644 --- a/lib/Backends/CPU/CPUBackend.h +++ b/lib/Backends/CPU/CPUBackend.h @@ -45,8 +45,14 @@ class CPUBackend : public BackendUsingGlowIR { std::unique_ptr compileIR(std::unique_ptr IR) const override; + std::unique_ptr + compileIRWithoutConstants(IRFunction *IR) const; + std::unique_ptr compile(Function *F) const override; + std::unique_ptr + compileWithoutConstants(Function *F) const override; + void save(Function *F, llvm::StringRef outputDir, llvm::StringRef networkName) const override; diff --git a/lib/Backends/CPU/CPUFunction.cpp b/lib/Backends/CPU/CPUFunction.cpp index 2bf9027531..16d00439a1 100644 --- a/lib/Backends/CPU/CPUFunction.cpp +++ b/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()) { diff --git a/lib/Backends/CPU/CPUFunction.h b/lib/Backends/CPU/CPUFunction.h index 7cdb48217d..c1ca4fdf4b 100644 --- a/lib/Backends/CPU/CPUFunction.h +++ b/lib/Backends/CPU/CPUFunction.h @@ -38,6 +38,10 @@ class CPUFunction final : public CompiledFunction { /// Ctor. CPUFunction(std::unique_ptr 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; diff --git a/lib/Backends/Interpreter/Interpreter.cpp b/lib/Backends/Interpreter/Interpreter.cpp index fedda97375..0c08e41015 100644 --- a/lib/Backends/Interpreter/Interpreter.cpp +++ b/lib/Backends/Interpreter/Interpreter.cpp @@ -29,8 +29,22 @@ std::unique_ptr Interpreter::compile(Function *F) const { return compileIR(std::move(IR)); } +std::unique_ptr +Interpreter::compileWithoutConstants(Function *F) const { + auto IR = generateAndOptimizeIR(F, shouldShareBuffers()); + return compileIRWithoutConstants(std::move(IR)); +} + std::unique_ptr Interpreter::compileIR(std::unique_ptr IR) const { + auto function = compileIRWithoutConstants(std::move(IR)); + auto IFunction = static_cast(function.get()); + IFunction->collectConstants(IFunction->getIR()); + return function; +} + +std::unique_ptr +Interpreter::compileIRWithoutConstants(std::unique_ptr IR) const { MemoryAllocator constantWeightsAllocator("ConstantWeights", 0); MemoryAllocator placeholderWeightsAllocator("PlaceholderWeights", 0); MemoryAllocator activationsAllocator("Activations", 0); diff --git a/lib/Backends/Interpreter/Interpreter.h b/lib/Backends/Interpreter/Interpreter.h index 6f1efe8d42..e134e95e6a 100644 --- a/lib/Backends/Interpreter/Interpreter.h +++ b/lib/Backends/Interpreter/Interpreter.h @@ -37,8 +37,14 @@ class Interpreter final : public BackendUsingGlowIR { std::unique_ptr compileIR(std::unique_ptr IR) const override; + std::unique_ptr + compileIRWithoutConstants(std::unique_ptr IR) const; + std::unique_ptr compile(Function *F) const override; + std::unique_ptr + compileWithoutConstants(Function *F) const override; + bool isOpSupported(Kinded::Kind opKind, ElemKind elementTy) const override; bool shouldLower(const Node *N) const override; diff --git a/lib/Backends/Interpreter/InterpreterFunction.cpp b/lib/Backends/Interpreter/InterpreterFunction.cpp index 66520e9637..c95b7386ad 100644 --- a/lib/Backends/Interpreter/InterpreterFunction.cpp +++ b/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()) { diff --git a/lib/Backends/Interpreter/InterpreterFunction.h b/lib/Backends/Interpreter/InterpreterFunction.h index 77ed5ecb7c..9fb49aa187 100644 --- a/lib/Backends/Interpreter/InterpreterFunction.h +++ b/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: diff --git a/lib/Backends/OpenCL/OpenCL.cpp b/lib/Backends/OpenCL/OpenCL.cpp index e10d61685f..7038824159 100644 --- a/lib/Backends/OpenCL/OpenCL.cpp +++ b/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 OCLBackend::compileIR(std::unique_ptr IR) const { + auto function = compileIRWithoutConstants(std::move(IR)); + auto OCLFunction = static_cast(function.get()); + OCLFunction->collectConstants(OCLFunction->getIR()); + return function; +} + +std::unique_ptr +OCLBackend::compileIRWithoutConstants(std::unique_ptr IR) const { MemoryAllocator allocator("GPU", 0xFFFFFFFF); runtime::RuntimeBundle bundle = generateRuntimeBundle(*IR, allocator, allocator, allocator); return llvm::make_unique(std::move(IR), bundle); } - std::unique_ptr OCLBackend::compile(Function *F) const { auto IR = generateAndOptimizeIR(F, shouldShareBuffers()); return compileIR(std::move(IR)); } + +std::unique_ptr +OCLBackend::compileWithoutConstants(Function *F) const { + auto IR = generateAndOptimizeIR(F, shouldShareBuffers()); + return compileIRWithoutConstants(std::move(IR)); +} diff --git a/lib/Backends/OpenCL/OpenCL.h b/lib/Backends/OpenCL/OpenCL.h index 20a32de9fa..5662a1e138 100644 --- a/lib/Backends/OpenCL/OpenCL.h +++ b/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 compileIR(std::unique_ptr IR) const override; + std::unique_ptr + compileIRWithoutConstants(std::unique_ptr IR) const; std::unique_ptr compile(Function *F) const override; + std::unique_ptr + compileWithoutConstants(Function *F) const override; bool transformPostLowering(Function *F, CompilationMode mode) const override; diff --git a/tests/unittests/BackendCorrectnessTest.cpp b/tests/unittests/BackendCorrectnessTest.cpp index b3a566e781..f96d8f68ff 100644 --- a/tests/unittests/BackendCorrectnessTest.cpp +++ b/tests/unittests/BackendCorrectnessTest.cpp @@ -239,6 +239,11 @@ class MockCPUBackend : public BackendUsingGlowIR { std::unique_ptr compile(Function *F) const override { return backend_->compile(F); } + + std::unique_ptr + compileWithoutConstants(Function *F) const override { + return backend_->compile(F); + } std::unique_ptr compileIR(std::unique_ptr IR) const override { return backend_->compileIR(std::move(IR)); diff --git a/tests/unittests/BackendTest.cpp b/tests/unittests/BackendTest.cpp index 6155d62e9c..b96dfc87ec 100644 --- a/tests/unittests/BackendTest.cpp +++ b/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(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. diff --git a/tests/unittests/BackendTestUtils.h b/tests/unittests/BackendTestUtils.h index d69473558d..9685374bbc 100644 --- a/tests/unittests/BackendTestUtils.h +++ b/tests/unittests/BackendTestUtils.h @@ -32,6 +32,10 @@ class MockBackend : public Backend { std::unique_ptr compile(Function *F) const override { return llvm::make_unique(); } + std::unique_ptr + compileWithoutConstants(Function *F) const override { + return llvm::make_unique(); + } bool isOpSupported(Kinded::Kind opKind, ElemKind elementTy) const override { return false; } diff --git a/tests/unittests/QuantizationTest.cpp b/tests/unittests/QuantizationTest.cpp index d0abb497a2..262604b980 100644 --- a/tests/unittests/QuantizationTest.cpp +++ b/tests/unittests/QuantizationTest.cpp @@ -837,6 +837,11 @@ class MockQuantBackend : public Backend { std::unique_ptr compile(Function *F) const override { return backend_->compile(F); } + + std::unique_ptr + 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) {