cpu: gemm: make sure to return out of memory if so

Some cases were still not returning proper status in case of out of
memory.

Author: aaraujom

src/cpu/gemm/f32/jit_avx512_common_gemm_f32.cpp

@@ -1927,13 +1927,22 @@ dnnl_status_t jit_avx512_common_gemm_f32(const char *transa, const char *transb,
         c_buffers = (float *)malloc(
                 nthr_m * nthr_n * (nthr_k - 1) * MB * NB * sizeof(float),
                 PAGE_4K);
+        if (!c_buffers) {
+            free(ompstatus_);
+            return dnnl_out_of_memory;
+        }
     }
 
     const size_t ws_elems_per_thr = (size_t)k * 48 + 64;
     const size_t ws_size_per_thr
             = rnd_up(ws_elems_per_thr * sizeof(float), PAGE_4K);
     if (k > STACK_K_CAPACITY) {
         ws_buffers = (float *)malloc(nthr_to_use * ws_size_per_thr, PAGE_4K);
+        if (!ws_buffers) {
+            free(ompstatus_);
+            free(c_buffers);
+            return dnnl_out_of_memory;
+        }
     }
 
     parallel(nthr_to_use, [&](int ithr, int nthr) {

src/cpu/gemm/f32/jit_avx_gemm_f32.cpp

@@ -2496,6 +2496,8 @@ dnnl_status_t jit_avx_gemm_f32(const char *transa, const char *transb,
     if (nthr_k > 1) {
         ompstatus_ = (unsigned char *)malloc(
                 nthr_to_use * CACHE_LINE_SIZE, CACHE_LINE_SIZE);
+        if (!ompstatus_) return dnnl_out_of_memory;
+
         ompstatus = (unsigned char volatile *)ompstatus_;
         assert(ompstatus);
 
@@ -2505,13 +2507,22 @@ dnnl_status_t jit_avx_gemm_f32(const char *transa, const char *transb,
         c_buffers = (float *)malloc(
                 nthr_m * nthr_n * (nthr_k - 1) * MB * NB * sizeof(float),
                 PAGE_4K);
+        if (!c_buffers) {
+            free(ompstatus_);
+            return dnnl_out_of_memory;
+        }
     }
 
     const size_t ws_elems_per_thr = (size_t)k * 16 + 64;
     const size_t ws_size_per_thr
             = rnd_up(ws_elems_per_thr * sizeof(float), PAGE_4K);
     if (k > STACK_K_CAPACITY) {
         ws_buffers = (float *)malloc(nthr_to_use * ws_size_per_thr, PAGE_4K);
+        if (!ws_buffers) {
+            free(ompstatus_);
+            free(c_buffers);
+            return dnnl_out_of_memory;
+        }
     }
 
     parallel(nthr_to_use, [&](int ithr, int nthr) {

src/cpu/gemm/gemm_driver.cpp

@@ -1598,6 +1598,11 @@ static dnnl_status_t gemm_threading_driver(
         c_local_storage = (c_type *)malloc(
                 sizeof(c_type) * c_local_stride * nthr_goal, PAGE_4K);
 
+        if (!c_local_storage) {
+            free(thread_arg);
+            return dnnl_out_of_memory;
+        }
+
         for (int ithr = 0; ithr < nthr_goal; ithr++) {
             thread_arg[ithr].c_local = c_local_storage + ithr * c_local_stride;
             thread_arg[ithr].ldc_local = ldc_local;

src/cpu/gemm/gemm_pack.cpp

@@ -165,6 +165,7 @@ dnnl_status_t sgemm_pack_get_size(const char *identifier, const char *transa,
     bool do_a = utils::one_of(*identifier, 'a', 'A');
     float alpha = 1.0f;
     gemm_pack_storage_shell_t shell {dnnl_get_max_threads()};
+    if (!shell.get()) return dnnl_out_of_memory;
 
     result = gemm_pack_driver<float, float, float>(identifier, transa, transb,
             M, N, K, &alpha, lda, ldb, nullptr, &shell, true);
@@ -204,6 +205,7 @@ dnnl_status_t gemm_bf16bf16f32_pack_get_size(const char *identifier,
 
     float alpha = 1.0f;
     gemm_pack_storage_shell_t shell {dnnl_get_max_threads()};
+    if (!shell.get()) return dnnl_out_of_memory;
 
     result = gemm_pack_driver<bfloat16_t, bfloat16_t, float>(identifier, transa,
             transb, &M_s32, &N_s32, &K_s32, &alpha, &lda_s32, &ldb_s32, nullptr,
@@ -244,6 +246,7 @@ dnnl_status_t gemm_x8x8s32_pack_get_size(const char *identifier,
     bool do_a = utils::one_of(*identifier, 'a', 'A');
     float alpha = 1.0f;
     gemm_pack_storage_shell_t shell {dnnl_get_max_threads(), do_a, !do_a};
+    if (!shell.get()) return dnnl_out_of_memory;
 
     if (!use_reference_igemm<a_dt, b_dt>()) {
         result = gemm_pack_driver<a_dt, b_dt, int32_t>(identifier, transa,