Remove static default allocator and create one on demand in the cpu backend

src/ngraph/runtime/aligned_buffer.cpp

@@ -32,29 +32,20 @@ runtime::AlignedBuffer::AlignedBuffer()
 {
 }
 
-runtime::AlignedBuffer::AlignedBuffer(size_t byte_size, size_t alignment)
-    : m_allocator(nullptr)
-    , m_byte_size(byte_size)
-{
-    m_byte_size = std::max<size_t>(1, byte_size);
-    size_t allocation_size = m_byte_size + alignment;
-    m_allocated_buffer = static_cast<char*>(malloc(allocation_size));
-    m_aligned_buffer = m_allocated_buffer;
-    size_t mod = size_t(m_aligned_buffer) % alignment;
-
-    if (mod != 0)
-    {
-        m_aligned_buffer += (alignment - mod);
-    }
-}
-
 runtime::AlignedBuffer::AlignedBuffer(size_t byte_size, size_t alignment, Allocator* allocator)
     : m_allocator(allocator)
     , m_byte_size(byte_size)
 {
     m_byte_size = std::max<size_t>(1, byte_size);
     size_t allocation_size = m_byte_size + alignment;
-    m_allocated_buffer = static_cast<char*>(m_allocator->malloc(allocation_size, alignment));
+    if (allocator)
+    {
+        m_allocated_buffer = static_cast<char*>(m_allocator->malloc(allocation_size, alignment));
+    }
+    else
+    {
+        m_allocated_buffer = static_cast<char*>(malloc(allocation_size));
+    }
     m_aligned_buffer = m_allocated_buffer;
     size_t mod = size_t(m_aligned_buffer) % alignment;
 

src/ngraph/runtime/aligned_buffer.hpp

@@ -34,11 +34,10 @@ namespace ngraph
 class ngraph::runtime::AlignedBuffer
 {
 public:
-    AlignedBuffer(size_t byte_size, size_t alignment);
     // Allocator objects and the allocation interfaces are owned by the
     // creators of AlignedBuffers. They need to ensure that the lifetime of
     // allocator exceeds the lifetime of this AlignedBuffer.
-    AlignedBuffer(size_t byte_size, size_t alignment, Allocator* allocator);
+    AlignedBuffer(size_t byte_size, size_t alignment, Allocator* allocator = nullptr);
 
     AlignedBuffer();
     ~AlignedBuffer();

src/ngraph/runtime/allocator.cpp

@@ -49,8 +49,7 @@ public:
     }
 };
 
-ngraph::runtime::Allocator* ngraph::runtime::get_default_allocator()
+std::unique_ptr<ngraph::runtime::Allocator> ngraph::runtime::create_default_allocator()
 {
-    static std::unique_ptr<DefaultAllocator> allocator(new DefaultAllocator());
-    return allocator.get();
+    return std::unique_ptr<DefaultAllocator>(new DefaultAllocator());
 }

src/ngraph/runtime/allocator.hpp

@@ -28,9 +28,9 @@ namespace ngraph
     {
         class Allocator;
         class DefaultAllocator;
-        /// \brief Returns a pointer to a statically allocated singleton
-        ///        allocator that calls into system allocation libraries
-        Allocator* get_default_allocator();
+        /// \brief Create a default allocator that calls into system
+        ///        allocation libraries
+        std::unique_ptr<Allocator> create_default_allocator();
     }
 }
 

src/ngraph/runtime/backend.hpp

@@ -143,11 +143,7 @@ public:
     virtual std::shared_ptr<ngraph::Node> get_backend_op(const std::string& op_name, ...);
 
     /// \brief Returns memory allocator used by backend for host allocations
-    virtual Allocator* get_host_memory_allocator()
-    {
-        return ngraph::runtime::get_default_allocator();
-    }
-
+    virtual Allocator* get_host_memory_allocator() { return nullptr; }
     /// \brief Set the host memory allocator to be used by the backend
     /// \param allocator is pointer to host memory allocator object
     virtual void set_host_memory_allocator(std::unique_ptr<Allocator> allocator) {}

src/ngraph/runtime/cpu/cpu_backend.cpp

@@ -59,6 +59,11 @@ namespace
     } s_cpu_static_init;
 }
 
+runtime::cpu::CPU_Backend::~CPU_Backend()
+{
+    m_exec_map.clear();
+}
+
 shared_ptr<runtime::cpu::CPU_CallFrame> runtime::cpu::CPU_Backend::make_call_frame(
     const shared_ptr<runtime::cpu::CPU_ExternalFunction>& external_function,
     ngraph::pass::PassConfig& pass_config,
@@ -159,19 +164,23 @@ void runtime::cpu::CPU_Backend::remove_compiled_function(shared_ptr<Executable>
 
 runtime::Allocator* runtime::cpu::CPU_Backend::get_host_memory_allocator()
 {
-    if (m_allocator)
-    {
-        return m_allocator.get();
-    }
-    else
+    if (!m_allocator)
     {
-        return runtime::get_default_allocator();
+        m_allocator = std::move(create_default_allocator());
     }
+    return m_allocator.get();
 }
 
 void runtime::cpu::CPU_Backend::set_host_memory_allocator(
     std::unique_ptr<runtime::Allocator> allocator)
 {
+    if (m_allocator)
+    {
+        // Resources allocated with the existing allocator might still be around and expect it
+        // to be available for freeing. We cannot switch to the new allocator
+        throw ngraph_error(
+            "Allocator already exists. Changing allocators mid-execution is not permitted.");
+    }
     m_allocator = std::move(allocator);
 }
 

src/ngraph/runtime/cpu/cpu_backend.hpp

@@ -36,6 +36,8 @@ namespace ngraph
             class CPU_BACKEND_API CPU_Backend : public runtime::Backend
             {
             public:
+                ~CPU_Backend() override;
+
                 std::shared_ptr<CPU_CallFrame>
                     make_call_frame(const std::shared_ptr<CPU_ExternalFunction>& external_function,
                                     ngraph::pass::PassConfig& pass_config,