Fix memory allocation bug

1 files changed, 7 insertions, 6 deletions

diff --git a/src/FbgemmI8Spmdm.cc b/src/FbgemmI8Spmdm.cc
index 0865bff..edcc4e8 100644
--- a/src/FbgemmI8Spmdm.cc
+++ b/src/FbgemmI8Spmdm.cc
@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 #include "fbgemm/FbgemmI8Spmdm.h"
+#include "fbgemm/Utils.h"
 
 #include <algorithm>
 #include <array>
@@ -150,15 +151,15 @@ void CompressedSparseColumn::SpMDM(
   t_start = std::chrono::high_resolution_clock::now();
 #endif
 
-  alignas(64) uint8_t* A_buffer = new uint8_t[K * 32];
-  alignas(64) int32_t* C_buffer = new int32_t[N * 32];
+  uint8_t* A_buffer = static_cast<uint8_t*>(ALIGNED_MALLOC(K * 32 * sizeof(uint8_t), 64));
+  int32_t* C_buffer = static_cast<int32_t*>(ALIGNED_MALLOC(N * 32 * sizeof(int32_t), 64));
 
   // Take 32 rows at a time
   int i_end = block.row_start + block.row_size;
   for (int i1 = block.row_start; i1 < i_end; i1 += 32) {
     // Transpose 32 x K submatrix of A
     if (i_end - i1 < 32) {
-      alignas(64) uint8_t* A_temp_buffer = new uint8_t[K * 32];
+      uint8_t* A_temp_buffer = static_cast<uint8_t*>(ALIGNED_MALLOC(K * 32 * sizeof(uint8_t), 64));
       for (int i2 = 0; i2 < (i_end - i1) / 8 * 8; i2 += 8) {
         transpose_8rows(K, A + (i1 + i2) * lda, lda, A_buffer + i2, 32);
       }
@@ -174,7 +175,7 @@ void CompressedSparseColumn::SpMDM(
       for (int i2 = (i_end - i1) / 8 * 8; i2 < 32; i2 += 8) {
         transpose_8rows(K, A_temp_buffer + i2 * K, K, A_buffer + i2, 32);
       }
-      delete[] A_temp_buffer;
+      FREE(A_temp_buffer);
     } else {
       for (int i2 = 0; i2 < 32; i2 += 8) {
         transpose_8rows(K, A + (i1 + i2) * lda, lda, A_buffer + i2, 32);
@@ -253,8 +254,8 @@ void CompressedSparseColumn::SpMDM(
   t_start = std::chrono::high_resolution_clock::now();
 #endif
 
-  delete[] A_buffer;
-  delete[] C_buffer;
+  FREE(A_buffer);
+  FREE(C_buffer);
 }
 
 void CompressedSparseColumn::SparseConv(