fix vmm free error

hipudding · Jul 11, 2024 · 0267af3 · 0267af3
1 parent b12a40f
commit 0267af3
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/ggml/src/ggml-cann.cpp b/ggml/src/ggml-cann.cpp
@@ -183,6 +183,7 @@ struct ggml_cann_pool_vmm : public ggml_cann_pool {
     size_t pool_size = 0;
     size_t granularity;
     std::vector<aclrtDrvMemHandle> handles;
+    std::vector<void*> map_offsets;
 
     explicit ggml_cann_pool_vmm(int device) :
         device(device),
@@ -191,7 +192,9 @@ struct ggml_cann_pool_vmm : public ggml_cann_pool {
 
     ~ggml_cann_pool_vmm() {
         if (pool_addr != 0) {
-            ACL_CHECK(aclrtUnmapMem(pool_addr));
+            for(auto& offset : map_offsets){
+                ACL_CHECK(aclrtUnmapMem(offset));
+            }
             for (auto& handle : handles) {
                 ACL_CHECK(aclrtFreePhysical(handle));
             }
@@ -233,6 +236,7 @@ struct ggml_cann_pool_vmm : public ggml_cann_pool {
             ACL_CHECK(aclrtMapMem((char*)pool_addr + pool_size, reserve_size, 0, handle, 0));
 
             handles.push_back(handle);
+            map_offsets.push_back((char*)pool_addr + pool_size);
 
             // add to the pool
             pool_size += reserve_size;