wip debug

src/ngraph/runtime/cpu/cpu_aligned_buffer.cpp

@@ -22,15 +22,14 @@
 
 using namespace ngraph;
 
-runtime::cpu::CPUAlignedBuffer::CPUAlignedBuffer(size_t byte_size, size_t alignment)
+runtime::cpu::CPUAlignedBuffer::CPUAlignedBuffer(size_t byte_size, size_t alignment, ngraph::runtime::cpu::CPUAllocator* cpu_allocator)
 {
     m_byte_size = byte_size;
-    AllocateFunc allocator = ngraph::runtime::cpu::CPUAllocator::framework_allocator;
+    m_cpu_allocator = cpu_allocator;
     if (m_byte_size > 0)
     {
         size_t allocation_size = m_byte_size + alignment;
-        m_allocated_buffer = static_cast<char*>(
-            ngraph::runtime::cpu::cpu_malloc(allocation_size, alignment, allocator));
+        m_allocated_buffer = static_cast<char*>(m_cpu_allocator->malloc(allocation_size));
         m_aligned_buffer = m_allocated_buffer;
         size_t mod = size_t(m_aligned_buffer) % alignment;
 
@@ -48,9 +47,8 @@ runtime::cpu::CPUAlignedBuffer::CPUAlignedBuffer(size_t byte_size, size_t alignm
 
 runtime::cpu::CPUAlignedBuffer::~CPUAlignedBuffer()
 {
-    DestroyFunc deallocator = ngraph::runtime::cpu::CPUAllocator::framework_deallocator;
     if (m_allocated_buffer != nullptr)
     {
-        ngraph::runtime::cpu::cpu_free(m_allocated_buffer, deallocator);
+        m_cpu_allocator->free(m_allocated_buffer);
     }
 }

src/ngraph/runtime/cpu/cpu_aligned_buffer.hpp

@@ -36,7 +36,7 @@ namespace ngraph
 class ngraph::runtime::cpu::CPUAlignedBuffer
 {
 public:
-    CPUAlignedBuffer(size_t byte_size, size_t alignment);
+    CPUAlignedBuffer(size_t byte_size, size_t alignment, ngraph::runtime::cpu::CPUAllocator* cpu_allocator);
     ~CPUAlignedBuffer();
 
     size_t size() const { return m_byte_size; }
@@ -47,7 +47,7 @@ private:
     CPUAlignedBuffer(CPUAlignedBuffer&&) = delete;
     CPUAlignedBuffer& operator=(const CPUAlignedBuffer&) = delete;
 
-    ngraph::runtime::cpu::CPUAllocator m_cpu_allocator;
+    ngraph::runtime::cpu::CPUAllocator* m_cpu_allocator;
     char* m_allocated_buffer;
     char* m_aligned_buffer;
     size_t m_byte_size;

src/ngraph/runtime/cpu/cpu_call_frame.cpp

@@ -146,13 +146,23 @@ void runtime::cpu::CPU_CallFrame::setup_runtime_context()
     size_t alignment = runtime::cpu::CPU_ExternalFunction::s_memory_pool_alignment;
 
     // assign the passed memory allocators
-    ngraph::runtime::cpu::CPUAllocator::framework_allocator = m_framework_allocator;
+    /*ngraph::runtime::cpu::CPUAllocator::framework_allocator = m_framework_allocator;
     ngraph::runtime::cpu::CPUAllocator::framework_deallocator = m_framework_deallocator;
     ngraph::runtime::cpu::CPUAllocator::alignment = alignment;
+    */
+    std::unique_ptr<ngraph::runtime::cpu::CPUAllocator> alloctor = nullptr;
+    if (m_framework_allocator && m_framework_deallocator)
+    {
+        alloctor = new ngraph::runtime::cpu::CPUAllocator(ngraph::runtime::FrameworkAllocator(m_framework_allocator, m_framework_deallocator, s_memory_pool_alignment));
+    }
+    else
+    {
+        allocator =  new ngraph::runtime::cpu::CPUAllocator(ngraph::runtime::SystemAllocator(s_memory_pool_alignment)); 
+    }
      
     for (auto buffer_size : m_external_function->get_memory_buffer_sizes())
     {
-        auto buffer = new CPUAlignedBuffer(buffer_size, alignment);
+        auto buffer = new CPUAlignedBuffer(buffer_size, alignment, allocator);
         ctx->memory_buffers.push_back(buffer);
     }
     const auto& mkldnn_emitter = m_external_function->get_mkldnn_emitter();

src/ngraph/runtime/cpu/cpu_mkl_allocator.cpp

@@ -14,64 +14,48 @@
 // limitations under the License.
 //*****************************************************************************
 
-#include "ngraph/runtime/cpu/cpu_mkl_allocator.hpp"
 #include <string>
-#include "ngraph/except.hpp"
 #include "ngraph/runtime/cpu/cpu_external_function.hpp"
+#include "ngraph/runtime/cpu/cpu_mkl_allocator.hpp"
 
 ngraph::runtime::cpu::CPUAllocator::CPUAllocator()
 {
 }
 
-AllocateFunc ngraph::runtime::cpu::CPUAllocator::framework_allocator = nullptr;
-DestroyFunc ngraph::runtime::cpu::CPUAllocator::framework_deallocator = nullptr;
-size_t ngraph::runtime::cpu::CPUAllocator::alignment =
-    ngraph::runtime::cpu::CPU_ExternalFunction::s_memory_pool_alignment;
-;
+ngraph::runtime::cpu::CPUAllocator::CPUAllocator(ngraph::runtime::Allocator* allocator)
+: m_allocator(allocator)
+{
+    
+}
 
-ngraph::runtime::cpu::CPUAllocator::CPUAllocator(AllocateFunc allocator,
-                                                 DestroyFunc deallocator,
-                                                 size_t alignment)
+ngraph::runtime::cpu::CPUAllocator::~CPUAllocator()
 {
-    mkl::i_malloc = MallocHook;
-    mkl::i_free = FreeHook;
 }
 
-void* ngraph::runtime::cpu::cpu_malloc(size_t size,
-                                       size_t alignment,
-                                       AllocateFunc framework_allocator)
+void* ngraph::runtime::cpu::cpuAllocator::malloc(size_t size)
 {
-    void* ptr;
-    if (framework_allocator != nullptr)
-    {
-        ptr = framework_allocator(nullptr, alignment, size);
-    }
-    else
-    {
-        ptr = malloc(size);
-    }
+    m_allocator->cpu_malloc(size);
+}
 
-    // check for exception
-    if (size != 0 && !ptr)
-    {
-        throw ngraph_error("malloc failed to allocate memory of size " + std::to_string(size));
-        throw std::bad_alloc();
-    }
-    return ptr;
+void ngraph::runtime::cpu::cpuAllocator::free(void* ptr)
+{
+     m_allocator->cpu_free(ptr);
 }
 
-void ngraph::runtime::cpu::cpu_free(void* ptr, DestroyFunc framework_deallocator)
+ngraph::runtime::SystemAllocator::SystemAllocator(size_t alignment)
+: m_alignment(alignment)
 {
-    if (framework_deallocator && ptr)
-    {
-        framework_deallocator(nullptr, ptr);
-    }
-    else if (ptr)
-    {
-        free(ptr);
-    }
+
 }
 
-ngraph::runtime::cpu::CPUAllocator::~CPUAllocator()
+ngraph::runtime::FrameworkAllocator::FrameworkAllocator(AllocateFunc allocator, DestroyFunc deallocator, size_t alignment)
+:   m_allocator(allocator)
+,   m_deallocator(deallocator)
+,   m_alignment(alignment)
 {
+
+
 }
+
+
+

src/ngraph/runtime/cpu/cpu_mkl_allocator.hpp

@@ -20,6 +20,7 @@
 #include <cstdint>
 #include "ngraph/runtime/aligned_buffer.hpp"
 #include "ngraph/util.hpp"
+#include "ngraph/except.hpp"
 
 using namespace ngraph;
 using AllocateFunc = void* (*)(void*, size_t, size_t);
@@ -44,6 +45,9 @@ namespace ngraph
 {
     namespace runtime
     {
+        class Allocator;
+        class SystemAllocator;
+        class FrameworkAllocator;
         namespace cpu
         {
             class CPUAllocator;
@@ -56,22 +60,94 @@ namespace ngraph
 class ngraph::runtime::cpu::CPUAllocator
 {
 public:
-    CPUAllocator(AllocateFunc allocator, DestroyFunc deallocator, size_t alignment);
+    CPUAllocator(ngraph::runtime::Allocator* alloctor);
     CPUAllocator();
     ~CPUAllocator();
 
-    static AllocateFunc framework_allocator;
-    static DestroyFunc framework_deallocator;
-    static size_t alignment;
+    void* malloc(size_t size);
+    void free(void* ptr);
 
 private:
-    static inline void* MallocHook(size_t size)
+
+    std::unique_ptr<ngraph::runtime::Allocator> m_allocator;
+    
+};
+
+
+// Abstarct class for the allocator 
+class  ngraph::runtime::Allocator
+{
+    public:
+        virtual void* cpu_malloc(void*, size_t size, size_t alignment) = 0;
+        virtual void cpu_free(void* ptr, void*) = 0;
+
+};
+
+
+class  ngraph::runtime::SystemAllocator : public  ngraph::runtime::Allocator
+{
+
+    public:
+        SystemAllocator(size_t alignment);
+        ~SystemAllocator();
+
+    void* cpu_malloc(void*, size_t size, size_t alignment) override 
+    {
+        void *ptr = malloc(size);
+         
+        // check for exception
+        if (size != 0 && !ptr)
         {
-        ngraph::runtime::cpu::cpu_malloc(size, alignment, framework_allocator);
+            throw ngraph_error("malloc failed to allocate memory of size " + std::to_string(size));
+            throw std::bad_alloc();
+        }
+        return ptr;
     }   
 
-    static inline void FreeHook(void* ptr)
+    void cpu_free(void* ptr, void*) override
+    {
+        if (ptr)
         {
-        ngraph::runtime::cpu::cpu_free(ptr, framework_deallocator);
+            free(ptr);
         }
+    }
+
+    private:
+        size_t m_alignment;
+};
+
+
+class  ngraph::runtime::FrameworkAllocator : public ngraph::runtime::Allocator
+{
+
+    public:
+    FrameworkAllocator(AllocateFunc allocator, DestroyFunc deallocator, size_t alignment);
+    ~FrameworkAllocator();
+    
+    void* cpu_malloc(void*, size_t size, size_t alignment) override 
+    {
+        void *ptr =  m_allocator(nullptr, alignment, size);
+         
+        // check for exception
+        if (size != 0 && !ptr)
+        {
+            throw ngraph_error("malloc failed to allocate memory of size " + std::to_string(size));
+            throw std::bad_alloc();
+        }
+        return ptr;
+    }   
+
+    void cpu_free(void* ptr, void*) override
+    {
+        if (ptr)
+        {
+            m_deallocator(nullptr, ptr);
+        }
+    }
+
+    private:
+    AllocateFunc m_allocator;
+    DestroyFunc m_deallocator;
+    size_t m_alignment;
+
 };