Skip to content

Commit

Permalink
Fix default CPU allocator memory alignment (apache#18885)
Browse files Browse the repository at this point in the history
* Replace std::malloc to aligned memory allocation in Pooled StorageManager

* Add test checking CPU memory alignment

* Fix memory allocation success check

* Fix sanity
  • Loading branch information
bgawrych authored Aug 14, 2020
1 parent 344587f commit ee80b77
Show file tree
Hide file tree
Showing 4 changed files with 63 additions and 15 deletions.
22 changes: 22 additions & 0 deletions src/common/utils.h
Original file line number Diff line number Diff line change
Expand Up @@ -950,6 +950,28 @@ inline int GetDefaultDtype(int dtype) {
mshadow::kFloat32;
}

inline bool AlignedMemAlloc(void** ptr, size_t size, size_t alignment) {
#if _MSC_VER
*ptr = _aligned_malloc(size, alignment);
if (*ptr == nullptr)
return false;
#else
int res = posix_memalign(ptr, alignment, size);
if (res != 0)
return false;
#endif
return true;
}

inline void AlignedMemFree(void* ptr) {
#if _MSC_VER
_aligned_free(ptr);
#else
free(ptr);
#endif
}


} // namespace common
} // namespace mxnet
#endif // MXNET_COMMON_UTILS_H_
15 changes: 3 additions & 12 deletions src/storage/cpu_device_storage.h
Original file line number Diff line number Diff line change
Expand Up @@ -60,21 +60,12 @@ class CPUDeviceStorage {
}; // class CPUDeviceStorage

inline void CPUDeviceStorage::Alloc(Storage::Handle* handle) {
#if _MSC_VER
handle->dptr = _aligned_malloc(handle->size, alignment_);
if (handle->dptr == nullptr) LOG(FATAL) << "Failed to allocate CPU Memory";
#else
int ret = posix_memalign(&handle->dptr, alignment_, handle->size);
if (ret != 0) LOG(FATAL) << "Failed to allocate CPU Memory";
#endif
bool success = mxnet::common::AlignedMemAlloc(&(handle->dptr), handle->size, alignment_);
if (!success) LOG(FATAL) << "Failed to allocate CPU Memory";
}

inline void CPUDeviceStorage::Free(Storage::Handle handle) {
#if _MSC_VER
_aligned_free(handle.dptr);
#else
free(handle.dptr);
#endif
mxnet::common::AlignedMemFree(handle.dptr);
}

} // namespace storage
Expand Down
18 changes: 15 additions & 3 deletions src/storage/storage_manager_helpers.h
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ typedef mxnet::common::cuda::DeviceStore CudaDeviceStore;
#endif // _WIN32

#include <tuple>
#include "../common/utils.h"

namespace mxnet {
namespace storage {
Expand Down Expand Up @@ -110,10 +111,22 @@ class ContextHelperCPU : public ContextHelper {
}

int Malloc(void **ppNtr, size_t size) const override {
return (*ppNtr = std::malloc(size))? 0 : -1;
bool success = mxnet::common::AlignedMemAlloc(ppNtr, size, alignment_);
return success ? 0 : -1;
}

void Free(void *dptr) const override { std::free(dptr); }
void Free(void *dptr) const override {
mxnet::common::AlignedMemFree(dptr);
}

private:
#if MXNET_USE_MKLDNN == 1
// MKLDNN requires special alignment. 64 is used by the MKLDNN library in
// memory allocation.
static constexpr size_t alignment_ = kMKLDNNAlign;
#else
static constexpr size_t alignment_ = 16;
#endif
};

#if MXNET_USE_CUDA
Expand Down Expand Up @@ -155,7 +168,6 @@ class ContextHelperPinned : public ContextHelperGPU {
#else
typedef ContextHelperCPU ContextHelperPinned;
#endif

} // namespace storage
} // namespace mxnet

Expand Down
23 changes: 23 additions & 0 deletions tests/cpp/storage/storage_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,29 @@ TEST(Storage, Basic_CPU) {
storage->Free(handle);
}

TEST(Storage, CPU_MemAlign) {
#if MXNET_USE_MKLDNN == 1
// MKLDNN requires special alignment. 64 is used by the MKLDNN library in
// memory allocation.
static constexpr size_t alignment_ = mxnet::kMKLDNNAlign;
#else
static constexpr size_t alignment_ = 16;
#endif

auto&& storage = mxnet::Storage::Get();
mxnet::Context context_cpu = mxnet::Context::CPU(0);

for (int i = 0; i < 5; ++i) {
const size_t kSize = (std::rand() % 1024) + 1;
auto&& handle = storage->Alloc(kSize, context_cpu);
EXPECT_EQ(handle.ctx, context_cpu);
EXPECT_EQ(handle.size, kSize);
EXPECT_EQ(reinterpret_cast<intptr_t>(handle.dptr) % alignment_, 0);
storage->Free(handle);
}
}


#if MXNET_USE_CUDA
TEST(Storage_GPU, Basic_GPU) {
if (mxnet::test::unitTestsWithCuda) {
Expand Down

0 comments on commit ee80b77

Please sign in to comment.