Add new memory handler for generated code

src/ngraph/runtime/cpu/external_function.cpp

@@ -74,6 +74,7 @@
 #include "ngraph/pass/memory_layout.hpp"
 #include "ngraph/pass/topological_sort.hpp"
 #include "ngraph/runtime/cpu/call_frame.hpp"
+#include "ngraph/runtime/cpu/cpu_backend.hpp"
 #include "ngraph/runtime/cpu/emitter.hpp"
 #include "ngraph/runtime/cpu/external_function.hpp"
 #include "ngraph/runtime/utils.hpp"
@@ -93,25 +94,6 @@ static StaticInitializers s_static_initializers;
 
 using ngraph::descriptor::layout::DenseTensorViewLayout;
 
-extern "C" void
-    allocate_aligned_buffer(size_t size, size_t alignment, char** allocated, char** aligned_ptr)
-{
-    size_t allocation_size = size + alignment;
-    *allocated = static_cast<char*>(malloc(allocation_size));
-    *aligned_ptr = *allocated;
-    size_t mod = size_t(*aligned_ptr) % alignment;
-
-    if (mod != 0)
-    {
-        (*aligned_ptr) += (alignment - mod);
-    }
-}
-
-extern "C" void free_aligned_buffer(void* allocated)
-{
-    free(allocated);
-}
-
 #define TI(x) type_index(typeid(x))
 
 static const OpMap dispatcher{
@@ -211,28 +193,18 @@ void ExternalFunction::compile()
 
 #include "ngraph/runtime/cpu/cpu_kernels.hpp"
 #include "ngraph/runtime/cpu/eigen_utils.hpp"
+#include "ngraph/runtime/cpu/memory_handler.hpp"
 
 using namespace ngraph::runtime::cpu::eigen;
 
-extern "C" void allocate_aligned_buffer(
-    size_t size,
-    size_t alignment,
-    char** allocated,
-    char** aligned_ptr);
-
-extern "C" void free_aligned_buffer(void* allocated);
-
 )";
 
-    TU << "// Declare any functions that are not main\n";
+    TU << "// Declare all functions\n";
     for (shared_ptr<Function> f : pass_manager.get_state().get_functions())
     {
-        if (f != m_function)
-        {
-            TU << "extern \"C\" void " << f->get_name() << "(\n";
-            TU << "    const std::vector<void*>& inputs,\n";
-            TU << "    const std::vector<void*>& outputs);\n";
-        }
+        TU << "extern \"C\" void " << f->get_name() << "(\n";
+        TU << "    const std::vector<void*>& inputs,\n";
+        TU << "    const std::vector<void*>& outputs);\n";
     }
     TU << "\n";
 
@@ -244,12 +216,10 @@ extern "C" void free_aligned_buffer(void* allocated);
         TU << "{\n";
 
         TU.indent++;
-        TU << "// Allocate the memory pool\n";
         size_t temp_pool_size = pass_manager.get_state().get_temporary_pool_size();
-        TU << "char* allocated_buffer_pool;\n";
-        TU << "char* aligned_buffer_pool;\n";
-        TU << "allocate_aligned_buffer(" << temp_pool_size << ", 64"
-           << ", &allocated_buffer_pool, &aligned_buffer_pool);\n";
+        TU << "// Allocate the memory pool\n";
+        TU << "ngraph::runtime::cpu::MemoryHandler memory_handler(" << temp_pool_size << ", "
+           << ngraph::runtime::cpu::alignment << ");\n";
         TU << "\n";
 
         TU << "// Define temporary tensors\n";
@@ -258,8 +228,8 @@ extern "C" void free_aligned_buffer(void* allocated);
             for (descriptor::Tensor* tensor : node->liveness_new_list)
             {
                 TU << tensor->get_element_type() << "* " << tensor->get_name() << " = ("
-                   << tensor->get_element_type() << "*)(aligned_buffer_pool + "
-                   << tensor->get_pool_offset() << ");\n";
+                   << tensor->get_element_type() << "*)(memory_handler.get_ptr("
+                   << tensor->get_pool_offset() << "));\n";
             }
         }
         TU << "\n";
@@ -321,8 +291,6 @@ extern "C" void free_aligned_buffer(void* allocated);
             handler->second(&emitter, node.get(), this, in, out);
         }
 
-        TU << "\nfree_aligned_buffer(allocated_buffer_pool);\n";
-
         TU.indent--;
 
         // End TU

src/ngraph/runtime/cpu/tensor_view.cpp

@@ -20,35 +20,47 @@
 using namespace ngraph;
 using namespace std;
 
-extern "C" void
-    allocate_aligned_buffer(size_t size, size_t alignment, char** allocated, char** aligned_ptr);
-extern "C" void free_aligned_buffer(void* allocated);
-
 runtime::cpu::CPUTensorView::CPUTensorView(const ngraph::element::Type& element_type,
                                            const Shape& shape)
     : runtime::TensorView(std::make_shared<ngraph::descriptor::PrimaryTensorView>(
           std::make_shared<ngraph::TensorViewType>(element_type, shape), "external", true, true))
+    , m_allocated_buffer_pool(nullptr)
+    , m_aligned_buffer_pool(nullptr)
+
 {
     m_descriptor->set_tensor_view_layout(
         std::make_shared<ngraph::descriptor::layout::DenseTensorViewLayout>(*m_descriptor));
 
     m_buffer_size = m_descriptor->get_tensor_view_layout()->get_size() * element_type.size();
-    allocate_aligned_buffer(m_buffer_size, runtime::cpu::alignment, &m_allocated, &m_buffer);
+    if (m_buffer_size > 0)
+    {
+        size_t allocation_size = m_buffer_size + runtime::cpu::alignment;
+        m_allocated_buffer_pool = static_cast<char*>(malloc(allocation_size));
+        m_aligned_buffer_pool = m_allocated_buffer_pool;
+        size_t mod = size_t(m_aligned_buffer_pool) % alignment;
+        if (mod != 0)
+        {
+            m_aligned_buffer_pool += (alignment - mod);
+        }
+    }
 }
 
 runtime::cpu::CPUTensorView::~CPUTensorView()
 {
-    free_aligned_buffer(m_allocated);
+    if (m_allocated_buffer_pool != nullptr)
+    {
+        free(m_allocated_buffer_pool);
+    }
 }
 
 char* runtime::cpu::CPUTensorView::get_data_ptr()
 {
-    return m_buffer;
+    return m_aligned_buffer_pool;
 }
 
 const char* runtime::cpu::CPUTensorView::get_data_ptr() const
 {
-    return m_buffer;
+    return m_aligned_buffer_pool;
 }
 
 void runtime::cpu::CPUTensorView::write(const void* source, size_t tensor_offset, size_t n)

src/ngraph/runtime/cpu/tensor_view.hpp

@@ -52,7 +52,7 @@ public:
     void read(void* p, size_t tensor_offset, size_t n) const override;
 
 private:
-    char* m_allocated;
-    char* m_buffer;
+    char* m_allocated_buffer_pool;
+    char* m_aligned_buffer_pool;
     size_t m_buffer_size;
 };