cpu: x64: gemm: fix memory leak

If shared library is loaded/unloaded multiple times memory will be leaked.
oneapi-src · Aug 19, 2021 · fd6d14c · fd6d14c
1 parent b7d40a0
commit fd6d14c
Show file tree

Hide file tree

Showing 4 changed files with 265 additions and 232 deletions.
diff --git a/src/cpu/x64/gemm/f32/jit_avx512_common_gemm_f32.cpp b/src/cpu/x64/gemm/f32/jit_avx512_common_gemm_f32.cpp
@@ -16,6 +16,7 @@
 
 #include <atomic>
 #include <cmath>
+#include <memory>
 #include <mutex>
 
 #include "common/dnnl_thread.hpp"
@@ -1740,7 +1741,7 @@ xbyak_gemm_t *get_xbyak_gemm(
     };
 
     // Kernel table [isTransA][isTransB][hasBias][beta (0, 1, other)]
-    static xbyak_gemm_t *kernel_table[2][2][2][3];
+    static std::unique_ptr<xbyak_gemm_t> kernel_table[2][2][2][3];
     static std::once_flag initialized;
     dnnl_status_t st = dnnl_success;
     std::call_once(initialized, [&] {
@@ -1753,8 +1754,8 @@ xbyak_gemm_t *get_xbyak_gemm(
                         auto &kern = kernel_table[isTransA][isTransB][hasBias]
                                                  [beta_idx(beta)];
 
-                        kern = new xbyak_gemm_t(
-                                isTransA, isTransB, beta, hasBias);
+                        kern.reset(new xbyak_gemm_t(
+                                isTransA, isTransB, beta, hasBias));
                         if (kern->create_kernel() != dnnl_success) {
                             st = dnnl_runtime_error;
                             return;
@@ -1763,7 +1764,7 @@ xbyak_gemm_t *get_xbyak_gemm(
     });
 
     return (st == dnnl_success)
-            ? kernel_table[isTransA][isTransB][hasBias][beta_idx(beta)]
+            ? kernel_table[isTransA][isTransB][hasBias][beta_idx(beta)].get()
             : nullptr;
 }
 

diff --git a/src/cpu/x64/gemm/f32/jit_avx512_core_gemm_smalln_tn_f32_kern.cpp b/src/cpu/x64/gemm/f32/jit_avx512_core_gemm_smalln_tn_f32_kern.cpp
@@ -15,6 +15,7 @@
 *******************************************************************************/
 
 #include <atomic>
+#include <memory>
 #include <mutex>
 
 #include "common/dnnl_thread.hpp"
@@ -723,7 +724,7 @@ dnnl_status_t sgemm_smalln_tn(const dim_t m, const dim_t n, const dim_t k,
         const dim_t ldb, const float beta, float *C, const dim_t ldc) {
     using namespace avx512_core_gemm_smalln_tn_f32;
 
-    static xbyak_gemm_smalln_tn_t *kernels[4][3][3];
+    static std::unique_ptr<xbyak_gemm_smalln_tn_t> kernels[4][3][3];
     static std::once_flag initialized;
 
     dnnl_status_t st = dnnl_success;
@@ -732,7 +733,7 @@ dnnl_status_t sgemm_smalln_tn(const dim_t m, const dim_t n, const dim_t k,
             for (float al : {0.0f, 1.0f, 2.0f}) {
                 for (float be : {0.0f, 1.0f, 2.0f}) {
                     auto &kern = kernels[N - 1][(dim_t)al][(dim_t)be];
-                    kern = new xbyak_gemm_smalln_tn_t(N, be, al);
+                    kern.reset(new xbyak_gemm_smalln_tn_t(N, be, al));
                     st = kern->create_kernel();
                     if (st != dnnl_success) return;
                 }

diff --git a/src/cpu/x64/gemm/f32/jit_avx_gemm_f32.cpp b/src/cpu/x64/gemm/f32/jit_avx_gemm_f32.cpp
@@ -16,6 +16,7 @@
 
 #include <atomic>
 #include <cmath>
+#include <memory>
 #include <mutex>
 
 #include "common/dnnl_thread.hpp"
@@ -2193,7 +2194,7 @@ xbyak_gemm_t *get_xbyak_gemm(
     };
 
     // Kernel table [isTransA][isTransB][hasBias][beta (0, 1, other)]
-    static xbyak_gemm_t *kernel_table[2][2][2][3];
+    static std::unique_ptr<xbyak_gemm_t> kernel_table[2][2][2][3];
     static std::once_flag initialized;
     dnnl_status_t st = dnnl_success;
     std::call_once(initialized, [&] {
@@ -2206,8 +2207,8 @@ xbyak_gemm_t *get_xbyak_gemm(
                         auto &kern = kernel_table[isTransA][isTransB][hasBias]
                                                  [beta_idx(beta)];
 
-                        kern = new xbyak_gemm_t(
-                                isTransA, isTransB, beta, hasBias);
+                        kern.reset(new xbyak_gemm_t(
+                                isTransA, isTransB, beta, hasBias));
                         if (kern->create_kernel() != dnnl_success) {
                             st = dnnl_runtime_error;
                             return;
@@ -2216,7 +2217,7 @@ xbyak_gemm_t *get_xbyak_gemm(
     });
 
     return (st == dnnl_success)
-            ? kernel_table[isTransA][isTransB][hasBias][beta_idx(beta)]
+            ? kernel_table[isTransA][isTransB][hasBias][beta_idx(beta)].get()
             : nullptr;
 }