add chunk allocator posix_memalign return value check (#60208) (#60495)

* fix chunk allocator posix_memalign return value check;test=develop * fix chunk allocator posix_memalign return value check;test=develop * fix chunk allocator posix_memalign return value check;test=develop
PaddlePaddle · Jan 2, 2024 · b065877 · b065877
1 parent 203754e
commit b065877
Showing 1 changed file with 11 additions and 3 deletions.
diff --git a/paddle/fluid/distributed/common/chunk_allocator.h b/paddle/fluid/distributed/common/chunk_allocator.h
@@ -14,6 +14,7 @@
 
 #pragma once
 #include <glog/logging.h>
+#include "paddle/fluid/platform/enforce.h"
 
 namespace paddle {
 namespace distributed {
@@ -77,9 +78,16 @@ class ChunkAllocator {
 
   void create_new_chunk() {
     Chunk* chunk;
-    posix_memalign(reinterpret_cast<void**>(&chunk),
-                   std::max<size_t>(sizeof(void*), alignof(Chunk)),
-                   sizeof(Chunk) + sizeof(Node) * _chunk_size);
+    size_t alloc_size = sizeof(Chunk) + sizeof(Node) * _chunk_size;
+    int error = posix_memalign(reinterpret_cast<void**>(&chunk),
+                               std::max<size_t>(sizeof(void*), alignof(Chunk)),
+                               alloc_size);
+    PADDLE_ENFORCE_EQ(error,
+                      0,
+                      paddle::platform::errors::ResourceExhausted(
+                          "Fail to alloc memory of %ld size, error code is %d.",
+                          alloc_size,
+                          error));
     chunk->next = _chunks;
     _chunks = chunk;