From d09097dd7e2c148e3b178935c39a6f1a935a268c Mon Sep 17 00:00:00 2001 From: Bartlomiej Gawrych Date: Thu, 6 Aug 2020 12:52:54 +0200 Subject: [PATCH 1/4] Replace std::malloc to aligned memory allocation in Pooled StorageManager --- src/common/utils.h | 22 ++++++++++++++++++++++ src/storage/cpu_device_storage.h | 15 +++------------ src/storage/storage_manager_helpers.h | 18 +++++++++++++++--- 3 files changed, 40 insertions(+), 15 deletions(-) diff --git a/src/common/utils.h b/src/common/utils.h index 9ea3329f2c24..22d695b0d727 100644 --- a/src/common/utils.h +++ b/src/common/utils.h @@ -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_ diff --git a/src/storage/cpu_device_storage.h b/src/storage/cpu_device_storage.h index b81cdbc17e25..eca7eaa057fc 100644 --- a/src/storage/cpu_device_storage.h +++ b/src/storage/cpu_device_storage.h @@ -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 diff --git a/src/storage/storage_manager_helpers.h b/src/storage/storage_manager_helpers.h index e144af2ab9a3..fdb32386a2dd 100644 --- a/src/storage/storage_manager_helpers.h +++ b/src/storage/storage_manager_helpers.h @@ -42,6 +42,7 @@ typedef mxnet::common::cuda::DeviceStore CudaDeviceStore; #endif // _WIN32 #include +#include "../common/utils.h" namespace mxnet { namespace storage { @@ -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 @@ -155,7 +168,6 @@ class ContextHelperPinned : public ContextHelperGPU { #else typedef ContextHelperCPU ContextHelperPinned; #endif - } // namespace storage } // namespace mxnet From af1aabade8a15638a378d7f41651b4c6d72be2dc Mon Sep 17 00:00:00 2001 From: Bartlomiej Gawrych Date: Fri, 7 Aug 2020 11:05:53 +0200 Subject: [PATCH 2/4] Add test checking CPU memory alignment --- tests/cpp/storage/storage_test.cc | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/cpp/storage/storage_test.cc b/tests/cpp/storage/storage_test.cc index d9e7d8bc0294..046eee5e9d26 100644 --- a/tests/cpp/storage/storage_test.cc +++ b/tests/cpp/storage/storage_test.cc @@ -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(handle.dptr) % alignment_, 0); + storage->Free(handle); + } +} + + #if MXNET_USE_CUDA TEST(Storage_GPU, Basic_GPU) { if (mxnet::test::unitTestsWithCuda) { From 711b9f627a6ffbaa517a8094e80fbffc26ba8026 Mon Sep 17 00:00:00 2001 From: Bartlomiej Gawrych Date: Mon, 10 Aug 2020 09:21:35 +0200 Subject: [PATCH 3/4] Fix memory allocation success check --- src/common/utils.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/common/utils.h b/src/common/utils.h index 22d695b0d727..f6d0aecbd207 100644 --- a/src/common/utils.h +++ b/src/common/utils.h @@ -953,7 +953,7 @@ inline int GetDefaultDtype(int dtype) { inline bool AlignedMemAlloc(void** ptr, size_t size, size_t alignment) { #if _MSC_VER *ptr = _aligned_malloc(size, alignment); - if(ptr == nullptr) + if(*ptr == nullptr) return false; #else int res = posix_memalign(ptr, alignment, size); From a3e9fec3bf22f6736f8c9f6b1a83d7d0816829c8 Mon Sep 17 00:00:00 2001 From: Bartlomiej Gawrych Date: Mon, 10 Aug 2020 11:47:48 +0200 Subject: [PATCH 4/4] Fix sanity --- src/common/utils.h | 2 +- src/storage/storage_manager_helpers.h | 2 +- tests/cpp/storage/storage_test.cc | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/common/utils.h b/src/common/utils.h index f6d0aecbd207..aa0cb6b1b454 100644 --- a/src/common/utils.h +++ b/src/common/utils.h @@ -953,7 +953,7 @@ inline int GetDefaultDtype(int dtype) { inline bool AlignedMemAlloc(void** ptr, size_t size, size_t alignment) { #if _MSC_VER *ptr = _aligned_malloc(size, alignment); - if(*ptr == nullptr) + if (*ptr == nullptr) return false; #else int res = posix_memalign(ptr, alignment, size); diff --git a/src/storage/storage_manager_helpers.h b/src/storage/storage_manager_helpers.h index fdb32386a2dd..1fccb5a08f45 100644 --- a/src/storage/storage_manager_helpers.h +++ b/src/storage/storage_manager_helpers.h @@ -117,7 +117,7 @@ class ContextHelperCPU : public ContextHelper { void Free(void *dptr) const override { mxnet::common::AlignedMemFree(dptr); - } + } private: #if MXNET_USE_MKLDNN == 1 diff --git a/tests/cpp/storage/storage_test.cc b/tests/cpp/storage/storage_test.cc index 046eee5e9d26..497af35c83cd 100644 --- a/tests/cpp/storage/storage_test.cc +++ b/tests/cpp/storage/storage_test.cc @@ -59,7 +59,7 @@ TEST(Storage, CPU_MemAlign) { auto&& storage = mxnet::Storage::Get(); mxnet::Context context_cpu = mxnet::Context::CPU(0); - for(int i=0; i < 5; ++i) { + 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);