Check if the pointer returned by malloc is nullptr. (#2347)

* Check if the pointer returned by malloc is nullptr.

* Use ngraph_error.

* Add more pointer checks.

* Replace malloc and free with ngraph_malloc and ngraph_free.

src/ngraph/runtime/aligned_buffer.cpp

@@ -14,10 +14,10 @@
 // limitations under the License.
 //*****************************************************************************
 
-#include <cstdlib> // llvm 8.1 gets confused about `malloc` otherwise
 #include <memory>
 
 #include "ngraph/runtime/aligned_buffer.hpp"
+#include "ngraph/util.hpp"
 
 using namespace ngraph;
 
@@ -34,7 +34,7 @@ runtime::AlignedBuffer::AlignedBuffer(size_t byte_size, size_t alignment)
     if (m_byte_size > 0)
     {
         size_t allocation_size = m_byte_size + alignment;
-        m_allocated_buffer = static_cast<char*>(malloc(allocation_size));
+        m_allocated_buffer = static_cast<char*>(ngraph_malloc(allocation_size));
         m_aligned_buffer = m_allocated_buffer;
         size_t mod = size_t(m_aligned_buffer) % alignment;
 
@@ -54,6 +54,6 @@ runtime::AlignedBuffer::~AlignedBuffer()
 {
     if (m_allocated_buffer != nullptr)
     {
-        free(m_allocated_buffer);
+        ngraph_free(m_allocated_buffer);
     }
 }

src/ngraph/runtime/cpu/cpu_tensor_view.cpp

@@ -24,6 +24,7 @@
 #include "ngraph/runtime/cpu/cpu_layout_descriptor.hpp"
 #include "ngraph/runtime/cpu/mkldnn_utils.hpp"
 #include "ngraph/shape.hpp"
+#include "ngraph/util.hpp"
 
 using namespace mkldnn;
 using namespace ngraph;
@@ -56,11 +57,7 @@ runtime::cpu::CPUTensorView::CPUTensorView(const ngraph::element::Type& element_
     else if (buffer_size > 0)
     {
         size_t allocation_size = buffer_size + BufferAlignment;
-        auto ptr = malloc(allocation_size);
-        if (!ptr)
-        {
-            throw ngraph_error("Error allocating CPU Tensor View memory");
-        }
+        auto ptr = ngraph_malloc(allocation_size);
         buffer = static_cast<char*>(ptr);
 
 // GCC major versions below 5 do not implement C++11 std::align
@@ -85,7 +82,7 @@ runtime::cpu::CPUTensorView::CPUTensorView(const ngraph::element::Type& element_
 
 runtime::cpu::CPUTensorView::~CPUTensorView()
 {
-    free(buffer);
+    ngraph_free(buffer);
 }
 
 char* runtime::cpu::CPUTensorView::get_data_ptr()

src/ngraph/runtime/cpu/mkldnn_emitter.hpp

@@ -40,6 +40,7 @@
 #include "ngraph/shape.hpp"
 #include "ngraph/strides.hpp"
 #include "ngraph/type/element_type.hpp"
+#include "ngraph/util.hpp"
 
 #define MKLDNN_DIMS(X) mkldnn::memory::dims(X.begin(), X.end())
 
@@ -55,8 +56,8 @@ namespace ngraph
             class MKLDNNWorkspace
             {
             public:
-                MKLDNNWorkspace(size_t size) { buf = reinterpret_cast<char*>(malloc(size)); }
-                ~MKLDNNWorkspace() { free(buf); }
+                MKLDNNWorkspace(size_t size) { buf = reinterpret_cast<char*>(ngraph_malloc(size)); }
+                ~MKLDNNWorkspace() { ngraph_free(buf); }
                 char* buf;
 
                 MKLDNNWorkspace(const MKLDNNWorkspace&) = delete;

src/ngraph/runtime/host_tensor.cpp

@@ -19,6 +19,7 @@
 
 #include "ngraph/descriptor/layout/dense_tensor_layout.hpp"
 #include "ngraph/runtime/host_tensor.hpp"
+#include "ngraph/util.hpp"
 
 using namespace ngraph;
 using namespace std;
@@ -46,7 +47,7 @@ runtime::HostTensor::HostTensor(const ngraph::element::Type& element_type,
     else if (m_buffer_size > 0)
     {
         size_t allocation_size = m_buffer_size + runtime::alignment;
-        m_allocated_buffer_pool = static_cast<char*>(malloc(allocation_size));
+        m_allocated_buffer_pool = static_cast<char*>(ngraph_malloc(allocation_size));
         m_aligned_buffer_pool = m_allocated_buffer_pool;
         size_t mod = size_t(m_aligned_buffer_pool) % alignment;
         if (mod != 0)
@@ -83,7 +84,7 @@ runtime::HostTensor::~HostTensor()
 {
     if (m_allocated_buffer_pool != nullptr)
     {
-        free(m_allocated_buffer_pool);
+        ngraph_free(m_allocated_buffer_pool);
     }
 }
 

src/ngraph/serializer.cpp

@@ -333,10 +333,10 @@ shared_ptr<ngraph::Function> ngraph::deserialize(istream& in)
                         {
                             if (info.get_name() == const_name)
                             {
-                                void* const_data = malloc(info.get_size());
+                                void* const_data = ngraph_malloc(info.get_size());
                                 reader.read(const_name, const_data, info.get_size());
                                 const_node = make_shared<op::Constant>(et, shape, const_data);
-                                free(const_data);
+                                ngraph_free(const_data);
                                 break;
                             }
                         }

src/ngraph/util.cpp

@@ -180,6 +180,25 @@ void ngraph::aligned_free(void* p)
 #endif
 }
 
+void* ngraph::ngraph_malloc(size_t size)
+{
+    auto ptr = malloc(size);
+    if (size != 0 && !ptr)
+    {
+        NGRAPH_ERR << "malloc failed to allocate memory of size " << size;
+        throw std::bad_alloc();
+    }
+    return ptr;
+}
+
+void ngraph::ngraph_free(void* ptr)
+{
+    if (ptr)
+    {
+        free(ptr);
+    }
+}
+
 size_t ngraph::round_up(size_t size, size_t alignment)
 {
     if (alignment == 0)

src/ngraph/util.hpp

@@ -18,6 +18,7 @@
 
 #include <chrono>
 #include <cmath>
+#include <cstdlib> // llvm 8.1 gets confused about `malloc` otherwise
 #include <functional>
 #include <iostream>
 #include <map>
@@ -182,6 +183,10 @@ namespace ngraph
 
     void* aligned_alloc(size_t alignment, size_t size);
     void aligned_free(void*);
+
+    void* ngraph_malloc(size_t size);
+    void ngraph_free(void*);
+
     size_t round_up(size_t size, size_t alignment);
     template <typename T>
     T apply_permutation(T input, ngraph::AxisVector order);