[nnc] Fixed memory leak and heap overflow (#2703)

* Fixed memory leak and heap overflow in softbackend elementwise op
* Collect caffe2 model problems before tensors loading

Signed-off-by: Ivan Vagin <ivan.vagin@samsung.com>

diff --git a/contrib/nnc/passes/caffe2_frontend/caffe2_importer.cpp b/contrib/nnc/passes/caffe2_frontend/caffe2_importer.cpp
index 3ed8086..1b874c3 100644 (file)
--- a/contrib/nnc/passes/caffe2_frontend/caffe2_importer.cpp
+++ b/contrib/nnc/passes/caffe2_frontend/caffe2_importer.cpp
@@ -74,9 +74,9 @@ void Caffe2Importer::import() {
     throw PassException("Could not load model: " + _initNet + "\n");
   _net->MergeFrom(*net2);
 
-  preloadAllTensors();
-
   collectUnsupportedOps();
+
+  preloadAllTensors();
 }
 
 mir::Graph* Caffe2Importer::createIR() {

diff --git a/contrib/nnc/passes/soft_backend/code_snippets/cpp_batchnorm.def b/contrib/nnc/passes/soft_backend/code_snippets/cpp_batchnorm.def
index 95bb817..08305ed 100644 (file)
--- a/contrib/nnc/passes/soft_backend/code_snippets/cpp_batchnorm.def
+++ b/contrib/nnc/passes/soft_backend/code_snippets/cpp_batchnorm.def
@@ -23,5 +23,5 @@ void batchNorm(Tensor &out, const char *params, const Tensor &in)
   (float)avgFraction;
   bool spatial = deserializeT<int32_t>(params);
   (void)spatial;
-  out.fillData(in.getData());
+  out.fillData(in.getData(), in.getShape().getNumElems());
 }

diff --git a/contrib/nnc/passes/soft_backend/code_snippets/cpp_dropout.def b/contrib/nnc/passes/soft_backend/code_snippets/cpp_dropout.def
index 3415c43..47b54f5 100644 (file)
--- a/contrib/nnc/passes/soft_backend/code_snippets/cpp_dropout.def
+++ b/contrib/nnc/passes/soft_backend/code_snippets/cpp_dropout.def
@@ -19,5 +19,5 @@ void dropout(Tensor &out, const char *params, const Tensor &in)
   out.reShape(in.getShape());
   float rate = deserializeT<float>(params);
   (void)rate;
-  out.fillData(in.getData());
+  out.fillData(in.getData(), in.getShape().getNumElems());
 }

diff --git a/contrib/nnc/passes/soft_backend/code_snippets/cpp_header_types.def b/contrib/nnc/passes/soft_backend/code_snippets/cpp_header_types.def
index 239e6eb..3c57921 100644 (file)
--- a/contrib/nnc/passes/soft_backend/code_snippets/cpp_header_types.def
+++ b/contrib/nnc/passes/soft_backend/code_snippets/cpp_header_types.def
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include <cstring>
 #include <initializer_list>
 #include <memory>
 #include <cassert>
@@ -137,11 +138,10 @@ public:
   }
 
   /** Copies data from external source into table*/
-  void fillData(const float *data)
+  void fillData(const float *data, const index_t num_elements)
   {
     assert(_managed);
-    for (index_t i = 0; i < _shape.getNumElems(); ++i)
-      _data[i] = data[i];
+    std::memcpy(_data, data, num_elements * sizeof(float));
   }
 
   Tensor &operator=(const Tensor &t)
@@ -150,7 +150,7 @@ public:
     if (this != &t)
     {
       reShape(t.getShape());
-      fillData(t._data);
+      fillData(t._data, t._shape.getNumElems());
     }
     return *this;
   }

