Removed the memory leaks directly in execution session (#1746)

* removed the memory leaks directly in execution session
Signed-off-by: Alexandre Eichenberger <[email protected]>
Co-authored-by: chentong319 <[email protected]>

docs/doc_example/CMakeLists.txt

@@ -82,7 +82,6 @@ add_onnx_mlir_executable(OMRuntimeTest
   LINK_LIBS PRIVATE
   OnnxMlirCompiler
   ExecutionSession
-  cruntime
   )
 
 

src/Runtime/ExecutionSession.cpp

@@ -25,6 +25,7 @@
 #include "llvm/Support/Path.h"
 
 #include "ExecutionSession.hpp"
+#include "OMTensorListHelper.hpp"
 
 namespace onnx_mlir {
 const std::string ExecutionSession::_queryEntryPointsName =
@@ -85,6 +86,13 @@ std::vector<OMTensorUniquePtr> ExecutionSession::run(
   auto *wrappedInput = omTensorListCreate(&omts[0], (int64_t)omts.size());
 
   auto *wrappedOutput = _entryPointFunc(wrappedInput);
+
+  // We created a wrapper for the input list, but the input list does not really
+  // own the tensor in the list, as they are coming as OMTensorUniquePtr. So we
+  // need to simply deallocate the list structure without touching the
+  // OMTensors.
+  omTensorListDestroyShallow(wrappedInput);
+
   if (!wrappedOutput)
     throw std::runtime_error(reportErrnoError());
   std::vector<OMTensorUniquePtr> outs;
@@ -93,6 +101,12 @@ std::vector<OMTensorUniquePtr> ExecutionSession::run(
     outs.emplace_back(OMTensorUniquePtr(
         omTensorListGetOmtByIndex(wrappedOutput, i), omTensorDestroy));
   }
+
+  // We created a wrapper for the output list, but the output list does not
+  // really own the tensor in the list, as they are returned in a vector of
+  // OMTensorUniquePtr. So we need to simply deallocate the list structure
+  // without touching the OMTensors.
+  omTensorListDestroyShallow(wrappedOutput);
   errno = 0; // No errors.
   return outs;
 }

src/Runtime/OMTensor.inc

@@ -284,8 +284,9 @@ OMTensor *omTensorCreateEmpty(
 void omTensorDestroy(OMTensor *tensor) {
   if (!tensor)
     return;
-  if (tensor->_owning)
+  if (tensor->_owning) {
     free(tensor->_allocatedPtr);
+  }
   free(tensor->_shape);
   free(tensor->_strides);
   free(tensor);

src/Runtime/OMTensorList.inc

@@ -103,8 +103,23 @@ void omTensorListDestroy(OMTensorList *list) {
     return;
   for (int64_t i = 0; i < list->_size; i++)
     omTensorDestroy(list->_omts[i]);
-  if (list->_owning)
+  if (list->_owning) {
     free(list->_omts);
+  }
+  free(list);
+}
+
+/* OMTensorList destroyer which does not destroy the tensors. Currently an
+ * unpublished call that is mainly used to handle the unique tensor scheme used
+ * in Execution Session.
+ */
+void omTensorListDestroyShallow(OMTensorList *list) {
+  if (!list)
+    return;
+  // Omit destruction of the OMTensors.
+  if (list->_owning) {
+    free(list->_omts);
+  }
   free(list);
 }
 

src/Runtime/OMTensorListHelper.hpp

@@ -0,0 +1,25 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+//===----- OMTensorListHelper.hpp - OMTensor List Helper Func header ------===//
+//
+// Copyright 2022 The IBM Research Authors.
+//
+// =============================================================================
+//
+// This file contains declaration of OMTensorList C++ helper functions.
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "OnnxMlirRuntime.h"
+
+/*
+ * Destroy the OMTensorList data structure (including the array of OMTensor when
+ * ownership is asserted), but not the OMTensor themselves. Assumed here is that
+ * their live range is managed explicitly or implicitly using a different
+ * mechanism.
+ */
+void omTensorListDestroyShallow(OMTensorList *list);

test/numerical/TestLoop.cpp

@@ -118,7 +118,7 @@ bool isOMLoopTheSameAsNaiveImplFor(std::string moduleIR,
   omTensorGetElem<bool>(condTensor.get(), {}) = true;
   inputs.emplace_back(move(condTensor));
 
-  auto *yInitShape = new int64_t[1]{1};
+  int64_t yInitShape[1] = {1};
   auto yInitTensor = OMTensorUniquePtr(
       omTensorCreateEmpty(&yInitShape[0], 1, OM_DATA_TYPE::ONNX_TYPE_INT64),
       omTensorDestroy);
@@ -135,7 +135,7 @@ bool isOMLoopTheSameAsNaiveImplFor(std::string moduleIR,
     return false;
   }
 
-  auto *yRefInitShape = new int64_t[1]{1};
+  int64_t yRefInitShape[1] = {1};
   auto vFinalRef = OMTensorUniquePtr(
       omTensorCreateEmpty(&yRefInitShape[0], 1, OM_DATA_TYPE::ONNX_TYPE_INT64),
       omTensorDestroy);