yb/fix mem leak (DeepLink-org#746)

* fix mem leak
xuq7410 · Dec 19, 2023 · be0bebb · be0bebb
1 parent 3792185
commit be0bebb
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/impl/ascend/common/debug.cpp b/impl/ascend/common/debug.cpp
@@ -98,6 +98,8 @@ void printContiguousTensor(diopiContextHandle_t ctx, const AscendTensor& at, cha
         CALL_ACLRT(aclrtMemcpyAsync(
             ptrHost, at.numel() * at.elemsize(), at.data(), at.numel() * at.elemsize(), ACL_MEMCPY_DEVICE_TO_HOST, reinterpret_cast<aclrtStream>(stream)));
         CALL_ACLRT(aclrtSynchronizeStream(reinterpret_cast<aclrtStream>(stream)));
+        CALL_ACLRT(aclrtFreeHost(ptrHost));
+
     } else {
         const void* ptrHostCopy;
         diopiGetTensorDataConst(at.tensorHandle(), &ptrHostCopy);

diff --git a/impl/ascend_npu/torch_npu/csrc/DIOPIAdapter.cpp b/impl/ascend_npu/torch_npu/csrc/DIOPIAdapter.cpp
@@ -1795,8 +1795,8 @@ class OpCommandImpl {
         size_t outputDataBuffArrLen = outputNum * sizeof(uintptr_t);
 
         size_t totalMemLen = inputTensorDescArrLen + inputDataBuffArrLen + outputTensorDescArrLen + outputDataBuffArrLen;
-
-        char* basePtr = static_cast<char*>(malloc(totalMemLen));
+        std::unique_ptr<char, void (*)(void*)> basePtrUnique((char*)malloc(totalMemLen), free);
+        char* basePtr = basePtrUnique.get();
         AT_ASSERT(basePtr != nullptr);
         const aclTensorDesc** aclTensorInputDescArr = reinterpret_cast<const aclTensorDesc**>(basePtr);
         basePtr += inputTensorDescArrLen;