add back aligned alloc

1 files changed, 10 insertions, 2 deletions

diff --git a/src/tensors/cpu/device.cpp b/src/tensors/cpu/device.cpp
index c3055cd2..c4710f5c 100644
--- a/src/tensors/cpu/device.cpp
+++ b/src/tensors/cpu/device.cpp
@@ -15,18 +15,26 @@ Device::~Device() {
   size_ = 0;
 }
 
+#ifdef _WIN32
+#define MALLOC(size) _aligned_alloc(size, alignment_)
+#elif __GNUC__
+#define MALLOC(size) aligned_alloc(alignment_, size)
+#else
+#define MALLOC(size) malloc(size)
+#endif
+
 void Device::reserve(size_t size) {
   size = align(size);
   ABORT_IF(size < size_ || size == 0,
            "New size must be larger than old size and larger than 0");
 
   if(data_) {
-    uint8_t *temp = static_cast<uint8_t*>(malloc(size));
+    uint8_t *temp = static_cast<uint8_t*>(MALLOC(size));
     std::copy(data_, data_ + size_, temp);
     free(data_);
     data_ = temp;
   } else {
-    data_ = static_cast<uint8_t*>(malloc(size));
+    data_ = static_cast<uint8_t*>(MALLOC(size));
   }
   size_ = size;
 }