diff --git a/contrib/nnc/passes/soft_backend/code_snippets/cpp_operations.def b/contrib/nnc/passes/soft_backend/code_snippets/cpp_operations.def
index d837a9b..29b5106 100644 (file)
--- a/contrib/nnc/passes/soft_backend/code_snippets/cpp_operations.def
+++ b/contrib/nnc/passes/soft_backend/code_snippets/cpp_operations.def
@@ -451,7 +451,7 @@ void biasAdd(Tensor &out, const char *params, const Tensor &in)
   Kernel bias = deserializeKernel(params);
 
   out.reShape(in.getShape());
-  out.fillData(in.getData());
+  out.fillData(in.getData(), in.getShape().getNumElems());
 
   AddBiasAndEvalActivationFunction(bias.data, bias.dims, out.getData(), shapeToDims(out.getShape()));
 }
@@ -516,37 +516,43 @@ void tanhActivation(Tensor &out, const char* params, const Tensor& in) {
 
 // These operations (add, mul, max) takes as input multiple tensors, at least 2, likely less then 7
 // parameter pack provides generalization for all possible number of inputs
-template <typename F, class ...Args>
-void ElementWise(Tensor &out,
-  const char *params, const Args &...inputs) {
-  const float *input[] = {inputs.getData()...};
-  RuntimeShape in_shapes[] = {shapeToRuntimeShape(inputs.getShape())...};
+template <typename F, typename ...Args>
+void ElementWise(Tensor &out, const char* params, const Args& ...inputs) {
+  static_assert(sizeof...(inputs) >= 2, "ElementWise op must have >= 2 inputs");
+
+  const float* input[] = {inputs.getData()...};
+  Shape in_shapes[] = {inputs.getShape()...};
+  RuntimeShape in_runtime_shapes[] = {shapeToRuntimeShape(inputs.getShape())...};
 
-  const int32_t num_inputs = sizeof(input) / sizeof(const float*);
+  const int32_t num_inputs = sizeof...(inputs);
   const bool needs_broadcast = (bool)deserializeT<int32_t>(params);
   const Shape out_shape = deserializeShape(params);
 
   out.reShape(out_shape);
-  out.fillData(input[0]);
+  out.fillData(input[0], in_shapes[0].getNumElems());
   const auto out_rt = shapeToRuntimeShape(out_shape);
   if (!needs_broadcast) {
     for (int32_t i = 1; i < num_inputs; ++i) {
       F::Call(out.getData(), out_rt,
-              input[i], in_shapes[i],
+              input[i], in_runtime_shapes[i],
               out.getData(), out_rt,
-              0);
+              false);
     }
   } else {
-    auto running_shape = RuntimeShape::ExtendedShape(4, in_shapes[0]);
+    auto running_shape = RuntimeShape::ExtendedShape(4, in_runtime_shapes[0]);
+    std::vector<float> inp_tmp(out_shape.getNumElems());
+
     for (int32_t i = 1; i < num_inputs; ++i) {
-      float* inp_tmp = new float[running_shape.FlatSize()];
-      memcpy(inp_tmp, out.getData(), (size_t)running_shape.FlatSize() * 4);
-      F::Call(inp_tmp, running_shape,
-              input[i], in_shapes[i],
+      assert(running_shape.FlatSize() <= out_shape.getNumElems());
+
+      std::copy(out.getData(), out.getData() + running_shape.FlatSize(), inp_tmp.begin());
+
+      F::Call(inp_tmp.data(), running_shape,
+              input[i], in_runtime_shapes[i],
               out.getData(), out_rt,
-              1);
+              true);
       // This modifies the running shape
-      running_shape.maxShape(RuntimeShape::ExtendedShape(4, in_shapes[i]));
+      running_shape.maxShape(RuntimeShape::ExtendedShape(4, in_runtime_shapes[i]));
     }
   }
 
@@ -558,7 +564,7 @@ void reshape(Tensor& out, const char* params, const Tensor& in) {
   assert(out_s.getNumElems() == in.getShape().getNumElems());
 
   out.reShape(out_s);
-  out.fillData(in.getData());
+  out.fillData(in.getData(), in.getShape().getNumElems());
 }
 
 void reduceMean(Tensor& out, const char* params, const Tensor& in) {