Use aligned_unique_ptr in more places to avoid leak (#5621)

Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/2578

Pull Request resolved: #5621

Reviewed By: henrylhtsang

Differential Revision: D100536982

Pulled By: q10

fbshipit-source-id: c0b1b9f5dbd462a123ffce359b5e9bc4092c2b12

Author: cyyever

fbgemm_gpu/codegen/training/backward/embedding_backward_split_cpu_template.cpp

@@ -111,8 +111,11 @@ for (const auto t : c10::irange(num_tables)) {
     int feature_begin = table_to_feature_offset[t];
 
     int num_non_zero_columns = cscs[t].num_non_zero_columns;
-    int* col_segment_ptr = cscs[t].column_segment_ptr;
-    int* col_segment_indices = cscs[t].column_segment_indices;
+    int* col_segment_ptr = cscs[t].column_segment_ptr.get();
+    int* col_segment_indices = cscs[t].column_segment_indices.get();
+    int* col_segment_ids = cscs[t].column_segment_ids.get();
+    int* row_indices = cscs[t].row_indices.get();
+    float* ind_weights = cscs[t].weights.get();
 
     const auto D_begin = D_offsets_data[feature_begin];
     const auto D =
@@ -134,7 +137,7 @@ for (const auto t : c10::irange(num_tables)) {
           /*IndexType=*/int32_t,
           /*OffsetType=*/int32_t>(
           D,
-          cscs[t].weights != nullptr,
+          ind_weights != nullptr,
           /*normalize_by_lengths=*/false,
           /*prefetch=*/16,
           /*is_weight_positional=*/false,
@@ -156,11 +159,11 @@ for (const auto t : c10::irange(num_tables)) {
               B,
               reinterpret_cast<const fbgemm_weight_t*>(
                   grad_output_data + D_begin),
-              cscs[t].row_indices + *offsets_begin_ptr,
+              row_indices + *offsets_begin_ptr,
               offsets_begin_ptr,
-              cscs[t].weights == nullptr
+              ind_weights == nullptr
                   ? nullptr
-                  : cscs[t].weights + *offsets_begin_ptr,
+                  : ind_weights + *offsets_begin_ptr,
               reinterpret_cast<float*>(grad_blocked_buffer));
 
           if (!success) {
@@ -170,7 +173,7 @@ for (const auto t : c10::irange(num_tables)) {
               c,
               c_block_end,
               col_segment_ptr,
-              cscs[t].row_indices,
+              row_indices,
               hash_size,
               /*allow_minus_one=*/false);
           }
@@ -218,14 +221,14 @@ for (const auto t : c10::irange(num_tables)) {
         for (int r = col_segment_ptr[c]; r < col_segment_ptr[c + 1]; ++r) {
           int D_offset = D_begin;
           if (is_shared_table) {
-            D_offset += cscs[t].column_segment_ids[r] * D;
+            D_offset += col_segment_ids[r] * D;
           }
-          int b = cscs[t].row_indices[r];
+          int b = row_indices[r];
 
           for (const auto d : c10::irange(D)) {
-            if (cscs[t].weights != nullptr) {
+            if (ind_weights != nullptr) {
               grad_buffer[d] += grad_output_data[b * grad_stride + D_offset + d] *
-                    cscs[t].weights[r];
+                    ind_weights[r];
             } else {
               grad_buffer[d] += grad_output_data[b * grad_stride + D_offset + d];
             }

fbgemm_gpu/codegen/training/forward/embedding_forward_split_cpu.cpp

@@ -414,12 +414,10 @@ void csr2csc_template_(
   if (nnz == 0) {
     return;
   }
-  csc.row_indices =
-      static_cast<int*>(fbgemm::fbgemmAlignedAlloc(64, nnz * sizeof(int)));
+  csc.row_indices = fbgemm::makeAlignedUniquePtr<int>(64, nnz);
   bool has_weights = csr_weights.data() != nullptr;
   if (IS_VALUE_PAIR) {
-    csc.weights = static_cast<float*>(
-        fbgemm::fbgemmAlignedAlloc(64, nnz * sizeof(float)));
+    csc.weights = fbgemm::makeAlignedUniquePtr<float>(64, nnz);
   }
 
   [[maybe_unused]] int column_ptr_curr = 0;
@@ -431,16 +429,15 @@ void csr2csc_template_(
   using pair_t = std::pair<int, scalar_t>;
   using value_t = typename std::conditional<IS_VALUE_PAIR, pair_t, int>::type;
 
-  csc.column_segment_ids =
-      static_cast<int*>(fbgemm::fbgemmAlignedAlloc(64, nnz * sizeof(int)));
-  int* tmpBufKeys =
-      static_cast<int*>(fbgemm::fbgemmAlignedAlloc(64, NS * sizeof(int)));
-  value_t* tmpBufValues = static_cast<value_t*>(
-      fbgemm::fbgemmAlignedAlloc(64, NS * sizeof(value_t)));
-  int* tmpBuf1Keys =
-      static_cast<int*>(fbgemm::fbgemmAlignedAlloc(64, NS * sizeof(int)));
-  value_t* tmpBuf1Values = static_cast<value_t*>(
-      fbgemm::fbgemmAlignedAlloc(64, NS * sizeof(value_t)));
+  csc.column_segment_ids = fbgemm::makeAlignedUniquePtr<int>(64, nnz);
+  auto tmpBufKeys = fbgemm::makeAlignedUniquePtr<int>(64, NS);
+  fbgemm::aligned_unique_ptr<value_t> tmpBufValues(
+      static_cast<value_t*>(
+          fbgemm::fbgemmAlignedAlloc(64, NS * sizeof(value_t))));
+  auto tmpBuf1Keys = fbgemm::makeAlignedUniquePtr<int>(64, NS);
+  fbgemm::aligned_unique_ptr<value_t> tmpBuf1Values(
+      static_cast<value_t*>(
+          fbgemm::fbgemmAlignedAlloc(64, NS * sizeof(value_t))));
 
   const auto FBo = csr_offsets[table_to_feature_offset[0] * B];
   for (int feature = table_to_feature_offset[0];
@@ -461,11 +458,11 @@ void csr2csc_template_(
           : 1.0;
       for (const auto p : c10::irange(pool_begin, pool_end)) {
         tmpBufKeys[p - FBo] = csr_indices[p];
-        if (IS_VALUE_PAIR) {
-          reinterpret_cast<pair_t*>(tmpBufValues)[p - FBo] = std::pair{
+        if constexpr (IS_VALUE_PAIR) {
+          tmpBufValues[p - FBo] = std::pair{
               FBs + b, scale_factor * (has_weights ? csr_weights[p] : 1.0f)};
         } else {
-          reinterpret_cast<int*>(tmpBufValues)[p - FBo] = FBs + b;
+          tmpBufValues[p - FBo] = FBs + b;
         }
       }
     }
@@ -475,10 +472,10 @@ void csr2csc_template_(
   value_t* sorted_col_row_index_values;
   std::tie(sorted_col_row_index_keys, sorted_col_row_index_values) =
       fbgemm::radix_sort_parallel(
-          tmpBufKeys,
-          tmpBufValues,
-          tmpBuf1Keys,
-          tmpBuf1Values,
+          tmpBufKeys.get(),
+          tmpBufValues.get(),
+          tmpBuf1Keys.get(),
+          tmpBuf1Values.get(),
           NS,
           num_embeddings);
 
@@ -509,10 +506,8 @@ void csr2csc_template_(
     U = num_uniq[max_thds - 1][0];
   }
 
-  csc.column_segment_ptr =
-      static_cast<int*>(fbgemm::fbgemmAlignedAlloc(64, (NS + 1) * sizeof(int)));
-  csc.column_segment_indices =
-      static_cast<int*>(fbgemm::fbgemmAlignedAlloc(64, NS * sizeof(int)));
+  csc.column_segment_ptr = fbgemm::makeAlignedUniquePtr<int>(64, NS + 1);
+  csc.column_segment_indices = fbgemm::makeAlignedUniquePtr<int>(64, NS);
   csc.column_segment_ptr[0] = 0;
   const pair_t* sorted_col_row_index_values_pair =
       reinterpret_cast<const pair_t*>(sorted_col_row_index_values);
@@ -528,26 +523,31 @@ void csr2csc_template_(
   }
   csc.column_segment_indices[0] = sorted_col_row_index_keys[0];
 
+  int* col_seg_indices = csc.column_segment_indices.get();
+  int* col_seg_ptr = csc.column_segment_ptr.get();
+
 #pragma omp parallel
   {
     int tid = omp_get_thread_num();
     int* tstart =
-        (tid == 0 ? csc.column_segment_indices + 1
-                  : csc.column_segment_indices + num_uniq[tid - 1][0]);
+        (tid == 0 ? col_seg_indices + 1
+                  : col_seg_indices + num_uniq[tid - 1][0]);
 
     int* t_offs =
-        (tid == 0 ? csc.column_segment_ptr + 1
-                  : csc.column_segment_ptr + num_uniq[tid - 1][0]);
+        (tid == 0 ? col_seg_ptr + 1 : col_seg_ptr + num_uniq[tid - 1][0]);
 
     if (!IS_VALUE_PAIR && !is_shared_table) {
       // For non shared table, no need for computing modulo.
       // As an optimization, pointer swap instead of copying.
 #pragma omp master
-      std::swap(
-          csc.row_indices,
-          *reinterpret_cast<int**>(
-              sorted_col_row_index_values == tmpBufValues ? &tmpBufValues
-                                                          : &tmpBuf1Values));
+      {
+        auto& buf = sorted_col_row_index_values == tmpBufValues.get()
+            ? tmpBufValues
+            : tmpBuf1Values;
+        int* tmp = csc.row_indices.release();
+        csc.row_indices.reset(reinterpret_cast<int*>(buf.release()));
+        buf.reset(reinterpret_cast<value_t*>(tmp));
+      }
     } else {
 #ifdef FBCODE_CAFFE2
       libdivide::divider<int> divisor(B);
@@ -582,7 +582,7 @@ void csr2csc_template_(
 
     if (at::get_num_threads() == 1 && tid == 0) {
       // Special handling of single thread case
-      U = t_offs - csc.column_segment_ptr;
+      U = t_offs - csc.column_segment_ptr.get();
     }
 
   } // omp parallel
@@ -591,11 +591,6 @@ void csr2csc_template_(
   csc.column_segment_ptr[U] = NS;
   column_ptr_curr += NS;
 
-  fbgemm::fbgemmAlignedFree(tmpBufKeys);
-  fbgemm::fbgemmAlignedFree(tmpBufValues);
-  fbgemm::fbgemmAlignedFree(tmpBuf1Keys);
-  fbgemm::fbgemmAlignedFree(tmpBuf1Values);
-
   assert(column_ptr_curr == nnz);
 }
 

fbgemm_gpu/include/fbgemm_gpu/embedding_forward_split_cpu.h

@@ -44,22 +44,15 @@ struct HyperCompressedSparseColumn {
   // For a shared table, a column can have multiple segments, each for a
   // feature sharing the table. In this case, the segments will have the
   // same column_segment_indices but different column_segment_ids.
-  int* column_segment_ptr = nullptr;
-  int* column_segment_indices = nullptr; // length num_non_zero_columns
-  int* column_segment_ids = nullptr; // length num_non_zero_columns
-  int* row_indices = nullptr; // length column_ptr[num_non_zero_columns]
-  float* weights = nullptr; // length column_ptr[num_non_zero_columns]
-  ~HyperCompressedSparseColumn() {
-    if (column_segment_ptr) {
-      fbgemm::fbgemmAlignedFree(column_segment_ptr);
-      fbgemm::fbgemmAlignedFree(column_segment_indices);
-      fbgemm::fbgemmAlignedFree(column_segment_ids);
-      fbgemm::fbgemmAlignedFree(row_indices);
-    }
-    if (weights) {
-      fbgemm::fbgemmAlignedFree(weights);
-    }
-  }
+  fbgemm::aligned_unique_ptr<int> column_segment_ptr;
+  fbgemm::aligned_unique_ptr<int>
+      column_segment_indices; // length num_non_zero_columns
+  fbgemm::aligned_unique_ptr<int>
+      column_segment_ids; // length num_non_zero_columns
+  fbgemm::aligned_unique_ptr<int>
+      row_indices; // length column_ptr[num_non_zero_columns]
+  fbgemm::aligned_unique_ptr<float>
+      weights; // length column_ptr[num_non_zero_columns]
 };
 
 template <typename index_t, typename scalar_t>

src/Utils.cc

@@ -369,15 +369,16 @@ void* fbgemmAlignedAlloc(
     size_t size,
     bool raiseException /*=false*/) {
   void* aligned_mem = nullptr;
-  int ret = 0;
 #ifdef _MSC_VER
   aligned_mem = _aligned_malloc(size, align);
-  ret = 0;
 #else
-  ret = posix_memalign(&aligned_mem, align, size);
+  int ret = posix_memalign(&aligned_mem, align, size);
+  if (ret != 0) {
+    aligned_mem = nullptr;
+  }
 #endif
   // Throw std::bad_alloc in the case of memory allocation failure.
-  if (raiseException && (ret || aligned_mem == nullptr)) {
+  if (raiseException && aligned_mem == nullptr) {
     throw std::bad_alloc();
   }
   return aligned_mem;