- add new class for cpu_aligned_buffers

src/ngraph/runtime/aligned_buffer.cpp

@@ -32,7 +32,6 @@ runtime::AlignedBuffer::AlignedBuffer()
 runtime::AlignedBuffer::AlignedBuffer(size_t byte_size, size_t alignment)
 {
     m_byte_size = byte_size;
-    auto cpu_allocator = ngraph::runtime::cpu::GetCPUAllocator();
     if (m_byte_size > 0)
     {
         size_t allocation_size = m_byte_size + alignment;

src/ngraph/runtime/cpu/CMakeLists.txt

@@ -17,6 +17,7 @@
 include(FindOpenMP)
 
 set(SRC
+    cpu_aligned_buffer.cpp
     cpu_backend.cpp
     cpu_builder.cpp
     cpu_call_frame.cpp

src/ngraph/runtime/cpu/cpu_aligned_buffer.cpp

+//*****************************************************************************
+// Copyright 2017-2019 Intel Corporation
+//
+// 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.
+//*****************************************************************************
+
+#include <memory>
+
+#include "ngraph/runtime/cpu/cpu_aligned_buffer.hpp"
+#include "ngraph/runtime/cpu/cpu_mkl_allocator.hpp"
+#include "ngraph/util.hpp"
+
+using namespace ngraph;
+
+runtime::cpu::CPUAlignedBuffer::CPUAlignedBuffer()
+    : m_cpu_allocator(nullptr, nullptr, 0)
+    , m_allocated_buffer(nullptr)
+    , m_aligned_buffer(nullptr)
+    , m_byte_size(0)
+{
+}
+
+runtime::cpu::CPUAlignedBuffer::CPUAlignedBuffer(size_t byte_size,
+                                                 size_t alignment,
+                                                 ngraph::runtime::cpu::CPUAllocator& cpu_allocator)
+{
+    m_byte_size = byte_size;
+    m_cpu_allocator = cpu_allocator;
+    if (m_byte_size > 0)
+    {
+        size_t allocation_size = m_byte_size + alignment;
+        m_allocated_buffer = static_cast<char*>(m_cpu_allocator.cpu_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);
+        }
+    }
+    else
+    {
+        m_allocated_buffer = nullptr;
+        m_aligned_buffer = nullptr;
+    }
+}
+
+runtime::cpu::CPUAlignedBuffer::~CPUAlignedBuffer()
+{
+    if (m_allocated_buffer != nullptr)
+    {
+        m_cpu_allocator.cpu_free(m_allocated_buffer);
+    }
+}

src/ngraph/runtime/cpu/cpu_aligned_buffer.hpp

+//*****************************************************************************
+// Copyright 2017-2019 Intel Corporation
+//
+// 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.
+//*****************************************************************************
+
+#pragma once
+
+#include <cstddef>
+#include "ngraph/runtime/cpu/cpu_mkl_allocator.hpp"
+
+namespace ngraph
+{
+    namespace runtime
+    {
+        namespace cpu
+        {
+            class CPUAlignedBuffer;
+        }
+    }
+}
+
+/// \brief Allocates a block of memory on the specified alignment. The actual size of the
+/// allocated memory is larger than the requested size by the alignment, so allocating 1 byte
+/// on 64 byte alignment will allocate 65 bytes.
+class ngraph::runtime::cpu::CPUAlignedBuffer
+{
+public:
+    CPUAlignedBuffer(size_t byte_size,
+                     size_t alignment,
+                     ngraph::runtime::cpu::CPUAllocator& cpu_allocator);
+    CPUAlignedBuffer();
+    ~CPUAlignedBuffer();
+
+    size_t size() const { return m_byte_size; }
+    void* get_ptr(size_t offset) const { return m_aligned_buffer + offset; }
+    void* get_ptr() const { return m_aligned_buffer; }
+private:
+    CPUAlignedBuffer(const CPUAlignedBuffer&) = delete;
+    CPUAlignedBuffer(CPUAlignedBuffer&&) = delete;
+    CPUAlignedBuffer& operator=(const CPUAlignedBuffer&) = delete;
+
+    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

@@ -17,6 +17,7 @@
 #include <algorithm>
 
 #include "ngraph/runtime/aligned_buffer.hpp"
+#include "ngraph/runtime/cpu/cpu_aligned_buffer.hpp"
 #include "ngraph/runtime/cpu/cpu_call_frame.hpp"
 #include "ngraph/runtime/cpu/cpu_external_function.hpp"
 #include "ngraph/runtime/cpu/cpu_tensor_view.hpp"
@@ -113,7 +114,6 @@ void runtime::cpu::CPU_CallFrame::propagate_layouts(
 void runtime::cpu::CPU_CallFrame::setup_runtime_context()
 {
     ctx = new CPURuntimeContext;
-
     ctx->pc = 0;
     ctx->op_durations = nullptr;
     if (runtime::cpu::IsTracingEnabled())
@@ -126,9 +126,10 @@ void runtime::cpu::CPU_CallFrame::setup_runtime_context()
 
     // Create temporary buffer pools
     size_t alignment = runtime::cpu::CPU_ExternalFunction::s_memory_pool_alignment;
+    CPUAllocator cpu_allocator(nullptr, nullptr, alignment);
     for (auto buffer_size : m_external_function->get_memory_buffer_sizes())
     {
-        auto buffer = new AlignedBuffer(buffer_size, alignment);
+        auto buffer = new CPUAlignedBuffer(buffer_size, alignment, cpu_allocator);
         ctx->memory_buffers.push_back(buffer);
     }
     const auto& mkldnn_emitter = m_external_function->get_mkldnn_emitter();

src/ngraph/runtime/cpu/cpu_external_function.cpp

@@ -132,6 +132,7 @@
 #include "ngraph/pass/reshape_sinking.hpp"
 #include "ngraph/pass/zero_dim_tensor_elimination.hpp"
 #include "ngraph/runtime/aligned_buffer.hpp"
+#include "ngraph/runtime/cpu/cpu_aligned_buffer.hpp"
 #include "ngraph/runtime/cpu/cpu_backend.hpp"
 #include "ngraph/runtime/cpu/cpu_builder.hpp"
 #include "ngraph/runtime/cpu/cpu_call_frame.hpp"

src/ngraph/runtime/cpu/cpu_mkl_allocator.cpp

@@ -18,6 +18,13 @@
 #include <string>
 #include "ngraph/except.hpp"
 
+ngraph::runtime::cpu::CPUAllocator::CPUAllocator()
+    : m_framework_allocator(nullptr)
+    , m_framework_deallocator(nullptr)
+    , m_alignment(0)
+{
+}
+
 ngraph::runtime::cpu::CPUAllocator::CPUAllocator(AllocateFunc allocator,
                                                  DestroyFunc deallocator,
                                                  size_t alignment)
@@ -25,8 +32,8 @@ ngraph::runtime::cpu::CPUAllocator::CPUAllocator(AllocateFunc allocator,
     , m_framework_deallocator(deallocator)
     , m_alignment(alignment)
 {
-    mkl::i_malloc = MallocHook;
-    mkl::i_free = FreeHook;
+    //    mkl::i_malloc = MallocHook;
+    //    mkl::i_free = FreeHook;
 }
 
 void* ngraph::runtime::cpu::CPUAllocator::cpu_malloc(size_t size)
@@ -62,8 +69,6 @@ void ngraph::runtime::cpu::CPUAllocator::cpu_free(void* ptr)
     }
 }
 
-ngraph::runtime::cpu::CPUAllocator& ngraph::runtime::cpu::GetCPUAllocator()
+ngraph::runtime::cpu::CPUAllocator::~CPUAllocator()
 {
-    static ngraph::runtime::cpu::CPUAllocator cpu_allocator(nullptr, nullptr, 4096);
-    return cpu_allocator;
 }

src/ngraph/runtime/cpu/cpu_mkl_allocator.hpp

@@ -47,7 +47,6 @@ namespace ngraph
         namespace cpu
         {
             class CPUAllocator;
-            extern CPUAllocator& GetCPUAllocator();
         }
     }
 }
@@ -56,17 +55,18 @@ class ngraph::runtime::cpu::CPUAllocator
 {
 public:
     CPUAllocator(AllocateFunc allocator, DestroyFunc deallocator, size_t alignment);
-    //~CPUAllocator();
+    CPUAllocator();
+    ~CPUAllocator();
 
     void* cpu_malloc(size_t size);
     void cpu_free(void* ptr);
 
 private:
+    size_t m_alignment;
     AllocateFunc m_framework_allocator;
     DestroyFunc m_framework_deallocator;
-    size_t m_alignment;
 
-    static inline void* MallocHook(size_t size)
+    /*static inline void* MallocHook(size_t size)
     {
         ngraph::runtime::cpu::GetCPUAllocator().cpu_malloc(size);
     }
@@ -74,5 +74,5 @@ private:
     static inline void FreeHook(void* ptr)
     {
         ngraph::runtime::cpu::GetCPUAllocator().cpu_free(ptr);
-    }
+    }*/
 };

src/ngraph/runtime/cpu/cpu_runtime_context.hpp

@@ -39,9 +39,11 @@ namespace ngraph
 {
     namespace runtime
     {
-        class AlignedBuffer;
+        namespace cpu
+        {
+            class CPUAlignedBuffer;
+        }
     }
-
     class State;
 }
 
@@ -62,7 +64,7 @@ namespace ngraph
                 bool* p_en;
                 bool first_iteration;
                 mkldnn::primitive* const* mkldnn_primitives;
-                std::vector<AlignedBuffer*> memory_buffers;
+                std::vector<CPUAlignedBuffer*> memory_buffers;
                 char* const* mkldnn_workspaces;
                 tbb::flow::graph* G;
                 tbb::global_control* c;