chore: import upstream snapshot with attribution

paddle/phi/core/memory/allocation/CMakeLists.txt

@@ -0,0 +1,58 @@
+set(ALLOCATOR_SRCS
+    allocator.cc
+    cpu_allocator.cc
+    aligned_allocator.cc
+    buffered_allocator.cc
+    best_fit_allocator.cc
+    naive_best_fit_allocator.cc
+    allocator_strategy.cc
+    allocator_facade.cc
+    auto_growth_best_fit_allocator.cc
+    auto_growth_best_fit_allocator_v2.cc
+    virtual_memory_auto_growth_best_fit_allocator.cc
+    retry_allocator.cc
+    memory_block.cc
+    memory_block_desc.cc
+    meta_cache.cc
+    buddy_allocator.cc
+    system_allocator.cc)
+
+if(WITH_GPU OR WITH_ROCM)
+  list(
+    APPEND
+    ALLOCATOR_SRCS
+    cuda_allocator.cc
+    cuda_managed_allocator.cc
+    cuda_malloc_async_allocator.cc
+    pinned_allocator.cc
+    stream_safe_cuda_allocator.cc
+    thread_local_allocator.cc)
+endif()
+
+if(CUDA_VERSION VERSION_GREATER_EQUAL 10.2)
+  list(APPEND ALLOCATOR_SRCS cuda_virtual_mem_allocator.cc
+       cuda_virtual_mem_allocator_v2.cc vmm_backing_map.cc
+       vmm_auto_growth_best_fit_allocator_v2.cc)
+endif()
+
+if(NOT WIN32)
+  list(APPEND ALLOCATOR_SRCS mmap_allocator.cc)
+  if(WITH_GPU)
+    list(APPEND ALLOCATOR_SRCS cuda_ipc_allocator.cc)
+  endif()
+  if(WITH_XPU)
+    list(APPEND ALLOCATOR_SRCS xpu_ipc_allocator.cc)
+  endif()
+endif()
+
+if(WITH_CUSTOM_DEVICE)
+  list(APPEND ALLOCATOR_SRCS custom_allocator.cc
+       stream_safe_custom_device_allocator.cc)
+endif()
+
+if(WITH_XPU)
+  list(APPEND ALLOCATOR_SRCS xpu_allocator.cc xpu_pinned_allocator.cc
+       stream_safe_xpu_allocator.cc)
+endif()
+
+collect_srcs(core_srcs SRCS ${ALLOCATOR_SRCS})

paddle/phi/core/memory/allocation/aligned_allocator.cc

@@ -0,0 +1,71 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/aligned_allocator.h"
+
+#include <utility>
+
+#include "paddle/common/macros.h"
+#include "paddle/phi/core/enforce.h"
+
+REGISTER_FILE_SYMBOLS(aligned_allocator);
+namespace paddle::memory::allocation {
+
+// For memory address alignment
+class AlignedAllocation : public Allocation {
+ public:
+  AlignedAllocation(DecoratedAllocationPtr underlying_allocation, size_t offset)
+      : Allocation(
+            reinterpret_cast<uint8_t*>(underlying_allocation->ptr()) + offset,
+            underlying_allocation->base_ptr(),
+            underlying_allocation->size() - offset,
+            underlying_allocation->place()),
+        underlying_allocation_(std::move(underlying_allocation)) {}
+
+ private:
+  DecoratedAllocationPtr underlying_allocation_;
+};
+
+AlignedAllocator::AlignedAllocator(
+    std::shared_ptr<Allocator> underlying_allocator, size_t alignment)
+    : underlying_allocator_(std::move(underlying_allocator)),
+      alignment_(alignment) {
+  PADDLE_ENFORCE_GT(
+      alignment_,
+      0,
+      common::errors::InvalidArgument(
+          "Alignment should be larger than 0, but got %d", alignment_));
+  if (alignment_ & (alignment_ - 1)) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Alignment should be power of 2 (2^N), but got %d", alignment_));
+  }
+}
+
+bool AlignedAllocator::IsAllocThreadSafe() const {
+  return underlying_allocator_->IsAllocThreadSafe();
+}
+
+phi::Allocation* AlignedAllocator::AllocateImpl(size_t size) {
+  auto raw_allocation = underlying_allocator_->Allocate(size + alignment_);
+  size_t offset = AlignedPtrOffset(raw_allocation->ptr(), alignment_);
+  auto* p = new AlignedAllocation(
+      static_unique_ptr_cast<Allocation>(std::move(raw_allocation)), offset);
+  return p;
+}
+
+void AlignedAllocator::FreeImpl(phi::Allocation* allocation) {
+  delete allocation;
+}
+
+}  // namespace paddle::memory::allocation

paddle/phi/core/memory/allocation/aligned_allocator.h

@@ -0,0 +1,44 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include <memory>
+#include <utility>
+
+#include "paddle/phi/core/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class PADDLE_API AlignedAllocator : public Allocator {
+ public:
+  AlignedAllocator(std::shared_ptr<Allocator> underlying_allocator,
+                   size_t alignment);
+
+  bool IsAllocThreadSafe() const override;
+
+ protected:
+  phi::Allocation* AllocateImpl(size_t size) override;
+
+  void FreeImpl(phi::Allocation* allocation) override;
+
+ private:
+  std::shared_ptr<Allocator> underlying_allocator_;
+  size_t alignment_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/allocator.cc

@@ -0,0 +1,57 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/backends/gpu/gpu_info.h"
+#include "paddle/phi/core/memory/stats.h"
+
+namespace paddle::memory::allocation {
+
+void Allocator::FreeImpl(phi::Allocation* allocation) {
+  static_cast<Allocation*>(allocation)
+      ->TopDecoratedAllocator()
+      ->Free(allocation);
+}
+
+void MultiScalePoolAllocator::RecordAlloc(uintptr_t allocator,
+                                          uint64_t id,
+                                          size_t size) {
+#if defined(PADDLE_WITH_CUDA)
+  std::lock_guard<SpinLock> lock(spinlock_);
+  const auto current_device_id = phi::backends::gpu::GetCurrentDeviceId();
+  const auto max_reserved =
+      paddle::memory::DeviceMemoryStatPeakValue("Reserved", current_device_id);
+  const auto cur_allocated = paddle::memory::DeviceMemoryStatCurrentValue(
+      "Allocated", current_device_id);
+  allocation_records_.emplace_back(
+      allocator, true, id, size, cur_allocated, max_reserved);
+#endif
+}
+
+void MultiScalePoolAllocator::RecordFree(uintptr_t allocator,
+                                         uint64_t id,
+                                         size_t size) {
+#if defined(PADDLE_WITH_CUDA)
+  std::lock_guard<SpinLock> lock(spinlock_);
+  const auto current_device_id = phi::backends::gpu::GetCurrentDeviceId();
+  const auto max_reserved =
+      paddle::memory::DeviceMemoryStatPeakValue("Reserved", current_device_id);
+  const auto cur_allocated = paddle::memory::DeviceMemoryStatCurrentValue(
+      "Allocated", current_device_id);
+  allocation_records_.emplace_back(
+      allocator, false, id, size, cur_allocated, max_reserved);
+#endif
+}
+
+}  // namespace paddle::memory::allocation

paddle/phi/core/memory/allocation/allocator.h

@@ -0,0 +1,338 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include <atomic>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <tuple>
+#include <type_traits>
+#include <utility>
+#include <vector>
+
+#include "paddle/common/enforce.h"
+#include "paddle/common/flags.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/allocator.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/memory/allocation/inlined_vector.h"
+#include "paddle/phi/core/memory/allocation/spin_lock.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_types.h"
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/cuda_stream.h"
+#endif
+
+#ifdef PADDLE_WITH_NCCL
+#include <nccl.h>
+#include "paddle/phi/backends/dynload/nccl.h"
+#endif
+
+COMMON_DECLARE_string(allocator_strategy);
+COMMON_DECLARE_bool(sync_after_alloc);
+COMMON_DECLARE_int64(alloc_fill_value);
+COMMON_DECLARE_bool(record_alloc_event);
+
+namespace paddle {
+namespace memory {
+class AllocatorVisitor;
+namespace allocation {
+
+// Exception when `Alloc`/`AllocShared` failed
+struct BadAlloc : public std::exception {
+  inline explicit BadAlloc(std::string err_msg, const char* file, int line)
+      : err_str_(phi::enforce::GetCompleteTraceBackString(
+            std::move(err_msg), file, line)) {}
+
+  const char* what() const noexcept override { return err_str_.c_str(); }
+
+  std::string err_str_;
+};
+
+class Allocator;
+
+// Allocation is the object holding the actually pointer. Use
+// `Allocation::ptr()` will returns the pointer that allocated.
+//
+// NOTE: this is the base class of Allocation. Each allocator can use its own
+//       allocation object.
+// NOTE: the `Allocation::ptr()` could be nullptr, if the allocation size is 0
+
+/**
+ * Allocation is returned by Allocator::Allocate() method.
+ *
+ * An allocator may be decorated by another allocator. For example, we can
+ * decorate a RetryAllocator to any allocator to perform allocation retry when
+ * first allocation request fails.
+ *
+ * Explanations of Allocator design are as follows:
+ *
+ * Suppose we have an allocator which is decorated by several allocators:
+ *
+ *   A(1) <- A(2) <- A(3) <- ... <- A(n)
+ *
+ * , and the public allocator is A(1).
+ *
+ * The allocation process would be:
+ *
+ *   A(n).Allocate() -> ... -> A(2).Allocate() -> A(1).Allocate()
+ *
+ * , and the free process would be:
+ *
+ *   A(1).Free() -> A(2).Free() -> ... -> A(n).Free()
+ *
+ * Therefore, we should record the allocator chain when allocating, so
+ * that we can free the allocation in the reverse order of allocator chain.
+ * The field `decorated_allocators_` is used to record this chain.
+ *
+ * Another example is that we want to add additional fields in Allocation,
+ * e.g., something what is done in AlignedAllocator, etc.
+ * In this case, we should declare a derived class of Allocation, which
+ * contains an underlying Allocation allocated by the underlying allocator.
+ * Therefore, `decorated_allocators_` of the new Allocation object
+ * would
+ * be a new chain, differing from the underlying Allocation object.
+ */
+class Allocation : public phi::Allocation {
+ public:
+  Allocation(void* ptr, size_t size, Place place)
+      : phi::Allocation(ptr, size, place), base_ptr_(ptr) {}
+  Allocation(void* ptr, void* base_ptr, size_t size, const Place& place)
+      : phi::Allocation(ptr, size, place), base_ptr_(base_ptr) {}
+
+  void* base_ptr() const { return base_ptr_; }
+
+ private:
+  inline void RegisterDecoratedAllocator(Allocator* allocator) {
+    decorated_allocators_.emplace_back(allocator);
+  }
+
+  inline void PopDecoratedAllocator() { decorated_allocators_.pop_back(); }
+
+  inline Allocator* TopDecoratedAllocator() {
+    return decorated_allocators_.back();
+  }
+
+ private:
+  void* base_ptr_;  // the point that directly requested from system
+
+  /**
+   * NOTE(zjl): Since decorated_allocators_ is usually a small vector.
+   * We reserve a small buffer to it to prevent frequent heap allocation
+   *
+   * Instead, we can use a std::vector<Allocator *> here, and reserve
+   * kReserveAllocatorNum in constructor of Allocation.
+   * But using std::vector<Allocator *> would make ocr recognition model
+   * fail in CE. The train duration is 8% slower than KPI.
+   */
+  static constexpr size_t kReserveAllocatorNum = 8;
+  using DecoratedAllocatorStack =
+      InlinedVector<Allocator*, kReserveAllocatorNum>;
+
+  DecoratedAllocatorStack decorated_allocators_;
+
+  friend class Allocator;
+  friend class MultiScalePoolAllocator;
+};
+
+using AllocationPtr = phi::Allocator::AllocationPtr;
+using DecoratedAllocationPtr =
+    std::unique_ptr<Allocation, phi::Allocator::DeleterType>;
+
+template <typename T>
+static T&& FillValue(T&& allocation) {
+#if defined(PADDLE_WITH_CUDA)
+  if (allocation != nullptr) {
+    if (FLAGS_sync_after_alloc || FLAGS_alloc_fill_value >= 0) {
+      bool need_sync = !phi::is_cpu_place(allocation->place());
+      if (need_sync) {
+        PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceSynchronize());
+      }
+      if (FLAGS_alloc_fill_value >= 0) {
+        PADDLE_ENFORCE_LE(FLAGS_alloc_fill_value,
+                          255,
+                          common::errors::InvalidArgument(
+                              "The value of FLAGS_alloc_fill_value must be in "
+                              "range [0, 255]. Expected 0 <= "
+                              "FLAGS_alloc_fill_value <= 255, but received "
+                              "FLAGS_alloc_fill_value = %ld.",
+                              FLAGS_alloc_fill_value));
+        const int fill_value = static_cast<int>(FLAGS_alloc_fill_value);
+        if (phi::is_gpu_place(allocation->place())) {
+          PADDLE_ENFORCE_GPU_SUCCESS(
+              cudaMemset(allocation->ptr(), fill_value, allocation->size()));
+        } else {
+          std::memset(allocation->ptr(), fill_value, allocation->size());
+        }
+        if (need_sync) {
+          PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceSynchronize());
+        }
+      }
+    }
+  }
+#endif
+  return std::forward<T>(allocation);
+}
+
+// Base interface class of memory Allocator.
+class PADDLE_API Allocator : public phi::Allocator {
+ public:
+  static void AllocationDeleter(phi::Allocation* allocation) {
+    Allocator* allocator =
+        static_cast<Allocation*>(allocation)->TopDecoratedAllocator();
+    allocator->Free(allocation);
+  }
+
+  // Allocate an allocation.
+  // size may be 0, but it would be too complex if we handle size == 0
+  // in each Allocator. So we handle size == 0 inside AllocatorFacade
+  // in our design.
+  AllocationPtr Allocate(size_t size) override {
+    auto* ptr = AllocateImpl(size);
+    static_cast<Allocation*>(ptr)->RegisterDecoratedAllocator(this);
+    return FillValue(AllocationPtr(ptr, AllocationDeleter));
+  }
+
+  void Free(phi::Allocation* allocation) override {
+    static_cast<Allocation*>(allocation)->PopDecoratedAllocator();
+    FreeImpl(allocation);
+  }
+
+  uint64_t Release(const Place& place) { return ReleaseImpl(place); }
+  size_t Compact(const Place& place) { return CompactImpl(place); }
+
+  virtual void Accept(AllocatorVisitor* visitor);
+
+ protected:
+  virtual phi::Allocation* AllocateImpl(size_t size) = 0;
+  virtual void FreeImpl(phi::Allocation* allocation);
+  virtual uint64_t ReleaseImpl(const Place& place UNUSED) { return 0; }
+  virtual size_t CompactImpl(const Place& place UNUSED) {
+    PADDLE_THROW(phi::errors::Unimplemented("Compact is not supported"));
+    return 0;
+  }
+};
+
+inline size_t AlignedSize(size_t size, size_t alignment) {
+  auto remaining = size % alignment;  // NOLINT
+  return remaining == 0 ? size : size + alignment - remaining;
+}
+
+inline size_t AlignedPtrOffset(const void* ptr, size_t alignment) {
+  auto ptr_addr = reinterpret_cast<uintptr_t>(ptr);
+  auto diff = ptr_addr % alignment;
+  return diff == 0 ? 0 : alignment - diff;
+}
+
+template <typename Derived, typename Base, typename BaseDel>
+decltype(auto) static_unique_ptr_cast(std::unique_ptr<Base, BaseDel>&& p) {
+  static_assert(std::is_base_of<Base, Derived>::value,
+                "Derived type must derive from Base.");
+  auto d = static_cast<Derived*>(p.release());
+  return std::unique_ptr<Derived, BaseDel>(d, p.get_deleter());
+}
+
+/**
+ * \brief MultiScalePoolAllocator is a decorator of Allocator.
+ * It allocates small request from small_allocator and large request from
+ * large_allocator.
+ */
+
+class PADDLE_API MultiScalePoolAllocator : public Allocator {
+ public:
+  MultiScalePoolAllocator(const std::shared_ptr<Allocator>& small_allocator,
+                          const std::shared_ptr<Allocator>& large_allocator,
+                          size_t alignment,
+                          const GPUPlace& place)
+      : small_allocator_(small_allocator),
+        large_allocator_(large_allocator),
+        alignment_(alignment),
+        place_(place) {}
+
+  // Allocate an allocation from small_allocator or large_allocator according to
+  // size.
+  AllocationPtr Allocate(size_t size) override {
+    auto allocation = IsSmallRequest(size) ? small_allocator_->Allocate(size)
+                                           : large_allocator_->Allocate(size);
+    static_cast<Allocation*>(allocation.get())
+        ->RegisterDecoratedAllocator(this);
+    if (FLAGS_record_alloc_event) {
+      uint64_t id = global_seq_counter_.fetch_add(1, std::memory_order_relaxed);
+      uintptr_t allocator_instance = reinterpret_cast<uintptr_t>(this);
+      RecordAlloc(allocator_instance, id, size);
+      allocation->set_id(id);
+    }
+    return allocation;
+  };
+  // Free an allocation from small_allocator or large_allocator.
+  void Free(phi::Allocation* allocation) override {
+    if (FLAGS_record_alloc_event) {
+      uint64_t id = allocation->id();
+      uintptr_t allocator_instance = reinterpret_cast<uintptr_t>(this);
+      RecordFree(allocator_instance, id, allocation->size());
+    }
+    auto* decorated_allocation = static_cast<Allocation*>(allocation);
+    decorated_allocation->PopDecoratedAllocator();
+    Allocator* underlying_allocator =
+        decorated_allocation->TopDecoratedAllocator();
+    PADDLE_ENFORCE_EQ(
+        underlying_allocator == small_allocator_.get() ||
+            underlying_allocator == large_allocator_.get(),
+        true,
+        common::errors::InvalidArgument(
+            "MultiScalePoolAllocator found an unexpected underlying "
+            "allocator when freeing allocation %p.",
+            allocation->ptr()));
+    underlying_allocator->Free(allocation);
+  };
+  // Get allocate event when start FLAGS_record_alloc_event.
+  std::vector<std::tuple<uintptr_t, bool, uint64_t, size_t, int64_t, int64_t>>
+  GetEvents() {
+    std::lock_guard<SpinLock> lock(spinlock_);
+    return allocation_records_;
+  }
+  // Get small_allocator_ and large_allocator_.
+  std::shared_ptr<Allocator>& GetSmallAllocator() { return small_allocator_; }
+  std::shared_ptr<Allocator>& GetLargeAllocator() { return large_allocator_; }
+  virtual bool IsSmallRequest(size_t size) = 0;
+
+ private:
+  phi::Allocation* AllocateImpl(size_t UNUSED) override { return nullptr; }
+  std::shared_ptr<Allocator> small_allocator_;
+  std::shared_ptr<Allocator> large_allocator_;
+  size_t alignment_;
+  Place place_;
+
+  // Record allocate event into `allocation_records_` when
+  // `FLAGS_record_alloc_event` is True.
+  void RecordAlloc(uintptr_t allocator, uint64_t id, size_t size);
+
+  // Record free event into `allocation_records_` when
+  // `FLAGS_record_alloc_event` is True.
+  void RecordFree(uintptr_t allocator, uint64_t id, size_t size);
+
+  // Return tuple is <allocator_instance, is_allocate, id, allocate_size,
+  // cur_allocated, max_reserved>, if more fields are added later, consider
+  // using a struct to combine them.
+  std::vector<std::tuple<uintptr_t, bool, uint64_t, size_t, int64_t, int64_t>>
+      allocation_records_;
+  SpinLock spinlock_;
+  static inline std::atomic<uint64_t> global_seq_counter_{0};
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/allocator_facade.h

@@ -0,0 +1,159 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include <memory>
+
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#ifdef PADDLE_WITH_CUDA
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+#endif
+#ifdef PADDLE_WITH_XPU
+#include "paddle/phi/core/platform/device/xpu/xpu_info.h"
+#endif
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/stream.h"
+
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+#include "paddle/phi/backends/device_manager.h"
+#include "paddle/phi/core/memory/allocation/custom_allocator.h"
+#endif
+
+namespace paddle {
+namespace memory {
+class AllocatorVisitor;
+namespace allocation {
+
+// Allocator Facade is the interface exposed to other modules.
+// All the configuration or dirty code under development should
+// be hidden behind this facade.
+//
+// NOTE(yy): This class is a singleton class.
+// NOTE(yy): To create a stable ABI and make compilation faster. Here we use
+// a Pimpl trick;
+class AllocatorFacadePrivate;
+class AllocatorFacade {
+ public:
+  using Allocation = phi::Allocation;
+  AllocatorFacade(const AllocatorFacade& o) = delete;
+  const AllocatorFacade& operator=(const AllocatorFacade& o) = delete;
+  ~AllocatorFacade();
+
+  PADDLE_API static AllocatorFacade& Instance();
+
+  AllocatorFacadePrivate* GetPrivate() const;
+
+  PADDLE_API const std::shared_ptr<Allocator>& GetAllocator(const Place& place);
+
+  PADDLE_API const std::shared_ptr<Allocator>& GetAutoGrowthAllocator(
+      const Place& place);
+
+  void* GetBasePtr(const std::shared_ptr<Allocation>& allocation);
+
+  PADDLE_API const std::shared_ptr<Allocator>& GetZeroAllocator(
+      const Place& place);
+
+  // Allocate a shared allocation.
+  std::shared_ptr<Allocation> AllocShared(const Place& place, size_t size);
+  // Allocate a unique allocation.
+  PADDLE_API AllocationPtr Alloc(const Place& place, size_t size);
+  // Release unused memory pool.
+  uint64_t Release(const Place& place);
+  // Compact memory of free blocks held by the VmmAllocator.
+  size_t Compact(const Place& place);
+
+  /**
+   * @brief Accepts an AllocatorVisitor and iterates over all nested Allocator
+   * instances associated with a specific memory location (Place), executing the
+   * visitor's corresponding Visit method for each one.
+   *
+   * This method facilitates the traversal of the Allocator hierarchy for the
+   * given memory Place, allowing the visitor to collect statistics or perform
+   * operations on all constituent allocators.
+   *
+   * @param place The memory location
+   * @param visitor A pointer to the AllocatorVisitor whose Visit methods will
+   * be executed against the nested allocators found at the specified Place.
+   */
+  void Accept(const Place& place, AllocatorVisitor* visitor);
+
+  std::shared_ptr<Allocation> AllocShared(const Place& place,
+                                          size_t size,
+                                          const phi::Stream& stream);
+
+  AllocationPtr Alloc(const Place& place,
+                      size_t size,
+                      const phi::Stream& stream);
+
+  bool InSameStream(const std::shared_ptr<Allocation>& allocation,
+                    const phi::Stream& stream);
+
+  PADDLE_API bool IsStreamSafeCUDAAllocatorUsed();
+  PADDLE_API bool IsCUDAMallocAsyncAllocatorUsed();
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  // TODO(zhiqiu): change gpuStream_t to phi::Stream if needed.
+  uint64_t Release(const GPUPlace& place, gpuStream_t stream);
+  bool RecordStream(std::shared_ptr<Allocation> allocation, gpuStream_t stream);
+  void EraseStream(std::shared_ptr<Allocation> allocation, gpuStream_t stream);
+
+  PADDLE_API const std::shared_ptr<Allocator>& GetAllocator(const Place& place,
+                                                            gpuStream_t stream);
+  gpuStream_t GetStream(const std::shared_ptr<Allocation>& allocation) const;
+  void SetDefaultStream(const GPUPlace& place, gpuStream_t stream);
+#elif defined(PADDLE_WITH_XPU)
+  PADDLE_API const std::shared_ptr<Allocator>& GetAllocator(const Place& place,
+                                                            XPUStream stream);
+  bool RecordStream(std::shared_ptr<Allocation> allocation, XPUStream stream);
+  void SetDefaultStream(const XPUPlace& place, XPUStream stream);
+  void EraseStream(std::shared_ptr<Allocation> allocation, XPUStream stream);
+#endif
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
+    defined(PADDLE_WITH_CUSTOM_DEVICE)
+  void PrepareMemoryPoolForCUDAGraph(int64_t id);
+  void RemoveMemoryPoolOfCUDAGraph(int64_t id);
+#endif
+#if defined(PADDLE_WITH_XPU)
+  void RemoveMemoryPoolOfXPUGraph(int64_t id);
+  void PrepareMemoryPoolForXPUGraph(int64_t id);
+#endif
+
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  uint64_t Release(const CustomPlace& place, phi::stream::stream_t stream);
+  bool RecordStream(std::shared_ptr<Allocation> allocation,
+                    phi::stream::stream_t stream);
+  void EraseStream(std::shared_ptr<Allocation> allocation,
+                   phi::stream::stream_t stream);
+  PADDLE_API const std::shared_ptr<Allocator>& GetAllocator(
+      const Place& place, phi::stream::stream_t stream);
+  phi::stream::stream_t GetStream(
+      const std::shared_ptr<Allocation>& allocation) const;
+  void SetDefaultStream(const CustomPlace& place, phi::stream::stream_t stream);
+#endif
+  // TODO(yy): Allocate a Copy-On-Write allocation?
+ private:
+  AllocatorFacade();
+  AllocatorFacadePrivate* m_;
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
+    defined(PADDLE_WITH_CUSTOM_DEVICE) || defined(PADDLE_WITH_XPU)
+  std::unordered_map<int64_t, std::unique_ptr<AllocatorFacadePrivate>>
+      cuda_graph_map_;
+  std::unordered_map<int64_t, int64_t> cuda_graph_ref_cnt_;
+#endif
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/allocator_strategy.cc

@@ -0,0 +1,53 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/allocator_strategy.h"
+
+#include "paddle/common/flags.h"
+#include "paddle/phi/core/enforce.h"
+
+COMMON_DECLARE_string(allocator_strategy);
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+static AllocatorStrategy GetStrategyFromFlag() {
+  if (FLAGS_allocator_strategy == "naive_best_fit") {
+    return AllocatorStrategy::kNaiveBestFit;
+  }
+
+  if (FLAGS_allocator_strategy == "auto_growth") {
+    return AllocatorStrategy::kAutoGrowth;
+  }
+
+  if (FLAGS_allocator_strategy == "thread_local") {
+    return AllocatorStrategy::kThreadLocal;
+  }
+
+  PADDLE_THROW(common::errors::InvalidArgument(
+      "Unsupported allocator strategy: %s, candidates are naive_best_fit, "
+      "auto_growth or thread_local.",
+      FLAGS_allocator_strategy));
+}
+
+AllocatorStrategy GetAllocatorStrategy() {
+  static AllocatorStrategy strategy = GetStrategyFromFlag();
+  return strategy;
+}
+
+void UseAllocatorStrategyGFlag() {}
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/allocator_strategy.h

@@ -0,0 +1,32 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include "paddle/common/macros.h"
+#include "paddle/utils/test_macros.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+enum class AllocatorStrategy { kNaiveBestFit, kAutoGrowth, kThreadLocal };
+
+extern AllocatorStrategy GetAllocatorStrategy();
+
+// Do nothing, just make sure linker do not prune this file.
+PADDLE_API void UseAllocatorStrategyGFlag();
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/auto_growth_best_fit_allocator.cc

@@ -0,0 +1,390 @@
+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/auto_growth_best_fit_allocator.h"
+
+#include <algorithm>
+#include <mutex>  // NOLINT
+#include <utility>
+
+#include "glog/logging.h"
+
+#include "paddle/common/flags.h"
+#include "paddle/phi/api/profiler/event_tracing.h"
+#include "paddle/phi/backends/device_manager.h"
+#include "paddle/phi/core/memory/allocation/aligned_allocator.h"
+
+PHI_DEFINE_EXPORTED_READONLY_bool(
+    free_idle_chunk,
+    false,
+    "Whether to free idle chunk when each allocation is freed. "
+    "If false, all freed allocation would be cached to speed up next "
+    "allocation request. If true, no allocation would be cached. This "
+    "flag only works when FLAGS_allocator_strategy=auto_growth.");
+
+PHI_DEFINE_EXPORTED_READONLY_bool(
+    free_when_no_cache_hit,
+    false,
+    "Whether to free idle chunks when no cache hit. If true, idle "
+    "chunk would be freed when no cache hit; if false, idle "
+    "chunk would be freed when out of memory occurs. This flag "
+    "only works when FLAGS_allocator_strategy=auto_growth.");
+
+PHI_DEFINE_EXPORTED_READONLY_bool(print_allocator_trace_info,
+                                  false,
+                                  "print trace memory info");
+
+PHI_DEFINE_EXPORTED_READONLY_bool(dump_chunk_info, false, "dump chunk info");
+PHI_DEFINE_EXPORTED_uint64(
+    alignment_size,
+    256,
+    "All sizes are rounded up to a multiple of this value. Default: 256.");
+PHI_DEFINE_EXPORTED_uint64(
+    small_pool_size_in_mb,
+    0,
+    "Threshold (MiB) separating the small and large pools. "
+    "0 disables the small pool and enables single-pool mode "
+    "(all requests go to the large pool). When > 0, requests "
+    "<= threshold use the small pool; larger requests use the "
+    "large pool. Default: 0.");
+PHI_DEFINE_EXPORTED_uint64(small_pool_auto_growth_chunk_size_in_mb,
+                           0,
+                           "The minimal chunk size for the small pool in MiB. "
+                           "If small_pool_size_in_mb > 0, this overrides "
+                           "the constructor-provided global growth size "
+                           "(FLAGS_auto_growth_chunk_size_in_mb).");
+PHI_DEFINE_EXPORTED_uint64(large_pool_auto_growth_chunk_size_in_mb,
+                           0,
+                           "The minimal chunk size for the large pool in MiB. "
+                           "If small_pool_size_in_mb > 0, this overrides "
+                           "the constructor-provided global growth size "
+                           "(FLAGS_auto_growth_chunk_size_in_mb).");
+PHI_DEFINE_EXPORTED_uint64(
+    large_pool_pre_alloc_in_mb,
+    0,
+    "Pre-reserve this many MiB in the large pool. 0 disables pre-allocation.");
+PHI_DEFINE_EXPORTED_uint64(
+    small_pool_pre_alloc_in_mb,
+    0,
+    "Pre-reserve this many MiB in the small pool. 0 disables pre-allocation.");
+
+namespace paddle::memory::allocation {
+
+AutoGrowthBestFitAllocator::AutoGrowthBestFitAllocator(
+    std::shared_ptr<Allocator> underlying_allocator,
+    size_t alignment,
+    size_t chunk_size,
+    bool allow_free_idle_chunk,
+    int extra_padding_size)
+    : underlying_allocator_(std::move(underlying_allocator)),
+      alignment_(alignment),
+      chunk_size_(std::max(AlignedSize(chunk_size, alignment), alignment)),
+      allow_free_idle_chunk_(allow_free_idle_chunk),
+      extra_padding_size_(extra_padding_size) {
+  total_alloc_times_ = 0;
+  total_alloc_size_ = 0;
+  total_free_times_ = 0;
+  total_free_size_ = 0;
+  VLOG(7) << "chunk_size_:" << chunk_size_;
+}
+
+void AutoGrowthBestFitAllocator::DumpInfo() const {
+  for (auto chunk_it = chunks_.begin(); chunk_it != chunks_.end(); ++chunk_it) {
+    std::cout << "Chunk\t";
+    std::ostringstream oss_used;
+    std::ostringstream oss_free;
+    size_t total = 0, free = 0, used = 0;
+    for (auto &b : chunk_it->blocks_) {
+      total += b.size_;
+      if (b.is_free_) {
+        free += b.size_;
+        oss_free << "(" << b.size_ << "," << b.ptr_ << ")";
+      } else {
+        used += b.size_;
+        oss_used << "(" << b.size_ << "," << b.ptr_ << ")";
+      }
+    }
+    std::cout << total << "\t" << used << "\t" << free << "\t";
+    std::cout << "[" << oss_used.str() << "]\t[" << oss_free.str() << "]"
+              << std::endl;
+  }
+}
+
+bool AutoGrowthBestFitAllocator::is_small_free_block(size_t size) {
+  auto small_pool_size = FLAGS_small_pool_size_in_mb << 20;
+  if (size <= small_pool_size) {
+    return true;
+  } else {
+    return false;
+  }
+}
+
+size_t AutoGrowthBestFitAllocator::auto_growth_size(bool is_small,
+                                                    size_t chunk_size) {
+  // fallback to single pool and use constructor-provided chunk_size.
+  if (FLAGS_small_pool_size_in_mb == 0) {
+    return chunk_size;
+  }
+
+  const uint64_t pool_auto_growth_chunk_size_mb =
+      is_small ? FLAGS_small_pool_auto_growth_chunk_size_in_mb
+               : FLAGS_large_pool_auto_growth_chunk_size_in_mb;
+  const size_t auto_growth_size =
+      pool_auto_growth_chunk_size_mb
+          ? (static_cast<size_t>(pool_auto_growth_chunk_size_mb) << 20)
+          : 0;
+
+  return AlignedSize(auto_growth_size, alignment_);
+}
+
+void AutoGrowthBestFitAllocator::PreAlloc() {
+  auto small_pool_pre_alloc = FLAGS_small_pool_pre_alloc_in_mb << 20;
+  auto large_pool_pre_alloc = FLAGS_large_pool_pre_alloc_in_mb << 20;
+  if (small_pool_pre_alloc > 0) {
+    VLOG(10) << "PreAlloc small_pool_pre_alloc_in_mb = "
+             << FLAGS_small_pool_pre_alloc_in_mb;
+    chunks_.emplace_back(static_unique_ptr_cast<Allocation>(
+        underlying_allocator_->Allocate(small_pool_pre_alloc)));
+    auto *chunk = &(*chunks_.rbegin());
+    uint8_t *p = reinterpret_cast<uint8_t *>(chunk->allocation_->ptr());
+    auto &blocks = chunk->blocks_;
+    blocks.emplace_back(
+        p, small_pool_pre_alloc, /*is_free=*/true, /*is_small=*/true, chunk);
+    small_free_blocks_.emplace(std::make_pair(small_pool_pre_alloc, p),
+                               --(blocks.end()));
+  }
+
+  if (large_pool_pre_alloc > 0) {
+    VLOG(10) << "PreAlloc large_pool_pre_alloc_in_mb = "
+             << FLAGS_large_pool_pre_alloc_in_mb;
+    chunks_.emplace_back(static_unique_ptr_cast<Allocation>(
+        underlying_allocator_->Allocate(large_pool_pre_alloc)));
+    auto *chunk = &(*chunks_.rbegin());
+    uint8_t *p = reinterpret_cast<uint8_t *>(chunk->allocation_->ptr());
+    auto &blocks = chunk->blocks_;
+    blocks.emplace_back(
+        p, large_pool_pre_alloc, /*is_free=*/true, /*is_small=*/false, chunk);
+    large_free_blocks_.emplace(std::make_pair(large_pool_pre_alloc, p),
+                               --(blocks.end()));
+  }
+}
+
+phi::Allocation *AutoGrowthBestFitAllocator::AllocateImpl(
+    size_t unaligned_size) {
+  phi::RecordEvent record("AutoGrowthBestFitAllocator::Allocate",
+                          phi::TracerEventType::UserDefined,
+                          9 /*level*/);
+
+  size_t size = AlignedSize(unaligned_size + extra_padding_size_, alignment_);
+
+  VLOG(10) << "Allocate " << unaligned_size << " bytes, aligned to " << size
+           << ", extra size " << extra_padding_size_;
+
+  std::lock_guard<SpinLock> guard(spinlock_);
+  bool is_small = is_small_free_block(size);
+  auto &free_blocks = is_small ? small_free_blocks_ : large_free_blocks_;
+  auto iter = free_blocks.lower_bound(std::make_pair(size, nullptr));
+  BlockIt block_it;
+  if (iter != free_blocks.end()) {
+    block_it = iter->second;
+    free_blocks.erase(iter);
+    auto *chunk = block_it->chunk_;
+    size_t remaining_size = block_it->size_ - size;
+    VLOG(10) << "Allocate " << size << " bytes from chunk size "
+             << block_it->size_ << ", remaining " << remaining_size;
+    if (remaining_size == 0) {
+      block_it->is_free_ = false;
+      block_it->is_small_ = is_small;
+    } else {
+      auto remaining_free_block = chunk->blocks_.insert(
+          block_it,
+          Block(block_it->ptr_, remaining_size, true, is_small, chunk));
+      free_blocks.emplace(std::make_pair(remaining_size, block_it->ptr_),
+                          remaining_free_block);
+      block_it->ptr_ =
+          reinterpret_cast<uint8_t *>(block_it->ptr_) + remaining_size;
+      block_it->size_ = size;
+      block_it->is_free_ = false;
+      block_it->is_small_ = is_small;
+    }
+  } else {
+    if (FLAGS_dump_chunk_info) {
+      std::cout << "MemDbg memory not enough growth chunk, need size = " << size
+                << std::endl;
+      DumpInfo();
+    }
+
+    if (FLAGS_free_when_no_cache_hit) {
+      FreeIdleChunks();
+    }
+    size_t realloc_size =
+        std::max(size, auto_growth_size(is_small, chunk_size_));
+
+    try {
+      chunks_.emplace_back(static_unique_ptr_cast<Allocation>(
+          underlying_allocator_->Allocate(realloc_size)));
+    } catch (BadAlloc &ex) {
+      if (FLAGS_dump_chunk_info) {
+        std::cout << "MemDbg OOM" << std::endl;
+        DumpInfo();
+      }
+      if (FLAGS_free_when_no_cache_hit) throw ex;
+      FreeIdleChunks();
+      chunks_.emplace_back(static_unique_ptr_cast<Allocation>(
+          underlying_allocator_->Allocate(realloc_size)));
+    }
+
+    auto *chunk = &(*chunks_.rbegin());
+    realloc_size = chunk->allocation_->size();
+    uint8_t *p = reinterpret_cast<uint8_t *>(chunk->allocation_->ptr());
+    auto &blocks = chunk->blocks_;
+
+    size_t remaining_size = realloc_size - size;
+    if (remaining_size > 0) {
+      blocks.emplace_back(p, remaining_size, true, is_small, chunk);
+      free_blocks.emplace(std::make_pair(remaining_size, p), --(blocks.end()));
+    }
+    blocks.emplace_back(p + remaining_size, size, false, is_small, chunk);
+    block_it = --(blocks.end());
+    VLOG(5) << "Not found and reallocate " << realloc_size << "("
+            << static_cast<void *>(p) << "), and remaining " << remaining_size;
+    if (FLAGS_dump_chunk_info) {
+      std::cout << "MemDbg memory after growth chunk, realloc_size = "
+                << realloc_size << std::endl;
+      DumpInfo();
+    }
+  }
+  ++total_alloc_times_;
+  total_alloc_size_ += size;
+  VLOG(10) << "Alloc " << block_it->size_ << " bytes, ptr = " << block_it->ptr_;
+  auto block_t = new BlockAllocation(block_it);
+  return block_t;
+}
+
+void AutoGrowthBestFitAllocator::FreeImpl(phi::Allocation *allocation) {
+  phi::RecordEvent record("AutoGrowthBestFitAllocator::Free",
+                          phi::TracerEventType::UserDefined,
+                          9 /*level*/);
+  VLOG(10) << "Free " << allocation->size()
+           << " bytes, ptr = " << allocation->ptr();
+  std::lock_guard<SpinLock> guard(spinlock_);
+  auto block_it = static_cast<BlockAllocation *>(allocation)->block_it_;
+  auto &blocks = block_it->chunk_->blocks_;
+  bool is_small = block_it->is_small_;
+  auto &free_blocks = is_small ? small_free_blocks_ : large_free_blocks_;
+
+  total_free_times_ += 1;
+  total_free_size_ += block_it->size_;
+
+  block_it->is_free_ = true;
+
+  if (block_it != blocks.begin()) {
+    auto prev_it = block_it;
+    --prev_it;
+
+    if (prev_it->is_free_) {
+      free_blocks.erase(std::make_pair(prev_it->size_, prev_it->ptr_));
+      prev_it->size_ += block_it->size_;
+      blocks.erase(block_it);
+      block_it = prev_it;
+    }
+  }
+
+  auto next_it = block_it;
+  ++next_it;
+
+  // It's weird that using `next_it == blocks.end()` will cause a judgment fail.
+  if (block_it != (--blocks.end()) && next_it->is_free_) {
+    free_blocks.erase(std::make_pair(next_it->size_, next_it->ptr_));
+    block_it->size_ += next_it->size_;
+    blocks.erase(next_it);
+  }
+
+  free_blocks.emplace(std::make_pair(block_it->size_, block_it->ptr_),
+                      block_it);
+
+  delete allocation;
+
+  if (FLAGS_free_idle_chunk) {
+    FreeIdleChunks();
+  }
+  if (FLAGS_dump_chunk_info) {
+    DumpInfo();
+  }
+}
+
+uint64_t AutoGrowthBestFitAllocator::FreeIdleChunks() {
+  if (FLAGS_dump_chunk_info) {
+    std::cout << "FreeIdleChunks called" << std::endl;
+  }
+  if (!allow_free_idle_chunk_) {
+    return 0;
+  }
+  uint64_t bytes = 0;
+  for (auto chunk_it = chunks_.begin(); chunk_it != chunks_.end();) {
+    auto &blocks = chunk_it->blocks_;
+    if (blocks.size() == 1 && blocks.begin()->is_free_) {
+      auto &block = *blocks.begin();
+      bool is_small = block.is_small_;
+      auto &free_blocks = is_small ? small_free_blocks_ : large_free_blocks_;
+      VLOG(2) << "Free chunk with size " << block.size_;
+      if (FLAGS_dump_chunk_info) {
+        std::cout << "FreeIdleChunks chunk is " << block.size_ << ", "
+                  << block.ptr_ << std::endl;
+      }
+      bytes += block.size_;
+      free_blocks.erase(std::make_pair(block.size_, block.ptr_));
+      chunk_it = chunks_.erase(chunk_it);
+    } else {
+      ++chunk_it;
+    }
+  }
+
+  if (FLAGS_print_allocator_trace_info) {
+    Trace();
+  }
+  return bytes;
+}
+
+void AutoGrowthBestFitAllocator::Trace() const {
+  size_t small_cur_idle_bytes = 0;
+  auto small_it = small_free_blocks_.begin();
+  for (; small_it != small_free_blocks_.end(); ++small_it) {
+    small_cur_idle_bytes += small_it->second->size_;
+  }
+  size_t large_cur_idle_bytes = 0;
+  auto large_it = large_free_blocks_.begin();
+  for (; large_it != large_free_blocks_.end(); ++large_it) {
+    large_cur_idle_bytes += large_it->second->size_;
+  }
+
+  VLOG(1) << "alloc:"
+          << total_alloc_size_ / static_cast<double>(1024 * 1024)  // NOLINT
+          << "m free:"
+          << total_free_size_ / static_cast<double>(1024 * 1024)  // NOLINT
+          << "m busy:"
+          << (total_alloc_size_ - total_free_size_) /  // NOLINT
+                 static_cast<double>(1024 * 1024)
+          << "m small idle:"
+          << small_cur_idle_bytes / static_cast<double>(1024 * 1024)  // NOLINT
+          << "m large idle:"
+          << large_cur_idle_bytes / static_cast<double>(1024 * 1024)  // NOLINT
+          << "m alloc_times:" << total_alloc_times_
+          << " free_times:" << total_free_times_
+          << " small free_blocks_num:" << small_free_blocks_.size()
+          << " large free_blocks_num:" << large_free_blocks_.size()
+          << " curr_chunks_num:" << chunks_.size();
+}
+
+}  // namespace paddle::memory::allocation

paddle/phi/core/memory/allocation/auto_growth_best_fit_allocator.h

@@ -0,0 +1,135 @@
+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <list>
+#include <map>
+#include <memory>
+#include <mutex>  // NOLINT
+#include <utility>
+
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/allocation/spin_lock.h"
+#include "paddle/phi/core/memory/mem_visitor.h"
+
+COMMON_DECLARE_bool(enable_auto_growth_allocator_add_lock);
+
+namespace paddle {
+namespace memory {
+class AllBlocksInfoVisitor;
+namespace allocation {
+
+class PADDLE_API AutoGrowthBestFitAllocator : public Allocator {
+ public:
+  AutoGrowthBestFitAllocator(std::shared_ptr<Allocator> underlying_allocator,
+                             size_t alignment,
+                             size_t chunk_size = 0,
+                             bool allow_free_idle_chunk = true,
+                             int extra_padding_size = 0);
+
+  bool IsAllocThreadSafe() const override { return true; }
+
+  void DumpInfo() const;
+
+  void PreAlloc() override;
+
+  void Accept(AllocatorVisitor *visitor) override { visitor->Visit(this); }
+
+ protected:
+  phi::Allocation *AllocateImpl(size_t size) override;
+
+  void FreeImpl(phi::Allocation *allocation) override;
+
+  bool is_small_free_block(size_t size);
+  size_t auto_growth_size(bool is_small, size_t chunk_size);
+
+  // Release the memory block which is not used in pool.
+  uint64_t ReleaseImpl(const Place &place) override {
+    // TODO(vivienfanghuagood): the next line may cause the process to deadlock.
+    if (FLAGS_enable_auto_growth_allocator_add_lock) {
+      std::lock_guard<SpinLock> guard(spinlock_);
+      return FreeIdleChunks();
+    }
+    return FreeIdleChunks();
+  }
+
+ protected:
+  uint64_t FreeIdleChunks();
+  void Trace() const;
+
+  template <typename T>
+  using List = std::list<T>;
+
+  struct Chunk;
+
+  struct Block {
+    Block(void *ptr, size_t size, bool is_free, bool is_small, Chunk *chunk)
+        : ptr_(ptr),
+          size_(size),
+          is_free_(is_free),
+          is_small_(is_small),
+          chunk_(chunk) {}
+
+    void *ptr_;
+    size_t size_;
+    bool is_free_;
+    bool is_small_;
+    Chunk *chunk_;  // which chunk it is from
+  };
+
+  struct Chunk {
+    explicit Chunk(DecoratedAllocationPtr allocation)
+        : allocation_(std::move(allocation)) {}
+
+    DecoratedAllocationPtr allocation_;
+    List<Block> blocks_;
+  };
+
+  struct BlockAllocation : public Allocation {
+    explicit BlockAllocation(const List<Block>::iterator &it)
+        : Allocation(it->ptr_,
+                     it->chunk_->allocation_->base_ptr(),
+                     it->size_,
+                     it->chunk_->allocation_->place()),
+          block_it_(it) {}
+
+    List<Block>::iterator block_it_;
+  };
+
+  using BlockIt = List<Block>::iterator;
+
+  std::shared_ptr<Allocator> underlying_allocator_;
+  std::map<std::pair<size_t, void *>, BlockIt> small_free_blocks_;
+  std::map<std::pair<size_t, void *>, BlockIt> large_free_blocks_;
+  std::list<Chunk> chunks_;
+  size_t alignment_;
+  size_t chunk_size_;
+  bool allow_free_idle_chunk_;
+  int extra_padding_size_;
+
+  // stat info
+  size_t total_alloc_times_;
+  size_t total_alloc_size_;
+  size_t total_free_times_;
+  size_t total_free_size_;
+
+  SpinLock spinlock_;
+
+  friend class paddle::memory::AllBlocksInfoVisitor;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/auto_growth_best_fit_allocator_v2.cc

@@ -0,0 +1,165 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#include "paddle/phi/core/memory/allocation/auto_growth_best_fit_allocator_v2.h"
+#include <algorithm>
+#include <mutex>  // NOLINT
+#include "glog/logging.h"
+#include "paddle/common/flags.h"
+#include "paddle/phi/api/profiler/event_tracing.h"
+#include "paddle/phi/backends/device_manager.h"
+#include "paddle/phi/core/memory/allocation/aligned_allocator.h"
+#include "paddle/phi/core/platform/cuda_device_guard.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+
+PD_DECLARE_bool(free_idle_chunk);
+PD_DECLARE_bool(free_when_no_cache_hit);
+
+namespace paddle::memory::allocation {
+
+AutoGrowthBestFitAllocatorV2::AutoGrowthBestFitAllocatorV2(
+    const std::shared_ptr<Allocator> &underlying_allocator,
+    size_t alignment,
+    GPUPlace place,
+    size_t chunk_size,
+    bool allow_free_idle_chunk,
+    int extra_padding_size)
+    : AutoGrowthBestFitAllocator(underlying_allocator,
+                                 alignment,
+                                 chunk_size,
+                                 true,
+                                 extra_padding_size),
+      place_(place) {}
+
+phi::Allocation *AutoGrowthBestFitAllocatorV2::AllocateImpl(
+    size_t unaligned_size) {
+  phi::RecordEvent record("AutoGrowthBestFitAllocatorV2::Allocate",
+                          phi::TracerEventType::UserDefined,
+                          9 /*level*/);
+
+  size_t size = AlignedSize(unaligned_size + extra_padding_size_, alignment_);
+
+  VLOG(10) << "Allocate " << unaligned_size << " bytes, aligned to " << size
+           << ", extra size " << extra_padding_size_;
+
+  std::lock_guard<SpinLock> guard(spinlock_);
+
+  BlockIt block_it;
+  if (AutoGrowthBestFitAllocatorV2State::GetInstance().IsWarmup()) {
+    auto iter = free_blocks_.lower_bound(std::make_pair(size, nullptr));
+    if (iter != free_blocks_.end() && iter->second->size_ >= unaligned_size &&
+        iter->second->size_ <= size) {
+      block_it = iter->second;
+      free_blocks_.erase(iter);
+      block_it->is_free_ = false;
+      VLOG(10) << "Allocate " << size << " bytes from chunk size "
+               << block_it->size_ << " by strict_matching_state.";
+    } else {
+      size_t actual_avail, actual_total;
+      {
+        platform::CUDADeviceGuard guard(place_.device);
+#ifdef PADDLE_WITH_HIP
+        auto result = hipMemGetInfo(&actual_avail, &actual_total);
+#else
+        auto result = cudaMemGetInfo(&actual_avail, &actual_total);
+#endif
+        if (result != gpuSuccess) {
+          actual_avail = 0;
+        }
+      }
+
+      if (actual_avail < size) {
+        FreeIdleChunks();
+      }
+
+      chunks_.emplace_back(static_unique_ptr_cast<Allocation>(
+          underlying_allocator_->Allocate(size)));
+
+      auto *chunk = &(*chunks_.rbegin());
+      size = chunk->allocation_->size();
+      uint8_t *p = reinterpret_cast<uint8_t *>(chunk->allocation_->ptr());
+      auto &blocks = chunk->blocks_;
+      blocks.emplace_back(p, size, false, true, chunk);
+      block_it = --(blocks.end());
+      VLOG(2) << "Not found and reallocate " << size << "("
+              << static_cast<void *>(p) << ") by strict_matching_state.";
+    }
+  } else {
+    if (is_first_switch_to_regular_) {
+      FreeIdleChunks();
+      is_first_switch_to_regular_ = false;
+    }
+    auto iter = free_blocks_.lower_bound(std::make_pair(size, nullptr));
+
+    if (iter != free_blocks_.end()) {
+      block_it = iter->second;
+      free_blocks_.erase(iter);
+      auto *chunk = block_it->chunk_;
+      size_t remaining_size = block_it->size_ - size;
+      VLOG(10) << "Allocate " << size << " bytes from chunk size "
+               << block_it->size_ << ", remaining " << remaining_size;
+      if (remaining_size == 0) {
+        block_it->is_free_ = false;
+      } else {
+        auto remaining_free_block = chunk->blocks_.insert(
+            block_it, Block(block_it->ptr_, remaining_size, true, true, chunk));
+        free_blocks_.emplace(std::make_pair(remaining_size, block_it->ptr_),
+                             remaining_free_block);
+        block_it->ptr_ =
+            reinterpret_cast<uint8_t *>(block_it->ptr_) + remaining_size;
+        block_it->size_ = size;
+        block_it->is_free_ = false;
+      }
+    } else {
+      if (FLAGS_free_when_no_cache_hit) {
+        FreeIdleChunks();
+      }
+      size_t realloc_size = std::max(size, chunk_size_);
+
+      try {
+        chunks_.emplace_back(static_unique_ptr_cast<Allocation>(
+            underlying_allocator_->Allocate(realloc_size)));
+      } catch (BadAlloc &ex) {
+        if (FLAGS_free_when_no_cache_hit) throw ex;
+        FreeIdleChunks();
+        chunks_.emplace_back(static_unique_ptr_cast<Allocation>(
+            underlying_allocator_->Allocate(realloc_size)));
+      }
+
+      auto *chunk = &(*chunks_.rbegin());
+      realloc_size = chunk->allocation_->size();
+      uint8_t *p = reinterpret_cast<uint8_t *>(chunk->allocation_->ptr());
+      auto &blocks = chunk->blocks_;
+
+      size_t remaining_size = realloc_size - size;
+      if (remaining_size > 0) {
+        blocks.emplace_back(p, remaining_size, true, true, chunk);
+        free_blocks_.emplace(std::make_pair(remaining_size, p),
+                             --(blocks.end()));
+      }
+      blocks.emplace_back(p + remaining_size, size, false, true, chunk);
+      block_it = --(blocks.end());
+      VLOG(2) << "Not found and reallocate " << realloc_size << "("
+              << static_cast<void *>(p) << "), and remaining "
+              << remaining_size;
+    }
+  }
+  ++total_alloc_times_;
+  total_alloc_size_ += size;
+  VLOG(10) << "Alloc " << block_it->size_ << " bytes, ptr = " << block_it->ptr_;
+  return new BlockAllocation(block_it);
+}
+
+}  // namespace paddle::memory::allocation
+#endif

paddle/phi/core/memory/allocation/auto_growth_best_fit_allocator_v2.h

@@ -0,0 +1,72 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#include <list>
+#include <map>
+#include <memory>
+#include <mutex>  // NOLINT
+#include <utility>
+
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/allocation/auto_growth_best_fit_allocator.h"
+#include "paddle/phi/core/memory/allocation/spin_lock.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class AutoGrowthBestFitAllocatorV2 : public AutoGrowthBestFitAllocator {
+ public:
+  AutoGrowthBestFitAllocatorV2(
+      const std::shared_ptr<Allocator> &underlying_allocator,
+      size_t alignment,
+      GPUPlace place,
+      size_t chunk_size = 0,
+      bool allow_free_idle_chunk = true,
+      int extra_padding_size = 0);
+
+ protected:
+  phi::Allocation *AllocateImpl(size_t size) override;
+
+ private:
+  GPUPlace place_;
+  bool is_first_switch_to_regular_{true};
+  std::map<std::pair<size_t, void *>, BlockIt> free_blocks_;
+};
+
+class AutoGrowthBestFitAllocatorV2State {
+ public:
+  AutoGrowthBestFitAllocatorV2State() = default;
+
+  ~AutoGrowthBestFitAllocatorV2State() {}
+
+  void SetWarmup(bool warmup) { is_warmup_ = warmup; }
+
+  bool IsWarmup() { return is_warmup_; }
+
+  static AutoGrowthBestFitAllocatorV2State &GetInstance() {
+    static AutoGrowthBestFitAllocatorV2State instance;
+    return instance;
+  }
+
+ private:
+  bool is_warmup_{true};
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle
+#endif

paddle/phi/core/memory/allocation/best_fit_allocator.cc

@@ -0,0 +1,188 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/best_fit_allocator.h"
+
+#include <cmath>
+#include <mutex>
+
+#include "paddle/common/macros.h"
+#include "paddle/phi/core/enforce.h"
+
+REGISTER_FILE_SYMBOLS(best_fit_allocator);
+
+namespace paddle::memory::allocation {
+
+static int HighestBitPos(size_t N) {
+  if (UNLIKELY(N == 0)) {
+    return 0;
+  } else {
+#ifdef __GNUCC__
+    return sizeof(unsigned int) * 8 - __builtin_clz(N);
+#else
+    return static_cast<int>(std::log2(N) + 1);
+#endif
+  }
+}
+
+BestFitAllocator::BestFitAllocator(phi::Allocation* allocation)
+    : allocation_(allocation) {
+  details::Chunk chunk;
+  chunk.size_ = allocation_->size();
+  chunk.offset_ = 0;
+  chunk.is_free = true;
+  chunks_.emplace_back(chunk);
+  free_chunks_[HighestBitPos(chunk.size_)].insert(
+      {chunk.size_, chunks_.begin()});
+}
+
+size_t BestFitAllocator::FreeSize() const {
+  size_t acc = 0;
+  for (auto& array_item : free_chunks_) {
+    for (auto& pair : array_item) {
+      acc += pair.second->size_;
+    }
+  }
+  return acc;
+}
+
+BestFitAllocator::ListIt BestFitAllocator::SplitChunk(size_t request_size,
+                                                      size_t free_chunk_offset,
+                                                      MapIt bin_iterator) {
+  auto to_split_it = bin_iterator->second;
+  free_chunks_[free_chunk_offset].erase(bin_iterator);
+
+  PADDLE_ENFORCE_EQ(to_split_it->is_free,
+                    true,
+                    common::errors::PreconditionNotMet(
+                        "The memory chunk to split is not free"));
+  PADDLE_ENFORCE_GE(to_split_it->size_,
+                    request_size,
+                    common::errors::PreconditionNotMet(
+                        "The size of memory chunk to split is "
+                        "not larger than size of request memory"));
+
+  auto remaining_size = to_split_it->size_ - request_size;
+  details::Chunk to_use;
+  details::Chunk remaining;
+  to_use.size_ = request_size;
+  to_use.is_free = false;
+  remaining.size_ = remaining_size;
+  remaining.is_free = true;
+
+  // calc offsets
+  to_use.offset_ = to_split_it->offset_;
+  remaining.offset_ = to_use.offset_ + to_use.size_;
+
+  // insert to chunk list
+  auto to_use_it = chunks_.insert(to_split_it, to_use);
+  if (remaining.size_ != 0) {
+    auto bit_size = static_cast<size_t>(HighestBitPos(remaining.size_));
+    free_chunks_[bit_size].insert(
+        {remaining.size_, chunks_.insert(to_split_it, remaining)});
+  }
+  chunks_.erase(to_split_it);
+  return to_use_it;
+}
+
+void BestFitAllocator::InsertFreeNode(const ListIt& it) {
+  auto pos = static_cast<size_t>(HighestBitPos(it->size_));
+  auto& free_map = free_chunks_[pos];
+  free_map.insert({it->size_, it});
+}
+void BestFitAllocator::EraseFreeNode(const ListIt& it) {
+  size_t pos = static_cast<size_t>(HighestBitPos(it->size_));
+  auto& free_map = free_chunks_[pos];
+  auto map_it = free_map.find(it->size_);
+  while (map_it->second != it && map_it != free_map.end()) {
+    ++map_it;
+  }
+  PADDLE_ENFORCE_NE(
+      map_it,
+      free_map.end(),
+      common::errors::NotFound("The node to erase is not found in map"));
+  free_map.erase(map_it);
+}
+size_t BestFitAllocator::NumFreeChunks() const {
+  size_t num = 0;
+  for (auto& array_item : free_chunks_) {
+    num += array_item.size();
+  }
+  return num;
+}
+void BestFitAllocator::FreeImpl(phi::Allocation* allocation) {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  auto* bf_allocation = dynamic_cast<BestFitAllocation*>(allocation);
+  PADDLE_ENFORCE_NOT_NULL(
+      bf_allocation,
+      common::errors::InvalidArgument(
+          "The input allocation is not type of BestFitAllocation."));
+  auto chunk_it = bf_allocation->ChunkIterator();
+  PADDLE_ENFORCE_EQ(chunk_it->is_free,
+                    false,
+                    common::errors::PreconditionNotMet(
+                        "The chunk of allocation to free is freed already"));
+  chunk_it->is_free = true;
+  if (chunk_it != chunks_.begin()) {
+    auto prev_it = chunk_it;
+    --prev_it;
+
+    if (prev_it->is_free) {
+      // Merge Left.
+      EraseFreeNode(prev_it);
+      prev_it->size_ += chunk_it->size_;
+      chunks_.erase(chunk_it);
+      chunk_it = prev_it;
+    }
+  }
+
+  auto next_it = chunk_it;
+  ++next_it;
+  if (next_it != chunks_.end() && next_it->is_free) {
+    EraseFreeNode(next_it);
+    chunk_it->size_ += next_it->size_;
+    chunks_.erase(next_it);
+  }
+
+  InsertFreeNode(chunk_it);
+  delete allocation;
+}
+phi::Allocation* BestFitAllocator::AllocateImpl(size_t size) {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  auto highest_set_bit = static_cast<size_t>(HighestBitPos(size));
+  MapIt map_it;
+  for (; highest_set_bit < free_chunks_.size(); ++highest_set_bit) {
+    map_it = free_chunks_[highest_set_bit].lower_bound(size);
+    if (map_it != free_chunks_[highest_set_bit].end()) {
+      break;
+    }
+  }
+  if (UNLIKELY(highest_set_bit == free_chunks_.size())) {
+    PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+        "Cannot allocate %d, All fragments size is %d.", size, FreeSize()));
+  }
+  auto chunk_it = SplitChunk(size, highest_set_bit, map_it);
+  return new BestFitAllocation(this, chunk_it);
+}
+
+BestFitAllocation::BestFitAllocation(
+    paddle::memory::allocation::BestFitAllocator* allocator,
+    typename details::ChunkList::iterator chunk_it)
+    : Allocation(reinterpret_cast<void*>(
+                     reinterpret_cast<uintptr_t>(allocator->BasePtr()) +
+                     chunk_it->offset_),  // NOLINT
+                 chunk_it->size_,
+                 allocator->GetPlace()),
+      chunk_it_(chunk_it) {}
+}  // namespace paddle::memory::allocation

paddle/phi/core/memory/allocation/best_fit_allocator.h

@@ -0,0 +1,141 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include <stdint.h>
+
+#include <array>
+#include <list>
+#include <map>
+
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/allocation/spin_lock.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+namespace details {
+struct Chunk {
+  bool is_free{true};
+  // Offset to the base allocation.
+  uintptr_t offset_;
+  size_t size_;
+};
+
+// Here we use std::list to maintain chunk list.
+// NOTE(yy): The traditional implementation of ChunkList is add `prev`/`next`
+// pointers in `Chunk`, and split the allocation as `ChunkHeader` and
+// `Payload`. Such as
+//   *-------*---------------*---------------*--------------*
+//   | Chunk | prev_ pointer | next_ pointer | payload .... |
+//   *-------*---------------*---------------*--------------*
+// This implementation can just return a raw pointer, and we can get the list
+// structure by the raw pointer. However, we cannot use the same code on GPU
+// since CPU cannot access GPU memory directly.
+//
+// So we choose to use `std::list` and return an allocation instance, which
+// contains the list node iterator, then we can unify CPU/GPU code.
+//
+// To return an allocation is not a bad idea, since Tensor/Vector should holds
+// an allocation instead of raw pointer directly.
+using ChunkList = std::list<Chunk>;
+
+// Here we use a multi-level map of free chunks.
+// the map is
+//      MSB offset --> size --> [ChunkList::iterator]
+//
+// The time complexities:
+//     find a free chunk:
+//          O(logN),
+//               where N is the number of free nodes with the same MSB offset.
+//     find the position of a chunk iterator:
+//          O(logN + K),
+//               where N is the number of free nodes with the same MSB offset.
+//               where K is the number of free nodes with the same size.
+//     insert a free chunk:
+//          O(logN),
+//               where N is the number of free nodes with the same MSB offset.
+//     erase a free chunk:
+//          O(1)
+using FreeChunkBin =
+    std::array<std::multimap<size_t, ChunkList::iterator>, sizeof(size_t) * 8>;
+}  // namespace details
+
+class BestFitAllocator;
+
+// The BestFitAllocation maintain the List Node iterator.
+class BestFitAllocation : public Allocation {
+ private:
+  using ListIt = typename details::ChunkList::iterator;
+
+ public:
+  BestFitAllocation(BestFitAllocator* allocator, ListIt chunk_it);
+
+  const ListIt& ChunkIterator() const { return chunk_it_; }
+
+ private:
+  typename details::ChunkList::iterator chunk_it_;
+};
+
+// TODO(yy): Current BestFitAllocator is not thread-safe. To make it thread
+// safe, we must wrap a locked_allocator. However, we can implement a thread
+// safe allocator by locking each bin and chunks list independently. It will
+// make BestFitAllocator faster in multi-thread situation.
+//
+// This allocator implements a best-fit allocator with merging the free nodes.
+//
+// To allocate a buffer, it will find the best-fit chunk. If the best-fit chunk
+// is larger than request size, the original block will be split into two
+// chunks. The first block will be used and the second block will be put into
+// free chunks.
+//
+// To free an allocation, it will set the chunk of allocation to free and merge
+// the prev-chunk and the next-chunk when possible.
+class PADDLE_API BestFitAllocator : public Allocator {
+ public:
+  explicit BestFitAllocator(phi::Allocation* allocation);
+
+  void* BasePtr() const { return allocation_->ptr(); }
+
+  const Place& GetPlace() const { return allocation_->place(); }
+
+  size_t NumFreeChunks() const;
+
+  bool IsAllocThreadSafe() const override { return true; }
+
+ private:
+  size_t FreeSize() const;
+  using MapIt = typename details::FreeChunkBin::value_type::iterator;
+  using ListIt = typename details::ChunkList::iterator;
+
+  ListIt SplitChunk(size_t request_size,
+                    size_t free_chunk_offset,
+                    MapIt bin_iterator);
+  void EraseFreeNode(const ListIt& it);
+  void InsertFreeNode(const ListIt& it);
+
+ protected:
+  void FreeImpl(phi::Allocation* allocation) override;
+  phi::Allocation* AllocateImpl(size_t size) override;
+
+ private:
+  phi::Allocation* allocation_;  // not owned
+  details::ChunkList chunks_;
+  details::FreeChunkBin free_chunks_;
+  SpinLock spinlock_;
+};
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/buddy_allocator.cc

@@ -0,0 +1,383 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/phi/core/memory/allocation/buddy_allocator.h"
+
+#include <algorithm>
+
+#include "glog/logging.h"
+#include "paddle/common/flags.h"
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#define USE_DEVICE
+COMMON_DECLARE_uint64(reallocate_gpu_memory_in_mb);
+#endif
+
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/platform/device/device_wrapper.h"
+
+namespace paddle::memory::detail {
+
+BuddyAllocator::BuddyAllocator(
+    std::unique_ptr<SystemAllocator> system_allocator,
+    size_t min_chunk_size,
+    size_t max_chunk_size,
+    size_t extra_padding_size,
+    const std::string dev_type)
+    : min_chunk_size_(min_chunk_size),
+      max_chunk_size_(max_chunk_size),
+      extra_padding_size_(extra_padding_size),
+      cache_(system_allocator->UseGpu()),
+      system_allocator_(std::move(system_allocator)) {
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  if (!dev_type.empty()) {
+    init_allocate_size_func_ = [dev_type]() {
+      return phi::DeviceManager::GetInitAllocSize(
+          phi::PlaceHelper::CreatePlace(dev_type));
+    };
+    re_allocate_size_func_ = [dev_type]() {
+      return phi::DeviceManager::GetReallocSize(
+          phi::PlaceHelper::CreatePlace(dev_type));
+    };
+    use_custom_device_ = true;
+  } else {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+    init_allocate_size_func_ = &platform::GpuInitAllocSize;
+    re_allocate_size_func_ = &platform::GpuReallocSize;
+#endif
+  }
+#endif
+
+  VLOG(1) << "min_chunk_size_: " << min_chunk_size_
+          << ", max_chunk_size_:" << max_chunk_size_
+          << ", extra_padding_size_: " << extra_padding_size_;
+}
+
+BuddyAllocator::~BuddyAllocator() {
+  VLOG(10) << "BuddyAllocator destructor makes sure that all of these "
+              "have actually been freed";
+  while (!pool_.empty()) {
+    auto block = static_cast<MemoryBlock*>(std::get<2>(*pool_.begin()));
+    auto desc = cache_.LoadDesc(block);
+    VLOG(10) << "Free from block (" << block << ", " << desc->get_total_size()
+             << ")";
+
+    system_allocator_->Free(block, desc->get_total_size(), desc->get_index());
+    cache_.Invalidate(block);
+    pool_.erase(pool_.begin());
+  }
+}
+
+inline size_t align(size_t size, size_t alignment) {
+  size_t remaining = size % alignment;
+  return remaining == 0 ? size : size + (alignment - remaining);
+}
+
+void* BuddyAllocator::Alloc(size_t unaligned_size) {
+  // adjust allocation alignment
+  size_t size =
+      align(unaligned_size + sizeof(MemoryBlock::Desc) + extra_padding_size_,
+            min_chunk_size_);
+  VLOG(10) << "alloc: " << unaligned_size
+           << ", padding for desc: " << sizeof(MemoryBlock::Desc)
+           << ", extra padding: " << extra_padding_size_
+           << ", alignment: " << min_chunk_size_;
+  // acquire the allocator lock
+  std::lock_guard<std::mutex> lock(mutex_);
+
+  VLOG(10) << "Allocate " << unaligned_size << " bytes from chunk size "
+           << size;
+
+  // if the allocation is huge, send directly to the system allocator
+  if (size > max_chunk_size_) {
+    VLOG(10) << "Allocate from system allocator.";
+    return SystemAlloc(size);
+  }
+
+  // query and allocate from the existing chunk
+  auto it = FindExistChunk(size);
+
+  // refill the pool if failure
+  if (it == pool_.end()) {
+    it = RefillPool(size);
+    // if still failure, fail fatally
+    if (it == pool_.end()) {
+      return nullptr;
+    }
+  } else {
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+    if (use_custom_device_) {
+      VLOG(10) << "Allocation from existing memory block " << std::get<2>(*it)
+               << " at address " << std::get<2>(*it);
+    } else {
+      VLOG(10) << "Allocation from existing memory block " << std::get<2>(*it)
+               << " at address "
+               << reinterpret_cast<MemoryBlock*>(std::get<2>(*it))->Data();
+    }
+#else
+    VLOG(10) << "Allocation from existing memory block " << std::get<2>(*it)
+             << " at address "
+             << reinterpret_cast<MemoryBlock*>(std::get<2>(*it))->Data();
+#endif
+  }
+
+  total_used_ += size;
+  total_free_ -= size;
+
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  if (use_custom_device_) {
+    return SplitToAlloc(it, size);
+  }
+#endif
+  // split the allocation and return data for use
+  return reinterpret_cast<MemoryBlock*>(SplitToAlloc(it, size))->Data();
+}
+
+void BuddyAllocator::Free(void* p) {
+  // Point back to metadata
+  auto block = static_cast<MemoryBlock*>(p)->Metadata();
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  if (use_custom_device_) {
+    block = static_cast<MemoryBlock*>(p);
+  }
+#endif
+  // Acquire the allocator lock
+  std::lock_guard<std::mutex> lock(mutex_);
+
+  VLOG(10) << "Free from address " << block;
+
+  auto* desc = cache_.LoadDesc(block);
+  if (desc->get_type() == MemoryBlock::HUGE_CHUNK) {
+    VLOG(10) << "Free directly from system allocator";
+    system_allocator_->Free(block, desc->get_total_size(), desc->get_index());
+
+    // Invalidate GPU allocation from cache
+    cache_.Invalidate(block);
+
+    return;
+  }
+
+  block->MarkAsFree(&cache_);
+
+  total_used_ -= desc->get_total_size();
+  total_free_ += desc->get_total_size();
+
+  // Trying to merge the right buddy
+  MemoryBlock* right_buddy = block->GetRightBuddy(&cache_);
+  if (right_buddy) {
+    VLOG(10) << "Merging this block " << block << " with its right buddy "
+             << right_buddy;
+
+    auto rb_desc = cache_.LoadDesc(right_buddy);
+    if (rb_desc->get_type() == MemoryBlock::FREE_CHUNK) {
+      // Take away right buddy from pool
+      pool_.erase(IndexSizeAddress(
+          rb_desc->get_index(), rb_desc->get_total_size(), right_buddy));
+
+      // merge its right buddy to the block
+      block->Merge(&cache_, right_buddy);
+    }
+  }
+
+  // Trying to merge the left buddy
+  MemoryBlock* left_buddy = block->GetLeftBuddy(&cache_);
+  if (left_buddy) {
+    VLOG(10) << "Merging this block " << block << " with its left buddy "
+             << left_buddy;
+
+    // auto left_buddy = block->left_buddy(cache_);
+    auto* lb_desc = cache_.LoadDesc(left_buddy);
+    if (lb_desc->get_type() == MemoryBlock::FREE_CHUNK) {
+      // Take away right buddy from pool
+      pool_.erase(IndexSizeAddress(
+          lb_desc->get_index(), lb_desc->get_total_size(), left_buddy));
+
+      // merge the block to its left buddy
+      left_buddy->Merge(&cache_, block);
+      block = left_buddy;
+      desc = lb_desc;
+    }
+  }
+
+  // Dumping this block into pool
+  VLOG(10) << "Inserting free block (" << block << ", "
+           << desc->get_total_size() << ")";
+  pool_.insert(
+      IndexSizeAddress(desc->get_index(), desc->get_total_size(), block));
+}
+
+uint64_t BuddyAllocator::Release() {
+  std::lock_guard<std::mutex> lock(mutex_);
+  int num = 0;
+  uint64_t bytes = 0;
+  for (auto iter = pool_.begin(); iter != pool_.end();) {
+    auto remain_size = std::get<1>(*iter);
+    auto remain_ptr = std::get<2>(*iter);
+    auto found = chunks_.find({remain_size, remain_ptr});
+    if (found != chunks_.end()) {
+      size_t index = found->second;
+      ++num;
+      bytes += remain_size;
+      total_free_ -= remain_size;
+      auto block = static_cast<MemoryBlock*>(remain_ptr);
+      system_allocator_->Free(remain_ptr, remain_size, index);
+      cache_.Invalidate(block);
+      iter = pool_.erase(iter);
+    } else {
+      iter++;
+    }
+  }
+  VLOG(10) << "Release " << num << " chunks, Free " << bytes << " bytes.";
+  return bytes;
+}
+
+size_t BuddyAllocator::Used() { return total_used_; }
+size_t BuddyAllocator::GetMinChunkSize() { return min_chunk_size_; }
+size_t BuddyAllocator::GetMaxChunkSize() { return max_chunk_size_; }
+
+void* BuddyAllocator::SystemAlloc(size_t size) {
+  size_t index = 0;
+  void* p = system_allocator_->Alloc(&index, size);
+
+  VLOG(8) << "Allocated " << p << " size " << size << " from system allocator.";
+
+  if (p == nullptr) return nullptr;
+
+  static_cast<MemoryBlock*>(p)->Init(
+      &cache_, MemoryBlock::HUGE_CHUNK, index, size, nullptr, nullptr);
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  if (use_custom_device_) {
+    return p;
+  }
+#endif
+  return static_cast<MemoryBlock*>(p)->Data();
+}
+
+BuddyAllocator::PoolSet::iterator BuddyAllocator::RefillPool(
+    size_t request_bytes) {
+  size_t allocate_bytes = max_chunk_size_;  // NOLINT
+  size_t index = 0;
+
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  allocate_bytes = DeviceAllocateSize(
+      init_allocate_size_func_, re_allocate_size_func_, request_bytes);
+#else
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  allocate_bytes = DeviceAllocateSize(
+      &platform::GpuInitAllocSize, &platform::GpuReallocSize, request_bytes);
+#endif
+#endif
+
+  // Allocate a new block
+  void* p = system_allocator_->Alloc(&index, allocate_bytes);
+
+  if (p == nullptr) return pool_.end();
+
+  VLOG(8) << "Creating and inserting new block " << p << " size "
+          << allocate_bytes << " from system allocator";
+
+  static_cast<MemoryBlock*>(p)->Init(&cache_,
+                                     MemoryBlock::FREE_CHUNK,
+                                     index,
+                                     allocate_bytes,
+                                     nullptr,
+                                     nullptr);
+
+  total_free_ += allocate_bytes;
+
+  // record the chunk.
+  chunks_.insert({{allocate_bytes, p}, index});
+
+  // dump the block into pool
+  return pool_.insert(IndexSizeAddress(index, allocate_bytes, p)).first;
+}
+
+BuddyAllocator::PoolSet::iterator BuddyAllocator::FindExistChunk(size_t size) {
+  size_t index = 0;
+
+  while (true) {
+    auto it = pool_.lower_bound(IndexSizeAddress(index, size, nullptr));
+
+    // no match chunk memory
+    if (it == pool_.end()) return it;
+
+    if (std::get<0>(*it) > index) {
+      // find suitable one
+      if (std::get<1>(*it) >= size) {
+        return it;
+      }
+      // update and continue
+      index = std::get<0>(*it);
+      continue;
+    }
+    return it;
+  }
+}
+
+void* BuddyAllocator::SplitToAlloc(BuddyAllocator::PoolSet::iterator it,
+                                   size_t size) {
+  auto block = static_cast<MemoryBlock*>(std::get<2>(*it));
+  auto desc = cache_.LoadDesc(block);
+  pool_.erase(it);
+
+  VLOG(10) << "Split block (" << block << ", " << desc->get_total_size()
+           << ") into";
+  block->Split(&cache_, size);
+
+  VLOG(10) << "Left block (" << block << ", " << desc->get_total_size() << ")";
+  desc->set_type(MemoryBlock::ARENA_CHUNK);
+
+  // the rest of memory if exist
+  MemoryBlock* right_buddy = block->GetRightBuddy(&cache_);
+  if (right_buddy) {
+    auto* rb_desc = cache_.LoadDesc(right_buddy);
+    if (rb_desc->get_type() == MemoryBlock::FREE_CHUNK) {
+      VLOG(10) << "Insert right block (" << right_buddy << ", "
+               << rb_desc->get_total_size() << ")";
+
+      pool_.insert(IndexSizeAddress(
+          rb_desc->get_index(), rb_desc->get_total_size(), right_buddy));
+    }
+  }
+
+  return block;
+}
+
+size_t BuddyAllocator::DeviceAllocateSize(
+    std::function<size_t()> init_allocate_size_func,
+    std::function<size_t()> re_allocate_size_func,
+    size_t request_bytes) {
+  size_t allocate_bytes = max_chunk_size_;
+#if defined(USE_DEVICE)
+  const bool use_gpu = system_allocator_->UseGpu();
+  VLOG(10) << "use_gpu " << use_gpu << ", total_used " << total_used_
+           << ", total_free " << total_free_;
+  if (use_gpu) {
+    if (total_used_ == 0 && total_free_ == 0) {
+      // Compute the allocation size for gpu for the first allocation.
+      allocate_bytes = std::max(init_allocate_size_func(), request_bytes);
+    } else {
+      // Compute the re-allocation size, we store the re-allocation size when
+      // user set FLAGS_reallocate_gpu_memory_in_mb to fix value.
+      if (realloc_size_ == 0 || FLAGS_reallocate_gpu_memory_in_mb == 0ul) {
+        realloc_size_ = re_allocate_size_func();
+      }
+      allocate_bytes = std::max(realloc_size_, request_bytes);
+    }
+  }
+#endif
+  return allocate_bytes;
+}
+
+}  // namespace paddle::memory::detail

paddle/phi/core/memory/allocation/buddy_allocator.h

@@ -0,0 +1,135 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#include <stdint.h>
+
+#include <functional>
+#include <map>
+#include <memory>
+#include <mutex>  // NOLINT
+#include <set>
+#include <tuple>
+#include <utility>
+#include <vector>
+
+#include "paddle/phi/backends/cpu/cpu_info.h"
+#include "paddle/phi/core/memory/allocation/memory_block.h"
+#include "paddle/phi/core/memory/allocation/system_allocator.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+
+namespace paddle {
+namespace memory {
+namespace detail {
+
+class BuddyAllocator {
+ public:
+  BuddyAllocator(std::unique_ptr<SystemAllocator> system_allocator,
+                 size_t min_chunk_size,
+                 size_t max_chunk_size,
+                 size_t extra_padding_size = 0,
+                 const std::string dev_type = "");
+
+  ~BuddyAllocator();
+
+ public:
+  void* Alloc(size_t unaligned_size);
+  void Free(void* ptr);
+  // Release the unused memory pool, a real free operation for the OS.
+  uint64_t Release();
+  size_t Used();
+  size_t GetMinChunkSize();
+  size_t GetMaxChunkSize();
+
+ public:
+  // Disable copy and assignment
+  BuddyAllocator(const BuddyAllocator&) = delete;
+  BuddyAllocator& operator=(const BuddyAllocator&) = delete;
+
+ private:
+  // Tuple (allocator index, memory size, memory address)
+  using IndexSizeAddress = std::tuple<size_t, size_t, void*>;
+  // Each element in PoolSet is a free allocation
+  using PoolSet = std::set<IndexSizeAddress>;
+  // Each element in PoolMap is an allocation record
+  // key: <size, ptr>, value: index
+  using PoolMap = std::map<std::pair<size_t, void*>, size_t>;
+
+  /*! \brief Allocate fixed-size memory from system */
+  void* SystemAlloc(size_t size);
+
+  /*! \brief If existing chunks are not suitable, refill pool */
+  PoolSet::iterator RefillPool(size_t request_bytes);
+
+  /**
+   *  \brief   Find the suitable chunk from existing pool and split
+   *           it to left and right buddies
+   *
+   *  \param   it     the iterator of pool list
+   *  \param   size   the size of allocation
+   *
+   *  \return  the left buddy address
+   */
+  void* SplitToAlloc(PoolSet::iterator it, size_t size);
+
+  /*! \brief Find the existing chunk which used to allocation */
+  PoolSet::iterator FindExistChunk(size_t size);
+
+  /*! \brief Allocate bytes from the device */
+  size_t DeviceAllocateSize(std::function<size_t()> init_allocate_size_func,
+                            std::function<size_t()> re_allocate_size_func,
+                            size_t request_bytes);
+
+ private:
+  size_t total_used_ = 0;  // the total size of used memory
+  size_t total_free_ = 0;  // the total size of free memory
+
+  size_t min_chunk_size_;  // the minimum size of each chunk
+  size_t max_chunk_size_;  // the maximum size of each chunk
+
+  size_t realloc_size_ = 0;        // the size of re-allocated chunk
+  size_t extra_padding_size_ = 0;  // the size of padding to the size of memory
+                                   // to alloc, especially used in NPU
+
+ private:
+  /**
+   * \brief A list of free allocation
+   *
+   * \note  Only store free chunk memory in pool
+   */
+  PoolSet pool_;
+
+  /**
+   * \brief Record the allocated chunks when Refill pool.
+   */
+  PoolMap chunks_;
+
+ private:
+  /*! Unify the metadata format between GPU and CPU allocations */
+  MetadataCache cache_;
+
+ private:
+  /*! Allocate CPU/GPU memory from system */
+  std::unique_ptr<SystemAllocator> system_allocator_;
+  std::mutex mutex_;
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  std::function<size_t()> init_allocate_size_func_, re_allocate_size_func_;
+  bool use_custom_device_ = false;
+#endif
+};
+
+}  // namespace detail
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/buffered_allocator.cc

@@ -0,0 +1,76 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/buffered_allocator.h"
+
+#include "paddle/common/macros.h"
+
+REGISTER_FILE_SYMBOLS(buffered_allocator);
+
+namespace paddle::memory::allocation {
+
+BufferedAllocator::BufferedAllocator(std::shared_ptr<Allocator> allocator)
+    : underlying_allocator_(std::move(allocator)) {
+  PADDLE_ENFORCE_NOT_NULL(
+      underlying_allocator_,
+      common::errors::InvalidArgument(
+          "Underlying allocator of BufferedAllocator is NULL"));
+  if (underlying_allocator_->IsAllocThreadSafe()) {
+    mtx_ = std::make_unique<std::mutex>();
+  }
+}
+
+BufferedAllocator::~BufferedAllocator() { FreeCache(-1UL); }
+
+void BufferedAllocator::FreeCache(size_t size) {
+  platform::LockGuardPtr<std::mutex> guard(mtx_);
+  if (UNLIKELY(size == 0)) return;
+  size_t cur = 0;
+  while (!allocations_.empty()) {  // free the largest
+    auto it = --allocations_.end();
+    cur += it->second->size();
+    underlying_allocator_->Free(it->second.release());
+    allocations_.erase(it);
+    if (cur >= size) return;
+  }
+}
+
+bool BufferedAllocator::IsAllocThreadSafe() const { return mtx_ != nullptr; }
+
+void BufferedAllocator::FreeImpl(phi::Allocation *allocation) {
+  platform::LockGuardPtr<std::mutex> guard(mtx_);
+  allocations_.emplace(allocation->size(),
+                       AllocationPtr(allocation, Allocator::AllocationDeleter));
+}
+
+phi::Allocation *BufferedAllocator::AllocateImpl(size_t size) {
+  {
+    platform::LockGuardPtr<std::mutex> guard(mtx_);
+    auto it = allocations_.lower_bound(size);
+    if (it != allocations_.end() && it->first < size * 2) {
+      AllocationPtr result(std::move(it->second));
+      allocations_.erase(it);
+      return result.release();
+    }
+  }
+
+  try {
+    return underlying_allocator_->Allocate(size).release();
+  } catch (BadAlloc &) {
+    FreeCache(size);
+    return underlying_allocator_->Allocate(size).release();
+  }
+}
+
+}  // namespace paddle::memory::allocation

paddle/phi/core/memory/allocation/buffered_allocator.h

@@ -0,0 +1,59 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <vector>
+
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/platform/lock_guard_ptr.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// NOTE(zjl): BufferedAllocator maintains a memory pool to accelerate
+// memory allocation and reuse memory.
+// BufferedAllocator provides the same thread-safety level as
+// underlying_allocator_
+class PADDLE_API BufferedAllocator : public Allocator {
+ public:
+  explicit BufferedAllocator(std::shared_ptr<Allocator> allocator);
+
+  ~BufferedAllocator();
+
+  bool IsAllocThreadSafe() const override;
+
+  // only used in unittest
+  inline void ClearCache() { FreeCache(-1UL); }
+
+ private:
+  void FreeCache(size_t size);
+
+ protected:
+  void FreeImpl(phi::Allocation *allocation) override;
+  phi::Allocation *AllocateImpl(size_t size) override;
+
+ private:
+  std::shared_ptr<Allocator> underlying_allocator_;
+  std::multimap<size_t, AllocationPtr> allocations_;
+  std::unique_ptr<std::mutex> mtx_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/cpu_allocator.cc

@@ -0,0 +1,53 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/cpu_allocator.h"
+
+#include <cstdlib>
+
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/memory/stats.h"
+
+namespace paddle::memory::allocation {
+
+bool CPUAllocator::IsAllocThreadSafe() const { return true; }
+
+void CPUAllocator::FreeImpl(phi::Allocation *allocation) {
+  auto size = allocation->size();
+  void *p = allocation->ptr();
+#ifdef _WIN32
+  _aligned_free(p);
+#else
+  free(p);  // NOLINT
+#endif
+  HOST_MEMORY_STAT_UPDATE(Reserved, 0, -size);
+  delete allocation;
+}
+
+phi::Allocation *CPUAllocator::AllocateImpl(size_t size) {
+  void *p = nullptr;
+#ifdef _WIN32
+  p = _aligned_malloc(size, kAlignment);
+#else
+  int error = posix_memalign(&p, kAlignment, size);
+  PADDLE_ENFORCE_EQ(
+      error,
+      0,
+      common::errors::ResourceExhausted(
+          "Fail to alloc memory of %ld size, error code is %d.", size, error));
+#endif
+  HOST_MEMORY_STAT_UPDATE(Reserved, 0, size);
+  return new Allocation(p, size, CPUPlace());
+}
+}  // namespace paddle::memory::allocation

paddle/phi/core/memory/allocation/cpu_allocator.h

@@ -0,0 +1,45 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include "paddle/phi/core/memory/allocation/allocator.h"
+
+#ifdef _WIN32
+#define posix_memalign_free _aligned_free
+#define posix_memalign(p, a, s) \
+  (((*(p)) = _aligned_malloc((s), (a))), *(p) ? 0 : errno)
+#endif
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+// CPU system allocator and allocation.
+//
+// NOTE(yy): Should we just use `malloc` here since there is an
+// aligned_allocator.
+//
+// NOTE(yy): It is no need to use `BestFitAllocator` in CPU. We can import
+// an open-sourced allocator into Paddle.
+class PADDLE_API CPUAllocator : public Allocator {
+ public:
+  constexpr static size_t kAlignment = 4096UL;
+  bool IsAllocThreadSafe() const override;
+
+ protected:
+  void FreeImpl(phi::Allocation* allocation) override;
+  phi::Allocation* AllocateImpl(size_t size) override;
+};
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/cuda_allocator.cc

@@ -0,0 +1,94 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/cuda_allocator.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/memory/stats.h"
+
+#ifdef PADDLE_WITH_CUDA
+#include <cuda.h>
+#include <cuda_runtime.h>
+#endif
+
+#ifdef PADDLE_WITH_HIP
+#include <hip/hip_runtime.h>
+#endif
+
+#include <string>
+
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/platform/cuda_device_guard.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+
+namespace paddle::memory::allocation {
+bool CUDAAllocator::IsAllocThreadSafe() const { return true; }
+void CUDAAllocator::FreeImpl(phi::Allocation* allocation) {
+  PADDLE_ENFORCE_EQ(
+      allocation->place(),
+      place_,
+      common::errors::PermissionDenied(
+          "GPU memory is freed in incorrect device. This may be a bug"));
+  platform::RecordedGpuFree(
+      allocation->ptr(), allocation->size(), place_.device);
+  delete allocation;
+}
+
+phi::Allocation* CUDAAllocator::AllocateImpl(size_t size) {
+  std::call_once(once_flag_, [this] { platform::SetDeviceId(place_.device); });
+
+  void* ptr;
+  auto result = platform::RecordedGpuMalloc(&ptr, size, place_.device);
+  if (LIKELY(result == gpuSuccess)) {
+    return new Allocation(ptr, size, Place(place_));
+  }
+
+  size_t avail, total, actual_avail, actual_total;
+  bool is_limited = platform::RecordedGpuMemGetInfo(
+      &avail, &total, &actual_avail, &actual_total, place_.device);
+  size_t allocated = total - avail;
+
+  std::string err_msg;
+  if (is_limited) {
+    auto limit_size = (total >> 20);
+    err_msg = string::Sprintf(
+        "Or set environment variable `FLAGS_gpu_memory_limit_mb` to a larger "
+        "value. Currently `FLAGS_gpu_memory_limit_mb` is %d, so the maximum "
+        "GPU memory usage is limited to %d MB.\n"
+        "   The command is `export FLAGS_gpu_memory_limit_mb=xxx`.",
+        limit_size,
+        limit_size);
+  }
+
+  size_t actual_allocated_memory =
+      paddle::memory::DeviceMemoryStatCurrentValue("Allocated", place_.device);
+
+  PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+      "\n\nOut of memory error on GPU %d. "
+      "Cannot allocate %s memory on GPU %d, %s memory has been "
+      "allocated(actual using allocated memory %s) and "
+      "available memory is only %s.\n\n"
+      "Please check whether there is any other process using GPU %d.\n"
+      "1. If yes, please stop them, or start PaddlePaddle on another GPU.\n"
+      "2. If no, please decrease the batch size of your model. %s\n",
+      place_.device,
+      string::HumanReadableSize(size),
+      place_.device,
+      string::HumanReadableSize(allocated),
+      string::HumanReadableSize(actual_allocated_memory),
+      string::HumanReadableSize(avail),
+      place_.device,
+      err_msg));
+}
+
+}  // namespace paddle::memory::allocation

paddle/phi/core/memory/allocation/cuda_allocator.h

@@ -0,0 +1,42 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include <mutex>  // NOLINT
+
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class PADDLE_API CUDAAllocator : public Allocator {
+ public:
+  explicit CUDAAllocator(const GPUPlace& place) : place_(place) {}
+
+  bool IsAllocThreadSafe() const override;
+
+ protected:
+  void FreeImpl(phi::Allocation* allocation) override;
+  phi::Allocation* AllocateImpl(size_t size) override;
+
+ private:
+  GPUPlace place_;
+  std::once_flag once_flag_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/cuda_device_context_allocator.h

@@ -0,0 +1,175 @@
+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <map>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "glog/logging.h"
+
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/platform/cuda_device_guard.h"
+#include "paddle/phi/core/platform/device_context.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+/**
+ * GPUContextAllocation is a wrapper of the underbeneath allocation.
+ * GPUContextAllocation adds a CUDA stream callback for the underbeneath
+ * allocation so that GPUContextAllocation can be used in a CUDA stream
+ * which deletes allocation in the callback.
+ */
+class GPUContextAllocation : public Allocation {
+ public:
+  explicit GPUContextAllocation(DecoratedAllocationPtr allocation)
+      : Allocation(allocation->ptr(),
+                   allocation->base_ptr(),
+                   allocation->size(),
+                   allocation->place()),
+        underlying_allocation_(std::move(allocation)) {}
+
+  ~GPUContextAllocation() {
+    PADDLE_WARN_NOT_NULL(
+        dev_ctx_,
+        common::errors::PreconditionNotMet(
+            "Device context is not set for GPUContextAllocation"));
+
+    auto *p_allocation = underlying_allocation_.release();
+    VLOG(4) << "Adding callback to delete GPUContextAllocation at "
+            << p_allocation;
+    dev_ctx_->AddStreamCallback([p_allocation] {
+      VLOG(4) << "Delete GPUContextAllocation at " << p_allocation;
+      Allocator::AllocationDeleter(p_allocation);
+    });
+  }
+
+  void SetGPUContext(const phi::GPUContext *dev_ctx) { dev_ctx_ = dev_ctx; }
+
+ private:
+  DecoratedAllocationPtr underlying_allocation_;
+  const phi::GPUContext *dev_ctx_{nullptr};
+};
+
+/**
+ * GPUContextAllocator will allocate a GPUContextAllocation
+ * after waiting for a self-created event on the default stream. It does so to
+ * let the non-default stream be able to allocate GPU memory which will be
+ * released by stream callback
+ */
+class GPUContextAllocator : public Allocator {
+ public:
+  explicit GPUContextAllocator(GPUPlace place, gpuStream_t default_stream)
+      : place_(place), default_stream_(default_stream) {
+    platform::CUDADeviceGuard guard(place_.device);
+#ifdef PADDLE_WITH_HIP
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        hipEventCreateWithFlags(&event_, hipEventDisableTiming));
+#else
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        cudaEventCreate(&event_, cudaEventDisableTiming));
+#endif
+  }
+
+  ~GPUContextAllocator() {
+    if (event_) {
+      platform::CUDADeviceGuard guard(place_.device);
+#ifdef PADDLE_WITH_HIP
+
+      PADDLE_WARN_GPU_SUCCESS(hipEventDestroy(event_));
+#else
+      PADDLE_WARN_GPU_SUCCESS(cudaEventDestroy(event_));
+#endif
+    }
+  }
+
+ protected:
+  phi::Allocation *AllocateImpl(size_t size) override {
+    PADDLE_ENFORCE_NOT_NULL(
+        default_stream_,
+        common::errors::PreconditionNotMet(
+            "Default stream is not set for GPUContextAllocator"));
+    platform::CUDADeviceGuard guard(place_.device);
+    auto allocation = new GPUContextAllocation(
+        static_unique_ptr_cast<Allocation>(memory::Alloc(place_, size)));
+// Wait for the event on stream
+#ifdef PADDLE_WITH_HIP
+    PADDLE_ENFORCE_GPU_SUCCESS(hipEventRecord(event_, default_stream_));
+    PADDLE_ENFORCE_GPU_SUCCESS(hipStreamWaitEvent(default_stream_, event_, 0));
+#else
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaEventRecord(event_, default_stream_));
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamWaitEvent(default_stream_, event_, 0));
+#endif
+    return allocation;
+  }
+
+  void FreeImpl(phi::Allocation *allocation) override { delete allocation; }
+
+ private:
+  GPUPlace place_;
+  gpuEvent_t event_{nullptr};
+  gpuStream_t default_stream_{nullptr};
+};
+
+/**
+ * GPUContextAllocatorPool is a singleton stores mapping from
+ * CUDAPlace(s) to std::shared_ptr<GPUContextAllocator>. When a
+ * phi::GPUContext's compute stream isn't default stream, it can call this
+ * class to allocate GPU memory which will be released by a callback after
+ * stream execution.
+ */
+class GPUContextAllocatorPool {
+ public:
+  static GPUContextAllocatorPool &Instance() {
+    static GPUContextAllocatorPool pool;
+    return pool;
+  }
+
+  AllocationPtr Alloc(const phi::GPUContext &dev_ctx, size_t size) {
+    auto iter = allocators_.find(GPUPlace(dev_ctx.GetPlace().GetDeviceId()));
+    PADDLE_ENFORCE_NE(
+        iter,
+        allocators_.end(),
+        common::errors::NotFound("No allocator found for CUDAPlace."));
+    auto &allocator = iter->second;
+    AllocationPtr allocation = allocator->Allocate(size);
+    static_cast<GPUContextAllocation *>(allocation.get())
+        ->SetGPUContext(&dev_ctx);
+    return allocation;
+  }
+
+ private:
+  GPUContextAllocatorPool() {
+    std::vector<int> devices = platform::GetSelectedDevices();
+    for (int i : devices) {
+      auto place = GPUPlace(i);
+      auto compute_stream =
+          phi::DeviceContextPool::Instance().GetByPlace(place)->stream();
+      auto allocator = std::shared_ptr<GPUContextAllocator>(
+          new GPUContextAllocator(place, compute_stream));
+      allocators_.insert(make_pair(place, allocator));
+    }
+  }
+
+  std::map<GPUPlace, std::shared_ptr<GPUContextAllocator>> allocators_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/cuda_ipc_allocator.cc

@@ -0,0 +1,97 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef _WIN32
+
+#include "paddle/phi/core/memory/allocation/cuda_ipc_allocator.h"
+
+#include <fcntl.h>
+#include <sys/mman.h>
+#include <cstdlib>
+
+#include <random>
+#include <string>
+
+#include "glog/logging.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/platform/cuda_device_guard.h"
+
+namespace paddle::memory::allocation {
+
+namespace {
+std::mutex ipc_mutex_;
+std::unordered_map<std::string, std::weak_ptr<void>> ipc_handle_to_baseptr_;
+}  // namespace
+
+std::shared_ptr<void> GetIpcBasePtr(std::string handle) {
+  std::lock_guard<std::mutex> lock(ipc_mutex_);
+
+  auto iter = ipc_handle_to_baseptr_.find(handle);
+  if (iter != ipc_handle_to_baseptr_.end()) {
+    auto baseptr = iter->second.lock();
+    if (baseptr) return baseptr;
+  }
+  // The IpcMemHandle can only open once for the same handle,
+  // so here we cache it here.
+  void *baseptr = nullptr;
+  auto ipc_handle = reinterpret_cast<const gpuIpcMemHandle_t *>(handle.c_str());
+  PADDLE_ENFORCE_GPU_SUCCESS(gpuIpcOpenMemHandle(
+      &baseptr, *ipc_handle, gpuIpcMemLazyEnablePeerAccess));
+  // Close ipc handle on the same device.
+  int device_id = platform::GetCurrentDeviceId();
+  // Add deleter to close ipc handle.
+  auto sp = std::shared_ptr<void>(baseptr, [handle, device_id](void *ptr) {
+    platform::CUDADeviceGuard guard(device_id);
+    std::lock_guard<std::mutex> lock(ipc_mutex_);
+    PADDLE_ENFORCE_GPU_SUCCESS(gpuIpcCloseMemHandle(ptr));
+    ipc_handle_to_baseptr_.erase(handle);
+    VLOG(6) << "cudaIpcCloseMemHandle for ptr:"
+            << "\t" << ptr;
+  });
+  std::weak_ptr<void> wp = sp;
+  ipc_handle_to_baseptr_.insert(iter, {handle, wp});
+
+  return sp;
+}
+
+void IpcCollect() {
+  std::lock_guard<std::mutex> lock(ipc_mutex_);
+  size_t before = ipc_handle_to_baseptr_.size();
+  VLOG(6) << "The number of IPC handles before collection:" << before;
+
+  for (auto it = ipc_handle_to_baseptr_.begin();
+       it != ipc_handle_to_baseptr_.end();) {
+    if (it->second.expired()) {
+      it = ipc_handle_to_baseptr_.erase(it);
+    } else {
+      VLOG(6) << " Valid ipc handle is not expired";
+      ++it;
+    }
+  }
+
+  size_t after = ipc_handle_to_baseptr_.size();
+  size_t collected = before - after;
+  VLOG(1) << "IpcCollect: collected " << collected << " expired IPC handles"
+          << "out of " << before << " total handles";
+}
+
+CudaIpcAllocation::~CudaIpcAllocation() {
+  shared_ptr_.reset();
+  VLOG(6) << "tensor deleted cudaIpcCloseMemHandle for ptr:"
+          << "\t" << this->ptr();
+}
+
+}  // namespace paddle::memory::allocation
+
+#endif

paddle/phi/core/memory/allocation/cuda_ipc_allocator.h

@@ -0,0 +1,60 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef _WIN32
+#pragma once
+
+#include <memory>
+#include <mutex>  // NOLINT
+#include <string>
+#include <unordered_set>
+#include <utility>
+
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/platform/cuda_device_guard.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+std::shared_ptr<void> GetIpcBasePtr(std::string handle);
+
+void IpcCollect();
+
+class CudaIpcAllocation : public Allocation {
+ public:
+  explicit CudaIpcAllocation(void *ptr,
+                             size_t size,
+                             int device_id,
+                             std::shared_ptr<void> shared_ptr)
+      : Allocation(ptr, size, GPUPlace(device_id)),
+        device_id_(std::move(device_id)),
+        shared_ptr_(std::move(shared_ptr)) {}
+
+  inline const int &device_id() const { return device_id_; }
+
+  ~CudaIpcAllocation() override;
+
+ private:
+  int device_id_;
+  std::shared_ptr<void> shared_ptr_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle
+
+#endif

paddle/phi/core/memory/allocation/cuda_malloc_async_allocator.cc

@@ -0,0 +1,399 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/cuda_malloc_async_allocator.h"
+#include <cstddef>
+#include <cstdint>
+#include <mutex>
+#include "paddle/common/flags.h"
+#include "paddle/common/macros.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/allocation/stream_safe_cuda_allocator.h"
+
+#include "glog/logging.h"
+
+#ifdef PADDLE_WITH_CUDA
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include "paddle/phi/backends/gpu/cuda/cuda_graph.h"
+#endif
+
+#include <string>
+
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/platform/cuda_device_guard.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+
+#include "paddle/utils/optional.h"
+
+#ifdef PADDLE_WITH_CUDA
+
+/*
+ * Note: [cuda_malloc_async_pool_memory_throttle_ratio]
+ * The primary purpose of the memory_throttle_ratio is to provide a
+ * threshold that determines when to initiate synchronization operations to
+ * deallocate memory. This mechanism helps in ensuring that the system does
+ * not exceed its memory capacity while also attempting to minimize performance
+ * degradation caused by frequent memory synchronization.
+ *
+ * ```
+ *   utilization = (allocated_size + pending_release_size) / total_memory_size
+ *   if(utilization > memory_throttle_ratio)
+ *      sync(free_stream, malloc_stream)
+ * ```
+ *
+ * When the utilization exceeds the memory_throttle_ratio, we
+ * initiate a stream synchronization operation before malloc.
+ *
+ * During synchronization, all memory deallocation requests in the free queue
+ * are processed, effectively lowering the memory utilization before
+ * any new memory allocation operations are going to proceed.
+ *
+ * [Impact on Performance and Memory Usage]
+ *
+ * - Lower memory_throttle_ratio Values
+ * the synchronization operation will be triggered more frequently.
+ * This can lead to better memory utilization but might result in decreased
+ * performance due to the increased number of synchronization operations.
+ *
+ * - Higher memory_throttle_ratio Values
+ * Conversely, setting a higher value allows for more memory to be allocated
+ * before triggering synchronization, which can enhance performance by reducing
+ * the number of sync operations. However, this increases the risk of reaching
+ * an OOM condition since more memory can be allocated without
+ * immediate deallocation.
+ */
+COMMON_DECLARE_double(cuda_malloc_async_pool_memory_throttle_ratio);
+
+namespace paddle::memory::allocation {
+
+thread_local std::once_flag CUDAMallocAsyncAllocation::once_flag_;
+
+inline void sync_streams(gpuStream_t to_record, gpuStream_t to_wait) {
+  cudaEvent_t event = nullptr;
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaEventCreateWithFlags(&event, cudaEventDisableTiming));
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaEventRecord(event, to_record));
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamWaitEvent(to_wait, event));
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaEventDestroy(event));
+}
+
+// CUDAMallocAsyncAllocation
+
+bool CUDAMallocAsyncAllocation::RecordStream(gpuStream_t stream) {
+  std::call_once(once_flag_,
+                 [this] { phi::backends::gpu::SetDeviceId(place_.device); });
+  std::lock_guard<SpinLock> lock_guard(recorded_streams_lock_);
+  if (malloc_stream_ == stream) {
+    // Called record_stream on tensor whose original malloc_stream matches the
+    // recorded stream. This should have no effect.
+    return false;
+  }
+  recorded_streams_.insert(stream);
+  return true;
+}
+
+void CUDAMallocAsyncAllocation::EraseStream(gpuStream_t stream) {
+  std::lock_guard<SpinLock> lock_guard(recorded_streams_lock_);
+  recorded_streams_.erase(stream);
+}
+
+size_t CUDAMallocAsyncAllocation::Free() {
+  if (recorded_streams_.empty()) {
+    platform::RecordedGpuFreeAsync(
+        ptr(), size(), place_.device, malloc_stream_);
+
+    if (UNLIKELY(phi::backends::gpu::CUDAGraph::IsThisThreadCapturing())) {
+      phi::backends::gpu::CUDAGraph::AddJoiningStreamDuringCapturing(
+          malloc_stream_);
+    }
+    return size();
+  } else {
+    sync_streams(malloc_stream_, free_stream_);
+
+    for (const auto& recorded_stream : recorded_streams_) {
+      sync_streams(recorded_stream, free_stream_);
+    }
+
+    platform::RecordedGpuFreeAsync(ptr(), size(), place_.device, free_stream_);
+
+    if (UNLIKELY(phi::backends::gpu::CUDAGraph::IsThisThreadCapturing())) {
+      phi::backends::gpu::CUDAGraph::AddJoiningStreamDuringCapturing(
+          free_stream_);
+    }
+    return 0;
+  }
+}
+
+// CUDAMallocAsyncAllocator
+
+CUDAMallocAsyncAllocator::CUDAMallocAsyncAllocator(
+    std::shared_ptr<Allocator> underlying_allocator,
+    const GPUPlace& place,
+    gpuStream_t default_stream)
+    : underlying_allocator_(std::move(underlying_allocator)),
+      place_(place),
+      default_stream_(default_stream),
+      current_allocated_size_(0),
+      pending_release_size_(0),
+      memory_throttle_ratio_(
+          FLAGS_cuda_malloc_async_pool_memory_throttle_ratio) {
+  // CUDA operations are not allowed here. The cuInit function must be called
+  // after a new fork, and since this constructor is typically initialized
+  // before cuInit, we should avoid calling any CUDA API here.
+  phi::backends::gpu::CUDAGraph::AddPreCaptureCallback([&]() {
+    VLOG(0) << "[Before capture callback] " << (this) << " "
+            << std::this_thread::get_id();
+    this->ClearFreeStream(true);
+  });
+}
+
+uint64_t CUDAMallocAsyncAllocator::ReleaseImpl(const Place& place) {
+  if (UNLIKELY(phi::backends::gpu::CUDAGraph::IsThisThreadCapturing())) {
+    VLOG(7) << "Memory release forbidden in CUDA Graph Captruing";
+    return 0;
+  }
+
+  uint64_t released_size = 0;
+  // we synchronize the event so all the block could be release.
+  if (underlying_allocator_)
+    released_size += underlying_allocator_->Release(place_);
+  VLOG(8) << "Release " << released_size << " bytes memory from all streams";
+  return released_size;
+}
+
+void CUDAMallocAsyncAllocator::ClearFreeStream(bool sync) {
+  LazyInitializeCudaFreeStream();
+
+  if (sync) {
+    VLOG(0) << "[CUDAMallocAsyncAllocator] " << (this)
+            << " synchronize the free stream to ensure all unrelesed blocks "
+            << "are freed";
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaStreamSynchronize(free_stream_));
+  } else {
+    sync_streams(free_stream_, default_stream_);
+  }
+  current_allocated_size_ -= pending_release_size_;
+  pending_release_size_ = 0;
+}
+
+void CUDAMallocAsyncAllocator::MallocThrottling() {
+  if (UNLIKELY(phi::backends::gpu::CUDAGraph::IsThisThreadCapturing())) {
+    // we disable MallocThrottling when capturing
+    return;
+  }
+  double allocated =
+      static_cast<double>(current_allocated_size_ + pending_release_size_);
+  double utilization = allocated / static_cast<double>(max_size_);
+
+  if (utilization > memory_throttle_ratio_) {
+    VLOG(10) << "utilization_ratio " << utilization
+             << " current_allocated_size "
+             << string::HumanReadableSize(current_allocated_size_)
+             << " pending_release_size "
+             << string::HumanReadableSize(pending_release_size_);
+    CUDAMallocAsyncAllocator::ClearFreeStream();
+  }
+}
+
+void CUDAMallocAsyncAllocator::FreeAllocation(
+    CUDAMallocAsyncAllocation* allocation) {
+  auto current_released_size = allocation->Free();
+  current_allocated_size_ -= current_released_size;
+  // The amount of pending release size (the space that has been queued to
+  // free_stream, that are going to be freed in the future)
+  pending_release_size_ += (allocation->size() - current_released_size);
+}
+
+/*
+ * There are four distinct scenarios involving `cudaMalloc`, `cudaFree`, and
+ * `cudaGraph`:
+ *
+ * 1. When both `cudaMalloc` and `cudaFree` occur within a graph.
+ * 2. When `cudaMalloc` happens within a graph, but `cudaFree` occurs outside
+ * the graph.
+ * 3. When `cudaMalloc` takes place outside a graph, but `cudaFree` happens
+ * within a graph.
+ * 4. When both `cudaMalloc` and `cudaFree` are executed outside any graph.
+ *
+ * For cases (1.) and (4.), the usage aligns with the typical pattern of
+ * `cudaMalloc`/`cudaFree`.
+ *
+ * In case (1.), `FreeImpl` removes the allocation from
+ * `graph_owned_allocations_`, followed by `FreeAllocation`. In case (2.), the
+ * callback within `AllocateImpl` would free the allocation after the graph is
+ * destroyed. In case (3.), `FreeImpl` releases the allocation after the CUDA
+ * graph has completed its capture. Finally, in case (4.), `FreeImpl` would call
+ * `FreeAllocation`, and the allocation would be freed.
+ */
+
+void CUDAMallocAsyncAllocator::FreeImpl(phi::Allocation* phi_allocation) {
+  auto* allocation = dynamic_cast<CUDAMallocAsyncAllocation*>(phi_allocation);
+  std::lock_guard<SpinLock> lock_guard(graph_owned_allocations_lock_);
+
+  // During graph capturing, only free the memory blocks owned by the graph;
+  // others are cached.
+  if (UNLIKELY(phi::backends::gpu::CUDAGraph::IsThisThreadCapturing())) {
+    // Handles scenario (3.)
+    if (graph_owned_allocations_.find(allocation) ==
+        graph_owned_allocations_.end()) {
+      // If the block is not owned by the graph, cache it for release after
+      // capturing.
+      phi::backends::gpu::CUDAGraph::AddPostCaptureCallbackDuringCapturing(
+          [=]() {
+            // Release this block after capturing
+            VLOG(0) << "[PostCaptureCallback] Releasing ptr = "
+                    << allocation->ptr() << " size = "
+                    << string::HumanReadableSize(allocation->size());
+            FreeAllocation(allocation);
+          });
+
+      return;
+    }
+    // Handles scenario (1.)
+    graph_owned_allocations_.erase(allocation);
+  } else {
+    // Handles scenario (2.)
+    if (graph_owned_allocations_.find(allocation) !=
+        graph_owned_allocations_.end()) {
+      auto graph = graph_owned_allocations_[allocation];
+      VLOG(0) << "[Rescheduled cudaFreeAsync] Allocation ptr = "
+              << allocation->ptr()
+              << " is allocated in a graph but freed outside the graph."
+              << " The allocation is rescheduled to be freed after the "
+              << "destruction of graph " << graph;
+      graph_owned_allocations_.erase(allocation);
+
+      // No need to free the allocation
+      return;
+    }
+  }
+
+  // Handles scenario (1.) and (4.)
+  FreeAllocation(allocation);
+}
+
+void CUDAMallocAsyncAllocator::LazyInitializeCudaFreeStream() {
+  std::call_once(once_flag_, [this] {
+    size_t avail, total, actual_avail, actual_total;
+    platform::RecordedGpuMemGetInfo(
+        &avail, &total, &actual_avail, &actual_total, place_.device);
+    max_size_ = total;
+
+    VLOG(0) << "[CUDAMallocAsyncAllocator] " << (this) << " place " << place_
+            << " max_size " << string::HumanReadableSize(max_size_)
+            << " memory_throttle_ratio " << memory_throttle_ratio_
+            << " tid = " << std::this_thread::get_id();
+
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        cudaStreamCreateWithPriority(&free_stream_, cudaStreamNonBlocking, 0));
+    cudaDeviceGetDefaultMemPool(&mempool_, place_.device);
+
+    platform::SetDeviceId(place_.device);
+  });
+}
+
+phi::Allocation* CUDAMallocAsyncAllocator::AllocateImpl(size_t size) {
+  LazyInitializeCudaFreeStream();
+
+  MallocThrottling();
+
+  void* ptr;
+  auto result = platform::RecordedGpuMallocAsync(
+      &ptr, size, place_.device, default_stream_);
+  if (LIKELY(result == gpuSuccess)) {
+    auto* allocation = new CUDAMallocAsyncAllocation(
+        ptr, size, Place(place_), default_stream_, free_stream_);
+    VLOG(10) << "Allocate " << allocation->ptr() << " with allocator "
+             << (this);
+
+    // If capturing, associate allocation with the current graph.
+    if (UNLIKELY(phi::backends::gpu::CUDAGraph::IsThisThreadCapturing())) {
+      std::lock_guard<SpinLock> lock_guard(graph_owned_allocations_lock_);
+      auto capturing_graph = phi::backends::gpu::CUDAGraph::CapturingID();
+      graph_owned_allocations_[allocation] = capturing_graph;
+
+      // Handles scenario (2.)
+      phi::backends::gpu::CUDAGraph::AddPostResetCallbackDuringCapturing(
+          [=](paddle::optional<const phi::backends::gpu::CUDAGraph&> graph) {
+            std::lock_guard<SpinLock> lock_guard_free(
+                graph_owned_allocations_lock_);
+
+            // Returns if the allocation is freed during capture.
+            if (graph_owned_allocations_.find(allocation) ==
+                graph_owned_allocations_.end())
+              return;
+
+            bool replayed = graph.get().IsReplayed();
+            if (replayed) {
+              VLOG(0) << "[Rescheduled cudaFreeAsync] Graph " << capturing_graph
+                      << " is destructed. Allocation = " << allocation->ptr()
+                      << " is freed.";
+              FreeAllocation(allocation);
+            } else {
+              VLOG(0) << "[Rescheduled cudaFreeAsync] Graph " << capturing_graph
+                      << " is destructed without any replay. Allocation = "
+                      << allocation->ptr()
+                      << " is not initialized and would not be freed.";
+            }
+          });
+    }
+    current_allocated_size_ += size;
+    return allocation;
+  }
+
+  size_t avail, total, actual_avail, actual_total;
+  bool is_limited = platform::RecordedGpuMemGetInfo(
+      &avail, &total, &actual_avail, &actual_total, place_.device);
+  size_t allocated = total - avail;
+
+  std::string err_msg;
+  if (is_limited) {
+    auto limit_size = (total >> 20);
+    err_msg = string::Sprintf(
+        "Or set environment variable `FLAGS_gpu_memory_limit_mb` to a larger "
+        "value. Currently `FLAGS_gpu_memory_limit_mb` is %d, so the maximum "
+        "GPU memory usage is limited to %d MB.\n"
+        "   The command is `export FLAGS_gpu_memory_limit_mb=xxx`.",
+        limit_size,
+        limit_size);
+  }
+
+  PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+      "\n\nOut of memory error on GPU %d. "
+      "Cannot allocate %s memory on GPU %d, %s memory has been allocated and "
+      "available memory is only %s.\n\n"
+      "Please check whether there is any other process using GPU %d.\n"
+      "1. If yes, please stop them, or start PaddlePaddle on another GPU.\n"
+      "2. If no, please decrease the batch size of your model. %s\n",
+      place_.device,
+      string::HumanReadableSize(size),
+      place_.device,
+      string::HumanReadableSize(allocated),
+      string::HumanReadableSize(avail),
+      place_.device,
+      err_msg));
+}
+
+gpuStream_t CUDAMallocAsyncAllocator::GetDefaultStream() const {
+  return default_stream_;
+}
+
+void CUDAMallocAsyncAllocator::SetDefaultStream(gpuStream_t stream) {
+  default_stream_ = stream;
+}
+
+}  // namespace paddle::memory::allocation
+
+#endif

paddle/phi/core/memory/allocation/cuda_malloc_async_allocator.h

@@ -0,0 +1,137 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include <mutex>  // NOLINT
+#include <unordered_set>
+
+#include "paddle/phi/backends/gpu/cuda/cuda_graph.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/allocation/spin_lock.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_types.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class CUDAMallocAsyncAllocation;
+class CUDAMallocAsyncAllocator;
+
+#ifdef PADDLE_WITH_CUDA
+
+// TODO(eee4017): It may be beneficial to introduce an abstract class named
+// `StreamAllocator` in future developments. This class would serve as a central
+// entity for methods specifically related to stream management, such as
+// `RecordStream` and `EraseStream`. The introduction of `StreamAllocator` would
+// enable both `StreamSafeCUDAAllocator` and `CUDAMallocAsyncAllocator` to
+// inherit directly from it,
+
+// The `CUDAMallocAsyncAllocation` class extends `Allocation` and is used for
+// managing memory allocations with CUDA async malloc. It includes methods to
+// handle stream associations and to query the owning stream of the allocation.
+class CUDAMallocAsyncAllocation : public Allocation {
+ public:
+  CUDAMallocAsyncAllocation(void* ptr,
+                            size_t size,
+                            Place place,
+                            gpuStream_t malloc_stream,
+                            gpuStream_t free_stream)
+      : Allocation(ptr, size, place),
+        malloc_stream_(malloc_stream),
+        free_stream_(free_stream) {}
+
+  gpuStream_t GetOwningStream() const { return malloc_stream_; }
+
+  bool RecordStream(gpuStream_t stream);
+  void EraseStream(gpuStream_t stream);
+  size_t Free();
+
+ private:
+  static thread_local std::once_flag once_flag_;
+  gpuStream_t malloc_stream_;
+  gpuStream_t free_stream_;
+
+  SpinLock recorded_streams_lock_;
+  std::unordered_set<gpuStream_t> recorded_streams_;
+};
+
+// The `CUDAMallocAsyncAllocator` class extends `Allocator` and is specialized
+// for asynchronous memory allocation in CUDA. It offers thread-safe allocation
+// and incorporates a default stream for memory operations.
+class CUDAMallocAsyncAllocator : public Allocator {
+ public:
+  explicit CUDAMallocAsyncAllocator(
+      std::shared_ptr<Allocator> underlying_allocator,
+      const GPUPlace& place,
+      gpuStream_t default_stream);
+
+  bool IsAllocThreadSafe() const override { return true; }
+  gpuStream_t GetDefaultStream() const;
+  void SetDefaultStream(gpuStream_t stream);
+  void ClearFreeStream(bool sync = false);
+
+ protected:
+  void FreeImpl(phi::Allocation* allocation) override;
+  phi::Allocation* AllocateImpl(size_t size) override;
+  uint64_t ReleaseImpl(const Place& place) override;
+
+ private:
+  void LazyInitializeCudaFreeStream();
+  void MallocThrottling();
+  void FreeAllocation(CUDAMallocAsyncAllocation* allocation);
+
+  std::shared_ptr<Allocator> underlying_allocator_;
+  GPUPlace place_;  // Specifies the CUDA device context.
+
+  cudaMemPool_t mempool_;
+  gpuStream_t default_stream_;  // Default stream for memory operations.
+
+  // we create a `free stream` for each allocator (each device should have a
+  // unique allocator) if an allocation is recorded on other stream than default
+  // stream, we release the allocation on `free stream`
+  gpuStream_t free_stream_;
+
+  size_t current_allocated_size_;
+  size_t pending_release_size_;
+  size_t max_size_;
+
+  double memory_throttle_ratio_;
+
+  std::once_flag once_flag_;
+
+  /*
+   * Life cycle management of graph_owned_allocations_:
+   *
+   * Each element within `graph_owned_allocations_` is initialized at
+   * `AllocateImpl`. However, there are two distinct ways of deconstruction.
+   *
+   * (A.) Deallocating occurs within `FreeImpl`.
+   * This implies that the allocation is initialized and disposed of during a
+   * graph capture, as in scenario (1.)
+   *
+   * (B.) Deallocation takes place in the callback after the graph is
+   * destructed. Meaning, the allocation is initialized during a graph capture
+   * but disposed of outside that context, as in scenario (2.)
+   */
+  std::unordered_map<CUDAMallocAsyncAllocation*, CUDAGraphID>
+      graph_owned_allocations_;
+  SpinLock graph_owned_allocations_lock_;
+};
+
+#endif
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/cuda_managed_allocator.cc

@@ -0,0 +1,95 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/cuda_managed_allocator.h"
+
+#ifdef PADDLE_WITH_CUDA
+#include <cuda.h>
+#include <cuda_runtime.h>
+#endif
+
+#ifdef PADDLE_WITH_HIP
+#include <hip/hip_runtime.h>
+#endif
+
+#include <string>
+
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/platform/cuda_device_guard.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+bool CUDAManagedAllocator::IsAllocThreadSafe() const { return true; }
+
+void CUDAManagedAllocator::FreeImpl(phi::Allocation* allocation) {
+  PADDLE_ENFORCE_EQ(
+      allocation->place(),
+      place_,
+      common::errors::PermissionDenied(
+          "GPU memory is freed in incorrect device. This may be a bug"));
+  platform::RecordedGpuFree(
+      allocation->ptr(), allocation->size(), place_.device);
+  delete allocation;
+}
+
+phi::Allocation* CUDAManagedAllocator::AllocateImpl(size_t size) {
+  std::call_once(once_flag_, [this] { platform::SetDeviceId(place_.device); });
+
+  int dev_id = place_.device;  // NOLINT
+  void* ptr;
+  auto result = platform::RecordedGpuMalloc(&ptr,
+                                            size,
+                                            dev_id,
+                                            /* malloc_managed_memory = */ true);
+  if (LIKELY(result == gpuSuccess)) {
+    return new Allocation(ptr, size, Place(place_));
+  }
+
+  uint64_t limit_size = platform::RecordedGpuLimitSize(dev_id);
+  uint64_t malloc_size = platform::RecordedGpuMallocSize(dev_id);
+  bool is_limited =
+      platform::IsGpuMallocRecorded(dev_id) && malloc_size + size > limit_size;
+
+  std::string err_msg;
+  if (UNLIKELY(is_limited)) {
+    int64_t limit_size_mb = limit_size >> 20;  // NOLINT
+    err_msg = string::Sprintf(
+        "Or set environment variable `FLAGS_gpu_memory_limit_mb` to a larger "
+        "value. Currently `FLAGS_gpu_memory_limit_mb` is %d, so the maximum "
+        "GPU memory usage is limited to %d MB.\n"
+        "   The command is `export FLAGS_gpu_memory_limit_mb=xxx`.",
+        limit_size_mb,
+        limit_size_mb);
+  }
+
+  PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+      "\n\nOut of memory error on GPU %d. "
+      "Cannot allocate %s CUDA managed memory on GPU %d, %s memory has been "
+      "allocated.\n\n"
+      "Please check whether there is any other process using GPU %d.\n"
+      "1. If yes, please stop them, or start PaddlePaddle on another GPU.\n"
+      "2. If no, please decrease the batch size of your model. %s\n\n",
+      dev_id,
+      string::HumanReadableSize(size),
+      dev_id,
+      string::HumanReadableSize(malloc_size),
+      dev_id,
+      err_msg));
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/cuda_managed_allocator.h

@@ -0,0 +1,40 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class CUDAManagedAllocator : public Allocator {
+ public:
+  explicit CUDAManagedAllocator(const GPUPlace& place) : place_(place) {}
+
+  bool IsAllocThreadSafe() const override;
+
+ protected:
+  void FreeImpl(phi::Allocation* allocation) override;
+  phi::Allocation* AllocateImpl(size_t size) override;
+
+ private:
+  GPUPlace place_;
+  std::once_flag once_flag_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/cuda_virtual_mem_allocator.cc

@@ -0,0 +1,265 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifdef PADDLE_WITH_CUDA
+#include <cuda.h>
+#include <cuda_runtime.h>
+#endif
+
+#include <string>
+
+#include "glog/logging.h"
+
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/memory/allocation/cuda_virtual_mem_allocator.h"
+#include "paddle/phi/core/memory/stats.h"
+
+#ifdef PADDLE_WITH_CUDA
+#include "paddle/phi/backends/dynload/cuda_driver.h"
+#include "paddle/phi/core/platform/cuda_device_guard.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+
+namespace paddle::memory::allocation {
+
+constexpr size_t kVirtualAddressSpaceSizeMultiplier = 2;
+
+std::mutex CUDAVirtualMemAllocator::base_ptr_handle_mu_;
+std::unordered_map<void*, CUmemGenericAllocationHandle>
+    CUDAVirtualMemAllocator::base_ptr_handle_map_;
+
+CUDAVirtualMemAllocator::CUDAVirtualMemAllocator(const GPUPlace& place)
+    : place_(place), virtual_mem_base_(0), prop_{} {
+  CUmemAllocationProp prop = {};
+
+  // Setup the properties common for all the chunks
+  // The allocations will be device pinned memory.
+  // This property structure describes the physical location where the memory
+  // will be allocated via cuMemCreate along with additional properties In this
+  // case, the allocation will be pinned device memory local to a given device.
+  prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
+  prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+  prop.location.id = place.device;  // NOLINT
+#if defined(_WIN32)
+  prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_NONE;
+#else
+  prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
+#endif
+  prop_ = prop;
+
+  // Prepare the access descriptor array indicating where and how the backings
+  // should be visible.
+  access_desc_.clear();
+  {
+    CUmemAccessDesc self = {};
+    self.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+    self.location.id = place.device;
+    self.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+    access_desc_.push_back(self);
+  }
+}
+
+void CUDAVirtualMemAllocator::InitOnce() {
+  std::call_once(init_flag_, [this] {
+    platform::SetDeviceId(place_.device);
+    paddle::platform::CUDADeviceGuard guard(place_.device);
+
+    // Get the minimum granularity.
+    size_t granularity = 0;
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cuMemGetAllocationGranularity(
+        &granularity, &prop_, CU_MEM_ALLOC_GRANULARITY_MINIMUM));
+    granularity_ = granularity;
+
+    // total size & VA size
+    size_t actual_avail, actual_total;
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaMemGetInfo(&actual_avail, &actual_total));
+    VLOG(1) << "VMM InitOnce dev " << place_.device << " actual_avail: "
+            << static_cast<double>(actual_avail) / (1 << 20) << " MB, "
+            << "actual_total: " << static_cast<double>(actual_total) / (1 << 20)
+            << " MB";
+
+    virtual_mem_size_ = AlignedSize(
+        actual_total * kVirtualAddressSpaceSizeMultiplier, granularity_);
+
+    // Reserve the required contiguous virtual address space for the allocations
+    // Reserving a larger VA range does not allocate physical memory up front.
+    // It leaves room for VMM pool fragmentation before physical memory is
+    // exhausted.
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cuMemAddressReserve(
+        &virtual_mem_base_, virtual_mem_size_, 0, 0, 0));
+
+    virtual_mem_alloced_offset_ = 0;
+  });
+}
+
+bool CUDAVirtualMemAllocator::IsAllocThreadSafe() const { return false; }
+
+void CUDAVirtualMemAllocator::FreeImpl(phi::Allocation* allocation) {
+  PADDLE_ENFORCE_EQ(
+      allocation->place(),
+      place_,
+      common::errors::PermissionDenied(
+          "GPU memory is freed in incorrect device. This may be a bug"));
+
+  auto iter = virtual_2_physical_map_.find(
+      reinterpret_cast<CUdeviceptr>(allocation->ptr()));
+  if (iter == virtual_2_physical_map_.end()) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Can not find virtual memory address at %s", allocation->ptr()));
+  }
+
+  int prev_id;
+  cudaGetDevice(&prev_id);
+  if (prev_id != place_.device) {
+    cudaSetDevice(place_.device);
+  }
+
+  auto result = phi::dynload::cuMemUnmap(iter->first, iter->second.second);
+  if (result != CUDA_ERROR_DEINITIALIZED) {
+    PADDLE_ENFORCE_GPU_SUCCESS(result);
+  }
+
+  if (result != CUDA_ERROR_DEINITIALIZED) {
+    PADDLE_ENFORCE_GPU_SUCCESS(platform::RecordedGpuMemRelease(
+        iter->second.first, iter->second.second, place_.device));
+  }
+
+  if (prev_id != place_.device) {
+    cudaSetDevice(prev_id);
+  }
+
+  UnregisterHandle(allocation->ptr());
+  virtual_2_physical_map_.erase(iter);
+
+  delete allocation;
+}
+
+phi::Allocation* CUDAVirtualMemAllocator::AllocateImpl(size_t size) {
+  InitOnce();
+  size = AlignedSize(size, granularity_);
+
+  CUdeviceptr ptr = virtual_mem_base_ + virtual_mem_alloced_offset_;
+
+  if (ptr + size > virtual_mem_base_ + virtual_mem_size_) {
+    PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+        "\n\nOut of memory error on GPU Virtual Memory %d. "
+        "Cannot allocate %s memory on GPU Virtual Memory %d, %s memory has "
+        "been allocated and "
+        "available memory is only %s.\n\n"
+        "Please decrease the batch size of your model.\n\n",
+        place_.device,
+        string::HumanReadableSize(size),
+        place_.device,
+        string::HumanReadableSize(virtual_mem_alloced_offset_),
+        string::HumanReadableSize(virtual_mem_size_ -
+                                  virtual_mem_alloced_offset_)));
+    return nullptr;
+  }
+
+  CUmemGenericAllocationHandle handle;
+
+  paddle::platform::CUDADeviceGuard guard(place_.device);
+
+  // Create physical memory backing allocation.
+  auto result =
+      platform::RecordedGpuMemCreate(&handle, size, &prop_, 0, place_.device);
+
+  if (result != CUDA_SUCCESS) {
+    if (result == CUDA_ERROR_OUT_OF_MEMORY) {
+      size_t actual_avail, actual_total;
+      PADDLE_ENFORCE_GPU_SUCCESS(cudaMemGetInfo(&actual_avail, &actual_total));
+      size_t actual_allocated = actual_total - actual_avail;
+      size_t actual_allocated_memory =
+          paddle::memory::DeviceMemoryStatCurrentValue("Allocated",
+                                                       place_.device);
+
+      PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+          "\n\nOut of memory error on GPU %d. "
+          "Cannot allocate %s memory on GPU %d, %s memory has been "
+          "allocated(actual using allocated memory %s) and "
+          "available memory is only %s.\n\n"
+          "Please check whether there is any other process using GPU %d.\n"
+          "1. If yes, please stop them, or start PaddlePaddle on another GPU.\n"
+          "2. If no, please decrease the batch size of your model.\n\n",
+          place_.device,
+          string::HumanReadableSize(size),
+          place_.device,
+          string::HumanReadableSize(actual_allocated),
+          string::HumanReadableSize(actual_allocated_memory),
+          string::HumanReadableSize(actual_avail),
+          place_.device));
+    } else {
+      PADDLE_ENFORCE_GPU_SUCCESS(result);
+    }
+    return nullptr;
+  }
+
+  // Assign the chunk to the appropriate VA range and release the handle.
+  // After mapping the memory, it can be referenced by virtual address.
+  // The allocation will be kept live until it is unmapped.
+  result = phi::dynload::cuMemMap(ptr, size, 0, handle, 0);
+
+  if (result != CUDA_SUCCESS) {
+    platform::RecordedGpuMemRelease(handle, size, place_.device);
+    PADDLE_ENFORCE_GPU_SUCCESS(result);
+    return nullptr;
+  }
+
+  // Apply the access descriptors to the whole VA range.
+  result = phi::dynload::cuMemSetAccess(
+      ptr, size, access_desc_.data(), access_desc_.size());
+
+  if (result != CUDA_SUCCESS) {
+    phi::dynload::cuMemUnmap(ptr, size);
+    platform::RecordedGpuMemRelease(handle, size, place_.device);
+    PADDLE_ENFORCE_GPU_SUCCESS(result);
+    return nullptr;
+  }
+
+  virtual_2_physical_map_.emplace(ptr, std::make_pair(handle, size));
+
+  virtual_mem_alloced_offset_ += size;
+  VLOG(10) << "AllocateImpl chunk handle: " << static_cast<int64_t>(handle)
+           << ", size=" << size
+           << ", device=" << static_cast<int>(place_.device);
+
+  RegisterHandle(reinterpret_cast<void*>(ptr), handle);
+
+  return new Allocation(
+      reinterpret_cast<void*>(ptr), size, Place(place_));  // NOLINT
+}
+
+CUmemGenericAllocationHandle CUDAVirtualMemAllocator::GetHandleFromBasePtr(
+    void* base_ptr) {
+  std::lock_guard<std::mutex> guard(base_ptr_handle_mu_);
+  auto it = base_ptr_handle_map_.find(base_ptr);
+  if (it == base_ptr_handle_map_.end()) {
+    return 0;
+  }
+  return it->second;
+}
+
+void CUDAVirtualMemAllocator::RegisterHandle(
+    void* base_ptr, CUmemGenericAllocationHandle handle) {
+  std::lock_guard<std::mutex> guard(base_ptr_handle_mu_);
+  base_ptr_handle_map_.emplace(base_ptr, handle);
+}
+
+void CUDAVirtualMemAllocator::UnregisterHandle(void* base_ptr) {
+  std::lock_guard<std::mutex> guard(base_ptr_handle_mu_);
+  base_ptr_handle_map_.erase(base_ptr);
+}
+
+}  // namespace paddle::memory::allocation
+
+#endif

paddle/phi/core/memory/allocation/cuda_virtual_mem_allocator.h

@@ -0,0 +1,78 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#ifdef PADDLE_WITH_CUDA
+#include <cuda_runtime.h>
+
+#include "paddle/phi/backends/dynload/cuda_driver.h"
+#include "paddle/phi/core/memory/allocation/vmm_ipc_allocation.h"
+#include "paddle/phi/core/platform/cuda_device_guard.h"
+#endif
+
+#include <mutex>  // NOLINT
+#include <unordered_map>
+#include <vector>
+
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+
+#ifdef PADDLE_WITH_CUDA
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// Allocate memory using NVIDIA's virtual memory management technology
+class CUDAVirtualMemAllocator : public Allocator {
+ public:
+  explicit CUDAVirtualMemAllocator(const GPUPlace& place);
+
+  bool IsAllocThreadSafe() const override;
+
+  static CUmemGenericAllocationHandle GetHandleFromBasePtr(void* base_ptr);
+
+ protected:
+  void FreeImpl(phi::Allocation* allocation) override;
+  phi::Allocation* AllocateImpl(size_t size) override;
+
+ private:
+  GPUPlace place_;
+  std::once_flag init_flag_;
+
+  CUdeviceptr virtual_mem_base_;
+  size_t virtual_mem_size_;
+  size_t virtual_mem_alloced_offset_;
+  size_t granularity_;
+
+  CUmemAllocationProp prop_;
+  std::vector<CUmemAccessDesc> access_desc_;
+
+  std::map<CUdeviceptr, std::pair<CUmemGenericAllocationHandle, size_t>>
+      virtual_2_physical_map_;
+  static std::mutex base_ptr_handle_mu_;
+  static std::unordered_map<void*, CUmemGenericAllocationHandle>
+      base_ptr_handle_map_;
+
+  void InitOnce();
+  static void RegisterHandle(void* base_ptr, CUmemGenericAllocationHandle h);
+  static void UnregisterHandle(void* base_ptr);
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle
+
+#endif

paddle/phi/core/memory/allocation/cuda_virtual_mem_allocator_v2.cc

@@ -0,0 +1,459 @@
+// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/cuda_virtual_mem_allocator_v2.h"
+
+#if defined(PADDLE_WITH_CUDA)
+
+#include <algorithm>
+#include <limits>
+#include <utility>
+
+#include "glog/logging.h"
+#include "paddle/phi/core/platform/cuda_device_guard.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+#include "paddle/phi/core/scope_guard.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+namespace {
+
+constexpr size_t kVMMSetAccessChunkSize = 64UL << 20;
+
+bool IsCudaDeinitialized(CUresult result) {
+  return result == CUDA_ERROR_DEINITIALIZED;
+}
+
+size_t GetPoolVAMultiplier(PoolType pool_type) {
+  switch (pool_type) {
+    case PoolType::kSmall:
+      return 1;
+    case PoolType::kLarge:
+      return 4;
+  }
+  return 1;
+}
+
+struct SetAccessResult {
+  CUresult status{CUDA_SUCCESS};
+  size_t failed_offset{0};
+  size_t failed_size{0};
+};
+
+SetAccessResult SetAccessInChunks(VMMDevicePtr ptr,
+                                  size_t size,
+                                  size_t handle_size,
+                                  const std::vector<CUmemAccessDesc>& desc) {
+  const size_t chunk_size =
+      std::max(handle_size, AlignedSize(kVMMSetAccessChunkSize, handle_size));
+  size_t offset = 0;
+  while (offset < size) {
+    const size_t remaining = size - offset;
+    const size_t current_size = std::min(chunk_size, remaining);
+    auto status = phi::dynload::cuMemSetAccess(
+        ptr + offset, current_size, desc.data(), desc.size());
+    if (status != CUDA_SUCCESS) {
+      return {status, offset, current_size};
+    }
+    offset += current_size;
+  }
+  return {};
+}
+
+template <typename Map, typename Key, typename Value>
+void EmplaceOrEnforce(Map* map,
+                      Key&& key,
+                      Value&& value,
+                      const char* map_name) {
+  const bool inserted =
+      map->try_emplace(std::forward<Key>(key), std::forward<Value>(value))
+          .second;
+  PADDLE_ENFORCE_EQ(
+      inserted,
+      true,
+      common::errors::AlreadyExists(
+          "Duplicate key inserted into %s, allocator state is inconsistent.",
+          map_name));
+}
+
+}  // namespace
+
+void CUDAVirtualMemAllocatorV2::AllocationLayoutRegistry::Add(
+    void* ptr, const HandleLayout& layout) {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  EmplaceOrEnforce(&layouts_, ptr, layout, "allocation_layout_map_");
+}
+
+bool CUDAVirtualMemAllocatorV2::AllocationLayoutRegistry::Lookup(
+    void* ptr, HandleLayout* layout) const {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  auto it = layouts_.find(ptr);
+  if (it == layouts_.end()) {
+    return false;
+  }
+  if (layout != nullptr) {
+    *layout = it->second;
+  }
+  return true;
+}
+
+void CUDAVirtualMemAllocatorV2::AllocationLayoutRegistry::Remove(void* ptr) {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  layouts_.erase(ptr);
+}
+
+CUDAVirtualMemAllocatorV2::CUDAVirtualMemAllocatorV2(const GPUPlace& place,
+                                                     size_t handle_size,
+                                                     PoolType pool)
+    : place_(place), handle_size_(handle_size), pool_type_(pool) {}
+
+bool CUDAVirtualMemAllocatorV2::IsAllocThreadSafe() const { return false; }
+
+void CUDAVirtualMemAllocatorV2::RollbackCreatedHandles(
+    const HandleLayout& layout) const {
+  for (const auto& meta : layout) {
+    if (meta == nullptr) {
+      continue;
+    }
+    phi::dynload::cuMemUnmap(meta->base(), meta->size());
+    platform::RecordedGpuMemRelease(
+        meta->handle(), meta->size(), place_.device);
+  }
+}
+
+void CUDAVirtualMemAllocatorV2::MarkLayoutMapped(const HandleLayout& layout) {
+  for (const auto& meta : layout) {
+    backing_map_.MarkMapped(meta->base(), meta, meta->size());
+  }
+}
+
+void CUDAVirtualMemAllocatorV2::InitOnce() {
+  std::call_once(init_flag_, [this] {
+    platform::CUDADeviceGuard guard(place_.device);
+    prop_.type = CU_MEM_ALLOCATION_TYPE_PINNED;
+    prop_.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+    prop_.location.id = place_.device;
+#if defined(_WIN32)
+    prop_.requestedHandleTypes = CU_MEM_HANDLE_TYPE_NONE;
+#else
+    prop_.requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
+#endif
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cuMemGetAllocationGranularity(
+        &granularity_, &prop_, CU_MEM_ALLOC_GRANULARITY_MINIMUM));
+    // V2 uses a per-pool fixed handle size. Unlike V1, the allocator rounds
+    // user input up to the device granularity so upper layers can treat every
+    // handle in one HandleLayout as a stable fixed-size building block.
+    handle_size_ =
+        AlignedSize(std::max(handle_size_, granularity_), granularity_);
+    size_t actual_avail = 0;
+    size_t actual_total = 0;
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaMemGetInfo(&actual_avail, &actual_total));
+    const size_t va_multiplier = GetPoolVAMultiplier(pool_type_);
+    PADDLE_ENFORCE_LE(va_multiplier,
+                      std::numeric_limits<size_t>::max() / actual_total,
+                      common::errors::InvalidArgument(
+                          "VA multiplier %d for pool %d overflows size_t.",
+                          va_multiplier,
+                          static_cast<int>(pool_type_)));
+    // Reserve VA by pool to leave room for later split and in-place reuse.
+    virtual_mem_size_ = AlignedSize(actual_total * va_multiplier, granularity_);
+    PADDLE_ENFORCE_GPU_SUCCESS(phi::dynload::cuMemAddressReserve(
+        &virtual_mem_base_, virtual_mem_size_, 0, 0, 0));
+    backing_map_.Configure(
+        virtual_mem_base_, virtual_mem_size_, handle_size_, place_.device);
+    CUmemAccessDesc self = {};
+    self.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
+    self.location.id = place_.device;
+    self.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+    access_desc_.push_back(self);
+  });
+}
+
+phi::Allocation* CUDAVirtualMemAllocatorV2::AllocateImpl(size_t size) {
+  return AppendWithLayout(size).allocation.release();
+}
+
+CUDAVirtualMemAllocatorV2::AllocationWithLayout
+CUDAVirtualMemAllocatorV2::AppendWithLayout(size_t size) {
+  InitOnce();
+  size_t aligned = AlignedSize(size, handle_size_);
+  PADDLE_ENFORCE_LE(
+      virtual_mem_alloced_offset_,
+      virtual_mem_size_,
+      common::errors::InvalidArgument(
+          "VMMAllocatorV2 tail offset exceeds reserved VA space."));
+  PADDLE_ENFORCE_LE(
+      aligned,
+      virtual_mem_size_ - virtual_mem_alloced_offset_,
+      common::errors::ResourceExhausted("VMMAllocatorV2 virtual address space "
+                                        "is exhausted for place %s.",
+                                        place_));
+  VMMDevicePtr ptr = virtual_mem_base_ + virtual_mem_alloced_offset_;
+
+  auto layout = CreateMappedHandleLayout(ptr, aligned, "AppendWithLayout");
+
+  MarkLayoutMapped(layout);
+  return WrapTrackedAllocation(ptr, aligned, std::move(layout), true);
+}
+
+CUDAVirtualMemAllocatorV2::AllocationWithBlock
+CUDAVirtualMemAllocatorV2::AppendWithBlock(size_t size) {
+  return BuildAllocationWithBlock(AppendWithLayout(size));
+}
+
+CUDAVirtualMemAllocatorV2::AllocationWithLayout
+CUDAVirtualMemAllocatorV2::PlaceAtVAWithLayout(VMMDevicePtr ptr, size_t size) {
+  InitOnce();
+  const size_t aligned = AlignedSize(size, handle_size_);
+  const size_t num_handles = aligned / handle_size_;
+  PADDLE_ENFORCE_EQ(virtual_mem_base_ + virtual_mem_size_ < virtual_mem_base_,
+                    false,
+                    common::errors::InvalidArgument(
+                        "VMMAllocatorV2 reserved VA range overflows."));
+  PADDLE_ENFORCE_GE(
+      ptr,
+      virtual_mem_base_,
+      common::errors::InvalidArgument(
+          "VMMAllocatorV2 PlaceAtVA ptr is before reserved VA range."));
+  PADDLE_ENFORCE_LT(
+      ptr,
+      virtual_mem_base_ + virtual_mem_size_,
+      common::errors::InvalidArgument(
+          "VMMAllocatorV2 PlaceAtVA ptr is outside reserved VA range."));
+  PADDLE_ENFORCE_EQ(
+      (ptr - virtual_mem_base_) % handle_size_,
+      0UL,
+      common::errors::InvalidArgument(
+          "VMMAllocatorV2 PlaceAtVA requires handle-aligned VA, ptr=%p.",
+          reinterpret_cast<void*>(ptr)));
+  PADDLE_ENFORCE_LE(
+      aligned,
+      virtual_mem_base_ + virtual_mem_size_ - ptr,
+      common::errors::ResourceExhausted(
+          "VMMAllocatorV2 PlaceAtVA range exceeds reserved VA space."));
+
+  VLOG(6) << "VMM V2 PlaceAtVA(AllocateAtVA) ptr="
+          << reinterpret_cast<void*>(ptr) << " requested=" << size
+          << " aligned=" << aligned << " handle_count=" << num_handles
+          << " tail_offset=" << virtual_mem_alloced_offset_;
+  auto layout = CreateMappedHandleLayout(ptr, aligned, "PlaceAtVAWithLayout");
+
+  MarkLayoutMapped(layout);
+  return WrapTrackedAllocation(ptr, aligned, std::move(layout), false);
+}
+
+CUDAVirtualMemAllocatorV2::AllocationWithBlock
+CUDAVirtualMemAllocatorV2::PlaceAtVAWithBlock(VMMDevicePtr ptr, size_t size) {
+  return BuildAllocationWithBlock(PlaceAtVAWithLayout(ptr, size));
+}
+
+HandleLayout CUDAVirtualMemAllocatorV2::CreateMappedHandleLayout(
+    VMMDevicePtr ptr, size_t aligned_size, const char* context) {
+  platform::CUDADeviceGuard guard(place_.device);
+  const size_t num_handles = aligned_size / handle_size_;
+  HandleLayout layout;
+  layout.reserve(num_handles);
+  for (size_t i = 0; i < num_handles; ++i) {
+    VMMAllocHandle handle;
+    auto ce = platform::RecordedGpuMemCreate(
+        &handle, handle_size_, &prop_, 0, place_.device);
+    if (ce != CUDA_SUCCESS) {
+      RollbackCreatedHandles(layout);
+      if (ce == CUDA_ERROR_OUT_OF_MEMORY) {
+        PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+            "%s cuMemCreate failed: out of GPU memory at handle %zu/%zu "
+            "(handle_size=%zu).",
+            context,
+            i,
+            num_handles,
+            handle_size_));
+      }
+      PADDLE_ENFORCE_GPU_SUCCESS(ce);
+    }
+
+    const VMMDevicePtr dst = ptr + i * handle_size_;
+    auto me = phi::dynload::cuMemMap(dst, handle_size_, 0, handle, 0);
+    if (me != CUDA_SUCCESS) {
+      platform::RecordedGpuMemRelease(handle, handle_size_, place_.device);
+      RollbackCreatedHandles(layout);
+      PADDLE_THROW(common::errors::External(
+          "%s cuMemMap failed at handle %zu/%zu.", context, i, num_handles));
+    }
+    layout.push_back(std::make_shared<VMMHandleMeta>(
+        VMMHandleMeta{dst, handle_size_, handle, place_.device}));
+  }
+  try {
+    SetAccessOrThrow(ptr, aligned_size, num_handles, context);
+  } catch (...) {
+    RollbackCreatedHandles(layout);
+    throw;
+  }
+  return layout;
+}
+
+void CUDAVirtualMemAllocatorV2::SetAccessOrThrow(VMMDevicePtr ptr,
+                                                 size_t aligned_size,
+                                                 size_t num_handles,
+                                                 const char* context) {
+  auto access_result =
+      SetAccessInChunks(ptr, aligned_size, handle_size_, access_desc_);
+  if (access_result.status == CUDA_SUCCESS) {
+    return;
+  }
+
+  size_t actual_avail = 0;
+  size_t actual_total = 0;
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaMemGetInfo(&actual_avail, &actual_total));
+  if (access_result.status == CUDA_ERROR_OUT_OF_MEMORY) {
+    PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+        "%s cuMemSetAccess failed: out of GPU memory at offset %zu/%zu "
+        "(failed_size=%zu, handle_size=%zu, handle_count=%zu, "
+        "actual_avail=%zu, actual_total=%zu).",
+        context,
+        access_result.failed_offset,
+        aligned_size,
+        access_result.failed_size,
+        handle_size_,
+        num_handles,
+        actual_avail,
+        actual_total));
+  }
+  PADDLE_THROW(common::errors::External(
+      "%s cuMemSetAccess failed at offset %zu/%zu (failed_size=%zu, "
+      "handle_size=%zu, handle_count=%zu, status=%d, actual_avail=%zu, "
+      "actual_total=%zu).",
+      context,
+      access_result.failed_offset,
+      aligned_size,
+      access_result.failed_size,
+      handle_size_,
+      num_handles,
+      static_cast<int>(access_result.status),
+      actual_avail,
+      actual_total));
+}
+
+bool CUDAVirtualMemAllocatorV2::CollectAllocationHandleLayout(
+    void* ptr, HandleLayout* layout) const {
+  return allocation_layouts_.Lookup(ptr, layout);
+}
+
+void CUDAVirtualMemAllocatorV2::FreeImpl(phi::Allocation* allocation) {
+  auto* ptr = allocation->ptr();
+  HandleLayout layout = RequireHandleLayout(ptr);
+
+  int prev_id = -1;
+  bool restore_device = false;
+  if (cudaGetDevice(&prev_id) == cudaSuccess && prev_id != place_.device) {
+    restore_device = cudaSetDevice(place_.device) == cudaSuccess;
+  }
+  DEFINE_PADDLE_SCOPE_GUARD([&] {
+    if (restore_device) {
+      cudaSetDevice(prev_id);
+    }
+  });
+
+  for (const auto& handle : layout) {
+    auto result = phi::dynload::cuMemUnmap(handle->base(), handle->size());
+    if (IsCudaDeinitialized(result)) {
+      continue;
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(result);
+    backing_map_.MarkUnmapped(handle->base(), handle->size());
+    result = platform::RecordedGpuMemRelease(
+        handle->handle(), handle->size(), place_.device);
+    if (IsCudaDeinitialized(result)) {
+      continue;
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(result);
+    backing_map_.MarkReleased(handle->base(), handle->handle(), handle->size());
+  }
+
+  UnregisterHandleLayout(ptr);
+  delete allocation;
+}
+
+CUDAVirtualMemAllocatorV2::AllocationWithLayout
+CUDAVirtualMemAllocatorV2::WrapTrackedAllocation(VMMDevicePtr ptr,
+                                                 size_t size,
+                                                 HandleLayout layout,
+                                                 bool advance_tail) {
+  if (advance_tail) {
+    AdvanceTailOffset(size);
+  }
+  AllocationWithLayout result;
+  auto* alloc = CreateTrackedAllocation(ptr, size, layout);
+  CUDAVirtualMemAllocatorV2* self = this;
+  result.layout = std::move(layout);
+  result.allocation = DecoratedAllocationPtr(alloc, [self](phi::Allocation* a) {
+    self->FreeImpl(static_cast<Allocation*>(a));
+  });
+  return result;
+}
+
+CUDAVirtualMemAllocatorV2::AllocationWithBlock
+CUDAVirtualMemAllocatorV2::BuildAllocationWithBlock(
+    AllocationWithLayout allocation_with_layout) {
+  AllocationWithBlock result;
+  result.block =
+      BlockV2::MakeMappedBlock(BlockType::kFree,
+                               allocation_with_layout.allocation->ptr(),
+                               allocation_with_layout.allocation->size(),
+                               pool_type_);
+  result.allocation = std::move(allocation_with_layout.allocation);
+  return result;
+}
+
+Allocation* CUDAVirtualMemAllocatorV2::CreateTrackedAllocation(
+    VMMDevicePtr ptr, size_t size, const HandleLayout& layout) {
+  RegisterHandleLayout(reinterpret_cast<void*>(ptr), layout);
+  return new Allocation(reinterpret_cast<void*>(ptr), size, place_);  // NOLINT
+}
+
+void CUDAVirtualMemAllocatorV2::RegisterHandleLayout(
+    void* ptr, const HandleLayout& layout) {
+  allocation_layouts_.Add(ptr, layout);
+  if (!backing_map_.ValidateLayout(layout, "RegisterHandleLayout")) {
+    VLOG(0) << "VMM V2 BackingMap validation failed while registering layout "
+            << ptr;
+  }
+}
+
+HandleLayout CUDAVirtualMemAllocatorV2::RequireHandleLayout(void* ptr) const {
+  HandleLayout layout;
+  const bool found = allocation_layouts_.Lookup(ptr, &layout);
+  PADDLE_ENFORCE_EQ(
+      found,
+      true,
+      common::errors::NotFound(
+          "No VMMAllocatorV2 handle layout found for allocation %p.", ptr));
+  return layout;
+}
+
+void CUDAVirtualMemAllocatorV2::UnregisterHandleLayout(void* ptr) {
+  allocation_layouts_.Remove(ptr);
+}
+
+bool CUDAVirtualMemAllocatorV2::IsRangeReleasable(VMMDevicePtr ptr,
+                                                  size_t size) const {
+  return backing_map_.IsRangeReleasable(ptr, size);
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle
+
+#endif

paddle/phi/core/memory/allocation/cuda_virtual_mem_allocator_v2.h

@@ -0,0 +1,141 @@
+// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#if defined(PADDLE_WITH_CUDA)
+
+#include <unordered_map>
+#include <vector>
+
+#include "paddle/phi/backends/dynload/cuda_driver.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/allocation/spin_lock.h"
+#include "paddle/phi/core/memory/allocation/vmm_allocator_v2_types.h"
+#include "paddle/phi/core/memory/allocation/vmm_backing_map.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// Compared with CUDAVirtualMemAllocator, V2 does not expose a single
+// VA<->handle mapping per allocation. It keeps the handle layout registered in
+// the bottom allocator and hands upper layers either allocation-level layout
+// snapshots or materialized mapped-free BlockV2 views.
+class CUDAVirtualMemAllocatorV2 : public Allocator {
+ public:
+  struct AllocationWithLayout {
+    DecoratedAllocationPtr allocation;
+    HandleLayout layout;
+  };
+
+  struct AllocationWithBlock {
+    bool HasAllocation() const { return allocation != nullptr; }
+    BlockV2 TakeBlock() { return std::move(block); }
+    DecoratedAllocationPtr TakeAllocation() { return std::move(allocation); }
+
+    DecoratedAllocationPtr allocation;
+    BlockV2 block;
+  };
+
+  struct AllocationLayoutRegistry {
+    void Add(void* ptr, const HandleLayout& layout);
+    bool Lookup(void* ptr, HandleLayout* layout) const;
+    void Remove(void* ptr);
+
+   private:
+    std::unordered_map<void*, HandleLayout> layouts_;
+    mutable SpinLock spinlock_;
+  };
+
+  // Standalone use defaults to the large pool. Upper layers may also choose
+  // explicit small/large pool types.
+  CUDAVirtualMemAllocatorV2(const GPUPlace& place,
+                            size_t handle_size,
+                            PoolType pool = PoolType::kLarge);
+
+  bool IsAllocThreadSafe() const override;
+
+  size_t handle_size() const { return handle_size_; }
+  PoolType pool_type() const { return pool_type_; }
+  VMMDevicePtr virtual_mem_base() const { return virtual_mem_base_; }
+  size_t virtual_mem_size() const { return virtual_mem_size_; }
+  size_t tail_offset() const { return virtual_mem_alloced_offset_; }
+  // Best-fit layers may consume VA from the reserved range incrementally. V2
+  // keeps this as an explicit cursor instead of reusing V1's
+  // virtual_2_physical_map_ bookkeeping.
+  void AdvanceTailOffset(size_t bytes) { virtual_mem_alloced_offset_ += bytes; }
+  // Retreat the tail cursor after upper layers release tail-end backing.
+  void SetTailOffset(size_t offset) { virtual_mem_alloced_offset_ = offset; }
+
+  const GPUPlace& place() const { return place_; }
+  AllocationWithBlock AppendWithBlock(size_t size);
+  // Create fresh physical backing and map it at an existing reserved VA range.
+  // This is used by upper layers to reuse unmapped-free VA space in place.
+  AllocationWithBlock PlaceAtVAWithBlock(VMMDevicePtr ptr, size_t size);
+  bool IsRangeReleasable(VMMDevicePtr ptr, size_t size) const;
+
+ protected:
+  phi::Allocation* AllocateImpl(size_t size) override;
+  void FreeImpl(phi::Allocation* allocation) override;
+
+ private:
+  void InitOnce();
+  void RollbackCreatedHandles(const HandleLayout& layout) const;
+  void MarkLayoutMapped(const HandleLayout& layout);
+  AllocationWithLayout AppendWithLayout(size_t size);
+  AllocationWithLayout PlaceAtVAWithLayout(VMMDevicePtr ptr, size_t size);
+  HandleLayout CreateMappedHandleLayout(VMMDevicePtr ptr,
+                                        size_t aligned_size,
+                                        const char* context);
+  void SetAccessOrThrow(VMMDevicePtr ptr,
+                        size_t aligned_size,
+                        size_t num_handles,
+                        const char* context);
+  bool CollectAllocationHandleLayout(void* ptr, HandleLayout* layout) const;
+  AllocationWithLayout WrapTrackedAllocation(VMMDevicePtr ptr,
+                                             size_t size,
+                                             HandleLayout layout,
+                                             bool advance_tail);
+  AllocationWithBlock BuildAllocationWithBlock(
+      AllocationWithLayout allocation_with_layout);
+  Allocation* CreateTrackedAllocation(VMMDevicePtr ptr,
+                                      size_t size,
+                                      const HandleLayout& layout);
+  void RegisterHandleLayout(void* ptr, const HandleLayout& layout);
+  HandleLayout RequireHandleLayout(void* ptr) const;
+  void UnregisterHandleLayout(void* ptr);
+
+  GPUPlace place_;
+  size_t handle_size_;
+  PoolType pool_type_;
+  std::once_flag init_flag_;
+
+  VMMDevicePtr virtual_mem_base_{0};
+  size_t virtual_mem_size_{0};
+  size_t virtual_mem_alloced_offset_{0};
+  size_t granularity_{0};
+  CUmemAllocationProp prop_{};
+  std::vector<CUmemAccessDesc> access_desc_;
+
+  AllocationLayoutRegistry allocation_layouts_;
+  VMMBackingMap backing_map_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle
+
+#endif

paddle/phi/core/memory/allocation/custom_allocator.cc

@@ -0,0 +1,86 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/custom_allocator.h"
+
+#include "paddle/phi/api/profiler/trace_event.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/memory/stats.h"
+#include "paddle/phi/core/platform/device/device_wrapper.h"
+#include "paddle/phi/core/platform/profiler.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+bool CustomAllocator::IsAllocThreadSafe() const { return true; }
+void CustomAllocator::FreeImpl(phi::Allocation* allocation) {
+  PADDLE_ENFORCE_EQ(
+      allocation->place(),
+      place_,
+      common::errors::PermissionDenied("CustomDevice memory is "
+                                       "freed in incorrect device. "
+                                       "This may be a bug"));
+  if (phi::DeviceManager::HasDeviceType(place_.GetDeviceType())) {
+    phi::DeviceManager::GetDeviceWithPlace(place_)->MemoryDeallocate(
+        allocation->ptr(), allocation->size());
+  }
+  DEVICE_MEMORY_STAT_UPDATE(
+      Reserved, place_.GetDeviceId(), -allocation->size());
+  platform::RecordMemEvent(allocation->ptr(),
+                           place_,
+                           allocation->size(),
+                           phi::TracerMemEventType::ReservedFree);
+  delete allocation;
+}
+
+phi::Allocation* CustomAllocator::AllocateImpl(size_t size) {
+  std::call_once(once_flag_, [this] { phi::DeviceManager::SetDevice(place_); });
+
+  void* ptr =
+      phi::DeviceManager::GetDeviceWithPlace(place_)->MemoryAllocate(size);
+  if (LIKELY(ptr)) {
+    DEVICE_MEMORY_STAT_UPDATE(Reserved, place_.GetDeviceId(), size);
+    platform::RecordMemEvent(
+        ptr, place_, size, phi::TracerMemEventType::ReservedAllocate);
+    return new Allocation(ptr, size, place_);
+  }
+
+  size_t avail, total;
+  phi::DeviceManager::MemoryStats(place_, &total, &avail);
+
+  auto dev_type = phi::PlaceHelper::GetDeviceType(place_);
+  auto dev_id = phi::PlaceHelper::GetDeviceId(place_);
+
+  PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+      "\n\nOut of memory error on %s:%d. "
+      "Cannot allocate %s memory on %s:%d, "
+      "available memory is only %s.\n\n"
+      "Please check whether there is any other process using %s:%d.\n"
+      "1. If yes, please stop them, or start PaddlePaddle on another %s.\n"
+      "2. If no, please decrease the batch size of your model.\n\n",
+      dev_type,
+      dev_id,
+      string::HumanReadableSize(size),
+      dev_type,
+      dev_id,
+      string::HumanReadableSize(avail),
+      dev_type,
+      dev_id,
+      dev_type));
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/custom_allocator.h

@@ -0,0 +1,42 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include <mutex>  // NOLINT
+
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class CustomAllocator : public Allocator {
+ public:
+  explicit CustomAllocator(const CustomPlace& place) : place_(place) {}
+
+  bool IsAllocThreadSafe() const override;
+
+ protected:
+  void FreeImpl(phi::Allocation* allocation) override;
+  phi::Allocation* AllocateImpl(size_t size) override;
+
+ private:
+  Place place_;
+  std::once_flag once_flag_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/inlined_vector.h

@@ -0,0 +1,72 @@
+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include <cstdint>
+#include <vector>
+
+#include "paddle/phi/core/enforce.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+template <typename T, size_t N>
+class InlinedVector {
+  static_assert(N > 0, "N must be larger than 0");
+
+ public:
+  inline InlinedVector() { len_ = 0; }
+
+  inline size_t size() const { return len_; }
+
+  inline T& operator[](size_t i) { return i < N ? head_[i] : tail_[i - N]; }
+
+  inline const T& operator[](size_t i) const {
+    return i < N ? head_[i] : tail_[i - N];
+  }
+
+  inline void emplace_back(const T& item) {
+    if (LIKELY(len_ < N)) {
+      head_[len_++] = item;
+    } else {
+      tail_.emplace_back(item);
+      ++len_;
+    }
+  }
+
+  inline void pop_back() {
+    if (UNLIKELY(len_ > N)) {
+      tail_.pop_back();
+    }
+    --len_;
+  }
+
+  inline T& back() {
+    if (LIKELY(len_ <= N)) {
+      return head_[len_ - 1];
+    } else {
+      return tail_.back();
+    }
+  }
+
+ private:
+  T head_[N];
+  size_t len_;
+  std::vector<T> tail_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/memory_block.cc

@@ -0,0 +1,155 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/phi/core/memory/allocation/memory_block.h"
+
+#include "paddle/phi/core/enforce.h"
+
+namespace paddle::memory::detail {
+
+void MemoryBlock::Init(MetadataCache* cache,
+                       Type t,
+                       size_t index,
+                       size_t size,
+                       void* left_buddy,
+                       void* right_buddy) {
+  cache->Save(this,
+              MemoryBlock::Desc(t,
+                                index,
+                                size - sizeof(MemoryBlock::Desc),
+                                size,
+                                static_cast<MemoryBlock*>(left_buddy),
+                                static_cast<MemoryBlock*>(right_buddy)));
+}
+
+MemoryBlock* MemoryBlock::GetLeftBuddy(MetadataCache* cache) {
+  return cache->LoadDesc(this)->left_buddy;
+}
+
+MemoryBlock* MemoryBlock::GetRightBuddy(MetadataCache* cache) {
+  return cache->LoadDesc(this)->right_buddy;
+}
+
+void MemoryBlock::Split(MetadataCache* cache,
+                        size_t size,
+                        size_t extra_padding_size) {
+  auto desc = cache->LoadDesc(this);
+  // make sure the split fits
+  PADDLE_ENFORCE_GE(desc->total_size,
+                    size,
+                    common::errors::InvalidArgument(
+                        "The size of memory block (%d) to split is "
+                        "not larger than size of request memory (%d)",
+                        desc->total_size,
+                        size));
+
+  size_t pay_load_size = sizeof(MemoryBlock::Desc) + extra_padding_size;
+
+  // bail out if there is no room for another partition
+  if (desc->total_size - size <= pay_load_size) {
+    return;
+  }
+
+  // find the position of the split
+  void* right_partition = reinterpret_cast<uint8_t*>(this) + size;
+
+  size_t remaining_size = desc->total_size - size;
+
+  // Add the new block as a buddy
+  // Write the metadata for the new block
+  auto new_block_right_buddy = desc->right_buddy;
+
+  cache->Save(static_cast<MemoryBlock*>(right_partition),
+              MemoryBlock::Desc(FREE_CHUNK,
+                                desc->index,
+                                remaining_size - pay_load_size,
+                                remaining_size,
+                                this,
+                                new_block_right_buddy));
+
+  desc->right_buddy = static_cast<MemoryBlock*>(right_partition);
+  desc->size = size - pay_load_size;
+  desc->total_size = size;
+
+  desc->UpdateGuards();
+
+  // Write metadata for the new block's right buddy
+  if (new_block_right_buddy != nullptr) {
+    auto buddy_desc = cache->LoadDesc(new_block_right_buddy);
+
+    buddy_desc->left_buddy = static_cast<MemoryBlock*>(right_partition);
+    buddy_desc->UpdateGuards();
+  }
+}
+
+void MemoryBlock::Merge(MetadataCache* cache, MemoryBlock* right_buddy) {
+  // only free blocks can be merged
+  auto desc = cache->LoadDesc(this);
+  auto rb_desc = cache->LoadDesc(right_buddy);
+  PADDLE_ENFORCE_EQ(desc->type,
+                    FREE_CHUNK,
+                    common::errors::PreconditionNotMet(
+                        "The destination chunk to merge is not free"));
+  PADDLE_ENFORCE_EQ(rb_desc->type,
+                    FREE_CHUNK,
+                    common::errors::PreconditionNotMet(
+                        "The source chunk to merge is not free"));
+
+  // link this->buddy's buddy
+  desc->right_buddy = rb_desc->right_buddy;
+
+  // link buddy's buddy -> this
+  if (desc->right_buddy != nullptr) {
+    auto buddy_metadata = cache->LoadDesc(desc->right_buddy);
+
+    buddy_metadata->left_buddy = this;
+    buddy_metadata->UpdateGuards();
+  }
+
+  desc->size += rb_desc->total_size;
+  desc->total_size += rb_desc->total_size;
+
+  desc->UpdateGuards();
+
+  cache->Save(right_buddy,
+              MemoryBlock::Desc(INVALID_CHUNK, 0, 0, 0, nullptr, nullptr));
+}
+
+void MemoryBlock::MarkAsFree(MetadataCache* cache) {
+  // check for double free or corruption
+  auto desc = cache->LoadDesc(this);
+  PADDLE_ENFORCE_NE(desc->type,
+                    FREE_CHUNK,
+                    common::errors::PreconditionNotMet(
+                        "The chunk to mark as free is free already"));
+  PADDLE_ENFORCE_NE(desc->type,
+                    INVALID_CHUNK,
+                    common::errors::PreconditionNotMet(
+                        "The chunk to mark as free is invalid"));
+  desc->type = FREE_CHUNK;
+  desc->UpdateGuards();
+}
+
+void* MemoryBlock::Data() const {
+  return const_cast<MemoryBlock::Desc*>(
+             reinterpret_cast<const MemoryBlock::Desc*>(this)) +
+         1;
+}
+
+MemoryBlock* MemoryBlock::Metadata() const {
+  return const_cast<MemoryBlock*>(reinterpret_cast<const MemoryBlock*>(
+      reinterpret_cast<const MemoryBlock::Desc*>(this) - 1));
+}
+
+}  // namespace paddle::memory::detail

paddle/phi/core/memory/allocation/memory_block.h

@@ -0,0 +1,145 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <unordered_map>
+
+namespace paddle {
+namespace memory {
+namespace detail {
+
+// Forward declaration.
+class MetadataCache;
+
+// MemoryBlock represents Each allocated memory block, which contains
+// MemoryBlock::Desc and the payload.
+struct MemoryBlock {
+  enum Type {
+    FREE_CHUNK,    // memory is free and idle
+    ARENA_CHUNK,   // memory is being occupied
+    HUGE_CHUNK,    // memory is out of management
+    INVALID_CHUNK  // memory is invalid
+  };
+
+  // init saves the MemoryBlock::Desc of the memory block in a MetadataCache.
+  // If it is a CPU memory block, the MetadataCache writes the
+  // MemoryBlock::Desc to the beginning of the block; or, if it is a GPU memory
+  // block, the MetadataCache writes the Metadata to a std::map in
+  // the CPU.
+  void Init(MetadataCache* cache,
+            Type t,
+            size_t index,
+            size_t size,
+            void* left_buddy,
+            void* right_buddy);
+
+  MemoryBlock* GetLeftBuddy(MetadataCache* cache);
+  MemoryBlock* GetRightBuddy(MetadataCache* cache);
+
+  // Split the allocation into left/right blocks.
+  void Split(MetadataCache* cache, size_t size, size_t extra_padding_size = 0);
+
+  // Merge left and right blocks together.
+  void Merge(MetadataCache* cache, MemoryBlock* right_buddy);
+
+  // Mark the allocation as free.
+  void MarkAsFree(MetadataCache* cache);
+
+  void* Data() const;
+  MemoryBlock* Metadata() const;
+
+  // MemoryBlock::Desc describes a MemoryBlock.
+  struct Desc {
+    Desc(MemoryBlock::Type t,
+         size_t i,
+         size_t s,
+         size_t ts,
+         MemoryBlock* l,
+         MemoryBlock* r);
+    Desc();
+
+    // mutator for type
+    inline void set_type(const MemoryBlock::Type& type) {
+      this->type = type;
+      this->UpdateGuards();
+    }
+
+    // accessor for type
+    inline const MemoryBlock::Type& get_type() const { return this->type; }
+
+    // accessor for index
+    inline const size_t& get_index() const { return this->index; }
+
+    // accessor for size
+    inline const size_t& get_size() const { return this->size; }
+
+    // accessor for total_size
+    inline const size_t& get_total_size() const { return this->total_size; }
+
+    // Updates guard_begin and guard_end by hashes of the Metadata object.
+    void UpdateGuards();
+
+    // Checks that guard_begin and guard_end are hashes of the Metadata object.
+    bool CheckGuards() const;
+
+    // TODO(gangliao): compress this
+    size_t guard_begin = 0;
+    MemoryBlock::Type type = MemoryBlock::INVALID_CHUNK;
+    size_t index = 0;
+    size_t size = 0;
+    size_t total_size = 0;
+    MemoryBlock* left_buddy = nullptr;
+    MemoryBlock* right_buddy = nullptr;
+    size_t guard_end = 0;
+  };
+};
+
+// A cache for accessing memory block meta-data that may be expensive
+// to access directly.  This class exists to unify the
+// MemoryBlock::Desc format between GPU and CPU allocations. It should
+// be removed when the CPU can access all GPU allocations directly via
+// UVM.
+class MetadataCache {
+ public:
+  explicit MetadataCache(bool uses_gpu);
+
+  // Disable copying and assignment.
+  MetadataCache(const MetadataCache&) = delete;
+  MetadataCache& operator=(const MetadataCache&) = delete;
+
+  // Returns the MemoryBlock::Desc for a memory block.  When MetadataCache is
+  // used to manage CPU memory, the MemoryBlock::Desc resides at the beginning
+  // of the memory block; when used to manage GPU memory, the
+  // Metadata resides in CPU memory indexed by cache_.
+  MemoryBlock::Desc* LoadDesc(MemoryBlock* memory_block);
+
+  // Saves the MemoryBlock::Desc of a memory block into the cache.  For CPU
+  // memory block, writes the MemoryBlock::Desc to the beginning of the memory
+  // block; whereas for GPU memory, writes it to cache_.
+  void Save(MemoryBlock* memory_block, const MemoryBlock::Desc& meta_data);
+
+  // For GPU memory block, erases its MemoryBlock::Desc from cache_.
+  void Invalidate(MemoryBlock* memory_block);
+
+ private:
+  typedef std::unordered_map<const MemoryBlock*, MemoryBlock::Desc> MetadataMap;
+  MetadataMap cache_;
+  bool uses_gpu_;
+};
+
+}  // namespace detail
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/memory_block_desc.cc

@@ -0,0 +1,75 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include <cstddef>
+#include <functional>
+
+#include "paddle/phi/core/memory/allocation/memory_block.h"
+
+namespace paddle::memory::detail {
+
+MemoryBlock::Desc::Desc(MemoryBlock::Type t,
+                        size_t i,
+                        size_t s,
+                        size_t ts,
+                        MemoryBlock* l,
+                        MemoryBlock* r)
+    : type(t),
+      index(i),
+      size(s),
+      total_size(ts),
+      left_buddy(l),
+      right_buddy(r) {}
+
+MemoryBlock::Desc::Desc()
+    : type(MemoryBlock::INVALID_CHUNK),
+      index(0),
+      size(0),
+      total_size(0),
+      left_buddy(nullptr),
+      right_buddy(nullptr) {}
+
+namespace {
+
+template <class T>
+inline void hash_combine(std::size_t* seed, const T& v) {
+  std::hash<T> hasher;
+  (*seed) ^= hasher(v) + 0x9e3779b9 + ((*seed) << 6) + ((*seed) >> 2);
+}
+
+inline size_t hash(const MemoryBlock::Desc& metadata, size_t initial_seed) {
+  size_t seed = initial_seed;
+
+  hash_combine(&seed, static_cast<size_t>(metadata.type));
+  hash_combine(&seed, metadata.index);
+  hash_combine(&seed, metadata.size);
+  hash_combine(&seed, metadata.total_size);
+  hash_combine(&seed, metadata.left_buddy);
+  hash_combine(&seed, metadata.right_buddy);
+
+  return seed;
+}
+
+}  // namespace
+
+void MemoryBlock::Desc::UpdateGuards() {
+  guard_begin = hash(*this, 1);
+  guard_end = hash(*this, 2);
+}
+
+bool MemoryBlock::Desc::CheckGuards() const {
+  return guard_begin == hash(*this, 1) && guard_end == hash(*this, 2);
+}
+
+}  // namespace paddle::memory::detail

paddle/phi/core/memory/allocation/meta_cache.cc

@@ -0,0 +1,65 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "glog/logging.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/memory/allocation/memory_block.h"
+
+namespace paddle::memory::detail {
+
+MetadataCache::MetadataCache(bool uses_gpu) : uses_gpu_(uses_gpu) {}
+
+MemoryBlock::Desc* MetadataCache::LoadDesc(MemoryBlock* block) {
+  if (uses_gpu_) {
+    auto iter = cache_.find(block);
+    PADDLE_ENFORCE_NE(
+        iter,
+        cache_.end(),
+        common::errors::NotFound("The memory block is not found in cache"));
+    auto* desc = &(iter->second);
+    PADDLE_ENFORCE_EQ(
+        desc->CheckGuards(),
+        true,
+        common::errors::InvalidArgument("Invalid CPU memory access"));
+    return desc;
+  } else {
+    auto* desc = reinterpret_cast<MemoryBlock::Desc*>(block);
+    VLOG(10) << "Load MemoryBlock::Desc type=" << desc->type;
+    PADDLE_ENFORCE_EQ(
+        desc->CheckGuards(),
+        true,
+        common::errors::InvalidArgument("Invalid CPU memory access"));
+    return reinterpret_cast<MemoryBlock::Desc*>(block);
+  }
+}
+
+void MetadataCache::Save(MemoryBlock* block,
+                         const MemoryBlock::Desc& original_desc) {
+  auto desc = original_desc;
+  desc.UpdateGuards();
+
+  if (uses_gpu_) {
+    cache_[block] = desc;
+  } else {
+    *reinterpret_cast<MemoryBlock::Desc*>(block) = desc;
+  }
+}
+
+void MetadataCache::Invalidate(MemoryBlock* block) {
+  if (uses_gpu_) {
+    cache_.erase(block);
+  }
+}
+
+}  // namespace paddle::memory::detail

paddle/phi/core/memory/allocation/mmap_allocator.cc

@@ -0,0 +1,420 @@
+// Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef _WIN32
+
+#include "paddle/phi/core/memory/allocation/mmap_allocator.h"
+
+#include <fcntl.h>
+#include <sys/mman.h>
+#include <cstdlib>
+
+#include <atomic>
+#include <random>
+#include <string>
+
+#include "glog/logging.h"
+#include "paddle/common/flags.h"
+#include "paddle/phi/core/enforce.h"
+
+COMMON_DECLARE_bool(use_shm_cache);
+
+namespace paddle::memory::allocation {
+
+std::string GetIPCName() {
+  static std::random_device rd;
+  static std::atomic<uint64_t> counter{0};
+  std::string handle = "/paddle_";
+#ifdef _WIN32
+  handle += std::to_string(GetCurrentProcessId());
+#else
+  handle += std::to_string(getpid());
+#endif
+  handle += "_";
+  handle += std::to_string(counter.fetch_add(1));
+  handle += "_";
+  handle += std::to_string(rd());
+  return handle;
+}
+
+struct CountInfo {
+  std::atomic<int> refcount;
+};
+
+void AllocateMemoryMap(std::string filename,
+                       int *shared_fd,
+                       int flags,
+                       size_t size,
+                       void **map_ptr_) {
+  // TODO(@ZHUI): support win32
+  int file_flags = 0;
+  int fd = *shared_fd;
+  if (flags & MAPPED_SHAREDMEM) {
+    file_flags = O_RDWR | O_CREAT;
+  } else {
+    file_flags = O_RDONLY;
+  }
+  if (flags & MAPPED_EXCLUSIVE) {
+    file_flags |= O_EXCL;
+  }
+  if (flags & MAPPED_NOCREATE) {
+    file_flags &= ~O_CREAT;
+  }
+
+  if (!(flags & MAPPED_FROMFD) && fd == -1) {
+    if (flags & MAPPED_SHAREDMEM) {
+      fd = shm_open(filename.c_str(), file_flags, (mode_t)0600);
+      PADDLE_ENFORCE_NE(
+          fd,
+          -1,
+          common::errors::Unavailable(
+              "File descriptor %s open failed, unable in read-write mode",
+              filename.c_str()));
+      VLOG(6) << "shm_open: " << filename;
+      MemoryMapFdSet::Instance().Insert(filename);
+    }
+  }
+
+  PADDLE_ENFORCE_EQ(ftruncate(fd, size),
+                    0,
+                    common::errors::Unavailable(
+                        "Truncate a file to a specified length failed!"));
+
+  if (flags & MAPPED_SHAREDMEM) {
+    *map_ptr_ = mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
+  } else {
+    *map_ptr_ = mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_PRIVATE, fd, 0);
+  }
+
+  if (flags & MAPPED_UNLINK) {
+    VLOG(6) << "shm_unlink: " << filename;
+    shm_unlink(filename.c_str());
+  }
+
+  PADDLE_ENFORCE_NE(*map_ptr_,
+                    MAP_FAILED,
+                    common::errors::Unavailable(
+                        "Memory map failed when create shared memory."));
+  if (flags & MAPPED_KEEPFD) {
+    *shared_fd = fd;
+    VLOG(6) << "keep fd: " << *shared_fd;
+  } else {
+    PADDLE_ENFORCE_NE(::close(fd),
+                      -1,
+                      common::errors::Unavailable(
+                          "Error closing memory mapped file %s", filename));
+
+    *shared_fd = -1;
+  }
+}
+
+std::shared_ptr<RefcountedMemoryMapAllocation>
+AllocateRefcountedMemoryMapAllocation(std::string filename,
+                                      int shared_fd,
+                                      int flags,
+                                      size_t size,
+                                      int buffer_id) {
+  int fd = shared_fd;
+  void *base_ptr = nullptr;
+  if (buffer_id == -1) {
+    AllocateMemoryMap(filename, &fd, flags, size + mmap_alignment, &base_ptr);
+    VLOG(4) << "Create and mmap a new shm: " << filename;
+  } else {
+    base_ptr = MemoryMapAllocationPool::Instance().GetById(buffer_id).mmap_ptr_;
+    VLOG(4) << "Get a cached shm " << filename;
+  }
+  void *aligned_base_ptr =
+      static_cast<void *>(static_cast<char *>(base_ptr) + mmap_alignment);
+  return std::make_shared<RefcountedMemoryMapAllocation>(
+      aligned_base_ptr, size, filename, fd, flags, buffer_id);
+}
+
+RefcountedMemoryMapAllocation::RefcountedMemoryMapAllocation(
+    void *ptr,
+    size_t size,
+    std::string ipc_name,
+    int fd,
+    int flags,
+    int buffer_id)
+    : MemoryMapAllocation(ptr, size, ipc_name, fd, flags) {
+  // must reset base ptr first.
+  buffer_id_ = buffer_id;
+  fd_ = fd;
+  flags_ = flags;
+  resetBaseptr();
+  initializeRefercount();
+}
+
+void MemoryMapAllocation::close() {
+  if (!closed_fd_) {
+    closed_fd_ = true;
+    if (flags_ & MAPPED_KEEPFD) {
+      PADDLE_ENFORCE_NE(::close(fd_),
+                        -1,
+                        common::errors::Unavailable(
+                            "Error closing file descriptor <%d>", fd_));
+    }
+  }
+  if (closed_) {
+    return;
+  }
+  closed_ = true;
+}
+
+MemoryMapAllocation::~MemoryMapAllocation() { close(); }  // NOLINT
+
+void RefcountedMemoryMapAllocation::incref() {
+  CountInfo *info = static_cast<CountInfo *>(map_ptr_);
+  ++info->refcount;
+}
+
+int RefcountedMemoryMapAllocation::decref() {
+  CountInfo *info = static_cast<CountInfo *>(map_ptr_);
+  return --info->refcount == 0;
+}
+
+void RefcountedMemoryMapAllocation::resetBaseptr() {
+  map_ptr_ =
+      static_cast<void *>(static_cast<char *>(map_ptr_) - mmap_alignment);
+  map_size_ = map_size_ + mmap_alignment;
+}
+
+void RefcountedMemoryMapAllocation::initializeRefercount() {
+  CountInfo *info = reinterpret_cast<CountInfo *>(map_ptr_);
+
+  if (flags_ & MAPPED_EXCLUSIVE) {
+    new (&info->refcount) std::atomic<int>(1);
+  } else {
+    info->refcount++;
+  }
+}
+
+void RefcountedMemoryMapAllocation::close() {
+  VLOG(4) << "Close a RefcountedMemoryMapAllocation: " << ipc_name_;
+  if (closed_) {
+    return;
+  }
+  closed_ = true;
+  void *data = map_ptr_;
+  CountInfo *info = reinterpret_cast<CountInfo *>(data);
+  --info->refcount;
+  if (flags_ & MAPPED_KEEPFD) {
+    closed_fd_ = true;
+    PADDLE_ENFORCE_NE(
+        ::close(fd_),
+        -1,
+        common::errors::Unavailable("Error closing file descriptor <%d>", fd_));
+    VLOG(6) << "close fd: " << fd_;
+  }
+
+  if (FLAGS_use_shm_cache && buffer_id_ != -1) {
+    return;
+  } else {
+    if (FLAGS_use_shm_cache &&
+        MemoryMapAllocationPool::Instance().BufferSize() <
+            static_cast<size_t>(
+                MemoryMapAllocationPool::Instance().MaxPoolSize())) {
+      MemoryMapAllocationPool::Instance().Insert(MemoryMapInfo(
+          flags_, map_size_ - mmap_alignment, ipc_name_, map_ptr_));
+    } else {
+      if (info->refcount == 0) {
+        shm_unlink(ipc_name_.c_str());
+        VLOG(6) << "shm_unlink file: " << ipc_name_;
+      }
+
+      PADDLE_ENFORCE_NE(munmap(map_ptr_, map_size_),
+                        -1,
+                        common::errors::Unavailable(
+                            "could not unmap the shared memory file: %s (%d)",
+                            strerror(errno),
+                            errno));
+    }
+  }
+}
+
+MemoryMapWriterAllocation::~MemoryMapWriterAllocation() {
+  if (munmap(this->ptr(), this->size()) == -1) {
+    common::errors::Unavailable("could not unmap the shared memory file %s",
+                                this->ipc_name());
+  }
+}
+
+MemoryMapReaderAllocation::~MemoryMapReaderAllocation() {
+  if (munmap(this->ptr(), this->size()) == -1) {
+    common::errors::Unavailable("could not unmap the shared memory file %s",
+                                this->ipc_name());
+  }
+
+  /* Here we do not pay attention to the result of shm_unlink,
+     because the memory mapped file may have been cleared due to the
+     MemoryMapFdSet::Clear() */
+
+  // Code of DataLoader subprocess:
+  //
+  //    core._array_to_share_memory_tensor(b)
+  //    out_queue.put((idx, tensor_list, structure))
+  //    core._remove_tensor_list_mmap_fds(tensor_list)
+
+  /* If the tensor in already in the send queue, the tensor will be
+   * deconstructed by the function. If the tensor not send yet, it
+   * will be cleared by MemoryMapFdSet::Clear().
+   * If the `_remove_tensor_list_mmap_fds` have be interrupted, the
+   * tensor will be cleared by both methods.
+   * */
+
+  shm_unlink(this->ipc_name().c_str());
+  MemoryMapFdSet::Instance().Remove(this->ipc_name());
+  VLOG(3) << "~MemoryMapReaderAllocation: " << this->ipc_name();
+}
+
+std::shared_ptr<MemoryMapWriterAllocation> AllocateMemoryMapWriterAllocation(
+    size_t size) {
+  const std::string &ipc_name = GetIPCName();
+  int flags = O_RDWR | O_CREAT;
+  int fd = shm_open(ipc_name.c_str(), flags, 0600);
+
+  PADDLE_ENFORCE_NE(fd,
+                    -1,
+                    common::errors::Unavailable(
+                        "File descriptor %s open failed", ipc_name.c_str()));
+  PADDLE_ENFORCE_EQ(ftruncate(fd, size),
+                    0,
+                    common::errors::Unavailable(
+                        "Truncate a file to a specified length failed!"));
+
+  void *ptr = mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
+  PADDLE_ENFORCE_NE(ptr,
+                    MAP_FAILED,
+                    common::errors::Unavailable(
+                        "Memory map failed when create shared memory."));
+  close(fd);
+
+  return std::make_shared<MemoryMapWriterAllocation>(ptr, size, ipc_name);
+}
+
+std::shared_ptr<MemoryMapReaderAllocation> RebuildMemoryMapReaderAllocation(
+    const std::string &ipc_name, size_t size) {
+  int flags = O_RDWR | O_CREAT;
+  flags &= ~O_CREAT;
+  int fd = shm_open(ipc_name.c_str(), flags, 0600);
+  PADDLE_ENFORCE_NE(fd,
+                    -1,
+                    common::errors::Unavailable(
+                        "File descriptor %s open failed", ipc_name.c_str()));
+  void *ptr = mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
+  PADDLE_ENFORCE_NE(ptr,
+                    MAP_FAILED,
+                    common::errors::Unavailable(
+                        "Memory map failed when rebuild shared memory."));
+  close(fd);
+  return std::make_shared<MemoryMapReaderAllocation>(ptr, size, ipc_name);
+}
+
+MemoryMapFdSet &MemoryMapFdSet::Instance() {  // NOLINT
+  static MemoryMapFdSet set;
+  return set;
+}
+
+void MemoryMapFdSet::Insert(const std::string &ipc_name) {
+  std::lock_guard<std::mutex> guard(mtx_);
+  fd_set_.emplace(ipc_name);
+  VLOG(3) << "PID: " << getpid() << ", MemoryMapFdSet: insert " << ipc_name
+          << ", set size: " << fd_set_.size();
+}
+
+void MemoryMapFdSet::Remove(const std::string &ipc_name) {
+  std::lock_guard<std::mutex> guard(mtx_);
+  fd_set_.erase(ipc_name);
+  VLOG(3) << "PID: " << getpid() << ", MemoryMapFdSet: erase " << ipc_name
+          << ", set size: " << fd_set_.size();
+}
+
+void MemoryMapFdSet::Clear() {
+  VLOG(7) << "PID: " << getpid() << ", MemoryMapFdSet: set size - "
+          << fd_set_.size();
+  std::lock_guard<std::mutex> guard(mtx_);
+  for (auto const &fd : fd_set_) {
+    int rlt = shm_unlink(fd.c_str());
+    if (rlt == 0) {
+      VLOG(7) << "PID: " << getpid() << ", MemoryMapFdSet: clear " << fd;
+    }
+  }
+  fd_set_.clear();
+}
+
+MemoryMapFdSet::~MemoryMapFdSet() { Clear(); }
+
+MemoryMapAllocationPool *MemoryMapAllocationPool::pool_ = nullptr;
+
+void MemoryMapAllocationPool::Insert(const MemoryMapInfo &memory_map) {
+  std::lock_guard<std::mutex> guard(mtx_);
+  memory_map_allocations_.push_back(memory_map);
+  VLOG(4) << this << "Insert a new shm: " << memory_map.file_name_;
+}
+
+int MemoryMapAllocationPool::FindFromCache(const int &flag,
+                                           const size_t &data_size,
+                                           const std::string &file_name,
+                                           bool check_refcount) {
+  std::lock_guard<std::mutex> guard(mtx_);
+  for (int idx = 0; idx < static_cast<int>(memory_map_allocations_.size());
+       idx++) {
+    if (memory_map_allocations_.at(idx).flags_ == flag &&
+        memory_map_allocations_.at(idx).data_size_ == data_size) {
+      if (file_name.empty() ||
+          memory_map_allocations_.at(idx).file_name_ == file_name) {
+        if (!check_refcount || reinterpret_cast<CountInfo *>(
+                                   memory_map_allocations_.at(idx).mmap_ptr_)
+                                       ->refcount == 0) {
+          VLOG(4) << "Match at: " << idx;
+          return idx;
+        }
+      }
+    }
+  }
+  return -1;
+}
+
+const MemoryMapInfo &MemoryMapAllocationPool::GetById(int id) {
+  std::lock_guard<std::mutex> guard(mtx_);
+  return memory_map_allocations_.at(id);
+}
+
+void MemoryMapAllocationPool::SetMaxPoolSize(const int &size) {
+  max_pool_size_ = size;
+  VLOG(4) << this << "Set max pool size is: " << max_pool_size_;
+}
+
+void MemoryMapAllocationPool::Clear() {
+  std::lock_guard<std::mutex> guard(mtx_);
+  for (auto const &mmap : memory_map_allocations_) {
+    int rlt = shm_unlink(mmap.file_name_.c_str());
+    if (rlt == 0) {
+      VLOG(4) << "MemoryMapAllocationPool: clear " << mmap.file_name_;
+    }
+    PADDLE_ENFORCE_NE(munmap(mmap.mmap_ptr_, mmap.data_size_ + mmap_alignment),
+                      -1,
+                      common::errors::Unavailable(
+                          "could not unmap the shared memory file: %s (%d)",
+                          strerror(errno),
+                          errno));
+  }
+  memory_map_allocations_.clear();
+}
+
+MemoryMapAllocationPool::~MemoryMapAllocationPool() { Clear(); }  // NOLINT
+
+}  // namespace paddle::memory::allocation
+
+#endif

paddle/phi/core/memory/allocation/mmap_allocator.h

@@ -0,0 +1,240 @@
+// Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#ifndef _WIN32
+
+#include <atomic>
+#include <memory>
+#include <mutex>
+#include <string>
+#include <unordered_set>
+#include <utility>
+
+#include "paddle/phi/core/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+std::string GetIPCName();
+
+static constexpr int64_t mmap_alignment = 64;
+
+enum MappedModes {
+  MAPPED_SHAREDMEM = 1,
+  MAPPED_EXCLUSIVE = 2,
+  MAPPED_NOCREATE = 4,
+  MAPPED_KEEPFD = 8,
+  MAPPED_FROMFD = 16,
+  MAPPED_UNLINK = 32
+};
+
+class MemoryMapAllocation : public Allocation {
+ public:
+  explicit MemoryMapAllocation(void *ptr,
+                               size_t size,
+                               std::string ipc_name,
+                               int fd)
+      : Allocation(ptr, size, CPUPlace()),
+        ipc_name_(std::move(ipc_name)),
+        fd_(fd),
+        map_ptr_(ptr),
+        map_size_(size) {}
+  explicit MemoryMapAllocation(
+      void *ptr, size_t size, std::string ipc_name, int fd, int flags)
+      : Allocation(ptr, size, CPUPlace()),
+        ipc_name_(std::move(ipc_name)),
+        fd_(fd),
+        flags_(flags),
+        map_ptr_(ptr),
+        map_size_(size) {}
+
+  inline const std::string &ipc_name() const { return ipc_name_; }
+  inline int shared_fd() const { return fd_; }
+
+  virtual void close();
+
+  ~MemoryMapAllocation() override;
+
+ protected:
+  std::string ipc_name_;
+  int fd_ = -1;
+  int flags_ = 0;
+  void *map_ptr_ = nullptr;
+  size_t map_size_ = 0;
+  bool closed_ = false;
+  bool closed_fd_ = false;
+};
+
+class RefcountedMemoryMapAllocation : public MemoryMapAllocation {
+ public:
+  RefcountedMemoryMapAllocation(void *ptr,
+                                size_t size,
+                                std::string ipc_name,
+                                int flags,
+                                int fd,
+                                int buffer_id = -1);
+
+  void incref();
+  int decref();
+  void close() override;
+  virtual ~RefcountedMemoryMapAllocation() { close(); }
+
+ protected:
+  int buffer_id_ = -1;
+  void initializeRefercount();
+  void resetBaseptr();
+};
+
+void AllocateMemoryMap(std::string filename,
+                       int *shared_fd,
+                       int flags,
+                       size_t size,
+                       void **base_ptr_);
+
+std::shared_ptr<RefcountedMemoryMapAllocation>
+AllocateRefcountedMemoryMapAllocation(std::string filename,
+                                      int shared_fd,
+                                      int flags,
+                                      size_t size,
+                                      int buffer_id = -1);
+
+class MemoryMapWriterAllocation : public Allocation {
+ public:
+  explicit MemoryMapWriterAllocation(void *ptr,
+                                     size_t size,
+                                     std::string ipc_name)
+      : Allocation(ptr, size, CPUPlace()), ipc_name_(std::move(ipc_name)) {}
+
+  inline const std::string &ipc_name() const { return ipc_name_; }
+  inline int shared_fd() const { return fd_; }
+
+  ~MemoryMapWriterAllocation() override;
+
+ private:
+  std::string ipc_name_;
+  int fd_ = -1;
+};
+
+class MemoryMapReaderAllocation : public Allocation {
+ public:
+  explicit MemoryMapReaderAllocation(void *ptr,
+                                     size_t size,
+                                     std::string ipc_name)
+      : Allocation(ptr, size, CPUPlace()), ipc_name_(std::move(ipc_name)) {}
+
+  inline const std::string &ipc_name() const { return ipc_name_; }
+  inline int shared_fd() const { return fd_; }
+
+  ~MemoryMapReaderAllocation() override;
+
+ private:
+  std::string ipc_name_;
+  int fd_ = -1;
+};
+
+std::shared_ptr<MemoryMapWriterAllocation> AllocateMemoryMapWriterAllocation(
+    size_t size);
+
+std::shared_ptr<MemoryMapReaderAllocation> RebuildMemoryMapReaderAllocation(
+    const std::string &ipc_name, size_t size);
+
+class MemoryMapFdSet {
+ public:
+  static MemoryMapFdSet &Instance();  // NOLINT
+
+  void Insert(const std::string &ipc_name);
+
+  void Remove(const std::string &ipc_name);
+
+  void Clear();
+
+  ~MemoryMapFdSet();
+
+ private:
+  MemoryMapFdSet() = default;
+
+  std::unordered_set<std::string> fd_set_;
+  std::mutex mtx_;
+};
+
+class MemoryMapInfo {
+ public:
+  explicit MemoryMapInfo(int flags,
+                         size_t data_size,
+                         std::string file_name,
+                         void *mmap_ptr)
+      : flags_(flags),
+        data_size_(data_size),
+        file_name_(file_name),
+        mmap_ptr_(mmap_ptr) {}
+
+  int flags_ = 0;
+  size_t data_size_ = 0;
+  std::string file_name_;
+  void *mmap_ptr_ = nullptr;
+};
+
+/* Note(zhangbo):
+MemoryMapAllocationPool is used to cache and reuse shm, thus reducing munmap in
+dataloader. The munmap(shm_mmap_ptr) instruction in
+RefcountedMemoryMapAllocation::close() function may block other threads of the
+process. Therefore, the logic of shm cache and reuse is designed: the shm
+created by the _share_filename process will be cached and reused according to
+the data_size of shm, thus eliminating the problem of munmap blocking other
+threads
+*/
+class MemoryMapAllocationPool {
+ public:
+  static MemoryMapAllocationPool &Instance() {
+    if (pool_ == nullptr) {
+      pool_ = new MemoryMapAllocationPool();
+    }
+    return *pool_;
+  }
+
+  void Insert(const MemoryMapInfo &memory_map);
+
+  int FindFromCache(const int &flag,
+                    const size_t &data_size,
+                    const std::string &file_name = "",
+                    bool check_refcount = true);
+
+  const MemoryMapInfo &GetById(int id);
+
+  size_t BufferSize() { return memory_map_allocations_.size(); }
+
+  void Clear();
+
+  void SetMaxPoolSize(const int &size);
+
+  int MaxPoolSize() { return max_pool_size_; }
+
+  ~MemoryMapAllocationPool();
+
+ private:
+  MemoryMapAllocationPool() = default;
+  static MemoryMapAllocationPool *pool_;
+  std::vector<MemoryMapInfo> memory_map_allocations_;
+  int max_pool_size_ = 0;
+  std::mutex mtx_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle
+
+#endif

paddle/phi/core/memory/allocation/naive_best_fit_allocator.cc

@@ -0,0 +1,716 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/naive_best_fit_allocator.h"
+
+#include <mutex>
+
+#include "glog/logging.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/memory/allocation/buddy_allocator.h"
+#include "paddle/phi/core/memory/allocation/system_allocator.h"
+#include "paddle/phi/core/platform/device/device_wrapper.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+#include "paddle/phi/core/platform/profiler.h"
+#include "paddle/phi/core/utils/visit_place.h"
+#include "paddle/utils/string/printf.h"
+#include "paddle/utils/string/split.h"
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#include "paddle/phi/core/platform/cuda_device_guard.h"
+#endif
+#include "paddle/common/flags.h"
+PHI_DEFINE_EXPORTED_bool(
+    init_allocated_mem,
+    false,
+    "It is a mistake that the values of the memory allocated by "
+    "BuddyAllocator are always zeroed in some op's implementation. "
+    "To find this error in time, we use init_allocated_mem to indicate "
+    "that initializing the allocated memory with a small value "
+    "during unit testing.");
+COMMON_DECLARE_double(fraction_of_gpu_memory_to_use);
+COMMON_DECLARE_uint64(initial_gpu_memory_in_mb);
+COMMON_DECLARE_uint64(reallocate_gpu_memory_in_mb);
+COMMON_DECLARE_bool(benchmark);
+
+namespace paddle::memory::legacy {
+template <typename Place>
+void *Alloc(const Place &place, size_t size);
+
+template <typename Place>
+void Free(const Place &place, void *p, size_t size);
+
+template <typename Place>
+uint64_t Release(const Place &place);
+
+template <typename Place>
+size_t Used(const Place &place);
+
+struct Usage {
+  size_t operator()(const CPUPlace &cpu) const;
+  size_t operator()(const GPUPlace &gpu) const;
+  size_t operator()(const GPUPinnedPlace &cuda_pinned) const;
+  size_t operator()(const XPUPlace &xpu) const;
+  size_t operator()(const XPUPinnedPlace &xpu_pinned) const;
+};
+
+size_t memory_usage(const Place &p);
+
+using BuddyAllocator = detail::BuddyAllocator;
+
+BuddyAllocator *GetCPUBuddyAllocator() {
+  // We tried thread_local for inference::RNN1 model, but that not works much
+  // for multi-thread test.
+  static std::once_flag init_flag;
+  static detail::BuddyAllocator *a = nullptr;
+
+  std::call_once(init_flag, []() {
+    a = new detail::BuddyAllocator(
+        std::unique_ptr<detail::SystemAllocator>(new detail::CPUAllocator),
+        phi::backends::cpu::CpuMinChunkSize(),
+        phi::backends::cpu::CpuMaxChunkSize());
+  });
+
+  return a;
+}
+
+template <>
+void *Alloc<CPUPlace>(const CPUPlace &place, size_t size) {
+  VLOG(10) << "Allocate " << size << " bytes on " << Place(place);
+  void *p = GetCPUBuddyAllocator()->Alloc(size);
+  if (FLAGS_init_allocated_mem) {
+    memset(p, 0xEF, size);
+  }
+  VLOG(10) << "  pointer=" << p;
+  return p;
+}
+
+template <>
+void Free<CPUPlace>(const CPUPlace &place, void *p, size_t size) {
+  VLOG(10) << "Free pointer=" << p << " on " << Place(place);
+  GetCPUBuddyAllocator()->Free(p);
+}
+
+template <>
+uint64_t Release<CPUPlace>(const CPUPlace &place) {
+  return GetCPUBuddyAllocator()->Release();
+}
+
+template <>
+size_t Used<CPUPlace>(const CPUPlace &place) {
+  return GetCPUBuddyAllocator()->Used();
+}
+
+// For Graphcore IPU
+template <>
+void *Alloc<IPUPlace>(const IPUPlace &place, size_t size) {
+  VLOG(10) << "Allocate " << size << " bytes on " << Place(place);
+  VLOG(10) << "IPUPlace, Allocate on cpu.";
+
+  void *p = GetCPUBuddyAllocator()->Alloc(size);
+  if (FLAGS_init_allocated_mem) {
+    memset(p, 0xEF, size);
+  }
+  VLOG(10) << "  pointer=" << p;
+  return p;
+}
+template <>
+void Free<IPUPlace>(const IPUPlace &place, void *p, size_t size) {
+  VLOG(10) << "Free pointer=" << p << " on " << Place(place);
+  GetCPUBuddyAllocator()->Free(p);
+}
+template <>
+uint64_t Release<IPUPlace>(const IPUPlace &place) {
+  return GetCPUBuddyAllocator()->Release();
+}
+template <>
+size_t Used<IPUPlace>(const IPUPlace &place) {
+  return GetCPUBuddyAllocator()->Used();
+}
+
+// For kunlun XPU
+template <>
+void *Alloc<XPUPlace>(const XPUPlace &place, size_t size) {
+#ifdef PADDLE_WITH_XPU
+  VLOG(10) << "Allocate " << size << " bytes on " << Place(place);
+  void *p = nullptr;
+
+  phi::backends::xpu::XPUDeviceGuard guard(place.device);
+  int ret = xpu_malloc(reinterpret_cast<void **>(&p), size);
+  if (ret != XPU_SUCCESS) {
+    VLOG(10) << "xpu memory malloc(" << size << ") failed, try again";
+    xpu_wait();
+    ret = xpu_malloc(reinterpret_cast<void **>(&p), size);
+  }
+  PADDLE_ENFORCE_EQ(
+      ret,
+      XPU_SUCCESS,
+      common::errors::External(
+          "XPU API return wrong value[%d], no enough memory", ret));
+  if (FLAGS_init_allocated_mem) {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "xpu memory FLAGS_init_allocated_mem is not implemented."));
+  }
+  VLOG(10) << "  pointer=" << p;
+  return p;
+#else
+  PADDLE_THROW(
+      common::errors::PermissionDenied("'XPUPlace' is not supported."));
+  return nullptr;
+#endif
+}
+
+template <>
+void Free<XPUPlace>(const XPUPlace &place, void *p, size_t size) {
+#ifdef PADDLE_WITH_XPU
+  VLOG(10) << "Free " << size << " bytes on " << Place(place);
+  VLOG(10) << "Free pointer=" << p << " on " << Place(place);
+
+  phi::backends::xpu::XPUDeviceGuard guard(place.device);
+  xpu_free(p);
+#else
+  PADDLE_THROW(
+      common::errors::PermissionDenied("'XPUPlace' is not supported."));
+#endif
+}
+
+template <>
+uint64_t Release<XPUPlace>(const XPUPlace &place) {
+#ifdef PADDLE_WITH_XPU
+  LOG(WARNING) << "Release XPU pool is not supported now, no action here.";
+#else
+  PADDLE_THROW(
+      common::errors::PermissionDenied("'XPUPlace' is not supported."));
+#endif
+  return -1;
+}
+
+template <>
+size_t Used<XPUPlace>(const XPUPlace &place) {
+#ifdef PADDLE_WITH_XPU
+  printf("Used func return 0 for XPUPlace\n");
+  return 0;
+#else
+  PADDLE_THROW(
+      common::errors::PermissionDenied("'XPUPlace' is not supported."));
+#endif
+}
+
+// For CUDA
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+class GPUBuddyAllocatorList {
+ private:
+  GPUBuddyAllocatorList()
+      : devices_(platform::GetSelectedDevices()), init_flags_(), allocators_() {
+    auto gpu_num = devices_.size();
+    allocators_.resize(gpu_num);
+    init_flags_.reserve(gpu_num);
+    for (size_t i = 0; i < gpu_num; ++i) {
+      init_flags_.emplace_back(new std::once_flag());
+    }
+  }
+
+  static GPUBuddyAllocatorList *CreateNewInstance() {
+    return new GPUBuddyAllocatorList();
+  }
+
+ public:
+  static GPUBuddyAllocatorList *Instance() {
+    static auto *instance = CreateNewInstance();
+    return instance;
+  }
+
+  BuddyAllocator *Get(int gpu_id) {
+    auto pos = std::distance(
+        devices_.begin(), std::find(devices_.begin(), devices_.end(), gpu_id));
+    PADDLE_ENFORCE_LT(pos,
+                      devices_.size(),
+                      common::errors::OutOfRange(
+                          "The index exceeds the size of devices, the size of "
+                          "devices is %d, the index is %d",
+                          devices_.size(),
+                          pos));
+
+    std::call_once(*init_flags_[pos], [this, pos] {
+      platform::SetDeviceId(devices_[pos]);
+      allocators_[pos] = std::make_unique<BuddyAllocator>(
+          std::unique_ptr<detail::SystemAllocator>(
+              new detail::GPUAllocator(devices_[pos])),
+          platform::GpuMinChunkSize(),
+          platform::GpuMaxChunkSize());
+      VLOG(10) << "\n\nNOTE:\n"
+               << "You can set GFlags environment variable "
+               << "'FLAGS_fraction_of_gpu_memory_to_use' "
+               << "or 'FLAGS_initial_gpu_memory_in_mb' "
+               << "or 'FLAGS_reallocate_gpu_memory_in_mb' "
+               << "to change the memory size for GPU usage.\n"
+               << "Current 'FLAGS_fraction_of_gpu_memory_to_use' value is "
+               << FLAGS_fraction_of_gpu_memory_to_use
+               << ". Current 'FLAGS_initial_gpu_memory_in_mb' value is "
+               << FLAGS_initial_gpu_memory_in_mb
+               << ". Current 'FLAGS_reallocate_gpu_memory_in_mb' value is "
+               << FLAGS_reallocate_gpu_memory_in_mb << "\n\n";
+    });
+
+    return allocators_[pos].get();
+  }
+
+ private:
+  std::vector<int> devices_;
+  std::vector<std::unique_ptr<std::once_flag>> init_flags_;
+  std::vector<std::unique_ptr<BuddyAllocator>> allocators_;
+};
+
+BuddyAllocator *GetGPUBuddyAllocator(int gpu_id) {
+  return GPUBuddyAllocatorList::Instance()->Get(gpu_id);
+}
+#endif
+
+template <>
+size_t Used<GPUPlace>(const GPUPlace &place) {
+#if (defined PADDLE_WITH_CUDA || defined PADDLE_WITH_HIP)
+  return GetGPUBuddyAllocator(place.device)->Used();
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'CUDAPlace' is not supported in CPU only device."));
+#endif
+}
+
+template <>
+void *Alloc<GPUPlace>(const GPUPlace &place, size_t size) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  auto *buddy_allocator = GetGPUBuddyAllocator(place.device);
+  auto *ptr = buddy_allocator->Alloc(size);
+  if (ptr == nullptr) {
+    platform::CUDADeviceGuard guard(place.device);
+    size_t avail, total;
+    platform::GpuMemoryUsage(&avail, &total);
+    PADDLE_THROW(common::errors::ResourceExhausted(
+        "Cannot allocate %s in GPU %d, available %s, total %s, GpuMinChunkSize "
+        "%s, GpuMaxChunkSize %s, GPU memory used: %s.",
+        string::HumanReadableSize(size),
+        place.device,
+        string::HumanReadableSize(avail),
+        string::HumanReadableSize(total),
+        string::HumanReadableSize(buddy_allocator->GetMinChunkSize()),
+        string::HumanReadableSize(buddy_allocator->GetMaxChunkSize()),
+        string::HumanReadableSize(Used<GPUPlace>(place))));
+  } else {
+    if (FLAGS_init_allocated_mem) {
+#ifdef PADDLE_WITH_HIP
+      hipMemset(ptr, 0xEF, size);
+#else
+      cudaMemset(ptr, 0xEF, size);
+#endif
+    }
+  }
+  return ptr;
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'CUDAPlace' is not supported in CPU only device."));
+#endif
+}
+
+template <>
+void Free<GPUPlace>(const GPUPlace &place, void *p, size_t size) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  GetGPUBuddyAllocator(place.device)->Free(p);
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'CUDAPlace' is not supported in CPU only device."));
+#endif
+}
+
+template <>
+uint64_t Release<GPUPlace>(const GPUPlace &place) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  return GetGPUBuddyAllocator(place.device)->Release();
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'CUDAPlace' is not supported in CPU only device."));
+#endif
+}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+BuddyAllocator *GetCUDAPinnedBuddyAllocator() {
+  static std::once_flag init_flag;
+  static BuddyAllocator *ba = nullptr;
+
+  std::call_once(init_flag, []() {
+    ba = new BuddyAllocator(std::unique_ptr<detail::SystemAllocator>(
+                                new detail::CUDAPinnedAllocator),
+                            phi::backends::cpu::CUDAPinnedMinChunkSize(),
+                            phi::backends::cpu::CUDAPinnedMaxChunkSize());
+  });
+
+  return ba;
+}
+#endif
+
+template <>
+size_t Used<GPUPinnedPlace>(const GPUPinnedPlace &place) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  return GetCUDAPinnedBuddyAllocator()->Used();
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'CUDAPinnedPlace' is not supported in CPU only device."));
+#endif
+}
+
+template <>
+void *Alloc<GPUPinnedPlace>(const GPUPinnedPlace &place, size_t size) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  VLOG(10) << "Allocate " << size << " bytes on " << Place(place);
+  auto *buddy_allocator = GetCUDAPinnedBuddyAllocator();
+  void *ptr = buddy_allocator->Alloc(size);
+
+  if (ptr == nullptr) {
+    LOG(WARNING) << "cudaHostAlloc Cannot allocate " << size
+                 << " bytes in CUDAPinnedPlace";
+  } else if (FLAGS_init_allocated_mem) {
+    memset(ptr, 0xEF, size);
+  }
+  return ptr;
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'CUDAPinnedPlace' is not supported in CPU only device."));
+#endif
+}
+
+template <>
+void Free<GPUPinnedPlace>(const GPUPinnedPlace &place, void *p, size_t size) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  VLOG(10) << "Free " << size << " bytes on " << Place(place);
+  GetCUDAPinnedBuddyAllocator()->Free(p);
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'CUDAPinnedPlace' is not supported in CPU only device."));
+#endif
+}
+
+template <>
+uint64_t Release<GPUPinnedPlace>(const GPUPinnedPlace &place) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  VLOG(10) << "Release on " << Place(place);
+  return GetCUDAPinnedBuddyAllocator()->Release();
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'CUDAPinnedPlace' is not supported in CPU only device."));
+#endif
+}
+
+// For XPUPinnedPlace
+#if defined(PADDLE_WITH_XPU)
+BuddyAllocator *GetXPUPinnedBuddyAllocator() {
+  static std::once_flag init_flag;
+  static BuddyAllocator *ba = nullptr;
+
+  std::call_once(init_flag, []() {
+    ba = new BuddyAllocator(std::unique_ptr<detail::SystemAllocator>(
+                                new detail::XPUPinnedAllocator),
+                            phi::backends::cpu::CUDAPinnedMinChunkSize(),
+                            phi::backends::cpu::CUDAPinnedMaxChunkSize());
+  });
+
+  return ba;
+}
+#endif
+
+template <>
+size_t Used<XPUPinnedPlace>(const XPUPinnedPlace &place) {
+#if defined(PADDLE_WITH_XPU)
+  return GetXPUPinnedBuddyAllocator()->Used();
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'XPUPinnedPlace' is not supported in CPU only device."));
+#endif
+}
+
+template <>
+void *Alloc<XPUPinnedPlace>(const XPUPinnedPlace &place, size_t size) {
+#if defined(PADDLE_WITH_XPU)
+  VLOG(10) << "Allocate " << size << " bytes on " << Place(place);
+  auto *buddy_allocator = GetXPUPinnedBuddyAllocator();
+  void *ptr = buddy_allocator->Alloc(size);
+
+  if (ptr == nullptr) {
+    LOG(WARNING) << "cudaHostAlloc Cannot allocate " << size
+                 << " bytes in XPUPinnedPlace";
+  } else if (FLAGS_init_allocated_mem) {
+    memset(ptr, 0xEF, size);
+  }
+  return ptr;
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'XPUPinnedPlace' is not supported in CPU only device."));
+#endif
+}
+
+template <>
+void Free<XPUPinnedPlace>(const XPUPinnedPlace &place, void *p, size_t size) {
+#if defined(PADDLE_WITH_XPU)
+  VLOG(10) << "Free " << size << " bytes on " << Place(place);
+  GetXPUPinnedBuddyAllocator()->Free(p);
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'XPUPinnedPlace' is not supported in CPU only device."));
+#endif
+}
+
+template <>
+uint64_t Release<XPUPinnedPlace>(const XPUPinnedPlace &place) {
+#if defined(PADDLE_WITH_XPU)
+  VLOG(10) << "Release on " << Place(place);
+  return GetXPUPinnedBuddyAllocator()->Release();
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'XPUPinnedPlace' is not supported in CPU only device."));
+#endif
+}
+
+// For CustomDevice
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+class BuddyAllocatorList {
+ private:
+  explicit BuddyAllocatorList(const std::string &device_type)
+      : device_type_(device_type) {
+    auto devices = phi::DeviceManager::GetSelectedDeviceList(device_type);
+    for (auto dev_id : devices) {
+      init_flags_[dev_id] = std::make_unique<std::once_flag>();
+    }
+  }
+
+  static BuddyAllocatorList *CreateNewInstance(const std::string &device_type) {
+    return new BuddyAllocatorList(device_type);
+  }
+
+ public:
+  static BuddyAllocatorList *Instance(const std::string &device_type) {
+    // DeviceType -> AllocatorList
+    static std::unordered_map<std::string, BuddyAllocatorList *> pool;
+    if (pool.find(device_type) == pool.end()) {
+      pool[device_type] = CreateNewInstance(device_type);
+    }
+    return pool[device_type];
+  }
+
+  BuddyAllocator *Get(int dev_id) {
+    PADDLE_ENFORCE_NE(init_flags_.find(dev_id),
+                      init_flags_.end(),
+                      common::errors::OutOfRange(
+                          "Cannot find %s %d, please check visible devices.",
+                          device_type_,
+                          dev_id));
+
+    std::call_once(*init_flags_[dev_id], [this, dev_id] {
+      phi::DeviceManager::SetDevice(device_type_, dev_id);
+      CustomPlace place(device_type_, dev_id);
+
+      VLOG(10) << "Init BuddyAllocator on " << place
+               << " with GetExtraPaddingSize "
+               << phi::DeviceManager::GetExtraPaddingSize(place);
+      allocators_[dev_id] = std::make_unique<BuddyAllocator>(
+          std::unique_ptr<detail::SystemAllocator>(
+              new detail::CustomAllocator(device_type_, dev_id)),
+          phi::DeviceManager::GetMinChunkSize(place),
+          phi::DeviceManager::GetMaxChunkSize(place),
+          phi::DeviceManager::GetExtraPaddingSize(place),
+          device_type_);
+    });
+
+    return allocators_[dev_id].get();
+  }
+
+ private:
+  std::string device_type_;
+  std::unordered_map<size_t, std::unique_ptr<std::once_flag>> init_flags_;
+  std::unordered_map<size_t, std::unique_ptr<BuddyAllocator>> allocators_;
+};
+
+BuddyAllocator *GetBuddyAllocator(const Place &place) {
+  VLOG(10) << "GetBuddyAllocator place = " << place;
+  if (phi::is_custom_place(place)) {
+    return BuddyAllocatorList::Instance(phi::PlaceHelper::GetDeviceType(place))
+        ->Get(phi::PlaceHelper::GetDeviceId(place));
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument("place must be CustomPlace"));
+  }
+}
+#endif
+
+template <>
+void *Alloc<CustomPlace>(const CustomPlace &place, size_t size) {
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  VLOG(10) << "Allocate " << size << " bytes on " << Place(place);
+  auto *buddy_allocator = GetBuddyAllocator(place);
+  auto *ptr = buddy_allocator->Alloc(size);
+
+  if (ptr == nullptr) {
+    phi::DeviceGuard guard(place);
+    size_t avail, total;
+    phi::DeviceManager::MemoryStats(place, &total, &avail);
+    PADDLE_THROW(common::errors::ResourceExhausted(
+        "Cannot allocate %s in %s:%d, available %s, total %s, used "
+        "%s. ",
+        string::HumanReadableSize(size),
+        place.GetDeviceType(),
+        place.device,
+        string::HumanReadableSize(avail),
+        string::HumanReadableSize(total),
+        string::HumanReadableSize(total - avail)));
+  } else {
+    if (FLAGS_init_allocated_mem) {
+      phi::DeviceManager::GetDeviceWithPlace(place)->MemorySet(ptr, 0xEF, size);
+    }
+  }
+  VLOG(10) << "  pointer=" << ptr;
+  return ptr;
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'CustomPlace' is not supported in CPU only device."));
+#endif
+}
+
+template <>
+void Free<CustomPlace>(const CustomPlace &place, void *p, size_t size) {
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  VLOG(10) << "Free pointer=" << p << " on " << Place(place);
+  if (phi::DeviceManager::HasDeviceType(place.GetDeviceType())) {
+    GetBuddyAllocator(place)->Free(p);
+  }
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'CustomPlace' is not supported in CPU only device."));
+#endif
+}
+
+template <>
+uint64_t Release<CustomPlace>(const CustomPlace &place) {
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  return GetBuddyAllocator(place)->Release();
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'CustomPlace' is not supported in CPU only device."));
+#endif
+}
+
+template <>
+size_t Used<CustomPlace>(const CustomPlace &place) {
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+  return GetBuddyAllocator(place)->Used();
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'CustomPlace' is not supported in CPU only device."));
+#endif
+}
+
+struct AllocVisitor {
+  using argument_type = const Place;
+  using result_type = void *;
+  inline explicit AllocVisitor(size_t size) : size_(size) {}
+
+  template <typename Place>
+  inline void *operator()(const Place &place) const {
+    return Alloc<Place>(place, size_);
+  }
+
+ private:
+  size_t size_;
+};
+
+struct FreeVisitor {
+  using argument_type = const Place;
+  using result_type = void;
+  inline explicit FreeVisitor(void *ptr, size_t size)
+      : ptr_(ptr), size_(size) {}
+
+  template <typename Place>
+  inline void operator()(const Place &place) const {
+    Free<Place>(place, ptr_, size_);
+  }
+
+ private:
+  void *ptr_;
+  size_t size_;
+};
+
+struct ReleaseVisitor {
+  using argument_type = const Place;
+  using result_type = uint64_t;
+  template <typename Place>
+  inline uint64_t operator()(const Place &place) const {
+    return Release<Place>(place);
+  }
+};
+
+size_t Usage::operator()(const CPUPlace &cpu) const { return Used(cpu); }
+
+size_t Usage::operator()(const GPUPlace &gpu) const {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  return Used(gpu);
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'CUDAPlace' is not supported in CPU only device."));
+#endif
+}
+
+size_t Usage::operator()(const GPUPinnedPlace &cuda_pinned) const {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  return Used(cuda_pinned);
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'CUDAPinnedPlace' is not supported in CPU only device."));
+#endif
+}
+
+size_t Usage::operator()(const XPUPlace &xpu) const {
+#if defined(PADDLE_WITH_XPU)
+  return Used(xpu);
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'XPUPlace' is not supported in CPU only device."));
+#endif
+}
+
+size_t Usage::operator()(const XPUPinnedPlace &xpu_pinned) const {
+#if defined(PADDLE_WITH_XPU)
+  return Used(xpu_pinned);
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'XPUPinnedPlace' is not supported in CPU only device."));
+#endif
+}
+
+}  // namespace paddle::memory::legacy
+
+namespace paddle::memory::allocation {
+
+phi::Allocation *NaiveBestFitAllocator::AllocateImpl(size_t size) {
+  VLOG(10) << "NaiveBestFitAllocator::AllocateImpl: place_ = " << place_;
+  void *ptr = phi::VisitPlace(place_, legacy::AllocVisitor(size));
+  auto *tmp_alloc = new Allocation(ptr, size, place_);
+  return tmp_alloc;
+}
+
+void NaiveBestFitAllocator::FreeImpl(phi::Allocation *allocation) {
+  phi::VisitPlace(allocation->place(),
+                  legacy::FreeVisitor(allocation->ptr(), allocation->size()));
+  delete allocation;
+}
+
+uint64_t NaiveBestFitAllocator::ReleaseImpl(const Place &place) {
+  return phi::VisitPlace(place, legacy::ReleaseVisitor());
+}
+
+}  // namespace paddle::memory::allocation

paddle/phi/core/memory/allocation/naive_best_fit_allocator.h

@@ -0,0 +1,48 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include <stdint.h>
+
+#include <algorithm>
+#include <mutex>  // NOLINT
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class PADDLE_API NaiveBestFitAllocator : public Allocator {
+ public:
+  explicit NaiveBestFitAllocator(const Place &p) : place_(p) {}
+
+  bool IsAllocThreadSafe() const override { return true; }
+
+ protected:
+  phi::Allocation *AllocateImpl(size_t size) override;
+  void FreeImpl(phi::Allocation *allocation) override;
+  uint64_t ReleaseImpl(const Place &place) override;
+
+ private:
+  Place place_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/pinned_allocator.cc

@@ -0,0 +1,48 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/pinned_allocator.h"
+
+#include "paddle/phi/core/memory/stats.h"
+#include "paddle/phi/core/platform/profiler/mem_tracing.h"
+namespace paddle::memory::allocation {
+bool CPUPinnedAllocator::IsAllocThreadSafe() const { return true; }
+void CPUPinnedAllocator::FreeImpl(phi::Allocation *allocation) {
+#ifdef PADDLE_WITH_HIP
+  PADDLE_ENFORCE_GPU_SUCCESS(hipHostFree(allocation->ptr()));
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaFreeHost(allocation->ptr()));
+#endif
+  VLOG(10) << "cudaFreeHost " << allocation->ptr();
+  HOST_MEMORY_STAT_UPDATE(Reserved, 0, -allocation->size());
+  platform::RecordMemEvent(allocation->ptr(),
+                           allocation->place(),
+                           allocation->size(),
+                           phi::TracerMemEventType::ReservedFree);
+  delete allocation;
+}
+phi::Allocation *CPUPinnedAllocator::AllocateImpl(size_t size) {
+  void *ptr;
+#ifdef PADDLE_WITH_HIP
+  PADDLE_ENFORCE_GPU_SUCCESS(hipHostMalloc(&ptr, size, hipHostMallocPortable));
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaHostAlloc(&ptr, size, cudaHostAllocPortable));
+#endif
+  VLOG(10) << "cudaHostAlloc " << size << " " << ptr;
+  HOST_MEMORY_STAT_UPDATE(Reserved, 0, size);
+  platform::RecordMemEvent(
+      ptr, GPUPinnedPlace(), size, phi::TracerMemEventType::ReservedAllocate);
+  return new Allocation(ptr, size, GPUPinnedPlace());
+}
+}  // namespace paddle::memory::allocation

paddle/phi/core/memory/allocation/pinned_allocator.h

@@ -0,0 +1,34 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include "paddle/phi/core/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// Allocator uses `cudaHostAlloc`
+class CPUPinnedAllocator : public Allocator {
+ public:
+  bool IsAllocThreadSafe() const override;
+
+ protected:
+  void FreeImpl(phi::Allocation *allocation) override;
+  phi::Allocation *AllocateImpl(size_t size) override;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/retry_allocator.cc

@@ -0,0 +1,119 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/retry_allocator.h"
+#include "paddle/common/flags.h"
+
+#include "glog/logging.h"
+
+COMMON_DECLARE_int64(offload_retry_times);
+
+namespace paddle::memory::allocation {
+
+static std::function<size_t(Place, size_t)> g_oom_callback;
+
+void RegisterOOMCallback(std::function<size_t(Place, size_t)> callback) {
+  g_oom_callback = std::move(callback);
+}
+
+class WaitedAllocateSizeGuard {
+ public:
+  WaitedAllocateSizeGuard(std::atomic<size_t>* waited_size,
+                          size_t requested_size)
+      : waited_size_(waited_size), requested_size_(requested_size) {
+    waited_size_->fetch_add(requested_size_, std::memory_order_relaxed);
+  }
+
+  ~WaitedAllocateSizeGuard() {
+    waited_size_->fetch_sub(requested_size_, std::memory_order_relaxed);
+  }
+
+ private:
+  std::atomic<size_t>* waited_size_;
+  size_t requested_size_;
+};
+
+void RetryAllocator::FreeImpl(phi::Allocation* allocation) {
+  // Delete underlying allocation first.
+  size_t size = allocation->size();
+  underlying_allocator_->Free(allocation);
+  if (UNLIKELY(waited_allocate_size_)) {
+    VLOG(10) << "Free " << size
+             << " bytes and notify all waited threads, "
+                "where waited_allocate_size_ = "
+             << waited_allocate_size_;
+    cv_.notify_all();
+  }
+}
+
+phi::Allocation* RetryAllocator::AllocateImpl(size_t size) {
+  auto alloc_func = [&, this]() {
+    return underlying_allocator_->Allocate(size).release();
+  };
+  // In fact, we can unify the code of allocation success and failure
+  // But it would add lock even when allocation success at the first time
+  try {
+    if (FLAGS_offload_retry_times <= 0 || g_oom_callback == nullptr) {
+      return alloc_func();
+    } else {
+      bool has_offloaded = true;
+      for (int64_t i = 0; i < FLAGS_offload_retry_times && has_offloaded; ++i) {
+        try {
+          return alloc_func();
+        } catch (BadAlloc&) {
+          VLOG(10) << "Allocation " << size << " on " << place_
+                   << " failed, try to run OOM callback " << i;
+          has_offloaded = (g_oom_callback(place_, size) > 0);
+        }
+      }
+      return alloc_func();
+    }
+  } catch (BadAlloc&) {
+    {
+      WaitedAllocateSizeGuard guard(&waited_allocate_size_, size);
+      VLOG(10) << "Allocation failed when allocating " << size
+               << " bytes, waited_allocate_size_ = " << waited_allocate_size_;
+      // We can just write allocation retry inside the predicate function of
+      // wait_until. But it needs to acquire the lock when executing predicate
+      // function. For better performance, we use loop here
+      auto end_time = std::chrono::high_resolution_clock::now() + retry_time_;
+      auto wait_until = [&, this] {
+        std::unique_lock<std::mutex> lock(mutex_);
+        return cv_.wait_until(lock, end_time);
+      };
+
+      size_t retry_time = 0;
+      while (wait_until() != std::cv_status::timeout) {
+        try {
+          return alloc_func();
+        } catch (BadAlloc&) {
+          // do nothing when it is not timeout
+          ++retry_time;
+          VLOG(10) << "Allocation failed when retrying " << retry_time
+                   << " times when allocating " << size
+                   << " bytes. Wait still.";
+        } catch (...) {
+          throw;
+        }
+      }
+    }
+    VLOG(10) << "Allocation failed because of timeout when allocating " << size
+             << " bytes.";
+    return alloc_func();  // If timeout, try last allocation request.
+  } catch (...) {
+    throw;
+  }
+}
+
+}  // namespace paddle::memory::allocation

paddle/phi/core/memory/allocation/retry_allocator.h

@@ -0,0 +1,82 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <atomic>              // NOLINT
+#include <chrono>              // NOLINT
+#include <condition_variable>  // NOLINT
+#include <memory>
+#include <mutex>  // NOLINT
+#include <utility>
+
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/mem_visitor.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+PADDLE_API void RegisterOOMCallback(
+    std::function<size_t(Place, size_t)> callback);
+
+class PADDLE_API RetryAllocator : public Allocator {
+ public:
+  RetryAllocator(std::shared_ptr<Allocator> allocator,
+                 Place place,
+                 size_t retry_ms)
+      : underlying_allocator_(std::move(allocator)),
+        place_(place),
+        retry_time_(retry_ms) {
+    PADDLE_ENFORCE_NOT_NULL(
+        underlying_allocator_,
+        common::errors::InvalidArgument(
+            "Underlying allocator of RetryAllocator is NULL"));
+    PADDLE_ENFORCE_EQ(
+        underlying_allocator_->IsAllocThreadSafe(),
+        true,
+        common::errors::PreconditionNotMet(
+            "Underlying allocator of RetryAllocator is not thread-safe"));
+  }
+  std::shared_ptr<Allocator>& GetUnderLyingAllocator() {
+    return underlying_allocator_;
+  }
+  void Accept(AllocatorVisitor* visitor) override { visitor->Visit(this); }
+
+  bool IsAllocThreadSafe() const override { return true; }
+
+ protected:
+  void FreeImpl(phi::Allocation* allocation) override;
+  phi::Allocation* AllocateImpl(size_t size) override;
+  uint64_t ReleaseImpl(const Place& place) override {
+    return underlying_allocator_->Release(place);
+  }
+  size_t CompactImpl(const Place& place) override {
+    return underlying_allocator_->Compact(place);
+  }
+
+ private:
+  std::shared_ptr<Allocator> underlying_allocator_;
+  Place place_;
+  std::chrono::milliseconds retry_time_;
+  std::mutex mutex_;
+  std::condition_variable cv_;
+
+  std::atomic<size_t> waited_allocate_size_{0};
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/spin_lock.h

@@ -0,0 +1,67 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <atomic>
+#if defined(_M_X64) || defined(__x86_64__) || defined(_M_IX86) || \
+    defined(__i386__)
+#define __PADDLE_x86__
+#include <immintrin.h>
+#endif
+#include <thread>
+
+#include "paddle/common/macros.h"
+
+namespace paddle {
+namespace memory {
+static inline void CpuRelax() {
+#if defined(__PADDLE_x86__)
+  _mm_pause();
+#endif
+}
+
+class SpinLock {
+ public:
+  SpinLock() : mlock_(false) {}
+
+  void lock() {
+    for (;;) {
+      if (!mlock_.exchange(true, std::memory_order_acquire)) {
+        break;
+      }
+      constexpr int kMaxLoop = 32;
+      for (int loop = 1; mlock_.load(std::memory_order_relaxed);) {
+        if (loop <= kMaxLoop) {
+          for (int i = 1; i <= loop; ++i) {
+            CpuRelax();
+          }
+          loop *= 2;
+        } else {
+          std::this_thread::yield();
+        }
+      }
+    }
+  }
+
+  void unlock() { mlock_.store(false, std::memory_order_release); }
+
+  DISABLE_COPY_AND_ASSIGN(SpinLock);
+
+ private:
+  std::atomic<bool> mlock_;
+};
+
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/stat_allocator.h

@@ -0,0 +1,87 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/mem_visitor.h"
+#include "paddle/phi/core/memory/stats.h"
+#include "paddle/phi/core/platform/profiler/mem_tracing.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class StatAllocator : public Allocator {
+ public:
+  explicit StatAllocator(std::shared_ptr<Allocator> underlying_allocator)
+      : underlying_allocator_(std::move(underlying_allocator)) {}
+
+  bool IsAllocThreadSafe() const override { return true; }
+  void Accept(AllocatorVisitor* visitor) override { visitor->Visit(this); }
+  std::shared_ptr<Allocator>& GetUnderLyingAllocator() {
+    return underlying_allocator_;
+  }
+
+ protected:
+  void FreeImpl(phi::Allocation* allocation) override {
+    if (phi::is_cpu_place(allocation->place()) ||
+        phi::is_pinned_place(allocation->place())) {
+      HOST_MEMORY_STAT_UPDATE(
+          Allocated, allocation->place().GetDeviceId(), -allocation->size());
+    } else {
+      DEVICE_MEMORY_STAT_UPDATE(
+          Allocated, allocation->place().GetDeviceId(), -allocation->size());
+    }
+    platform::RecordMemEvent(allocation->ptr(),
+                             allocation->place(),
+                             allocation->size(),
+                             phi::TracerMemEventType::Free);
+    underlying_allocator_->Free(allocation);
+  }
+
+  phi::Allocation* AllocateImpl(size_t size) override {
+    phi::Allocator::AllocationPtr allocation =
+        underlying_allocator_->Allocate(size);
+
+    const Place& place = allocation->place();
+    if (phi::is_cpu_place(place) || phi::is_pinned_place(place)) {
+      HOST_MEMORY_STAT_UPDATE(
+          Allocated, place.GetDeviceId(), allocation->size());
+    } else {
+      DEVICE_MEMORY_STAT_UPDATE(
+          Allocated, place.GetDeviceId(), allocation->size());
+    }
+    platform::RecordMemEvent(allocation->ptr(),
+                             allocation->place(),
+                             allocation->size(),
+                             phi::TracerMemEventType::Allocate);
+    return allocation.release();
+  }
+
+  uint64_t ReleaseImpl(const Place& place) override {
+    return underlying_allocator_->Release(place);
+  }
+
+  size_t CompactImpl(const Place& place) override {
+    return underlying_allocator_->Compact(place);
+  }
+
+ private:
+  std::shared_ptr<Allocator> underlying_allocator_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/stream_safe_cuda_allocator.cc

@@ -0,0 +1,304 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/stream_safe_cuda_allocator.h"
+#include <thread>
+#include "glog/logging.h"
+
+#include "paddle/phi/api/profiler/event_tracing.h"
+#include "paddle/phi/backends/gpu/gpu_info.h"
+
+#if defined(PADDLE_WITH_CUDA)
+#include "paddle/phi/backends/gpu/cuda/cuda_graph.h"
+#elif defined(PADDLE_WITH_HIP)
+#include "paddle/phi/backends/gpu/rocm/hip_graph.h"
+#endif
+
+namespace paddle::memory::allocation {
+
+StreamSafeCUDAAllocation::StreamSafeCUDAAllocation(
+    DecoratedAllocationPtr underlying_allocation,
+    gpuStream_t owning_stream,
+    StreamSafeCUDAAllocator* allocator)
+    : Allocation(underlying_allocation->ptr(),
+                 underlying_allocation->base_ptr(),
+                 underlying_allocation->size(),
+                 underlying_allocation->place()),
+      underlying_allocation_(std::move(underlying_allocation)),
+      owning_stream_(owning_stream),
+      allocator_(allocator->shared_from_this()) {}
+
+bool StreamSafeCUDAAllocation::RecordStream(gpuStream_t stream) {
+  VLOG(8) << "Try record stream " << stream << " for address " << ptr();
+  if (stream == owning_stream_) {
+    return false;
+  }
+
+  std::call_once(once_flag_,
+                 [this] { phi::backends::gpu::SetDeviceId(place_.device); });
+
+  std::lock_guard<SpinLock> lock_guard(outstanding_event_map_lock_);
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  if (UNLIKELY(phi::backends::gpu::CUDAGraph::IsThisThreadCapturing())) {
+    graph_capturing_stream_set_.insert(stream);
+    return true;
+  }
+#endif
+
+  RecordStreamWithNoGraphCapturing(stream);
+  RecordGraphCapturingStreams();
+  return true;
+}
+
+void StreamSafeCUDAAllocation::EraseStream(gpuStream_t stream) {
+  VLOG(8) << "Try remove stream " << stream << " for address " << ptr();
+  std::lock_guard<SpinLock> lock_guard(outstanding_event_map_lock_);
+  auto it = outstanding_event_map_.find(stream);
+  if (it == outstanding_event_map_.end()) {
+    return;
+  }
+
+#ifdef PADDLE_WITH_CUDA
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaEventDestroy(it->second));
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(hipEventDestroy(it->second));
+#endif
+  outstanding_event_map_.erase(it);
+}
+
+bool StreamSafeCUDAAllocation::CanBeFreed() {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  if (UNLIKELY(phi::backends::gpu::CUDAGraph::IsThisThreadCapturing())) {
+    return graph_capturing_stream_set_.empty() &&
+           outstanding_event_map_.empty();
+  }
+#endif
+
+  std::call_once(once_flag_,
+                 [this] { phi::backends::gpu::SetDeviceId(place_.device); });
+
+  RecordGraphCapturingStreams();
+
+  for (auto it = outstanding_event_map_.begin();
+       it != outstanding_event_map_.end();
+       ++it) {
+    gpuEvent_t& event = it->second;
+#ifdef PADDLE_WITH_CUDA
+    gpuError_t err = cudaEventQuery(event);
+    if (err == cudaErrorNotReady) {
+      VLOG(9) << "Event " << event << " for " << ptr() << " is not completed";
+      // Erase the completed event before "it"
+      outstanding_event_map_.erase(outstanding_event_map_.begin(), it);
+      return false;
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(err);
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaEventDestroy(event));
+#else
+    gpuError_t err = hipEventQuery(event);
+    if (err == hipErrorNotReady) {
+      VLOG(9) << "Event " << event << " for " << ptr() << " is not completed";
+      // Erase the completed event before "it"
+      outstanding_event_map_.erase(outstanding_event_map_.begin(), it);
+      return false;
+    }
+    PADDLE_ENFORCE_GPU_SUCCESS(err);
+    PADDLE_ENFORCE_GPU_SUCCESS(hipEventDestroy(event));
+#endif
+    VLOG(8) << "Destroy event " << event;
+  }
+  return true;
+}
+
+gpuStream_t StreamSafeCUDAAllocation::GetOwningStream() const {
+  return owning_stream_;
+}
+
+void StreamSafeCUDAAllocation::RecordGraphCapturingStreams() {
+  for (gpuStream_t stream : graph_capturing_stream_set_) {
+    RecordStreamWithNoGraphCapturing(stream);
+  }
+  graph_capturing_stream_set_.clear();
+}
+
+void StreamSafeCUDAAllocation::RecordStreamWithNoGraphCapturing(
+    gpuStream_t stream) {
+  gpuEvent_t record_event;
+  auto it = outstanding_event_map_.find(stream);
+  if (it == outstanding_event_map_.end()) {
+    gpuEvent_t new_event;
+#ifdef PADDLE_WITH_CUDA
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        cudaEventCreateWithFlags(&new_event, cudaEventDisableTiming));
+#else
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        hipEventCreateWithFlags(&new_event, hipEventDisableTiming));
+#endif
+    outstanding_event_map_[stream] = new_event;
+    record_event = new_event;
+    VLOG(9) << "Create a new event " << new_event;
+  } else {
+    record_event = it->second;
+    VLOG(9) << "Reuse event " << record_event;
+  }
+
+#ifdef PADDLE_WITH_CUDA
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaEventRecord(record_event, stream));
+#else
+  PADDLE_ENFORCE_GPU_SUCCESS(hipEventRecord(record_event, stream));
+#endif
+  VLOG(8) << "Record event " << record_event << " to stream " << stream;
+}
+
+StreamSafeCUDAAllocator::StreamSafeCUDAAllocator(
+    std::shared_ptr<Allocator> underlying_allocator,
+    GPUPlace place,
+    gpuStream_t default_stream,
+    bool in_cuda_graph_capturing)
+    : underlying_allocator_(std::move(underlying_allocator)),
+      place_(place),
+      default_stream_(default_stream),
+      in_cuda_graph_capturing_(in_cuda_graph_capturing) {
+  if (LIKELY(!in_cuda_graph_capturing)) {
+    std::lock_guard<SpinLock> lock_guard(allocator_map_lock_);
+    allocator_map_[place].emplace_back(this);
+  }
+}
+
+StreamSafeCUDAAllocator::~StreamSafeCUDAAllocator() {
+  if (LIKELY(!in_cuda_graph_capturing_)) {
+    std::lock_guard<SpinLock> lock_guard(allocator_map_lock_);
+    std::vector<StreamSafeCUDAAllocator*>& allocators = allocator_map_[place_];
+    allocators.erase(std::remove(allocators.begin(), allocators.end(), this),
+                     allocators.end());
+  }
+}
+
+bool StreamSafeCUDAAllocator::IsAllocThreadSafe() const { return true; }
+
+gpuStream_t StreamSafeCUDAAllocator::GetDefaultStream() const {
+  return default_stream_;
+}
+
+void StreamSafeCUDAAllocator::SetDefaultStream(gpuStream_t stream) {
+  default_stream_ = stream;
+}
+
+phi::Allocation* StreamSafeCUDAAllocator::AllocateImpl(size_t size) {
+  phi::RecordEvent record("StreamSafeCUDAAllocator::Allocate",
+                          phi::TracerEventType::UserDefined,
+                          9 /*level*/);
+  ProcessUnfreedAllocations();
+  VLOG(8) << "Try allocate " << size << " bytes";
+  AllocationPtr underlying_allocation;
+  try {
+    underlying_allocation = underlying_allocator_->Allocate(size);
+  } catch (BadAlloc&) {
+    VLOG(4) << "Allocation failed when allocating " << size << " bytes";
+    ReleaseImpl(place_);
+    try {
+      underlying_allocation = underlying_allocator_->Allocate(size);
+    } catch (...) {
+      VLOG(3)
+          << "Still allocation failed after release memory from all streams";
+      throw;
+    }
+  } catch (...) {
+    throw;
+  }
+  StreamSafeCUDAAllocation* allocation = new StreamSafeCUDAAllocation(
+      static_unique_ptr_cast<Allocation>(std::move(underlying_allocation)),
+      default_stream_,
+      this);
+  VLOG(8) << "Thread " << std::this_thread::get_id() << " Allocate "
+          << allocation->size() << " bytes at address " << allocation->ptr()
+          << "  , stream: " << default_stream_;
+  return allocation;
+}
+
+void StreamSafeCUDAAllocator::FreeImpl(phi::Allocation* allocation) {
+  phi::RecordEvent record("StreamSafeCUDAAllocator::Free",
+                          phi::TracerEventType::UserDefined,
+                          9 /*level*/);
+  StreamSafeCUDAAllocation* stream_safe_cuda_allocation =
+      static_cast<StreamSafeCUDAAllocation*>(allocation);
+
+  VLOG(8) << "Try free allocation " << stream_safe_cuda_allocation->ptr();
+  if (stream_safe_cuda_allocation->CanBeFreed()) {
+    VLOG(9) << "Directly delete allocation";
+    delete stream_safe_cuda_allocation;
+  } else {
+    VLOG(9) << "Put into unfreed_allocation list";
+    std::lock_guard<SpinLock> lock_guard(unfreed_allocation_lock_);
+    unfreed_allocations_.emplace_back(stream_safe_cuda_allocation);
+  }
+}
+
+uint64_t StreamSafeCUDAAllocator::ReleaseImpl(const Place& place) {
+  if (UNLIKELY(in_cuda_graph_capturing_)) {
+    VLOG(7) << "Memory release forbidden in CUDA Graph Capturing";
+    return 0;
+  }
+
+  std::lock_guard<SpinLock> lock_guard(allocator_map_lock_);
+  std::vector<StreamSafeCUDAAllocator*>& allocators = allocator_map_[place];
+  uint64_t released_size = 0;
+  for (StreamSafeCUDAAllocator* allocator : allocators) {
+    released_size += allocator->ProcessUnfreedAllocationsAndRelease();
+  }
+  VLOG(8) << "Release " << released_size << " bytes memory from all streams";
+  return released_size;
+}
+
+size_t StreamSafeCUDAAllocator::CompactImpl(const Place& place) {
+  std::lock_guard<SpinLock> lock_guard(allocator_map_lock_);
+  VLOG(4) << "enter StreamSafeCUDAAllocator compact!!";
+  std::vector<StreamSafeCUDAAllocator*>& allocators = allocator_map_[place];
+  size_t compact_free_size = 0;
+  for (StreamSafeCUDAAllocator* allocator : allocators) {
+    compact_free_size += allocator->underlying_allocator_->Compact(place_);
+  }
+  return compact_free_size;
+}
+
+void StreamSafeCUDAAllocator::ProcessUnfreedAllocations() {
+  // NOTE(Ruibiao): This condition is to reduce lock completion. It does not
+  // need to be thread-safe since here occasional misjudgments are permissible.
+  if (unfreed_allocations_.empty()) {
+    return;
+  }
+
+  std::lock_guard<SpinLock> lock_guard(unfreed_allocation_lock_);
+  for (auto it = unfreed_allocations_.begin();
+       it != unfreed_allocations_.end();) {
+    if ((*it)->CanBeFreed()) {
+      delete *it;
+      it = unfreed_allocations_.erase(it);
+    } else {
+      ++it;
+    }
+  }
+}
+
+uint64_t StreamSafeCUDAAllocator::ProcessUnfreedAllocationsAndRelease() {
+  ProcessUnfreedAllocations();
+  return underlying_allocator_->Release(place_);
+}
+
+thread_local std::once_flag StreamSafeCUDAAllocation::once_flag_;
+
+std::map<Place, std::vector<StreamSafeCUDAAllocator*>>
+    StreamSafeCUDAAllocator::allocator_map_;
+SpinLock StreamSafeCUDAAllocator::allocator_map_lock_;
+
+}  // namespace paddle::memory::allocation

paddle/phi/core/memory/allocation/stream_safe_cuda_allocator.h

@@ -0,0 +1,115 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <list>
+#include <map>
+#include <set>
+
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/allocation/spin_lock.h"
+#include "paddle/phi/core/memory/mem_visitor.h"
+
+#ifdef PADDLE_WITH_CUDA
+#include <cuda_runtime.h>
+#else
+#include <hip/hip_runtime.h>
+#endif
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class StreamSafeCUDAAllocator;
+
+class StreamSafeCUDAAllocation : public Allocation {
+ public:
+  StreamSafeCUDAAllocation(DecoratedAllocationPtr underlying_allocation,
+                           gpuStream_t owning_stream,
+                           StreamSafeCUDAAllocator *allocator);
+
+  bool RecordStream(gpuStream_t stream);
+  void EraseStream(gpuStream_t stream);
+  bool CanBeFreed();
+  gpuStream_t GetOwningStream() const;
+  void *ptr() const noexcept override { return underlying_allocation_->ptr(); }
+  size_t size() const noexcept override {
+    return underlying_allocation_->size();
+  }
+  const Place &place() const noexcept override {
+    return underlying_allocation_->place();
+  }
+
+ private:
+  thread_local static std::once_flag once_flag_;
+  void RecordGraphCapturingStreams();
+  void RecordStreamWithNoGraphCapturing(gpuStream_t stream);
+  DecoratedAllocationPtr underlying_allocation_;
+  std::set<gpuStream_t> graph_capturing_stream_set_;
+  std::map<gpuStream_t, gpuEvent_t> outstanding_event_map_;
+  gpuStream_t owning_stream_;
+  SpinLock outstanding_event_map_lock_;
+  // To compatible with CUDA Graph, hold the allocator shared_ptr so that
+  // Allocator will not deconstruct before Allocation
+  std::shared_ptr<Allocator> allocator_;
+};
+
+class StreamSafeCUDAAllocator
+    : public Allocator,
+      public std::enable_shared_from_this<StreamSafeCUDAAllocator> {
+ public:
+  StreamSafeCUDAAllocator(std::shared_ptr<Allocator> underlying_allocator,
+                          GPUPlace place,
+                          gpuStream_t default_stream,
+                          bool in_cuda_graph_capturing = false);
+  ~StreamSafeCUDAAllocator();
+
+  std::shared_ptr<Allocator> &GetUnderLyingAllocator() {
+    return underlying_allocator_;
+  }
+  std::vector<StreamSafeCUDAAllocator *> &GetAllocatorByPlace() {
+    return allocator_map_[place_];
+  }
+  bool IsAllocThreadSafe() const override;
+  gpuStream_t GetDefaultStream() const;
+  void SetDefaultStream(gpuStream_t stream);
+  void Accept(AllocatorVisitor *visitor) override { visitor->Visit(this); }
+
+ protected:
+  phi::Allocation *AllocateImpl(size_t size) override;
+  void FreeImpl(phi::Allocation *allocation) override;
+  uint64_t ReleaseImpl(const Place &place) override;
+  size_t CompactImpl(const Place &place) override;
+
+ private:
+  void ProcessUnfreedAllocations();
+  uint64_t ProcessUnfreedAllocationsAndRelease();
+
+  static std::map<Place, std::vector<StreamSafeCUDAAllocator *>> allocator_map_;
+  static SpinLock allocator_map_lock_;
+
+  std::shared_ptr<Allocator> underlying_allocator_;
+  GPUPlace place_;
+  gpuStream_t default_stream_;
+  std::list<StreamSafeCUDAAllocation *> unfreed_allocations_;
+  SpinLock unfreed_allocation_lock_;
+
+  bool in_cuda_graph_capturing_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/stream_safe_custom_device_allocator.cc

@@ -0,0 +1,239 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <thread>
+#include "glog/logging.h"
+
+#include "paddle/phi/api/profiler/event_tracing.h"
+#include "paddle/phi/backends/context_pool.h"
+#include "paddle/phi/core/memory/allocation/stream_safe_custom_device_allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+StreamSafeCustomDeviceAllocation::StreamSafeCustomDeviceAllocation(
+    DecoratedAllocationPtr underlying_allocation,
+    phi::stream::stream_t owning_stream,
+    StreamSafeCustomDeviceAllocator* allocator)
+    : Allocation(underlying_allocation->ptr(),
+                 underlying_allocation->base_ptr(),
+                 underlying_allocation->size(),
+                 underlying_allocation->place()),
+      underlying_allocation_(std::move(underlying_allocation)),
+      owning_stream_(std::move(owning_stream)),
+      allocator_(allocator->shared_from_this()) {}
+
+bool StreamSafeCustomDeviceAllocation::RecordStream(
+    phi::stream::stream_t stream) {
+  VLOG(8) << "Try record stream " << stream << " for address " << ptr();
+  if (stream == owning_stream_) {
+    return false;
+  }
+  std::call_once(once_flag_, [this] { phi::DeviceManager::SetDevice(place_); });
+  std::lock_guard<SpinLock> lock_guard(outstanding_event_map_lock_);
+
+  auto it = outstanding_event_map_.find(stream);
+  if (it == outstanding_event_map_.end()) {
+    outstanding_event_map_.insert(
+        {stream, std::make_shared<phi::event::Event>()});
+    outstanding_event_map_[stream]->Init(place());
+    VLOG(9) << "Create a new event "
+            << outstanding_event_map_[stream]->raw_event();
+  }
+  auto stream_wrapper = phi::stream::Stream(place(), stream);
+  VLOG(8) << "Record event " << outstanding_event_map_[stream]->raw_event()
+          << " to stream " << stream;
+  outstanding_event_map_[stream]->Record(&stream_wrapper);
+  return true;
+}
+
+void StreamSafeCustomDeviceAllocation::EraseStream(
+    phi::stream::stream_t stream) {
+  VLOG(8) << "Try remove stream " << stream << " for address " << ptr();
+  std::lock_guard<SpinLock> lock_guard(outstanding_event_map_lock_);
+  auto it = outstanding_event_map_.find(stream);
+  if (it == outstanding_event_map_.end()) {
+    return;
+  }
+  it->second->Destroy();
+  outstanding_event_map_.erase(it);
+}
+
+bool StreamSafeCustomDeviceAllocation::CanBeFreed() {
+  std::lock_guard<SpinLock> lock_guard(outstanding_event_map_lock_);
+  if (!phi::DeviceManager::HasDeviceType(place_.GetDeviceType())) {
+    return true;
+  }
+  std::call_once(once_flag_, [this] { phi::DeviceManager::SetDevice(place_); });
+  for (auto it = outstanding_event_map_.begin();
+       it != outstanding_event_map_.end();) {
+    auto& event = it->second;
+    if (!event->Query()) {
+      VLOG(9) << "Event " << event->raw_event() << " for " << ptr()
+              << " is not completed";
+      return false;
+    }
+    VLOG(8) << "Destroy event " << event->raw_event();
+    event->Destroy();
+    it = outstanding_event_map_.erase(it);
+  }
+  outstanding_event_map_.clear();
+  return true;
+}
+
+phi::stream::stream_t StreamSafeCustomDeviceAllocation::GetOwningStream()
+    const {
+  return owning_stream_;
+}
+
+void StreamSafeCustomDeviceAllocation::SetOwningStream(
+    phi::stream::stream_t s) {
+  owning_stream_ = s;
+}
+
+StreamSafeCustomDeviceAllocator::StreamSafeCustomDeviceAllocator(
+    std::shared_ptr<Allocator> underlying_allocator,
+    CustomPlace place,
+    phi::stream::stream_t default_stream)
+    : underlying_allocator_(std::move(underlying_allocator)),
+      place_(std::move(place)),
+      default_stream_(std::move(default_stream)) {
+  std::lock_guard<SpinLock> lock_guard(allocator_map_lock_);
+  allocator_map_[place_].emplace_back(this);
+}
+
+StreamSafeCustomDeviceAllocator::~StreamSafeCustomDeviceAllocator() {
+  std::lock_guard<SpinLock> lock_guard(allocator_map_lock_);
+  std::vector<StreamSafeCustomDeviceAllocator*>& allocators =
+      allocator_map_[place_];
+  allocators.erase(std::remove(allocators.begin(), allocators.end(), this),
+                   allocators.end());
+}
+
+phi::stream::stream_t StreamSafeCustomDeviceAllocator::GetDefaultStream()
+    const {
+  return default_stream_;
+}
+
+void StreamSafeCustomDeviceAllocator::SetDefaultStream(
+    phi::stream::stream_t stream) {
+  default_stream_ = stream;
+}
+
+phi::Allocation* StreamSafeCustomDeviceAllocator::AllocateImpl(size_t size) {
+  phi::RecordEvent record("StreamSafeCustomDeviceAllocator::Allocate",
+                          phi::TracerEventType::UserDefined,
+                          9 /*level*/);
+  ProcessUnfreedAllocations();
+  VLOG(8) << "Try allocate " << size << " bytes";
+  AllocationPtr underlying_allocation;
+  try {
+    underlying_allocation = underlying_allocator_->Allocate(size);
+  } catch (BadAlloc&) {
+    VLOG(4) << "Allocation failed when allocating " << size << " bytes";
+    ReleaseImpl(place_);
+    try {
+      underlying_allocation = underlying_allocator_->Allocate(size);
+    } catch (...) {
+      VLOG(3)
+          << "Still allocation failed after release memory from all streams";
+      throw;
+    }
+  } catch (...) {
+    throw;
+  }
+  StreamSafeCustomDeviceAllocation* allocation =
+      new StreamSafeCustomDeviceAllocation(
+          static_unique_ptr_cast<Allocation>(std::move(underlying_allocation)),
+          default_stream_,
+          this);
+  VLOG(8) << "Thread " << std::this_thread::get_id() << " Allocate "
+          << allocation->size() << " bytes at address " << allocation->ptr()
+          << "  , stream: " << default_stream_;
+  return allocation;
+}
+
+void StreamSafeCustomDeviceAllocator::FreeImpl(phi::Allocation* allocation) {
+  phi::RecordEvent record("StreamSafeCustomDeviceAllocator::Free",
+                          phi::TracerEventType::UserDefined,
+                          9 /*level*/);
+  StreamSafeCustomDeviceAllocation* stream_safe_cuda_allocation =
+      static_cast<StreamSafeCustomDeviceAllocation*>(allocation);
+
+  VLOG(8) << "Try free allocation " << stream_safe_cuda_allocation->ptr();
+  if (!stream_safe_cuda_allocation->GetOwningStream()) {
+    stream_safe_cuda_allocation->SetOwningStream(
+        default_stream_ ? default_stream_
+                        : reinterpret_cast<phi::CustomContext*>(
+                              phi::DeviceContextPool::Instance().Get(place_))
+                              ->stream());
+  }
+  if (stream_safe_cuda_allocation->CanBeFreed()) {
+    VLOG(9) << "Directly delete allocation";
+    delete stream_safe_cuda_allocation;
+  } else {
+    VLOG(9) << "Put into unfreed_allocation list";
+    std::lock_guard<SpinLock> lock_guard(unfreed_allocation_lock_);
+    unfreed_allocations_.emplace_back(stream_safe_cuda_allocation);
+  }
+}
+
+uint64_t StreamSafeCustomDeviceAllocator::ReleaseImpl(const Place& place) {
+  std::lock_guard<SpinLock> lock_guard(allocator_map_lock_);
+  std::vector<StreamSafeCustomDeviceAllocator*>& allocators =
+      allocator_map_[place];
+  uint64_t released_size = 0;
+  for (StreamSafeCustomDeviceAllocator* allocator : allocators) {
+    released_size += allocator->ProcessUnfreedAllocationsAndRelease();
+  }
+  VLOG(8) << "Release " << released_size << " bytes memory from all streams";
+  return released_size;
+}
+
+void StreamSafeCustomDeviceAllocator::ProcessUnfreedAllocations() {
+  // NOTE(Ruibiao): This condition is to reduce lock completion. It does not
+  // need to be thread-safe since here occasional misjudgments are
+  // permissible.
+  if (unfreed_allocations_.empty()) {
+    return;
+  }
+
+  std::lock_guard<SpinLock> lock_guard(unfreed_allocation_lock_);
+  for (auto it = unfreed_allocations_.begin();
+       it != unfreed_allocations_.end();) {
+    if ((*it)->CanBeFreed()) {
+      delete *it;
+      it = unfreed_allocations_.erase(it);
+    } else {
+      ++it;
+    }
+  }
+}
+
+uint64_t
+StreamSafeCustomDeviceAllocator::ProcessUnfreedAllocationsAndRelease() {
+  ProcessUnfreedAllocations();
+  return underlying_allocator_->Release(place_);
+}
+
+thread_local std::once_flag StreamSafeCustomDeviceAllocation::once_flag_;
+
+std::map<Place, std::vector<StreamSafeCustomDeviceAllocator*>>
+    StreamSafeCustomDeviceAllocator::allocator_map_;
+SpinLock StreamSafeCustomDeviceAllocator::allocator_map_lock_;
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/stream_safe_custom_device_allocator.h

@@ -0,0 +1,91 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <list>
+#include <map>
+#include <set>
+
+#include "paddle/phi/backends/device_manager.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/allocation/spin_lock.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class StreamSafeCustomDeviceAllocator;
+
+class StreamSafeCustomDeviceAllocation : public Allocation {
+ public:
+  StreamSafeCustomDeviceAllocation(DecoratedAllocationPtr underlying_allocation,
+                                   phi::stream::stream_t owning_stream,
+                                   StreamSafeCustomDeviceAllocator *allocator);
+
+  bool RecordStream(phi::stream::stream_t stream);
+  void EraseStream(phi::stream::stream_t stream);
+  bool CanBeFreed();
+  phi::stream::stream_t GetOwningStream() const;
+  void SetOwningStream(phi::stream::stream_t s);
+
+ private:
+  thread_local static std::once_flag once_flag_;
+  DecoratedAllocationPtr underlying_allocation_;
+  std::map<phi::stream::stream_t, std::shared_ptr<phi::event::Event>>
+      outstanding_event_map_;
+  phi::stream::stream_t owning_stream_;
+  SpinLock outstanding_event_map_lock_;
+  std::shared_ptr<Allocator> allocator_;
+  bool will_be_freed_{false};
+};
+
+class StreamSafeCustomDeviceAllocator
+    : public Allocator,
+      public std::enable_shared_from_this<StreamSafeCustomDeviceAllocator> {
+ public:
+  StreamSafeCustomDeviceAllocator(
+      std::shared_ptr<Allocator> underlying_allocator,
+      CustomPlace place,
+      phi::stream::stream_t default_stream);
+  ~StreamSafeCustomDeviceAllocator();
+
+  bool IsAllocThreadSafe() const override { return true; }
+  phi::stream::stream_t GetDefaultStream() const;
+  void SetDefaultStream(phi::stream::stream_t stream);
+
+ protected:
+  phi::Allocation *AllocateImpl(size_t size) override;
+  void FreeImpl(phi::Allocation *allocation) override;
+  uint64_t ReleaseImpl(const Place &place) override;
+
+ private:
+  void ProcessUnfreedAllocations();
+  uint64_t ProcessUnfreedAllocationsAndRelease();
+
+  static std::map<Place, std::vector<StreamSafeCustomDeviceAllocator *>>
+      allocator_map_;
+  static SpinLock allocator_map_lock_;
+
+  std::shared_ptr<Allocator> underlying_allocator_;
+  CustomPlace place_;
+  phi::stream::stream_t default_stream_;
+  std::list<StreamSafeCustomDeviceAllocation *> unfreed_allocations_;
+  SpinLock unfreed_allocation_lock_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/stream_safe_xpu_allocator.cc

@@ -0,0 +1,219 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/stream_safe_xpu_allocator.h"
+#include <thread>
+#include "glog/logging.h"
+
+#include "paddle/phi/api/profiler/event_tracing.h"
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/backends/xpu/xpu_info.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+StreamSafeXPUAllocation::StreamSafeXPUAllocation(
+    DecoratedAllocationPtr underlying_allocation,
+    XPUStream owning_stream,
+    StreamSafeXPUAllocator* allocator)
+    : Allocation(underlying_allocation->ptr(),
+                 underlying_allocation->base_ptr(),
+                 underlying_allocation->size(),
+                 underlying_allocation->place()),
+      underlying_allocation_(std::move(underlying_allocation)),
+      owning_stream_(std::move(owning_stream)),
+      allocator_(allocator->shared_from_this()) {}
+
+bool StreamSafeXPUAllocation::RecordStream(XPUStream stream) {
+  VLOG(8) << "Try record stream " << stream << " for address " << ptr();
+  if (stream == owning_stream_) {
+    VLOG(8) << "stream " << stream << " is the same as owning stream "
+            << owning_stream_;
+    VLOG(8) << "Skip recording the same stream " << stream << " for address "
+            << ptr();
+    return false;
+  }
+
+  std::call_once(once_flag_,
+                 [this] { phi::backends::xpu::SetXPUDeviceId(place_.device); });
+
+  std::lock_guard<SpinLock> lock_guard(outstanding_event_map_lock_);
+
+  RecordStreamPrivate(stream);
+  return true;
+}
+
+bool StreamSafeXPUAllocation::CanBeFreed() {
+  std::call_once(once_flag_,
+                 [this] { phi::backends::xpu::SetXPUDeviceId(place_.device); });
+  for (auto it = outstanding_event_map_.begin();
+       it != outstanding_event_map_.end();
+       ++it) {
+    XPUEvent& event = it->second;
+    if (xpu_event_query(event) == XPU_SUCCESS) {
+      PADDLE_ENFORCE_XRE_SUCCESS(xpu_event_destroy(event));
+      VLOG(8) << "Destroy event " << event;
+    } else {
+      outstanding_event_map_.erase(outstanding_event_map_.begin(), it);
+      return false;
+    }
+  }
+  return true;
+}
+
+XPUStream StreamSafeXPUAllocation::GetOwningStream() const {
+  return owning_stream_;
+}
+
+void StreamSafeXPUAllocation::RecordStreamPrivate(XPUStream stream) {
+  XPUEvent record_event;
+  auto it = outstanding_event_map_.find(stream);
+  if (it == outstanding_event_map_.end()) {
+    XPUEvent new_event;
+    PADDLE_ENFORCE_XRE_SUCCESS(xpu_event_create(&new_event));
+    outstanding_event_map_[stream] = new_event;
+    record_event = new_event;
+    VLOG(9) << "Create a new event " << new_event;
+  } else {
+    record_event = it->second;
+    VLOG(9) << "Reuse event " << record_event;
+  }
+
+  PADDLE_ENFORCE_XRE_SUCCESS(xpu_event_record(record_event, stream));
+  VLOG(8) << "Record event " << record_event << " to stream " << stream;
+}
+
+StreamSafeXPUAllocator::StreamSafeXPUAllocator(
+    std::shared_ptr<Allocator> underlying_allocator,
+    XPUPlace place,
+    XPUStream default_stream)
+    : underlying_allocator_(std::move(underlying_allocator)),
+      place_(std::move(place)),
+      default_stream_(std::move(default_stream)) {
+  std::lock_guard<SpinLock> lock_guard(allocator_map_lock_);
+  allocator_map_[place].emplace_back(this);
+}
+
+StreamSafeXPUAllocator::~StreamSafeXPUAllocator() {
+  std::lock_guard<SpinLock> lock_guard(allocator_map_lock_);
+  std::vector<StreamSafeXPUAllocator*>& allocators = allocator_map_[place_];
+  allocators.erase(std::remove(allocators.begin(), allocators.end(), this),
+                   allocators.end());
+}
+
+bool StreamSafeXPUAllocator::IsAllocThreadSafe() const { return true; }
+
+XPUStream StreamSafeXPUAllocator::GetDefaultStream() const {
+  return default_stream_;
+}
+
+void StreamSafeXPUAllocator::SetDefaultStream(XPUStream stream) {
+  default_stream_ = stream;
+}
+
+phi::Allocation* StreamSafeXPUAllocator::AllocateImpl(size_t size) {
+  phi::RecordEvent record("StreamSafeXPUAllocator::Allocate",
+                          phi::TracerEventType::UserDefined,
+                          9 /*level*/);
+  ProcessUnfreedAllocations();
+  VLOG(8) << "Try allocate " << size << " bytes";
+  AllocationPtr underlying_allocation;
+  try {
+    underlying_allocation = underlying_allocator_->Allocate(size);
+  } catch (BadAlloc&) {
+    VLOG(4) << "Allocation failed when allocating " << size << " bytes";
+    ReleaseImpl(place_);
+    try {
+      underlying_allocation = underlying_allocator_->Allocate(size);
+    } catch (...) {
+      VLOG(3)
+          << "Still allocation failed after release memory from all streams";
+      throw;
+    }
+  } catch (...) {
+    throw;
+  }
+  StreamSafeXPUAllocation* allocation = new StreamSafeXPUAllocation(
+      static_unique_ptr_cast<Allocation>(std::move(underlying_allocation)),
+      default_stream_,
+      this);
+  VLOG(8) << "Thread " << std::this_thread::get_id() << " Allocate "
+          << allocation->size() << " bytes at address " << allocation->ptr()
+          << "  , stream: " << default_stream_;
+  return allocation;
+}
+
+void StreamSafeXPUAllocator::FreeImpl(phi::Allocation* allocation) {
+  phi::RecordEvent record("StreamSafeXPUAllocator::Free",
+                          phi::TracerEventType::UserDefined,
+                          9 /*level*/);
+  StreamSafeXPUAllocation* stream_safe_xpu_allocation =
+      static_cast<StreamSafeXPUAllocation*>(allocation);
+
+  VLOG(8) << "Try free allocation " << stream_safe_xpu_allocation->ptr();
+  if (stream_safe_xpu_allocation->CanBeFreed()) {
+    VLOG(9) << "Directly delete allocation";
+    delete stream_safe_xpu_allocation;
+  } else {
+    VLOG(9) << "Put into unfreed_allocation list";
+    std::lock_guard<SpinLock> lock_guard(unfreed_allocation_lock_);
+    unfreed_allocations_.emplace_back(stream_safe_xpu_allocation);
+  }
+}
+
+uint64_t StreamSafeXPUAllocator::ReleaseImpl(const Place& place) {
+  std::lock_guard<SpinLock> lock_guard(allocator_map_lock_);
+  std::vector<StreamSafeXPUAllocator*>& allocators = allocator_map_[place];
+  uint64_t released_size = 0;
+  for (StreamSafeXPUAllocator* allocator : allocators) {
+    released_size += allocator->ProcessUnfreedAllocationsAndRelease();
+  }
+  VLOG(8) << "Release " << released_size << " bytes memory from all streams";
+  return released_size;
+}
+
+void StreamSafeXPUAllocator::ProcessUnfreedAllocations() {
+  // NOTE(Ruibiao): This condition is to reduce lock completion. It does not
+  // need to be thread-safe since here occasional misjudgments are permissible.
+  if (unfreed_allocations_.empty()) {
+    return;
+  }
+
+  std::lock_guard<SpinLock> lock_guard(unfreed_allocation_lock_);
+  for (auto it = unfreed_allocations_.begin();
+       it != unfreed_allocations_.end();) {
+    if ((*it)->CanBeFreed()) {
+      delete *it;
+      it = unfreed_allocations_.erase(it);
+    } else {
+      ++it;
+    }
+  }
+}
+
+uint64_t StreamSafeXPUAllocator::ProcessUnfreedAllocationsAndRelease() {
+  ProcessUnfreedAllocations();
+  return underlying_allocator_->Release(place_);
+}
+
+thread_local std::once_flag StreamSafeXPUAllocation::once_flag_;
+
+std::map<Place, std::vector<StreamSafeXPUAllocator*>>
+    StreamSafeXPUAllocator::allocator_map_;
+SpinLock StreamSafeXPUAllocator::allocator_map_lock_;
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/stream_safe_xpu_allocator.h

@@ -0,0 +1,90 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <list>
+#include <map>
+#include <set>
+
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/allocation/spin_lock.h"
+
+#include "paddle/phi/backends/xpu/xpu_context.h"
+#include "paddle/phi/core/platform/device/xpu/xpu_info.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class StreamSafeXPUAllocator;
+
+class StreamSafeXPUAllocation : public Allocation {
+ public:
+  StreamSafeXPUAllocation(DecoratedAllocationPtr underlying_allocation,
+                          XPUStream owning_stream,
+                          StreamSafeXPUAllocator *allocator);
+
+  bool RecordStream(XPUStream stream);
+  void EraseStream(XPUStream stream);
+  bool CanBeFreed();
+  XPUStream GetOwningStream() const;
+
+ private:
+  thread_local static std::once_flag once_flag_;
+  void RecordStreamPrivate(XPUStream stream);
+  DecoratedAllocationPtr underlying_allocation_;
+
+  std::map<XPUStream, XPUEvent> outstanding_event_map_;
+  XPUStream owning_stream_;
+  SpinLock outstanding_event_map_lock_;
+  std::shared_ptr<Allocator> allocator_;
+};
+
+class StreamSafeXPUAllocator
+    : public Allocator,
+      public std::enable_shared_from_this<StreamSafeXPUAllocator> {
+ public:
+  StreamSafeXPUAllocator(std::shared_ptr<Allocator> underlying_allocator,
+                         XPUPlace place,
+                         XPUStream default_stream);
+  ~StreamSafeXPUAllocator();
+
+  bool IsAllocThreadSafe() const override;
+  XPUStream GetDefaultStream() const;
+  void SetDefaultStream(XPUStream stream);
+
+ protected:
+  phi::Allocation *AllocateImpl(size_t size) override;
+  void FreeImpl(phi::Allocation *allocation) override;
+  uint64_t ReleaseImpl(const Place &place) override;
+
+ private:
+  void ProcessUnfreedAllocations();
+  uint64_t ProcessUnfreedAllocationsAndRelease();
+
+  static std::map<Place, std::vector<StreamSafeXPUAllocator *>> allocator_map_;
+  static SpinLock allocator_map_lock_;
+
+  std::shared_ptr<Allocator> underlying_allocator_;
+  XPUPlace place_;
+  XPUStream default_stream_;
+  std::list<StreamSafeXPUAllocation *> unfreed_allocations_;
+  SpinLock unfreed_allocation_lock_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/system_allocator.cc

@@ -0,0 +1,442 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/phi/core/memory/allocation/system_allocator.h"
+
+#include "paddle/phi/core/memory/stats.h"
+
+#ifdef _WIN32
+#include <malloc.h>
+#ifndef NOMINMAX
+#define NOMINMAX  // msvc max/min macro conflict with std::min/max
+#endif
+#include <windows.h>  // VirtualLock/VirtualUnlock
+#else
+#include <sys/mman.h>  // for mlock and munlock
+#endif
+
+#include "paddle/common/flags.h"
+#include "paddle/phi/backends/cpu/cpu_info.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#include "paddle/phi/core/platform/cuda_device_guard.h"
+#endif
+
+#ifdef PADDLE_WITH_XPU
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <chrono>
+#include <ctime>
+#include <iomanip>
+#include <iostream>
+#include <sstream>
+#endif
+
+#include "paddle/phi/core/platform/device/device_wrapper.h"
+#include "paddle/phi/core/platform/profiler/mem_tracing.h"
+
+COMMON_DECLARE_bool(use_pinned_memory);
+COMMON_DECLARE_double(fraction_of_gpu_memory_to_use);
+COMMON_DECLARE_uint64(initial_gpu_memory_in_mb);
+COMMON_DECLARE_uint64(reallocate_gpu_memory_in_mb);
+
+namespace paddle::memory::detail {
+
+void* AlignedMalloc(size_t size) {
+  void* p = nullptr;
+  size_t alignment = 32ul;
+#ifdef PADDLE_WITH_DNNL
+  // refer to https://github.com/uxlfoundation/oneDNN/blob/main/include/dnnl.hpp
+  // memory alignment
+  alignment = 4096ul;
+#endif
+#ifdef _WIN32
+  p = _aligned_malloc(size, alignment);
+#else
+  int error = posix_memalign(&p, alignment, size);
+  PADDLE_ENFORCE_EQ(
+      error,
+      0,
+      common::errors::ResourceExhausted(
+          "Fail to alloc memory of %ld size, error code is %d.", size, error));
+#endif
+  PADDLE_ENFORCE_NOT_NULL(p,
+                          common::errors::ResourceExhausted(
+                              "Fail to alloc memory of %ld size.", size));
+  return p;
+}
+
+void* CPUAllocator::Alloc(size_t* index, size_t size) {
+  // According to http://www.cplusplus.com/reference/cstdlib/malloc/,
+  // malloc might not return nullptr if size is zero, but the returned
+  // pointer shall not be dereferenced -- so we make it nullptr.
+  if (size <= 0) return nullptr;
+
+  *index = 0;  // unlock memory
+
+  void* p = AlignedMalloc(size);
+
+  if (p != nullptr) {
+    if (FLAGS_use_pinned_memory) {
+      *index = 1;
+#ifdef _WIN32
+      VirtualLock(p, size);
+#else
+      mlock(p, size);  // lock memory
+#endif
+    }
+  }
+
+  HOST_MEMORY_STAT_UPDATE(Reserved, 0, size);
+  platform::RecordMemEvent(
+      p, CPUPlace(), size, phi::TracerMemEventType::ReservedAllocate);
+  return p;
+}
+
+void CPUAllocator::Free(void* p, size_t size, size_t index) {
+  if (p != nullptr && index == 1) {
+#ifdef _WIN32
+    VirtualUnlock(p, size);
+#else
+    munlock(p, size);
+#endif
+  }
+  HOST_MEMORY_STAT_UPDATE(Reserved, 0, -size);
+  platform::RecordMemEvent(
+      p, CPUPlace(), size, phi::TracerMemEventType::ReservedFree);
+#ifdef _WIN32
+  _aligned_free(p);
+#else
+  free(p);  // NOLINT
+#endif
+}
+
+bool CPUAllocator::UseGpu() const { return false; }
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+
+void* GPUAllocator::Alloc(size_t* index, size_t size) {
+  // CUDA documentation doesn't explain if cudaMalloc returns nullptr
+  // if size is 0.  We just make sure it does.
+  if (size <= 0) return nullptr;
+
+  void* p;
+  auto result = platform::RecordedGpuMalloc(&p, size, gpu_id_);
+
+  if (result == gpuSuccess) {
+    *index = 0;
+    gpu_alloc_size_ += size;
+    return p;
+  } else {
+    size_t avail, total, actual_avail, actual_total;
+    bool is_limited = platform::RecordedGpuMemGetInfo(
+        &avail, &total, &actual_avail, &actual_total, gpu_id_);
+    size_t allocated = total - avail;
+
+    std::string err_msg;
+    if (is_limited) {
+      auto limit_size = (total >> 20);
+      err_msg = string::Sprintf(
+          "\n   3) Set environment variable `FLAGS_gpu_memory_limit_mb` to a "
+          "larger value. Currently `FLAGS_gpu_memory_limit_mb` is %d, so the "
+          "maximum GPU memory usage is limited to %d MB.\n"
+          "      The command is `export FLAGS_gpu_memory_limit_mb=xxx`.",
+          limit_size,
+          limit_size);
+    }
+
+    PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+        "\n\nOut of memory error on GPU %d. "
+        "Cannot allocate %s memory on GPU %d, %s memory has been allocated and "
+        "available memory is only %s.\n\n"
+        "Please check whether there is any other process using GPU %d.\n"
+        "1. If yes, please stop them, or start PaddlePaddle on another GPU.\n"
+        "2. If no, please try one of the following suggestions:\n"
+        "   1) Decrease the batch size of your model.\n"
+        "   2) FLAGS_fraction_of_gpu_memory_to_use is %.2lf now, "
+        "please set it to a higher value but less than 1.0.\n"
+        "      The command is "
+        "`export FLAGS_fraction_of_gpu_memory_to_use=xxx`.%s\n\n",
+        gpu_id_,
+        string::HumanReadableSize(size),
+        gpu_id_,
+        string::HumanReadableSize(allocated),
+        string::HumanReadableSize(avail),
+        gpu_id_,
+        FLAGS_fraction_of_gpu_memory_to_use,
+        err_msg));
+  }
+}
+
+void GPUAllocator::Free(void* p, size_t size, size_t index) {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The index should be 0, index is %d", index));
+  PADDLE_ENFORCE_GE(gpu_alloc_size_,
+                    size,
+                    common::errors::InvalidArgument(
+                        "The size of memory (%d) to free exceeds the size of "
+                        "allocated gpu memory (%d)",
+                        size,
+                        gpu_alloc_size_));
+  gpu_alloc_size_ -= size;
+
+  platform::RecordedGpuFree(p, size, gpu_id_);
+}
+
+bool GPUAllocator::UseGpu() const { return true; }
+
+// PINNED memory allows direct DMA transfers by the GPU to and from system
+// memory. It's locked to a physical address.
+void* CUDAPinnedAllocator::Alloc(size_t* index, size_t size) {
+  if (size <= 0) return nullptr;
+
+  // NOTE: here, we use CUDAPinnedMaxAllocSize as the maximum memory size
+  // of host pinned allocation. Allocates too much would reduce
+  // the amount of memory available to the underlying system for paging.
+  size_t usable =
+      phi::backends::cpu::CUDAPinnedMaxAllocSize() - cuda_pinnd_alloc_size_;
+
+  if (size > usable) {
+    LOG(WARNING) << "Cannot malloc " << size / 1024.0 / 1024.0
+                 << " MB pinned memory."
+                 << ", available " << usable / 1024.0 / 1024.0
+                 << " MB";  // NOLINT
+    return nullptr;
+  }
+
+  void* p;
+// PINNED memory is visible to all CUDA contexts.
+#ifdef PADDLE_WITH_HIP
+  hipError_t result = hipHostMalloc(&p, size, hipHostMallocPortable);
+#else
+  cudaError_t result = cudaHostAlloc(&p, size, cudaHostAllocPortable);
+#endif
+
+  if (result == gpuSuccess) {
+    *index = 1;  // PINNED memory
+    cuda_pinnd_alloc_size_ += size;
+    HOST_MEMORY_STAT_UPDATE(Reserved, 0, size);
+    platform::RecordMemEvent(
+        p, CPUPlace(), size, phi::TracerMemEventType::ReservedAllocate);
+    return p;
+  } else {
+    LOG(WARNING) << "cudaHostAlloc failed.";
+    return nullptr;
+  }
+
+  return nullptr;
+}
+
+void CUDAPinnedAllocator::Free(void* p, size_t size, size_t index) {
+  gpuError_t err;
+  PADDLE_ENFORCE_EQ(index,
+                    1,
+                    common::errors::InvalidArgument(
+                        "The index should be 1, but got %d", index));
+
+  PADDLE_ENFORCE_GE(cuda_pinnd_alloc_size_,
+                    size,
+                    common::errors::InvalidArgument(
+                        "The size of memory (%d) to free exceeds the size of "
+                        "allocated cuda pinned memory (%d)",
+                        size,
+                        cuda_pinnd_alloc_size_));
+  cuda_pinnd_alloc_size_ -= size;
+#ifdef PADDLE_WITH_HIP
+  err = hipHostFree(p);
+  if (err != hipErrorDeinitialized) {
+    PADDLE_ENFORCE_EQ(
+        err,
+        hipSuccess,
+        common::errors::Fatal(
+            "hipFreeHost failed in GPUPinnedAllocator, error code is %d", err));
+  }
+#else
+  err = cudaFreeHost(p);
+
+  // Purposefully allow cudaErrorCudartUnloading, because
+  // that is returned if you ever call cudaFreeHost after the
+  // driver has already shutdown. This happens only if the
+  // process is terminating, in which case we don't care if
+  // cudaFreeHost succeeds.
+  if (err != cudaErrorCudartUnloading) {
+    PADDLE_ENFORCE_EQ(
+        err,
+        0,
+        common::errors::Fatal(
+            "cudaFreeHost failed in GPUPinnedAllocator, error code is %d",
+            err));
+  }
+#endif
+  HOST_MEMORY_STAT_UPDATE(Reserved, 0, -size);
+  platform::RecordMemEvent(
+      p, CPUPlace(), size, phi::TracerMemEventType::ReservedFree);
+}
+
+bool CUDAPinnedAllocator::UseGpu() const { return false; }
+
+#endif
+
+#ifdef PADDLE_WITH_XPU
+
+// XPU PINNED memory allows direct DMA transfers by the XPU to and from system
+// memory. It’s locked to a physical address.
+void* XPUPinnedAllocator::Alloc(size_t* index, size_t size) {
+  VLOG(6) << "Alloc start, requested size: " << size << " bytes";
+  if (size <= 0) {
+    VLOG(6) << "Requested size <= 0, returning nullptr";
+    return nullptr;
+  }
+
+  // NOTE: here, we use CUDAPinnedMaxAllocSize as the maximum memory size
+  // of host pinned allocation. Allocates too much would reduce
+  // the amount of memory available to the underlying system for paging.
+  size_t usable =
+      phi::backends::cpu::CUDAPinnedMaxAllocSize() - xpu_pinned_alloc_size_;
+  VLOG(6) << "Usable pinned memory: " << usable << " bytes";
+
+  if (size > usable) {
+    VLOG(6) << "Requested size (" << size
+            << " bytes) exceeds usable pinned memory (" << usable << " bytes)";
+    LOG(WARNING) << "Cannot malloc " << size / 1024.0 / 1024.0
+                 << " MB pinned memory."
+                 << ", available " << usable / 1024.0 / 1024.0
+                 << " MB";  // NOLINT
+    return nullptr;
+  }
+
+  void* p = nullptr;
+  VLOG(6) << "Calling cudaHostAlloc for " << size << " bytes";
+  // PINNED memory is visible to all CUDA contexts.
+  // FIXME(yangjianbang): XPU does not support cudaHostAllocPortable with
+  // multiple cuda contexts yet.
+  cudaError_t result = cudaHostAlloc(&p, size, cudaHostAllocDefault);
+  VLOG(6) << "cudaHostAlloc returned: " << result;
+
+  if (result == cudaSuccess) {
+    *index = 1;  // PINNED memory
+    xpu_pinned_alloc_size_ += size;
+    HOST_MEMORY_STAT_UPDATE(Reserved, 0, size);
+    platform::RecordMemEvent(
+        p, CPUPlace(), size, phi::TracerMemEventType::ReservedAllocate);
+    VLOG(6) << "cudaHostAlloc succeeded. Allocated pointer: " << p;
+    return p;
+  } else {
+    VLOG(6) << "cudaHostAlloc failed.";
+    LOG(WARNING) << "cudaHostAlloc failed.";
+    return nullptr;
+  }
+}
+
+void XPUPinnedAllocator::Free(void* p, size_t size, size_t index) {
+  cudaError_t err;
+  PADDLE_ENFORCE_EQ(index,
+                    1,
+                    common::errors::InvalidArgument(
+                        "The index should be 1, but got %d", index));
+
+  PADDLE_ENFORCE_GE(xpu_pinned_alloc_size_,
+                    size,
+                    common::errors::InvalidArgument(
+                        "The size of memory (%d) to free exceeds the size of "
+                        "allocated cuda pinned memory (%d)",
+                        size,
+                        xpu_pinned_alloc_size_));
+  xpu_pinned_alloc_size_ -= size;
+  err = cudaFreeHost(p);
+
+  // Purposefully allow cudaErrorCudartUnloading, because
+  // that is returned if you ever call cudaFreeHost after the
+  // driver has already shutdown. This happens only if the
+  // process is terminating, in which case we don't care if
+  // cudaFreeHost succeeds.
+  if (err != cudaErrorCudartUnloading) {
+    PADDLE_ENFORCE_EQ(
+        err,
+        0,
+        common::errors::Fatal(
+            "cudaFreeHost failed in XPUPinnedAllocator, error code is %d",
+            err));
+  }
+  HOST_MEMORY_STAT_UPDATE(Reserved, 0, -size);
+  platform::RecordMemEvent(
+      p, CPUPlace(), size, phi::TracerMemEventType::ReservedFree);
+}
+
+bool XPUPinnedAllocator::UseGpu() const { return false; }
+#endif
+
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+void* CustomAllocator::Alloc(size_t* index, size_t size) {
+  if (size <= 0) return nullptr;
+
+  void* p;
+  auto place = CustomPlace(dev_type_, dev_id_);
+  auto device = phi::DeviceManager::GetDeviceWithPlace(place);
+  p = device->MemoryAllocate(size);
+  if (LIKELY(p)) {
+    VLOG(4) << "CustomAllocator::Alloc " << p << " size " << size;
+    *index = 0;
+    plug_alloc_size += size;
+    DEVICE_MEMORY_STAT_UPDATE(Reserved, dev_id_, size);
+    platform::RecordMemEvent(
+        p, place, size, phi::TracerMemEventType::ReservedAllocate);
+  } else {
+    size_t avail, total;
+
+    phi::DeviceManager::MemoryStats(place, &total, &avail);
+    PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+        "\n\nOut of memory error on %s %d. "
+        "total memory is %s, used memory is %s, "
+        "available memory is only %s.\n\n",
+        dev_type_,
+        dev_id_,
+        string::HumanReadableSize(total),
+        string::HumanReadableSize(total - avail),
+        string::HumanReadableSize(avail)));
+  }
+  return p;
+}
+
+void CustomAllocator::Free(void* p, size_t size, size_t index) {
+  VLOG(4) << "CustomAllocator::Free " << p << " size " << size;
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The index should be 0, index is %d", index));
+  PADDLE_ENFORCE_GE(plug_alloc_size,
+                    size,
+                    common::errors::InvalidArgument(
+                        "The size of memory (%d) to free exceeds the size of "
+                        "allocated gpu memory (%d)",
+                        size,
+                        plug_alloc_size));
+  plug_alloc_size -= size;
+  auto place = CustomPlace(dev_type_, dev_id_);
+  auto device = phi::DeviceManager::GetDeviceWithPlace(place);
+  device->MemoryDeallocate(p, size);
+  DEVICE_MEMORY_STAT_UPDATE(Reserved, dev_id_, size);
+  platform::RecordMemEvent(
+      p, place, size, phi::TracerMemEventType::ReservedFree);
+}
+
+bool CustomAllocator::UseGpu() const { return true; }
+#endif
+
+}  // namespace paddle::memory::detail

paddle/phi/core/memory/allocation/system_allocator.h

@@ -0,0 +1,103 @@
+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#include <stddef.h>  // for size_t
+
+#include <string>
+#include "paddle/common/macros.h"
+
+namespace paddle {
+namespace memory {
+namespace detail {
+
+/**
+ * \brief SystemAllocator is the parent class of CPUAllocator,
+ *        CUDAPinnedAllocator and GPUAllocator. A BuddyAllocator
+ *        object uses a SystemAllocator* pointing to the
+ *        underlying system allocator.
+ */
+class SystemAllocator {
+ public:
+  virtual ~SystemAllocator() {}
+  virtual void* Alloc(size_t* index, size_t size) = 0;
+  virtual void Free(void* p, size_t size, size_t index) = 0;
+  virtual bool UseGpu() const = 0;
+};
+
+class PADDLE_API CPUAllocator : public SystemAllocator {
+ public:
+  virtual void* Alloc(size_t* index, size_t size);
+  virtual void Free(void* p, size_t size, size_t index);
+  virtual bool UseGpu() const;
+};
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+class PADDLE_API GPUAllocator : public SystemAllocator {
+ public:
+  explicit GPUAllocator(int gpu_id) : gpu_id_(gpu_id) {}
+
+  virtual void* Alloc(size_t* index, size_t size);
+  virtual void Free(void* p, size_t size, size_t index);
+  virtual bool UseGpu() const;
+
+ private:
+  size_t gpu_alloc_size_ = 0;
+  int gpu_id_;
+};
+
+class PADDLE_API CUDAPinnedAllocator : public SystemAllocator {
+ public:
+  virtual void* Alloc(size_t* index, size_t size);
+  virtual void Free(void* p, size_t size, size_t index);
+  virtual bool UseGpu() const;
+
+ private:
+  size_t cuda_pinnd_alloc_size_ = 0;
+};
+#endif
+
+#if defined(PADDLE_WITH_XPU)
+class XPUPinnedAllocator : public SystemAllocator {
+ public:
+  virtual void* Alloc(size_t* index, size_t size);
+  virtual void Free(void* p, size_t size, size_t index);
+  virtual bool UseGpu() const;
+
+ private:
+  size_t xpu_pinned_alloc_size_ = 0;
+};
+#endif
+
+#ifdef PADDLE_WITH_CUSTOM_DEVICE
+class CustomAllocator : public SystemAllocator {
+ public:
+  explicit CustomAllocator(const std::string& device_type, size_t dev_id)
+      : dev_type_(device_type), dev_id_(dev_id) {}
+
+  virtual void* Alloc(size_t* index, size_t size);
+  virtual void Free(void* p, size_t size, size_t index);
+  virtual bool UseGpu() const;
+
+ private:
+  size_t plug_alloc_size = 0;
+  std::string dev_type_;
+  size_t dev_id_;
+};
+#endif
+
+}  // namespace detail
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/thread_local_allocator.cc

@@ -0,0 +1,78 @@
+// Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/thread_local_allocator.h"
+
+#include "glog/logging.h"
+
+namespace paddle::memory::allocation {
+
+ThreadLocalAllocatorImpl::ThreadLocalAllocatorImpl(const Place& p) : place_(p) {
+  if (phi::is_gpu_place(place_)) {
+    buddy_allocator_ = std::make_unique<memory::detail::BuddyAllocator>(
+        std::unique_ptr<memory::detail::SystemAllocator>(
+            new memory::detail::GPUAllocator(place_.device)),
+        platform::GpuMinChunkSize(),
+        platform::GpuMaxChunkSize());
+  } else {
+    PADDLE_THROW(common::errors::Unavailable(
+        "Thread local allocator only supports CUDAPlace now."));
+  }
+}
+
+std::shared_ptr<ThreadLocalAllocatorImpl> ThreadLocalCUDAAllocatorPool::Get(
+    int gpu_id) {
+  auto pos = std::distance(devices_.begin(),
+                           std::find(devices_.begin(), devices_.end(), gpu_id));
+  PADDLE_ENFORCE_LT(
+      pos,
+      devices_.size(),
+      common::errors::InvalidArgument(
+          "The position of device should be less than the size of devices."));
+  std::call_once(*init_flags_[pos], [this, pos, gpu_id] {
+    platform::SetDeviceId(devices_[pos]);
+    allocators_[pos].reset(new ThreadLocalAllocatorImpl(GPUPlace(gpu_id)));
+  });
+  return allocators_[pos];
+}
+
+ThreadLocalCUDAAllocatorPool::ThreadLocalCUDAAllocatorPool()
+    : devices_(platform::GetSelectedDevices()) {
+  auto gpu_num = devices_.size();
+  allocators_.resize(gpu_num);
+  init_flags_.reserve(gpu_num);
+  for (size_t i = 0; i < gpu_num; ++i) {
+    init_flags_.emplace_back(new std::once_flag());
+  }
+}
+
+ThreadLocalAllocation* ThreadLocalAllocatorImpl::AllocateImpl(size_t size) {
+  VLOG(10) << "ThreadLocalAllocatorImpl::AllocateImpl " << size;
+  void* ptr = buddy_allocator_->Alloc(size);
+  auto* tl_allocation = new ThreadLocalAllocation(ptr, size, place_);
+  tl_allocation->SetThreadLocalAllocatorImpl(shared_from_this());
+  return tl_allocation;
+}
+
+void ThreadLocalAllocatorImpl::FreeImpl(ThreadLocalAllocation* allocation) {
+  VLOG(10) << "ThreadLocalAllocatorImpl::FreeImpl " << allocation;
+  buddy_allocator_->Free(allocation->ptr());
+  delete allocation;
+}
+
+uint64_t ThreadLocalAllocatorImpl::ReleaseImpl() {
+  return buddy_allocator_->Release();
+}
+
+}  // namespace paddle::memory::allocation

paddle/phi/core/memory/allocation/thread_local_allocator.h

@@ -0,0 +1,104 @@
+// Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <memory>
+#include <vector>
+
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/allocation/buddy_allocator.h"
+#include "paddle/phi/core/memory/allocation/system_allocator.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class ThreadLocalAllocatorImpl;
+
+class ThreadLocalAllocation : public Allocation {
+ public:
+  ThreadLocalAllocation(void* ptr, size_t size, Place place)
+      : Allocation(ptr, size, place) {}
+
+  void SetThreadLocalAllocatorImpl(
+      std::shared_ptr<ThreadLocalAllocatorImpl> allocator) {
+    allocator_ = allocator;
+  }
+
+  std::shared_ptr<ThreadLocalAllocatorImpl> GetAllocator() {
+    return allocator_;
+  }
+
+ private:
+  std::shared_ptr<ThreadLocalAllocatorImpl> allocator_;
+};
+
+class ThreadLocalAllocatorImpl
+    : public std::enable_shared_from_this<ThreadLocalAllocatorImpl> {
+ public:
+  explicit ThreadLocalAllocatorImpl(const Place& p);
+  ThreadLocalAllocation* AllocateImpl(size_t size);
+  void FreeImpl(ThreadLocalAllocation* allocation);
+  uint64_t ReleaseImpl();
+
+ private:
+  std::unique_ptr<memory::detail::BuddyAllocator> buddy_allocator_;
+  Place place_;
+};
+
+class ThreadLocalCUDAAllocatorPool {
+ public:
+  static ThreadLocalCUDAAllocatorPool& Instance() {
+    static thread_local ThreadLocalCUDAAllocatorPool pool;
+    return pool;
+  }
+
+  PADDLE_API std::shared_ptr<ThreadLocalAllocatorImpl> Get(int gpu_id);
+
+ private:
+  PADDLE_API ThreadLocalCUDAAllocatorPool();
+  std::vector<int> devices_;
+  std::vector<std::unique_ptr<std::once_flag>> init_flags_;
+  std::vector<std::shared_ptr<ThreadLocalAllocatorImpl>> allocators_;
+};
+
+class ThreadLocalCUDAAllocator : public Allocator {
+ public:
+  explicit ThreadLocalCUDAAllocator(const GPUPlace& p) : gpu_id_(p.device) {}
+
+  bool IsAllocThreadSafe() const override { return true; }
+
+ protected:
+  phi::Allocation* AllocateImpl(size_t size) override {
+    return ThreadLocalCUDAAllocatorPool::Instance().Get(gpu_id_)->AllocateImpl(
+        size);
+  }
+  void FreeImpl(phi::Allocation* allocation) override {
+    auto* tl_allocation = static_cast<ThreadLocalAllocation*>(allocation);
+    auto allocator_impl = tl_allocation->GetAllocator();
+    allocator_impl->FreeImpl(tl_allocation);
+  }
+  uint64_t ReleaseImpl(const Place& p) override {
+    return ThreadLocalCUDAAllocatorPool::Instance().Get(gpu_id_)->ReleaseImpl();
+  }
+
+ private:
+  int gpu_id_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/virtual_memory_auto_growth_best_fit_allocator.cc

@@ -0,0 +1,584 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/virtual_memory_auto_growth_best_fit_allocator.h"
+#include <limits>
+#include <mutex>
+#include "glog/logging.h"
+#include "paddle/common/flags.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/memory/allocation/aligned_allocator.h"
+#include "paddle/phi/core/memory/allocation/cuda_virtual_mem_allocator.h"
+
+PHI_DEFINE_EXPORTED_uint64(
+    vmm_small_pool_size_in_mb,
+    1,
+    "Threshold (MiB) separating the small and large pools. "
+    "0 disables the small pool and enables single-pool mode "
+    "(all requests go to the large pool). When > 0, requests "
+    "<= threshold use the small pool; larger requests use the "
+    "large pool. Default: 0.");
+PHI_DEFINE_EXPORTED_uint64(vmm_small_pool_min_growth_size_in_mb,
+                           0,
+                           "The minimal chunk size for the small pool in MiB. "
+                           "If small_pool_size_in_mb > 0, this overrides "
+                           "the constructor-provided global growth size "
+                           "(FLAGS_auto_growth_chunk_size_in_mb).");
+PHI_DEFINE_EXPORTED_uint64(vmm_large_pool_min_growth_size_in_mb,
+                           0,
+                           "The minimal chunk size for the large pool in MiB. "
+                           "If small_pool_size_in_mb > 0, this overrides "
+                           "the constructor-provided global growth size "
+                           "(FLAGS_auto_growth_chunk_size_in_mb).");
+PHI_DEFINE_EXPORTED_uint64(
+    vmm_large_pool_pre_alloc_in_mb,
+    0,
+    "Pre-reserve this many MiB in the large pool. 0 disables pre-allocation.");
+PHI_DEFINE_EXPORTED_uint64(
+    vmm_small_pool_pre_alloc_in_mb,
+    0,
+    "Pre-reserve this many MiB in the small pool. 0 disables pre-allocation.");
+PHI_DEFINE_EXPORTED_uint64(
+    vmm_pre_alloc_in_mb,
+    0,
+    "Pre-reserve this many MiB in the small pool. 0 disables pre-allocation.");
+PHI_DEFINE_EXPORTED_bool(
+    dump_vmm_allocation_info,
+    false,
+    "dump VirtualMemoryAutoGrowthBestFitAllocator's allocation info");
+PHI_DEFINE_EXPORTED_bool(native_compact,
+                         false,
+                         "native_compact means compact memory after OOM, The "
+                         "algorithm still needs to be upgraded.");
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+bool NeedSplit(size_t block_size, size_t alignment, size_t alloc_size) {
+  return block_size > (alloc_size * 2) || (block_size - alloc_size) > alignment;
+}
+
+static BlockPart MakeBlockPart(void *ptr, size_t size, int device) {
+  auto chunk = std::make_shared<VmmChunkMeta>();
+  chunk->base = reinterpret_cast<VmmDevicePtr>(ptr);
+  chunk->size = size;
+#ifdef PADDLE_WITH_CUDA
+  auto handle = CUDAVirtualMemAllocator::GetHandleFromBasePtr(ptr);
+  PADDLE_ENFORCE_NE(
+      handle,
+      0,
+      common::errors::InvalidArgument(
+          "Allocation returned by underlying allocator is not VMM allocation"));
+  chunk->handle = handle;
+#else
+  PADDLE_THROW(common::errors::Unavailable(
+      "Virtual memory auto-growth allocator requires CUDA support."));
+#endif
+  chunk->device = device;
+  return BlockPart{chunk, 0, size};
+}
+
+VirtualMemoryAutoGrowthBestFitAllocator::
+    VirtualMemoryAutoGrowthBestFitAllocator(
+        const std::shared_ptr<Allocator> &underlying_allocator,
+        size_t alignment,
+        const GPUPlace &place)
+    : underlying_allocator_(
+          std::make_shared<AlignedAllocator>(underlying_allocator, alignment)),
+      alignment_(alignment),
+      place_(place) {
+  // NOTE(liujinnan): Only support TotalMemoryCompactor strategy for now.
+  memory_compactor_ = std::make_unique<TotalMemoryCompactor>();
+}
+
+phi::Allocation *VirtualMemoryAutoGrowthBestFitAllocator::AllocateImpl(
+    size_t size) {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  size = AlignedSize(size, alignment_);
+  auto result = AllocFromFreeBlocks(size);
+
+  if (!result) {
+    ExtendOrCompact(size);
+    result = AllocFromFreeBlocks(size);
+  }
+
+  return result;
+}
+
+void VirtualMemoryAutoGrowthBestFitAllocator::FreeImpl(
+    phi::Allocation *allocation) {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  auto block_it = static_cast<BlockAllocation *>(allocation)->block_it_;
+  TryMergeBlock2Blocks(block_it);
+  delete allocation;
+}
+
+bool VirtualMemoryAutoGrowthBestFitAllocator::CollectTensorParts(
+    void *ptr, size_t size, std::vector<BlockPart> *parts) {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  auto target_begin = reinterpret_cast<uintptr_t>(ptr);
+  PADDLE_ENFORCE_LE(
+      size,
+      std::numeric_limits<uintptr_t>::max() - target_begin,
+      common::errors::InvalidArgument(
+          "Invalid VMM tensor range: ptr %p plus size %zu overflows.",
+          ptr,
+          size));
+  auto target_end = target_begin + size;
+  for (const auto &block : all_blocks_) {
+    if (block.is_free_) {
+      continue;
+    }
+    auto block_begin = reinterpret_cast<uintptr_t>(block.ptr_);
+    auto block_end = block_begin + block.size_;
+    if (target_begin >= block_begin && target_end <= block_end) {
+      if (parts) {
+        *parts = SliceBlockPartsForRange(
+            block.parts_, target_begin - block_begin, size);
+      }
+      return true;
+    }
+  }
+  return false;
+}
+
+void VirtualMemoryAutoGrowthBestFitAllocator::TryMergeBlock2Blocks(
+    std::list<Block>::iterator block) {
+  if (block->ptr_ == all_blocks_.front().ptr_ &&
+      block->ptr_ == all_blocks_.back().ptr_) {
+    block->is_free_ = true;
+    free_blocks_.emplace(std::make_pair(block->size_, block->ptr_), block);
+  } else if (block->ptr_ == all_blocks_.front().ptr_) {
+    auto next = std::next(block);
+    if (next->is_free_ &&
+        reinterpret_cast<uint8_t *>(block->ptr_) + block->size_ == next->ptr_) {
+      // merge with next
+      AppendBlockPartsTail(&block->parts_, &next->parts_);
+      block->size_ += next->size_;
+      block->is_free_ = true;
+      free_blocks_.erase(std::make_pair(next->size_, next->ptr_));
+      all_blocks_.erase(next);
+      free_blocks_.emplace(std::make_pair(block->size_, block->ptr_), block);
+    } else {
+      block->is_free_ = true;
+      free_blocks_.emplace(std::make_pair(block->size_, block->ptr_), block);
+    }
+  } else if (block->ptr_ == all_blocks_.back().ptr_) {
+    auto pre = std::prev(block);
+    if (pre->is_free_ &&
+        reinterpret_cast<uint8_t *>(pre->ptr_) + pre->size_ == block->ptr_) {
+      // merge with pre
+      free_blocks_.erase(std::make_pair(pre->size_, pre->ptr_));
+      AppendBlockPartsTail(&pre->parts_, &block->parts_);
+      pre->size_ += block->size_;
+      all_blocks_.erase(block);
+      free_blocks_.emplace(std::make_pair(pre->size_, pre->ptr_), pre);
+    } else {
+      block->is_free_ = true;
+      free_blocks_.emplace(std::make_pair(block->size_, block->ptr_), block);
+    }
+  } else {
+    auto pre = std::prev(block);
+    auto next = std::next(block);
+    if (pre->is_free_ &&
+        reinterpret_cast<uint8_t *>(pre->ptr_) + pre->size_ == block->ptr_ &&
+        !(next->is_free_ &&
+          reinterpret_cast<uint8_t *>(block->ptr_) + block->size_ ==
+              next->ptr_)) {
+      // merge with pre
+      free_blocks_.erase(std::make_pair(pre->size_, pre->ptr_));
+      AppendBlockPartsTail(&pre->parts_, &block->parts_);
+      pre->size_ += block->size_;
+      all_blocks_.erase(block);
+      free_blocks_.emplace(std::make_pair(pre->size_, pre->ptr_), pre);
+    } else if (next->is_free_ &&
+               reinterpret_cast<uint8_t *>(block->ptr_) + block->size_ ==
+                   next->ptr_ &&
+               !(pre->is_free_ &&
+                 reinterpret_cast<uint8_t *>(pre->ptr_) + pre->size_ ==
+                     block->ptr_)) {
+      // merge with next
+      block->size_ += next->size_;
+      block->is_free_ = true;
+      AppendBlockPartsTail(&block->parts_, &next->parts_);
+      free_blocks_.erase(std::make_pair(next->size_, next->ptr_));
+      all_blocks_.erase(next);
+      free_blocks_.emplace(std::make_pair(block->size_, block->ptr_), block);
+    } else if (pre->is_free_ &&
+               reinterpret_cast<uint8_t *>(pre->ptr_) + pre->size_ ==
+                   block->ptr_ &&
+               next->is_free_ &&
+               reinterpret_cast<uint8_t *>(block->ptr_) + block->size_ ==
+                   next->ptr_) {
+      // merge with pre and next
+      free_blocks_.erase(std::make_pair(pre->size_, pre->ptr_));
+      free_blocks_.erase(std::make_pair(next->size_, next->ptr_));
+      AppendBlockPartsTail(&pre->parts_, &block->parts_);
+      AppendBlockPartsTail(&pre->parts_, &next->parts_);
+      pre->size_ += (block->size_ + next->size_);
+      all_blocks_.erase(block);
+      all_blocks_.erase(next);
+      free_blocks_.emplace(std::make_pair(pre->size_, pre->ptr_), pre);
+    } else {
+      block->is_free_ = true;
+      free_blocks_.emplace(std::make_pair(block->size_, block->ptr_), block);
+    }
+  }
+}
+
+std::optional<AllocationPtr>
+VirtualMemoryAutoGrowthBestFitAllocator::AllocateOrCompact(size_t size) {
+  AllocationPtr allocateptr = nullptr;
+  // Just Allocate, no compact.
+  if (!FLAGS_native_compact) {
+    if (all_blocks_.empty()) {
+      allocateptr = std::move(underlying_allocator_->Allocate(size));
+    } else {
+      auto free_block = std::prev(all_blocks_.end());
+      if (free_block->is_free_) {
+        assert(free_block->size_ < size);
+        auto remain_size = size - free_block->size_;
+        VLOG(4) << " Tail free block size {" << free_block->size_
+                << "} is smaller than allocate size {" << size
+                << "} after compact, re-alloc {" << remain_size << "}";
+        allocateptr = std::move(underlying_allocator_->Allocate(remain_size));
+      } else {
+        VLOG(4) << "Tail block is not free, just allocate {" << size << "}";
+        allocateptr = std::move(underlying_allocator_->Allocate(size));
+      }
+    }
+    return allocateptr;
+  }
+  // Compact branch, try allocate and compact.
+  try {
+    allocateptr = std::move(underlying_allocator_->Allocate(size));
+  } catch (const paddle::memory::allocation::BadAlloc &e) {
+    VLOG(4) << "Do Memory Compact allocate size and compact " << size;
+    size_t compact_free_size = memory_compactor_->Compact(
+        all_blocks_, all_blocks_.front().ptr_, all_blocks_.back().ptr_);
+    VLOG(4) << "Memory Compacted Size: " << compact_free_size;
+    auto free_block = std::prev(all_blocks_.end());
+    if (free_block->is_free_ && free_block->size_ < size) {
+      auto realloc_size = size - free_block->size_;
+      VLOG(4) << "Free block size {" << free_block->size_
+              << "} is smaller than allocate size {" << size
+              << "} after compact, re-alloc {" << realloc_size << "}";
+      try {
+        auto realloc_ptr =
+            underlying_allocator_->Allocate(size - free_block->size_);
+        VLOG(4) << "Re-alloc size {" << realloc_ptr->size() << "} success";
+        std::vector<BlockPart> realloc_parts;
+        realloc_parts.emplace_back(MakeBlockPart(
+            realloc_ptr->ptr(), realloc_ptr->size(), place_.device));
+        AppendBlockPartsTail(&free_block->parts_, &realloc_parts);
+        free_block->size_ += realloc_ptr->size();
+        allocations_.push_back(std::move(realloc_ptr));  // hold allocation
+      } catch (const paddle::memory::allocation::BadAlloc &e) {
+        VLOG(4) << "Re-alloc size {" << realloc_size << "} failed";
+        throw;
+      }
+    }
+    return std::nullopt;
+  }
+  return allocateptr;
+}
+
+void VirtualMemoryAutoGrowthBestFitAllocator::ExtendOrCompact(size_t size) {
+  void *alloc_ptr = nullptr;
+  size_t alloc_size = 0;
+  if (FLAGS_dump_vmm_allocation_info) {
+    DumpInfo("===== Before ExtendOrCompact ===== request size: " +
+             std::to_string(size));
+  }
+
+  std::optional<AllocationPtr> allocate_result;
+  try {
+    allocate_result = AllocateOrCompact(size);
+  } catch (const paddle::memory::allocation::BadAlloc &e) {
+    auto stats = SumLargestFreeBlockSizes(free_blocks_.size());
+    PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+        "%s\n"
+        "VMM allocator stats (pool): total_free=%s, max_free=%s.\n",
+        e.what(),
+        string::HumanReadableSize(stats.second),
+        string::HumanReadableSize(stats.first)));
+  }
+
+  AllocationPtr allocateptr = nullptr;
+  if (allocate_result.has_value()) {
+    allocateptr = std::move(allocate_result.value());
+  }
+  if (!allocateptr) {
+    // Allocate failed and Compact success branch.
+    free_blocks_.clear();
+    auto free_block = std::prev(all_blocks_.end());
+    if (free_block->is_free_) {
+      free_blocks_.emplace(std::make_pair(free_block->size_, free_block->ptr_),
+                           free_block);
+    } else {
+      LOG(INFO) << "Dont have free block after memory compact";
+    }
+    if (FLAGS_dump_vmm_allocation_info) {
+      DumpInfo("===== After ExtendOrCompact do compact =====");
+    }
+    // After compact, Merge is not needed. just return.
+    return;
+  }
+
+  alloc_ptr = allocateptr->ptr();
+  alloc_size = allocateptr->size();
+  allocations_.push_back(std::move(allocateptr));  // hold allocation
+
+  std::vector<BlockPart> new_parts;
+  new_parts.emplace_back(MakeBlockPart(alloc_ptr, alloc_size, place_.device));
+
+  if (all_blocks_.empty()) {
+    all_blocks_.emplace_back(alloc_ptr, alloc_size, true);
+    auto it = all_blocks_.begin();
+    it->parts_ = std::move(new_parts);
+    free_blocks_.emplace(std::make_pair(alloc_size, alloc_ptr), it);
+    return;
+  }
+
+  // insert to back
+  auto block_it = all_blocks_.end();
+  block_it--;
+  if (block_it->is_free_ &&
+      reinterpret_cast<uint8_t *>(block_it->ptr_) + block_it->size_ ==
+          alloc_ptr) {
+    // merge with pre
+    free_blocks_.erase(std::make_pair(block_it->size_, block_it->ptr_));
+    block_it->size_ += alloc_size;
+    AppendBlockPartsTail(&block_it->parts_, &new_parts);
+    free_blocks_.emplace(std::make_pair(block_it->size_, block_it->ptr_),
+                         block_it);
+  } else {
+    // do not merge
+    all_blocks_.emplace_back(alloc_ptr, alloc_size, true);
+    auto block_it = all_blocks_.end();
+    block_it--;
+    block_it->parts_ = std::move(new_parts);
+    free_blocks_.emplace(std::make_pair(alloc_size, alloc_ptr), block_it);
+  }
+  if (FLAGS_dump_vmm_allocation_info) {
+    DumpInfo("===== After ExtendOrCompact =====  request size: " +
+             std::to_string(size) +
+             " alloc size: " + std::to_string(alloc_size));
+  }
+}
+
+phi::Allocation *VirtualMemoryAutoGrowthBestFitAllocator::AllocFromFreeBlocks(
+    size_t size) {
+  auto iter = free_blocks_.lower_bound(std::make_pair(size, nullptr));
+  if (iter != free_blocks_.end()) {
+    std::list<Block>::iterator block_it = iter->second;
+    free_blocks_.erase(iter);
+    if (NeedSplit(block_it->size_, alignment_, size)) {
+      void *remaining_ptr = reinterpret_cast<uint8_t *>(block_it->ptr_) + size;
+      size_t remaining_size = block_it->size_ - size;
+
+      std::vector<BlockPart> alloc_parts =
+          SliceBlockPartsForRange(block_it->parts_, 0, size);
+      std::vector<BlockPart> remaining_parts =
+          SliceBlockPartsForRange(block_it->parts_, size, remaining_size);
+
+      block_it->size_ = size;
+      block_it->is_free_ = false;
+      block_it->parts_.swap(alloc_parts);
+
+      auto remaining_free_block = all_blocks_.insert(
+          std::next(block_it), Block(remaining_ptr, remaining_size, true));
+      remaining_free_block->parts_ = std::move(remaining_parts);
+      free_blocks_.emplace(std::make_pair(remaining_size, remaining_ptr),
+                           remaining_free_block);
+    } else {
+      block_it->is_free_ = false;
+    }
+    return new BlockAllocation(block_it, place_);
+  }
+  return nullptr;
+}
+
+size_t VirtualMemoryAutoGrowthBestFitAllocator::CompactImpl(
+    const Place &place) {
+  VLOG(1) << "Do Memory Compact Manual";
+  size_t compact_free_size = memory_compactor_->Compact(
+      all_blocks_, all_blocks_.front().ptr_, all_blocks_.back().ptr_);
+  VLOG(1) << "Memory Compact Manual Finish Compact size: " << compact_free_size;
+
+  if (compact_free_size > 0) {
+    auto free_block = std::prev(all_blocks_.end());
+    assert(free_block->is_free_);
+    // remove all old free blocks and put new free block into free_blocks_.
+    free_blocks_.clear();
+    free_blocks_.emplace(std::make_pair(free_block->size_, free_block->ptr_),
+                         free_block);
+  }
+  return compact_free_size;
+}
+
+bool VirtualMemoryAutoGrowthBestFitAllocator::TryAllocateBatch(
+    const std::vector<size_t> &sizes) {
+  auto SimulateAlloc =
+      [&](size_t size,
+          std::map<std::pair<size_t, void *>, size_t> &shadow_blocks) {
+        auto iter = shadow_blocks.lower_bound(std::make_pair(size, nullptr));
+        if (iter != shadow_blocks.end()) {
+          size_t block_size = iter->first.first;
+          void *block_ptr = iter->first.second;
+          shadow_blocks.erase(iter);
+          if (NeedSplit(block_size, alignment_, size)) {
+            size_t remaining_size = block_size - size;
+            void *remaining_ptr = reinterpret_cast<uint8_t *>(block_ptr) + size;
+            shadow_blocks.emplace(std::make_pair(remaining_size, remaining_ptr),
+                                  remaining_size);
+          }
+          return true;
+        }
+        return false;
+      };
+
+  std::lock_guard<SpinLock> guard(spinlock_);
+
+  // copy large N free_blocks_ to shadow_blocks_.
+  std::map<std::pair<size_t, void *>, size_t> shadow_blocks;
+  auto it = free_blocks_.rbegin();
+  for (int i = 0; i < sizes.size() && it != free_blocks_.rend(); ++i, ++it) {
+    shadow_blocks.emplace(it->first, it->first.first);
+  }
+  for (size_t size : sizes) {
+    size_t aligned_size = AlignedSize(size, alignment_);
+    if (!SimulateAlloc(aligned_size, shadow_blocks)) return false;
+  }
+  return true;
+}
+
+std::pair<size_t, size_t>
+VirtualMemoryAutoGrowthBestFitAllocator::SumLargestFreeBlockSizes(
+    size_t n) const {
+  if (n == 0 || free_blocks_.empty()) return std::make_pair(0, 0);
+
+  size_t large_size = free_blocks_.rbegin()->first.first;
+  size_t total_size = 0;
+  size_t count = 0;
+
+  for (auto it = free_blocks_.rbegin(); it != free_blocks_.rend() && count < n;
+       ++it, ++count) {
+    total_size += it->first.first;
+  }
+
+  return std::make_pair(large_size, total_size);
+}
+
+void VirtualMemoryAutoGrowthBestFitAllocator::DumpInfo(
+    std::string phase) const {
+  size_t total = 0, free = 0, used = 0;
+  std::cout << phase << std::endl;
+  std::cout << "All_blocks_:" << std::endl;
+  for (auto block = all_blocks_.begin(); block != all_blocks_.end(); ++block) {
+    std::ostringstream oss_used;
+    std::ostringstream oss_free;
+
+    if (block->is_free_) {
+      free += block->size_;
+      oss_free << "(" << block->size_ << "," << block->ptr_ << ")";
+    } else {
+      used += block->size_;
+      oss_used << "(" << block->size_ << "," << block->ptr_ << ","
+               << block->allocation_->ptr() << ")";
+    }
+
+    std::cout << "is_free? " << block->is_free_ << "[" << oss_used.str()
+              << "]\t[" << oss_free.str() << "]" << std::endl;
+  }
+  std::cout << total << "\t" << used << "\t" << free << std::endl;
+  std::cout << "Free_blocks_:" << std::endl;
+  for (const auto &[key, list_iter] : free_blocks_) {
+    auto [size, ptr] = key;
+    std::cout << "Size: " << size << ", Ptr: " << ptr << "\t" << list_iter->ptr_
+              << std::endl;
+  }
+}
+
+void VirtualMemoryAutoGrowthBestFitAllocator::PreAlloc() {
+  auto pre_alloc_size = FLAGS_vmm_pre_alloc_in_mb << 20;
+  VLOG(4)
+      << "Begin PreAllocate in VirtualMemoryAutoGrowthBestFitAllocator size "
+      << pre_alloc_size;
+  PreAllocate(pre_alloc_size);
+  VLOG(4)
+      << "Finish PreAllocate in VirtualMemoryAutoGrowthBestFitAllocator size "
+      << pre_alloc_size;
+}
+
+void VirtualMemoryAutoGrowthBestFitAllocator::PreAllocate(size_t size) {
+  if (size <= 0) return;
+  ExtendOrCompact(size);
+}
+
+bool VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocator::IsSmallRequest(
+    size_t size) {
+  const size_t routed_size = AlignedSize(size, alignment_);
+  const size_t small_pool_size = FLAGS_vmm_small_pool_size_in_mb << 20;
+  return routed_size < small_pool_size;
+}
+
+void VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocator::PreAlloc() {
+  auto small_allocator =
+      std::dynamic_pointer_cast<VirtualMemoryAutoGrowthBestFitAllocator>(
+          GetSmallAllocator());
+  auto large_allocator =
+      std::dynamic_pointer_cast<VirtualMemoryAutoGrowthBestFitAllocator>(
+          GetLargeAllocator());
+
+  auto vmm_small_pool_pre_alloc = FLAGS_vmm_small_pool_pre_alloc_in_mb << 20;
+  auto vmm_large_pool_pre_alloc = FLAGS_vmm_large_pool_pre_alloc_in_mb << 20;
+
+  if (vmm_small_pool_pre_alloc > 0 && small_allocator) {
+    VLOG(4) << "Begin Small Pool PreAllocate in "
+               "VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocator size "
+            << vmm_small_pool_pre_alloc;
+    small_allocator->PreAllocate(vmm_small_pool_pre_alloc);
+    VLOG(4) << "Finish Small Pool PreAllocate in "
+               "VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocator size "
+            << vmm_small_pool_pre_alloc;
+  }
+  if (vmm_large_pool_pre_alloc > 0 && large_allocator) {
+    VLOG(4) << "Begin Large Pool PreAllocate in "
+               "VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocator size "
+            << vmm_large_pool_pre_alloc;
+    large_allocator->PreAllocate(vmm_large_pool_pre_alloc);
+    VLOG(4) << "Finish Large Pool PreAllocate in "
+               "VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocator size "
+            << vmm_large_pool_pre_alloc;
+  }
+}
+
+size_t VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocator::CompactImpl(
+    const Place &place) {
+  auto large_allocator =
+      std::dynamic_pointer_cast<VirtualMemoryAutoGrowthBestFitAllocator>(
+          GetLargeAllocator());
+  VLOG(1) << "Do Memory Compact Large Pool Manual";
+  size_t compact_free_size = large_allocator->Compact(place);
+  VLOG(1) << "Memory Compact Large Pool Manual Finish Compact size: "
+          << compact_free_size;
+  compact_size_.emplace_back(compact_free_size);
+  return compact_free_size;
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/virtual_memory_auto_growth_best_fit_allocator.h

@@ -0,0 +1,135 @@
+// Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <list>
+#include <map>
+#include <memory>
+#include <optional>
+#include <set>
+#include <vector>
+
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/allocation/spin_lock.h"
+#include "paddle/phi/core/memory/allocation/vmm_ipc_allocation.h"
+#include "paddle/phi/core/memory/mem_utils.h"
+#include "paddle/phi/core/memory/mem_visitor.h"
+
+namespace paddle {
+namespace memory {
+
+namespace allocation {
+/**
+ * Like AutoGrowthBestFitAllocator, VirtualMemoryAutoGrowthBestFitAllocator will
+ * gradually apply to GPU for video memory as the model uses more video memory.
+ * However, the difference is that VirtualMemoryAutoGrowthBestFitAllocator uses
+ * NVIDIA's virtual memory management technology and obtains the virtual memory
+ * address. If the video memory applied for twice is continuous, we can combine
+ * the two video memories later. This combination can greatly reduce
+ * fragmentation.
+ */
+
+class VirtualMemoryAutoGrowthBestFitAllocator : public Allocator {
+ public:
+  VirtualMemoryAutoGrowthBestFitAllocator(
+      const std::shared_ptr<Allocator> &underlying_allocator,
+      size_t alignment,
+      const GPUPlace &place);
+
+  std::shared_ptr<Allocator> &GetUnderLyingAllocator() {
+    return underlying_allocator_;
+  }
+  const std::map<std::pair<size_t, void *>, std::list<Block>::iterator>
+      &GetFreeBlocks() const {
+    return free_blocks_;
+  }
+
+  const std::list<Block> &GetAllBlocks() const { return all_blocks_; }
+
+  std::pair<size_t, size_t> SumLargestFreeBlockSizes(size_t n) const;
+  void Accept(AllocatorVisitor *visitor) override { visitor->Visit(this); }
+
+  bool IsAllocThreadSafe() const override { return true; }
+  void PreAlloc() override;
+  void PreAllocate(size_t size);
+  // Try to simulate an allocation, simulating a request for vector<size>.
+
+  bool TryAllocateBatch(const std::vector<size_t> &sizes);
+
+  bool CollectTensorParts(void *ptr,
+                          size_t size,
+                          std::vector<BlockPart> *parts);
+
+ protected:
+  phi::Allocation *AllocateImpl(size_t size) override;
+  size_t CompactImpl(const Place &place) override;
+  void FreeImpl(phi::Allocation *allocation) override;
+
+ private:
+  // AllocateOrCompact will try to allocate memory from free blocks first, if
+  // OOM happens, it will try to compact memory.
+  std::optional<AllocationPtr> AllocateOrCompact(size_t size);
+  phi::Allocation *AllocFromFreeBlocks(size_t size);
+  void ExtendOrCompact(size_t size);
+  void TryMergeBlock2Blocks(std::list<Block>::iterator iter);
+  void DumpInfo(std::string phase) const;
+
+  std::shared_ptr<Allocator> underlying_allocator_;
+  std::unique_ptr<MemoryCompactionStrategy> memory_compactor_;
+  size_t alignment_;
+
+  std::map<std::pair<size_t, void *>, std::list<Block>::iterator> free_blocks_;
+  std::list<Block> all_blocks_;
+  std::list<AllocationPtr> allocations_;
+  Place place_;
+  SpinLock spinlock_;
+};
+
+/**
+ * VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocator is a multi-scale
+ * allocator that combines the virtual memory management technology of
+ * VirtualMemoryAutoGrowthBestFitAllocator and the multi-scale pooling strategy
+ * of MultiScalePoolAllocator.
+ */
+class VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocator
+    : public MultiScalePoolAllocator {
+ public:
+  VirtualMemoryAutoGrowthBestFitMultiScalePoolAllocator(
+      const std::shared_ptr<VirtualMemoryAutoGrowthBestFitAllocator>
+          &small_allocator,
+      const std::shared_ptr<VirtualMemoryAutoGrowthBestFitAllocator>
+          &large_allocator,
+      size_t alignment,
+      const GPUPlace &place)
+      : MultiScalePoolAllocator(
+            small_allocator, large_allocator, alignment, place),
+        alignment_(alignment) {}
+  bool IsAllocThreadSafe() const override { return true; }
+  void PreAlloc() override;
+  void Accept(AllocatorVisitor *visitor) override { visitor->Visit(this); }
+  bool IsSmallRequest(size_t size) override;
+  std::vector<size_t> GetCompactSize() const { return compact_size_; }
+
+ protected:
+  size_t CompactImpl(const Place &place) override;
+
+ private:
+  size_t alignment_;
+  std::vector<size_t> compact_size_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/vmm_allocator_v2_types.h

@@ -0,0 +1,159 @@
+// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#if defined(PADDLE_WITH_CUDA)
+#include "paddle/phi/backends/dynload/cuda_driver.h"
+#endif
+#include "paddle/phi/core/enforce.h"
+
+#if defined(PADDLE_WITH_CUDA)
+#include "paddle/phi/core/platform/device/gpu/gpu_types.h"
+#endif
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+#if defined(PADDLE_WITH_CUDA)
+using VMMDevicePtr = CUdeviceptr;
+using VMMAllocHandle = CUmemGenericAllocationHandle;
+#else
+using VMMDevicePtr = uintptr_t;
+using VMMAllocHandle = uint64_t;
+#endif
+
+// V2 keeps the bottom-layer shared types independent from the best-fit layer
+// so that CUDAVirtualMemAllocatorV2 can be reviewed and compiled separately.
+enum class PoolType : uint8_t {
+  kSmall = 0,
+  kLarge = 1,
+};
+
+// Fixed-size handle metadata returned by the bottom VMM provider. Upper layers
+// may later reference these handles from block-level views.
+struct VMMHandleMeta {
+  VMMHandleMeta() = default;
+
+  VMMHandleMeta(VMMDevicePtr base,
+                size_t size,
+                VMMAllocHandle handle,
+                int device)
+      : base_(base), size_(size), handle_(handle), device_(device) {}
+
+  VMMDevicePtr base() const { return base_; }
+  size_t size() const { return size_; }
+  VMMAllocHandle handle() const { return handle_; }
+  int device() const { return device_; }
+
+ private:
+  VMMDevicePtr base_{0};
+  size_t size_{0};
+  VMMAllocHandle handle_{0};
+  int device_{0};
+};
+
+// HandleLayout is a lightweight allocation-level handle list returned by the
+// bottom VMM provider. It is used to bootstrap upper-layer block state.
+using HandleLayout = std::vector<std::shared_ptr<VMMHandleMeta>>;
+
+enum class BlockType : uint8_t {
+  kActive = 0,
+  kFree = 1,
+  kUnmappedFree = 2,
+};
+
+struct BlockV2 {
+  static BlockV2 MakeMappedBlock(BlockType type,
+                                 void* ptr,
+                                 size_t size,
+                                 PoolType pool_type) {
+    BlockV2 block;
+    block.Reset(ptr, size, type, pool_type);
+    return block;
+  }
+
+  static BlockV2 MakeUnmappedFreeBlock(void* ptr,
+                                       size_t size,
+                                       PoolType pool_type) {
+    BlockV2 block;
+    block.Reset(ptr, size, BlockType::kUnmappedFree, pool_type);
+    return block;
+  }
+
+  bool IsActive() const { return type_ == BlockType::kActive; }
+  bool IsFree() const { return type_ == BlockType::kFree; }
+  bool IsMappedFree() const { return IsFree(); }
+  bool IsUnmappedFree() const { return type_ == BlockType::kUnmappedFree; }
+  void* ptr() const { return ptr_; }
+  size_t size() const { return size_; }
+  uint8_t* begin_ptr() const { return reinterpret_cast<uint8_t*>(ptr_); }
+  uint8_t* end_ptr() const { return begin_ptr() + size_; }
+  VMMDevicePtr begin_va() const {
+    return reinterpret_cast<VMMDevicePtr>(begin_ptr());
+  }
+  VMMDevicePtr end_va() const { return begin_va() + size_; }
+  bool IsAdjacentBefore(const BlockV2& next) const {
+    return end_ptr() == next.begin_ptr();
+  }
+  bool CanMergeAdjacentFreeBlock(const BlockV2& next) const {
+    return IsFree() && next.IsFree() && IsAdjacentBefore(next);
+  }
+  bool CanMergeAdjacentUnmappedFreeBlock(const BlockV2& next) const {
+    return IsUnmappedFree() && next.IsUnmappedFree() && IsAdjacentBefore(next);
+  }
+  BlockV2 MakeMappedFreeSubBlock(size_t offset, size_t len) const {
+    return MakeMappedBlock(
+        BlockType::kFree, begin_ptr() + offset, len, pool_type_);
+  }
+  BlockV2 MakeMappedActiveSubBlock(size_t offset, size_t len) const {
+    return MakeMappedBlock(
+        BlockType::kActive, begin_ptr() + offset, len, pool_type_);
+  }
+  BlockV2 MakeUnmappedFreeSubBlock(size_t offset, size_t len) const {
+    return MakeUnmappedFreeBlock(begin_ptr() + offset, len, pool_type_);
+  }
+  void MarkActive() { type_ = BlockType::kActive; }
+  void MarkFree() { type_ = BlockType::kFree; }
+  void Reset(void* ptr, size_t size, BlockType type, PoolType pool_type) {
+    ptr_ = ptr;
+    size_ = size;
+    type_ = type;
+    pool_type_ = pool_type;
+  }
+  void TrimToPrefix(size_t keep) { size_ = keep; }
+  void TrimToSuffix(size_t trim, size_t keep) {
+    ptr_ = reinterpret_cast<uint8_t*>(ptr_) + trim;
+    size_ = keep;
+  }
+  void MergeAdjacentBlock(const BlockV2& src) { size_ += src.size_; }
+  void MergeAdjacentUnmappedFreeBlock(const BlockV2& src) {
+    size_ += src.size_;
+  }
+
+  void* ptr_{nullptr};
+  size_t size_{0};
+  BlockType type_{BlockType::kUnmappedFree};
+
+  PoolType pool_type_{PoolType::kLarge};
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/vmm_auto_growth_best_fit_allocator_v2.cc

@@ -0,0 +1,668 @@
+// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/vmm_auto_growth_best_fit_allocator_v2.h"
+
+#if defined(PADDLE_WITH_CUDA)
+
+#include <exception>
+
+#include "glog/logging.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/platform/cuda_device_guard.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+namespace {
+
+template <typename Map, typename Key, typename Value>
+void EmplaceOrEnforce(Map* map,
+                      Key&& key,
+                      Value&& value,
+                      const char* map_name) {
+  const bool inserted =
+      map->try_emplace(std::forward<Key>(key), std::forward<Value>(value))
+          .second;
+  PADDLE_ENFORCE_EQ(
+      inserted,
+      true,
+      common::errors::AlreadyExists(
+          "Duplicate key inserted into %s, allocator state is inconsistent.",
+          map_name));
+}
+
+}  // namespace
+
+void VMMAutoGrowthBestFitAllocatorV2::UnderlyingAllocationRegistry::Add(
+    DecoratedAllocationPtr allocation) {
+  allocations_.emplace_back(std::move(allocation));
+  auto it = std::prev(allocations_.end());
+  auto* begin = Begin(*it);
+  PADDLE_ENFORCE_EQ(
+      allocations_by_ptr_.emplace(begin, it).second,
+      true,
+      common::errors::AlreadyExists(
+          "Duplicate underlying allocation base %p in VMM V2 registry.",
+          begin));
+}
+
+namespace {
+
+bool RangesOverlap(void* lhs_ptr,
+                   size_t lhs_size,
+                   void* rhs_ptr,
+                   size_t rhs_size) {
+  const auto* lhs_begin = reinterpret_cast<const uint8_t*>(lhs_ptr);
+  const auto* lhs_end = lhs_begin + lhs_size;
+  const auto* rhs_begin = reinterpret_cast<const uint8_t*>(rhs_ptr);
+  const auto* rhs_end = rhs_begin + rhs_size;
+  return lhs_end > rhs_begin && rhs_end > lhs_begin;
+}
+
+}  // namespace
+
+uint8_t* VMMAutoGrowthBestFitAllocatorV2::UnderlyingAllocationRegistry::Begin(
+    const DecoratedAllocationPtr& allocation) {
+  return reinterpret_cast<uint8_t*>(allocation->ptr());
+}
+
+uint8_t* VMMAutoGrowthBestFitAllocatorV2::UnderlyingAllocationRegistry::End(
+    const DecoratedAllocationPtr& allocation) {
+  return Begin(allocation) + allocation->size();
+}
+
+bool VMMAutoGrowthBestFitAllocatorV2::UnderlyingAllocationRegistry::HasOverlap(
+    void* ptr, size_t size) const {
+  auto* begin = reinterpret_cast<uint8_t*>(ptr);
+  auto* end = begin + size;
+  auto it = allocations_by_ptr_.lower_bound(begin);
+  if (it != allocations_by_ptr_.begin()) {
+    auto prev = std::prev(it);
+    if (End(*prev->second) > begin) {
+      return true;
+    }
+  }
+  return it != allocations_by_ptr_.end() && it->first < end;
+}
+
+bool VMMAutoGrowthBestFitAllocatorV2::UnderlyingAllocationRegistry::Overlaps(
+    void* ptr, size_t size) const {
+  return HasOverlap(ptr, size);
+}
+
+VMMAutoGrowthBestFitAllocatorV2::UnderlyingAllocationRegistry::iterator
+VMMAutoGrowthBestFitAllocatorV2::UnderlyingAllocationRegistry::Erase(
+    iterator it) {
+  allocations_by_ptr_.erase(Begin(*it));
+  return allocations_.erase(it);
+}
+
+VMMAutoGrowthBestFitAllocatorV2::VMMAutoGrowthBestFitAllocatorV2(
+    const std::shared_ptr<CUDAVirtualMemAllocatorV2>& underlying_allocator,
+    size_t alignment,
+    const GPUPlace& place,
+    PoolType pool_type)
+    : underlying_allocator_(underlying_allocator),
+      alignment_(alignment),
+      place_(place),
+      pool_type_(pool_type) {}
+
+phi::Allocation* VMMAutoGrowthBestFitAllocatorV2::AllocateImpl(size_t size) {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  const size_t requested_size = AlignedSize(size, alignment_);
+  if (auto* allocation = AllocFromFreeBlocks(requested_size)) {
+    return allocation;
+  }
+  if (auto* allocation = AllocFromUnmappedFreeBlocks(requested_size)) {
+    return allocation;
+  }
+
+  // Tail reuse: if the last block in the address space is FREE, detach it
+  // and only request the difference from the underlying allocator. The
+  // underlying VMM provider maps new handles at a monotonically increasing
+  // VA cursor, so the new allocation is guaranteed to be contiguous with
+  // the tail FREE block.
+  bool has_tail_reuse = false;
+  size_t tail_reuse_size = 0;
+  BlockV2 combined_free_block;
+  if (!all_blocks_.empty()) {
+    auto tail_it = std::prev(all_blocks_.end());
+    if (CanIndexFreeBlock(*tail_it)) {
+      has_tail_reuse = true;
+      tail_reuse_size = tail_it->size_;
+      EraseFreeBlock(tail_it);
+      combined_free_block = std::move(*tail_it);
+      all_blocks_.erase(tail_it);
+    }
+  }
+
+  const size_t grow_size = (requested_size > tail_reuse_size)
+                               ? (requested_size - tail_reuse_size)
+                               : 0;
+  auto restore_tail_free_block = [&] {
+    if (has_tail_reuse) {
+      auto restored_it =
+          all_blocks_.insert(all_blocks_.end(), std::move(combined_free_block));
+      InsertFreeBlock(restored_it);
+    }
+  };
+
+  // Grow: obtain a new raw allocation from the bottom VMM provider.
+  // If cuMemCreate fails due to physical memory exhaustion (CU error 2),
+  // the driver-level allocator throws EnforceNotMet. Convert it to BadAlloc so
+  // the outer retry path can handle it.
+  CUDAVirtualMemAllocatorV2::AllocationWithBlock grow_alloc;
+  if (grow_size > 0) {
+    try {
+      grow_alloc = underlying_allocator_->AppendWithBlock(grow_size);
+    } catch (const BadAlloc& bad_alloc) {
+      restore_tail_free_block();
+      PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+          "VMM V2 best-fit allocator (pool %d) failed to grow by %zu bytes.\n"
+          "Underlying VMM allocation failure:\n%s",
+          static_cast<int>(pool_type_),
+          grow_size,
+          bad_alloc.what()));
+    } catch (const std::exception& e) {
+      restore_tail_free_block();
+      PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+          "VMM V2 best-fit allocator (pool %d) failed to grow by %zu bytes.\n"
+          "Underlying VMM allocation exception:\n%s",
+          static_cast<int>(pool_type_),
+          grow_size,
+          e.what()));
+    } catch (...) {
+      restore_tail_free_block();
+      PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+          "VMM V2 best-fit allocator (pool %d) failed to grow by %zu bytes "
+          "with an unknown underlying VMM allocation exception.",
+          static_cast<int>(pool_type_),
+          grow_size));
+    }
+  }
+
+  size_t total_new_size = tail_reuse_size;
+
+  if (grow_alloc.HasAllocation()) {
+    BlockV2 grow_block = AdoptBackingBlock(&grow_alloc);
+    total_new_size += grow_block.size_;
+    if (has_tail_reuse) {
+      combined_free_block.MergeAdjacentBlock(grow_block);
+    } else {
+      combined_free_block = std::move(grow_block);
+    }
+  }
+
+  const size_t remaining_size = total_new_size - requested_size;
+
+  BlockV2 block =
+      combined_free_block.MakeMappedActiveSubBlock(0, requested_size);
+  auto it = all_blocks_.insert(all_blocks_.end(), std::move(block));
+
+  if (remaining_size > 0) {
+    BlockV2 remaining_block = combined_free_block.MakeMappedFreeSubBlock(
+        requested_size, remaining_size);
+    auto remain_it =
+        all_blocks_.insert(std::next(it), std::move(remaining_block));
+    InsertFreeBlock(remain_it);
+  }
+
+  return new VMMAutoGrowthBestFitBlockAllocationV2(it, place_, this);  // NOLINT
+}
+
+void VMMAutoGrowthBestFitAllocatorV2::FreeImpl(phi::Allocation* allocation) {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  auto* wrapped_allocation =
+      static_cast<VMMAutoGrowthBestFitBlockAllocationV2*>(allocation);
+  auto it = wrapped_allocation->block_it();
+  PADDLE_ENFORCE_NE(
+      it,
+      all_blocks_.end(),
+      common::errors::NotFound("Can not find active block for allocation %p in "
+                               "VMMAutoGrowthBestFitAllocatorV2.",
+                               allocation->ptr()));
+  it->MarkFree();
+  TryMerge(it);
+  delete allocation;
+}
+
+phi::Allocation* VMMAutoGrowthBestFitAllocatorV2::AllocFromFreeBlocks(
+    size_t size) {
+  auto it = free_blocks_.lower_bound({size, nullptr});
+  if (it == free_blocks_.end()) {
+    return nullptr;
+  }
+
+  auto block_it = it->second;
+  PADDLE_ENFORCE_EQ(
+      CanIndexFreeBlock(*block_it),
+      true,
+      common::errors::PreconditionNotMet(
+          "VMM V2 free block index points to a non-reusable block."));
+  free_blocks_.erase(it);
+
+  if (block_it->size_ > size) {
+    const size_t remaining_size = block_it->size_ - size;
+    BlockV2 remaining_block =
+        block_it->MakeMappedFreeSubBlock(size, remaining_size);
+    block_it->TrimToPrefix(size);
+    auto remain_it =
+        all_blocks_.insert(std::next(block_it), std::move(remaining_block));
+    InsertFreeBlock(remain_it);
+  }
+
+  block_it->MarkActive();
+  return new  // NOLINT
+      VMMAutoGrowthBestFitBlockAllocationV2(block_it, place_, this);
+}
+
+phi::Allocation* VMMAutoGrowthBestFitAllocatorV2::AllocFromUnmappedFreeBlocks(
+    size_t size) {
+  const size_t backing_size =
+      AlignedSize(size, underlying_allocator_->handle_size());
+  BlockListIt best = all_blocks_.end();
+  for (auto iter = unmapped_free_blocks_.lower_bound({backing_size, nullptr});
+       iter != unmapped_free_blocks_.end();) {
+    auto it = iter->second;
+    if (!it->IsUnmappedFree()) {
+      iter = unmapped_free_blocks_.erase(iter);
+      continue;
+    }
+    if (RangeOverlapsUnderlying(it->ptr_, backing_size)) {
+      VLOG(6) << "VMM V2 AllocFromUnmappedFreeBlocks skip ownership-overlapped "
+                 "unmapped-free ptr="
+              << it->ptr_ << " backing_size=" << backing_size
+              << " block_size=" << it->size_;
+      ++iter;
+      continue;
+    }
+    best = it;
+    break;
+  }
+  if (best == all_blocks_.end()) {
+    return nullptr;
+  }
+
+  const auto unmapped_free_ptr = best->begin_va();
+  VLOG(6) << "VMM V2 AllocFromUnmappedFreeBlocks ptr="
+          << reinterpret_cast<void*>(unmapped_free_ptr) << " requested=" << size
+          << " backing_size=" << backing_size
+          << " original_unmapped_free_size=" << best->size_
+          << " tail_offset=" << underlying_allocator_->tail_offset();
+  CUDAVirtualMemAllocatorV2::AllocationWithBlock unmapped_free_alloc;
+  try {
+    unmapped_free_alloc = underlying_allocator_->PlaceAtVAWithBlock(
+        unmapped_free_ptr, backing_size);
+  } catch (const BadAlloc&) {
+    // Do not mutate the allocation view if backing cannot be created in this
+    // unmapped-free range due to physical memory pressure. The normal grow
+    // path will surface the allocation failure if needed. Other exceptions
+    // indicate allocator state bugs and must not be hidden as a cache miss.
+    return nullptr;
+  }
+
+  BlockV2 mapped_block = AdoptBackingBlock(&unmapped_free_alloc);
+  PADDLE_ENFORCE_EQ(
+      mapped_block.size_,
+      backing_size,
+      common::errors::InvalidArgument(
+          "Unexpected unmapped-free backing size: got %zu, expected %zu.",
+          mapped_block.size_,
+          backing_size));
+
+  const size_t original_unmapped_free_size = best->size_;
+  const PoolType original_pool_type = best->pool_type_;
+
+  EraseUnmappedFreeBlock(best);
+  *best = mapped_block.MakeMappedActiveSubBlock(0, size);
+
+  auto insert_pos = std::next(best);
+  if (backing_size > size) {
+    BlockV2 mapped_remain =
+        mapped_block.MakeMappedFreeSubBlock(size, backing_size - size);
+    auto free_it = all_blocks_.insert(insert_pos, std::move(mapped_remain));
+    InsertFreeBlock(free_it);
+    insert_pos = std::next(free_it);
+  }
+
+  if (original_unmapped_free_size > backing_size) {
+    BlockV2 tail_unmapped_free = BlockV2::MakeUnmappedFreeBlock(
+        reinterpret_cast<uint8_t*>(best->ptr_) + backing_size,
+        original_unmapped_free_size - backing_size,
+        original_pool_type);
+    auto tail_it =
+        all_blocks_.insert(insert_pos, std::move(tail_unmapped_free));
+    InsertUnmappedFreeBlock(tail_it);
+  }
+
+  return new  // NOLINT
+      VMMAutoGrowthBestFitBlockAllocationV2(best, place_, this);
+}
+
+void VMMAutoGrowthBestFitAllocatorV2::TrackUnderlyingAllocation(
+    DecoratedAllocationPtr allocation) {
+  underlying_allocations_.Add(std::move(allocation));
+}
+
+BlockV2 VMMAutoGrowthBestFitAllocatorV2::AdoptBackingBlock(
+    CUDAVirtualMemAllocatorV2::AllocationWithBlock* allocation_with_block) {
+  PADDLE_ENFORCE_NOT_NULL(
+      allocation_with_block,
+      common::errors::InvalidArgument(
+          "AllocationWithBlock must not be null when adopting block."));
+  BlockV2 block = allocation_with_block->TakeBlock();
+  auto allocation = static_unique_ptr_cast<Allocation>(
+      allocation_with_block->TakeAllocation());
+  TrackUnderlyingAllocation(std::move(allocation));
+  return block;
+}
+
+bool VMMAutoGrowthBestFitAllocatorV2::RangeOverlapsUnderlying(
+    void* ptr, size_t size) const {
+  return underlying_allocations_.Overlaps(ptr, size);
+}
+
+bool VMMAutoGrowthBestFitAllocatorV2::HasReleasableIdleUnderlying() const {
+  for (const auto& allocation : underlying_allocations_) {
+    auto* base = reinterpret_cast<uint8_t*>(allocation->ptr());
+    if (CanReleaseIdleUnderlying(base, allocation->size())) {
+      return true;
+    }
+  }
+  return false;
+}
+
+bool VMMAutoGrowthBestFitAllocatorV2::CanReleaseIdleUnderlying(
+    uint8_t* base, size_t size) const {
+  if (!IsRangeEntirelyFree(base, size)) {
+    return false;
+  }
+  return underlying_allocator_->IsRangeReleasable(
+      reinterpret_cast<VMMDevicePtr>(base), size);
+}
+
+bool VMMAutoGrowthBestFitAllocatorV2::TryReleaseIdleUnderlying(
+    UnderlyingAllocationRegistry::iterator* alloc_it, uint64_t* released) {
+  auto* allocation = (**alloc_it).get();
+  auto* base = reinterpret_cast<uint8_t*>(allocation->ptr());
+  const size_t alloc_size = allocation->size();
+  if (!CanReleaseIdleUnderlying(base, alloc_size)) {
+    return false;
+  }
+
+  ReplaceRangeWithUnmappedFree(base, alloc_size);
+  *released += alloc_size;
+  VLOG(5) << "VMM V2 pool " << static_cast<int>(pool_type_)
+          << " released idle chunk: " << alloc_size << " bytes";
+  *alloc_it = underlying_allocations_.Erase(*alloc_it);
+  return true;
+}
+
+bool VMMAutoGrowthBestFitAllocatorV2::CanIndexFreeBlock(
+    const BlockV2& block) const {
+  return block.IsMappedFree();
+}
+
+void VMMAutoGrowthBestFitAllocatorV2::InsertFreeBlock(BlockListIt it) {
+  if (!CanIndexFreeBlock(*it)) {
+    return;
+  }
+  EmplaceOrEnforce(
+      &free_blocks_, std::make_pair(it->size_, it->ptr_), it, "free_blocks_");
+}
+
+void VMMAutoGrowthBestFitAllocatorV2::EraseFreeBlock(BlockListIt it) {
+  free_blocks_.erase({it->size_, it->ptr_});
+}
+
+void VMMAutoGrowthBestFitAllocatorV2::InsertUnmappedFreeBlock(BlockListIt it) {
+  if (!it->IsUnmappedFree()) {
+    return;
+  }
+  EmplaceOrEnforce(&unmapped_free_blocks_,
+                   std::make_pair(it->size_, it->ptr_),
+                   it,
+                   "unmapped_free_blocks_");
+}
+
+void VMMAutoGrowthBestFitAllocatorV2::EraseUnmappedFreeBlock(BlockListIt it) {
+  unmapped_free_blocks_.erase({it->size_, it->ptr_});
+}
+
+void VMMAutoGrowthBestFitAllocatorV2::TryMerge(BlockListIt it) {
+  // Only adjacent FREE blocks are merged here. ACTIVE blocks are never touched,
+  // and unmapped-free blocks remain as explicit holes for later reuse.
+  // all_blocks_ is the full VA-ordered block list, so adjacency is checked
+  // against neighboring entries in that list.
+  if (it != all_blocks_.begin()) {
+    auto prev = std::prev(it);
+    if (prev->CanMergeAdjacentFreeBlock(*it)) {
+      EraseFreeBlock(prev);
+      prev->MergeAdjacentBlock(*it);
+      all_blocks_.erase(it);
+      it = prev;
+    }
+  }
+
+  auto next = std::next(it);
+  if (next != all_blocks_.end() && it->CanMergeAdjacentFreeBlock(*next)) {
+    EraseFreeBlock(next);
+    it->MergeAdjacentBlock(*next);
+    all_blocks_.erase(next);
+  }
+
+  InsertFreeBlock(it);
+}
+
+void VMMAutoGrowthBestFitAllocatorV2::TryMergeUnmappedFree(BlockListIt it) {
+  if (it == all_blocks_.end() || !it->IsUnmappedFree()) {
+    return;
+  }
+
+  if (it != all_blocks_.begin()) {
+    auto prev = std::prev(it);
+    if (prev->CanMergeAdjacentUnmappedFreeBlock(*it)) {
+      EraseUnmappedFreeBlock(prev);
+      EraseUnmappedFreeBlock(it);
+      prev->MergeAdjacentUnmappedFreeBlock(*it);
+      all_blocks_.erase(it);
+      it = prev;
+      InsertUnmappedFreeBlock(it);
+    }
+  }
+
+  auto next = std::next(it);
+  if (next != all_blocks_.end() &&
+      it->CanMergeAdjacentUnmappedFreeBlock(*next)) {
+    EraseUnmappedFreeBlock(it);
+    EraseUnmappedFreeBlock(next);
+    it->MergeAdjacentUnmappedFreeBlock(*next);
+    all_blocks_.erase(next);
+    InsertUnmappedFreeBlock(it);
+  }
+}
+
+// ---------------------------------------------------------------------------
+// ReleaseImpl / FreeIdleChunks: release underlying allocations whose entire
+// VA range is covered by FREE blocks back to the CUDA VMM driver.
+//
+// Because TryMerge may have merged FREE blocks across allocation boundaries,
+// we must split the spanning block at the allocation edges, release the
+// backing, and keep the released VA range as explicit unmapped-free space for
+// later reuse.
+// ---------------------------------------------------------------------------
+
+uint64_t VMMAutoGrowthBestFitAllocatorV2::ReleaseImpl(
+    const Place& place UNUSED) {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  if (!HasReleasableIdleUnderlying()) {
+    return 0;
+  }
+  // FreeIdleChunks may release CUDA VMM mappings and physical handles. Those
+  // driver calls are not ordered by the stream-safe wrapper, so wait before
+  // making any previously returned VA range invalid.
+  platform::CUDADeviceGuard device_guard(place_.device);
+  PADDLE_ENFORCE_GPU_SUCCESS(cudaDeviceSynchronize());
+  return FreeIdleChunks();
+}
+
+uint64_t VMMAutoGrowthBestFitAllocatorV2::FreeIdleChunks() {
+  uint64_t released = 0;
+
+  for (auto alloc_it = underlying_allocations_.begin();
+       alloc_it != underlying_allocations_.end();) {
+    if (!TryReleaseIdleUnderlying(&alloc_it, &released)) {
+      ++alloc_it;
+    }
+  }
+
+  TrimTrailingUnmappedFreeBlocks();
+  underlying_allocator_->SetTailOffset(ComputeTailOffset());
+  return released;
+}
+
+void VMMAutoGrowthBestFitAllocatorV2::TrimTrailingUnmappedFreeBlocks() {
+  while (!all_blocks_.empty()) {
+    auto tail_it = std::prev(all_blocks_.end());
+    if (!tail_it->IsUnmappedFree() ||
+        underlying_allocations_.Overlaps(tail_it->ptr_, tail_it->size_)) {
+      break;
+    }
+    EraseUnmappedFreeBlock(tail_it);
+    all_blocks_.erase(tail_it);
+  }
+}
+
+size_t VMMAutoGrowthBestFitAllocatorV2::ComputeTailOffset() const {
+  for (auto it = all_blocks_.rbegin(); it != all_blocks_.rend(); ++it) {
+    if (it->IsUnmappedFree() &&
+        !underlying_allocations_.Overlaps(it->ptr_, it->size_)) {
+      continue;
+    }
+    return static_cast<size_t>(it->end_va() -
+                               underlying_allocator_->virtual_mem_base());
+  }
+  return 0;
+}
+
+bool VMMAutoGrowthBestFitAllocatorV2::IsRangeEntirelyFree(uint8_t* base,
+                                                          size_t size) const {
+  auto* end = base + size;
+  for (const auto& block : all_blocks_) {
+    auto* bptr = block.begin_ptr();
+    auto* bend = block.end_ptr();
+    if (bend <= base) continue;
+    if (bptr >= end) break;
+    if (block.IsActive()) {
+      return false;
+    }
+  }
+  // Return true when the range contains only FREE/unmapped-free blocks, or
+  // when blocks have already been removed by a prior FreeIdleChunks pass.
+  return true;
+}
+
+void VMMAutoGrowthBestFitAllocatorV2::ReplaceRangeWithUnmappedFree(
+    uint8_t* base, size_t size) {
+  auto* end = base + size;
+  auto erase_free_index = [this](BlockList::iterator it) {
+    if (it->IsUnmappedFree()) {
+      EraseUnmappedFreeBlock(it);
+    } else {
+      EraseFreeBlock(it);
+    }
+  };
+  auto insert_free_index = [this](BlockList::iterator it) {
+    if (it->IsUnmappedFree()) {
+      InsertUnmappedFreeBlock(it);
+    } else {
+      InsertFreeBlock(it);
+    }
+  };
+
+  for (auto it = all_blocks_.begin(); it != all_blocks_.end();) {
+    auto* bptr = it->begin_ptr();
+    auto* bend = it->end_ptr();
+
+    if (bend <= base) {
+      ++it;
+      continue;
+    }
+    if (bptr >= end) break;
+
+    // Case 1: block entirely within [base, end): remove it.
+    if (bptr >= base && bend <= end) {
+      erase_free_index(it);
+      it = all_blocks_.erase(it);
+      continue;
+    }
+
+    // Case 2: block straddles left boundary only: keep left remnant.
+    if (bptr < base && bend <= end) {
+      const size_t keep = static_cast<size_t>(base - bptr);
+      erase_free_index(it);
+      it->TrimToPrefix(keep);
+      insert_free_index(it);
+      ++it;
+      continue;
+    }
+
+    // Case 3: block straddles right boundary only: keep right remnant.
+    if (bptr >= base && bend > end) {
+      const size_t trim = static_cast<size_t>(end - bptr);
+      const size_t keep = it->size_ - trim;
+      erase_free_index(it);
+      it->TrimToSuffix(trim, keep);
+      insert_free_index(it);
+      break;  // nothing more in range
+    }
+
+    // Case 4: block fully encompasses [base, end): split into two.
+    if (bptr < base && bend > end) {
+      const size_t left_size = static_cast<size_t>(base - bptr);
+      const size_t right_offset = static_cast<size_t>(end - bptr);
+      const size_t right_size = it->size_ - right_offset;
+      BlockV2 right =
+          it->IsUnmappedFree()
+              ? it->MakeUnmappedFreeSubBlock(right_offset, right_size)
+              : it->MakeMappedFreeSubBlock(right_offset, right_size);
+
+      erase_free_index(it);
+      it->TrimToPrefix(left_size);
+      insert_free_index(it);
+      auto right_it = all_blocks_.insert(std::next(it), std::move(right));
+      insert_free_index(right_it);
+      break;  // done
+    }
+
+    ++it;
+  }
+
+  auto insert_pos = all_blocks_.begin();
+  while (insert_pos != all_blocks_.end() && insert_pos->begin_ptr() < base) {
+    ++insert_pos;
+  }
+  auto unmapped_it = all_blocks_.insert(
+      insert_pos, BlockV2::MakeUnmappedFreeBlock(base, size, pool_type_));
+  InsertUnmappedFreeBlock(unmapped_it);
+  TryMergeUnmappedFree(unmapped_it);
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle
+
+#endif

paddle/phi/core/memory/allocation/vmm_auto_growth_best_fit_allocator_v2.h

@@ -0,0 +1,142 @@
+// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <functional>
+#include <list>
+#include <map>
+#include <memory>
+
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/allocation/cuda_virtual_mem_allocator_v2.h"
+#include "paddle/phi/core/memory/allocation/spin_lock.h"
+#include "paddle/phi/core/memory/allocation/vmm_allocator_v2_types.h"
+#include "paddle/phi/core/memory/mem_visitor.h"
+
+#if defined(PADDLE_WITH_CUDA)
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+using BlockList = std::list<BlockV2>;
+using BlockListIt = BlockList::iterator;
+
+class VMMAutoGrowthBestFitAllocatorV2;
+
+class VMMAutoGrowthBestFitBlockAllocationV2 : public Allocation {
+ public:
+  VMMAutoGrowthBestFitBlockAllocationV2(BlockListIt block_it,
+                                        const Place& place,
+                                        VMMAutoGrowthBestFitAllocatorV2* owner)
+      : Allocation(block_it->ptr_, block_it->ptr_, block_it->size_, place),
+        block_it_(block_it),
+        owner_(owner) {}
+
+  BlockListIt block_it() const { return block_it_; }
+
+ private:
+  BlockListIt block_it_;
+  VMMAutoGrowthBestFitAllocatorV2* owner_;
+};
+
+class VMMAutoGrowthBestFitAllocatorV2 : public Allocator {
+ public:
+  VMMAutoGrowthBestFitAllocatorV2(
+      const std::shared_ptr<CUDAVirtualMemAllocatorV2>& underlying_allocator,
+      size_t alignment,
+      const GPUPlace& place,
+      PoolType pool_type);
+
+  bool IsAllocThreadSafe() const override { return true; }
+  void Accept(AllocatorVisitor* visitor) override { visitor->Visit(this); }
+
+  const BlockList& all_blocks() const { return all_blocks_; }
+  PoolType pool_type() const { return pool_type_; }
+  size_t alignment() const { return alignment_; }
+
+ protected:
+  phi::Allocation* AllocateImpl(size_t size) override;
+  void FreeImpl(phi::Allocation* allocation) override;
+  uint64_t ReleaseImpl(const Place& place) override;
+
+ private:
+  struct UnderlyingAllocationRegistry {
+    using List = std::list<DecoratedAllocationPtr>;
+    using iterator = List::iterator;
+    using OverlapPredicate = std::function<bool(const DecoratedAllocationPtr&)>;
+
+    void Add(DecoratedAllocationPtr allocation);
+    bool Overlaps(void* ptr, size_t size) const;
+    iterator begin() { return allocations_.begin(); }
+    iterator end() { return allocations_.end(); }
+    List::const_iterator begin() const { return allocations_.begin(); }
+    List::const_iterator end() const { return allocations_.end(); }
+    iterator Erase(iterator it);
+
+   private:
+    using Index = std::map<uint8_t*, iterator>;
+    static uint8_t* Begin(const DecoratedAllocationPtr& allocation);
+    static uint8_t* End(const DecoratedAllocationPtr& allocation);
+    bool HasOverlap(void* ptr, size_t size) const;
+
+    List allocations_;
+    Index allocations_by_ptr_;
+  };
+
+  phi::Allocation* AllocFromFreeBlocks(size_t size);
+  phi::Allocation* AllocFromUnmappedFreeBlocks(size_t size);
+  BlockV2 AdoptBackingBlock(
+      CUDAVirtualMemAllocatorV2::AllocationWithBlock* allocation_with_block);
+  void TrackUnderlyingAllocation(DecoratedAllocationPtr allocation);
+  bool RangeOverlapsUnderlying(void* ptr, size_t size) const;
+  bool HasReleasableIdleUnderlying() const;
+  bool CanReleaseIdleUnderlying(uint8_t* base, size_t size) const;
+  bool TryReleaseIdleUnderlying(
+      UnderlyingAllocationRegistry::iterator* alloc_it, uint64_t* released);
+  bool CanIndexFreeBlock(const BlockV2& block) const;
+  void InsertFreeBlock(BlockListIt it);
+  void EraseFreeBlock(BlockListIt it);
+  void InsertUnmappedFreeBlock(BlockListIt it);
+  void EraseUnmappedFreeBlock(BlockListIt it);
+  void TryMerge(BlockListIt it);
+  void TryMergeUnmappedFree(BlockListIt it);
+  uint64_t FreeIdleChunks();
+  void TrimTrailingUnmappedFreeBlocks();
+  size_t ComputeTailOffset() const;
+  bool IsRangeEntirelyFree(uint8_t* base, size_t size) const;
+  void ReplaceRangeWithUnmappedFree(uint8_t* base, size_t size);
+
+  // Best-fit V2 only grows from the fixed-handle CUDA VMM provider. The
+  // bottom allocator returns mapped-free BlockV2 views, while best-fit owns
+  // allocation/free-list policy over those block views.
+  std::shared_ptr<CUDAVirtualMemAllocatorV2> underlying_allocator_;
+  size_t alignment_;
+  GPUPlace place_;
+  PoolType pool_type_;
+  UnderlyingAllocationRegistry underlying_allocations_;
+  // Full block list ordered by VA address. This is the source of truth and
+  // contains ACTIVE/FREE/UNMAPPED-FREE blocks together.
+  BlockList all_blocks_;
+  std::map<std::pair<size_t, void*>, BlockListIt> free_blocks_;
+  std::map<std::pair<size_t, void*>, BlockListIt> unmapped_free_blocks_;
+  mutable SpinLock spinlock_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle
+
+#endif

paddle/phi/core/memory/allocation/vmm_backing_map.cc

@@ -0,0 +1,330 @@
+// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/vmm_backing_map.h"
+
+#if defined(PADDLE_WITH_CUDA)
+
+#include <algorithm>
+#include <mutex>
+
+#include "glog/logging.h"
+#include "paddle/phi/core/enforce.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+namespace {
+
+bool AddOverflow(VMMDevicePtr base, size_t size) { return base + size < base; }
+
+bool ComputeOverlappedPages(VMMDevicePtr base,
+                            size_t backing_size,
+                            size_t page_size,
+                            VMMDevicePtr va,
+                            size_t size,
+                            const char* context,
+                            size_t* start,
+                            size_t* count) {
+  if (size == 0 || page_size == 0 || AddOverflow(base, backing_size) ||
+      va < base || va + size < va || va + size > base + backing_size) {
+    VLOG(0) << "VMM V2 BackingMap invalid overlap range in " << context
+            << ": va=" << reinterpret_cast<void*>(va) << " size=" << size
+            << " base=" << reinterpret_cast<void*>(base)
+            << " backing_size=" << backing_size << " page_size=" << page_size;
+    return false;
+  }
+  const size_t begin_offset = va - base;
+  const size_t end_offset = va + size - base;
+  *start = begin_offset / page_size;
+  const size_t end_page = (end_offset + page_size - 1) / page_size;
+  *count = end_page - *start;
+  return true;
+}
+
+}  // namespace
+
+void VMMBackingMap::Configure(VMMDevicePtr base,
+                              size_t size,
+                              size_t page_size,
+                              int device) {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  PADDLE_ENFORCE_GT(page_size,
+                    0UL,
+                    common::errors::InvalidArgument(
+                        "VMM V2 BackingMap page_size must be positive."));
+  PADDLE_ENFORCE_EQ(
+      size % page_size,
+      0UL,
+      common::errors::InvalidArgument(
+          "VMM V2 BackingMap size %zu must be page-aligned by page_size %zu.",
+          size,
+          page_size));
+  if (configured_) {
+    if (base_ != base || size_ != size || page_size_ != page_size ||
+        device_ != device) {
+      VLOG(0) << "VMM V2 BackingMap reconfigure mismatch: old_base="
+              << reinterpret_cast<void*>(base_)
+              << " new_base=" << reinterpret_cast<void*>(base)
+              << " old_size=" << size_ << " new_size=" << size
+              << " old_page_size=" << page_size_
+              << " new_page_size=" << page_size << " old_device=" << device_
+              << " new_device=" << device;
+    }
+    return;
+  }
+
+  base_ = base;
+  size_ = size;
+  page_size_ = page_size;
+  device_ = device;
+  configured_ = true;
+  pages_.resize(size_ / page_size_);
+  mapped_page_count_ = 0;
+}
+
+bool VMMBackingMap::CheckRangeLocked(VMMDevicePtr va,
+                                     size_t size,
+                                     const char* context,
+                                     size_t* start,
+                                     size_t* count) const {
+  if (!configured_) {
+    VLOG(0) << "VMM V2 BackingMap " << context
+            << " before Configure, va=" << reinterpret_cast<void*>(va)
+            << " size=" << size;
+    return false;
+  }
+  if (size == 0 || page_size_ == 0 || size % page_size_ != 0 ||
+      AddOverflow(base_, size_) || va < base_ || va + size < va ||
+      va + size > base_ + size_ || (va - base_) % page_size_ != 0) {
+    VLOG(0) << "VMM V2 BackingMap invalid range in " << context
+            << ": va=" << reinterpret_cast<void*>(va) << " size=" << size
+            << " base=" << reinterpret_cast<void*>(base_)
+            << " backing_size=" << size_ << " page_size=" << page_size_;
+    return false;
+  }
+  *start = (va - base_) / page_size_;
+  *count = size / page_size_;
+  return true;
+}
+
+void VMMBackingMap::MarkPageMappedLocked(
+    Page* page,
+    VMMDevicePtr page_va,
+    VMMAllocHandle handle,
+    const std::shared_ptr<VMMHandleMeta>& meta) {
+  PADDLE_ENFORCE_EQ(
+      page->mapped && handle != 0 && page->handle != handle,
+      false,
+      common::errors::PreconditionNotMet(
+          "VMM V2 BackingMap cannot overwrite mapped page at %p from "
+          "handle %p to %p.",
+          reinterpret_cast<void*>(page_va),
+          reinterpret_cast<void*>(page->handle),
+          reinterpret_cast<void*>(handle)));
+  if (!page->mapped) {
+    mapped_page_count_++;
+  }
+  page->handle = handle;
+  page->meta = meta;
+  page->mapped = true;
+  page->epoch++;
+}
+
+void VMMBackingMap::ResetPageToUnmappedLocked(Page* page) {
+  if (page->mapped && mapped_page_count_ > 0) {
+    mapped_page_count_--;
+  }
+  page->handle = 0;
+  page->meta.reset();
+  page->mapped = false;
+  page->epoch++;
+}
+
+void VMMBackingMap::MarkMapped(VMMDevicePtr va,
+                               VMMAllocHandle handle,
+                               size_t size) {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  size_t start = 0;
+  size_t count = 0;
+  if (!CheckRangeLocked(va, size, "MarkMapped(handle)", &start, &count)) {
+    return;
+  }
+  for (size_t i = 0; i < count; ++i) {
+    auto& page = pages_[start + i];
+    MarkPageMappedLocked(
+        &page, va + i * page_size_, handle, std::shared_ptr<VMMHandleMeta>());
+  }
+}
+
+void VMMBackingMap::MarkMapped(VMMDevicePtr va,
+                               const std::shared_ptr<VMMHandleMeta>& meta,
+                               size_t size) {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  size_t start = 0;
+  size_t count = 0;
+  if (!CheckRangeLocked(va, size, "MarkMapped", &start, &count)) {
+    return;
+  }
+  const VMMAllocHandle handle =
+      meta == nullptr ? static_cast<VMMAllocHandle>(0) : meta->handle();
+  for (size_t i = 0; i < count; ++i) {
+    auto& page = pages_[start + i];
+    MarkPageMappedLocked(&page, va + i * page_size_, handle, meta);
+  }
+}
+
+void VMMBackingMap::MarkUnmapped(VMMDevicePtr va, size_t size) {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  size_t start = 0;
+  size_t count = 0;
+  if (!CheckRangeLocked(va, size, "MarkUnmapped", &start, &count)) {
+    return;
+  }
+  for (size_t i = 0; i < count; ++i) {
+    auto& page = pages_[start + i];
+    if (!page.mapped) {
+      VLOG(5) << "VMM V2 BackingMap unmapping already-unmapped page at "
+              << reinterpret_cast<void*>(va + i * page_size_);
+    }
+    ResetPageToUnmappedLocked(&page);
+  }
+}
+
+void VMMBackingMap::MarkReleased(VMMDevicePtr va,
+                                 VMMAllocHandle handle,
+                                 size_t size) {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  size_t start = 0;
+  size_t count = 0;
+  if (!CheckRangeLocked(va, size, "MarkReleased", &start, &count)) {
+    return;
+  }
+  for (size_t i = 0; i < count; ++i) {
+    auto& page = pages_[start + i];
+    if (handle != 0 && page.handle != 0 && page.handle != handle) {
+      VLOG(0) << "VMM V2 BackingMap release handle mismatch at "
+              << reinterpret_cast<void*>(va + i * page_size_)
+              << " tracked=" << reinterpret_cast<void*>(page.handle)
+              << " released=" << reinterpret_cast<void*>(handle);
+    }
+    ResetPageToUnmappedLocked(&page);
+  }
+}
+
+bool VMMBackingMap::ValidateLayout(const HandleLayout& layout,
+                                   const char* context) const {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  bool ok = true;
+  for (const auto& meta : layout) {
+    size_t start = 0;
+    size_t count = 0;
+    if (!CheckRangeLocked(
+            meta->base(), meta->size(), context, &start, &count)) {
+      ok = false;
+      continue;
+    }
+    for (size_t i = 0; i < count; ++i) {
+      const auto& page = pages_[start + i];
+      if (!page.mapped) {
+        VLOG(0) << "VMM V2 BackingMap mapped-state mismatch in " << context
+                << " va="
+                << reinterpret_cast<void*>(meta->base() + i * page_size_)
+                << " tracked_mapped=" << page.mapped;
+        ok = false;
+      }
+      if (page.mapped && page.handle != meta->handle()) {
+        VLOG(0) << "VMM V2 BackingMap handle mismatch in " << context << " va="
+                << reinterpret_cast<void*>(meta->base() + i * page_size_)
+                << " tracked=" << reinterpret_cast<void*>(page.handle)
+                << " meta=" << reinterpret_cast<void*>(meta->handle());
+        ok = false;
+      }
+    }
+  }
+  return ok;
+}
+
+bool VMMBackingMap::IsRangeMapped(VMMDevicePtr va, size_t size) const {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  size_t start = 0;
+  size_t count = 0;
+  if (!CheckRangeLocked(va, size, "IsRangeMapped", &start, &count)) {
+    return false;
+  }
+  for (size_t i = 0; i < count; ++i) {
+    if (!pages_[start + i].mapped) {
+      return false;
+    }
+  }
+  return true;
+}
+
+bool VMMBackingMap::IsRangeUnmapped(VMMDevicePtr va, size_t size) const {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  size_t start = 0;
+  size_t count = 0;
+  if (!CheckRangeLocked(va, size, "IsRangeUnmapped", &start, &count)) {
+    return false;
+  }
+  for (size_t i = 0; i < count; ++i) {
+    if (pages_[start + i].mapped) {
+      return false;
+    }
+  }
+  return true;
+}
+
+bool VMMBackingMap::IsRangeReleasable(VMMDevicePtr va, size_t size) const {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  size_t start = 0;
+  size_t count = 0;
+  if (!ComputeOverlappedPages(base_,
+                              size_,
+                              page_size_,
+                              va,
+                              size,
+                              "IsRangeReleasable",
+                              &start,
+                              &count)) {
+    return false;
+  }
+  for (size_t i = 0; i < count; ++i) {
+    if (!PageCanUseBackingLocked(&pages_[start + i], "IsRangeReleasable")) {
+      return false;
+    }
+  }
+  return true;
+}
+
+bool VMMBackingMap::PageCanUseBackingLocked(Page* page,
+                                            const char* context) const {
+  if (page == nullptr || !page->mapped || page->meta == nullptr) {
+    VLOG(6) << "VMM V2 BackingMap page cannot use backing in " << context;
+    return false;
+  }
+  return true;
+}
+
+size_t VMMBackingMap::total_mapped_bytes() const {
+  std::lock_guard<SpinLock> guard(spinlock_);
+  return mapped_page_count_ * page_size_;
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle
+
+#endif

paddle/phi/core/memory/allocation/vmm_backing_map.h

@@ -0,0 +1,96 @@
+// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#if defined(PADDLE_WITH_CUDA)
+
+#include <cstddef>
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "paddle/phi/core/memory/allocation/spin_lock.h"
+#include "paddle/phi/core/memory/allocation/vmm_allocator_v2_types.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// Page-granular backing state for VMM V2. Allocation blocks keep only logical
+// VA layout; backing ownership and release safety are decided from this map.
+class VMMBackingMap {
+ public:
+  struct MappedPage {
+    VMMDevicePtr va{0};
+    VMMAllocHandle handle{0};
+    std::shared_ptr<VMMHandleMeta> meta;
+    uint64_t epoch{0};
+  };
+  struct UnmappedPage {
+    VMMDevicePtr va{0};
+    uint64_t epoch{0};
+  };
+  void Configure(VMMDevicePtr base, size_t size, size_t page_size, int device);
+
+  bool IsConfigured() const { return configured_; }
+
+  void MarkMapped(VMMDevicePtr va, VMMAllocHandle handle, size_t size);
+  void MarkMapped(VMMDevicePtr va,
+                  const std::shared_ptr<VMMHandleMeta>& meta,
+                  size_t size);
+  void MarkUnmapped(VMMDevicePtr va, size_t size);
+  void MarkReleased(VMMDevicePtr va, VMMAllocHandle handle, size_t size);
+
+  bool ValidateLayout(const HandleLayout& layout, const char* context) const;
+  bool IsRangeMapped(VMMDevicePtr va, size_t size) const;
+  bool IsRangeUnmapped(VMMDevicePtr va, size_t size) const;
+  bool IsRangeReleasable(VMMDevicePtr va, size_t size) const;
+  size_t total_mapped_bytes() const;
+
+ private:
+  struct Page {
+    VMMAllocHandle handle{0};
+    std::shared_ptr<VMMHandleMeta> meta;
+    bool mapped{false};
+    uint64_t epoch{0};
+  };
+
+  bool CheckRangeLocked(VMMDevicePtr va,
+                        size_t size,
+                        const char* context,
+                        size_t* start,
+                        size_t* count) const;
+  void MarkPageMappedLocked(Page* page,
+                            VMMDevicePtr page_va,
+                            VMMAllocHandle handle,
+                            const std::shared_ptr<VMMHandleMeta>& meta);
+  void ResetPageToUnmappedLocked(Page* page);
+  bool PageCanUseBackingLocked(Page* page, const char* context) const;
+
+  VMMDevicePtr base_{0};
+  size_t size_{0};
+  size_t page_size_{0};
+  int device_{-1};
+  bool configured_{false};
+  mutable std::vector<Page> pages_;
+  size_t mapped_page_count_{0};
+  mutable SpinLock spinlock_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle
+
+#endif

paddle/phi/core/memory/allocation/vmm_ipc_allocation.h

@@ -0,0 +1,221 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <algorithm>
+#include <cstdint>
+#include <iterator>
+#include <limits>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#if defined(PADDLE_WITH_CUDA)
+#include "paddle/phi/backends/dynload/cuda_driver.h"
+using VmmDevicePtr = CUdeviceptr;
+using VmmAllocHandle = CUmemGenericAllocationHandle;
+#else
+using VmmDevicePtr = uintptr_t;
+using VmmAllocHandle = uint64_t;
+#endif
+
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+struct ImportedVmmMulti {
+  VmmDevicePtr base{0};
+  size_t reserved_size{0};
+  std::vector<VmmAllocHandle> hs;
+#if defined(PADDLE_WITH_CUDA)
+  ~ImportedVmmMulti() {
+    if (base && reserved_size) {
+      phi::dynload::cuMemUnmap(base, reserved_size);
+    }
+    for (auto h : hs) {
+      if (h) phi::dynload::cuMemRelease(h);
+    }
+    if (base && reserved_size) {
+      phi::dynload::cuMemAddressFree(base, reserved_size);
+    }
+  }
+#else
+  ~ImportedVmmMulti() = default;
+#endif
+};
+
+class VmmImportedAllocation : public phi::Allocation {
+ public:
+  VmmImportedAllocation(void* ptr,
+                        size_t bytes,
+                        Place place,
+                        std::shared_ptr<ImportedVmmMulti> keep)
+      : Allocation(ptr, bytes, place), keep_(std::move(keep)) {}
+
+ private:
+  std::shared_ptr<ImportedVmmMulti> keep_;
+};
+
+struct VmmChunkMeta {
+  VmmDevicePtr base;
+  size_t size;
+  VmmAllocHandle handle;
+  int device;
+};
+
+struct BlockPart {
+  std::shared_ptr<VmmChunkMeta> chunk;
+  size_t chunk_rel_off;
+  size_t len;
+};
+
+inline bool TryConcatAdjacentBlockPart(BlockPart* a, const BlockPart& b) {
+  if (!a) return false;
+  if (a->chunk.get() != b.chunk.get()) return false;
+  if (a->chunk_rel_off + a->len != b.chunk_rel_off) return false;
+  a->len += b.len;
+  return true;
+}
+
+inline std::vector<BlockPart> SliceBlockPartsForRange(
+    const std::vector<BlockPart>& parts,
+    size_t range_offset,
+    size_t range_len) {
+  // parts describes one logical block as an ordered list of VMM chunk slices.
+  // The target range is expressed in that logical block address space.
+  std::vector<BlockPart> sliced_parts;
+  if (range_len == 0 || parts.empty()) {
+    return sliced_parts;
+  }
+
+  PADDLE_ENFORCE_LE(
+      range_offset,
+      std::numeric_limits<size_t>::max() - range_len,
+      common::errors::InvalidArgument(
+          "Invalid VMM block-part slice range: offset %zu plus length %zu "
+          "overflows.",
+          range_offset,
+          range_len));
+
+  if (parts.size() == 1) {
+    const auto& part = parts.front();
+    PADDLE_ENFORCE_LE(
+        range_offset,
+        part.len,
+        common::errors::InvalidArgument(
+            "Invalid VMM block-part slice offset %zu for part length %zu.",
+            range_offset,
+            part.len));
+    PADDLE_ENFORCE_LE(
+        range_len,
+        part.len - range_offset,
+        common::errors::InvalidArgument(
+            "Invalid VMM block-part slice length %zu at offset %zu for part "
+            "length %zu.",
+            range_len,
+            range_offset,
+            part.len));
+    return {
+        BlockPart{part.chunk, part.chunk_rel_off + range_offset, range_len}};
+  }
+
+  sliced_parts.reserve(parts.size());
+  const size_t range_end = range_offset + range_len;
+  size_t cursor = 0;
+  size_t sliced_len = 0;
+
+  for (const auto& part : parts) {
+    const size_t part_block_begin = cursor;
+    const size_t part_block_end = cursor + part.len;
+    cursor = part_block_end;
+
+    if (part_block_end <= range_offset) {
+      continue;
+    }
+    if (part_block_begin >= range_end) {
+      break;
+    }
+
+    const size_t slice_begin = std::max(part_block_begin, range_offset);
+    const size_t slice_end = std::min(part_block_end, range_end);
+    BlockPart slice{part.chunk,
+                    part.chunk_rel_off + (slice_begin - part_block_begin),
+                    slice_end - slice_begin};
+
+    if (!sliced_parts.empty() &&
+        TryConcatAdjacentBlockPart(&sliced_parts.back(), slice)) {
+      sliced_len += slice.len;
+      continue;
+    }
+    sliced_parts.push_back(std::move(slice));
+    sliced_len += sliced_parts.back().len;
+  }
+
+  PADDLE_ENFORCE_EQ(
+      sliced_len,
+      range_len,
+      common::errors::InvalidArgument(
+          "Invalid VMM block-part slice range: requested %zu bytes at offset "
+          "%zu, but only sliced %zu bytes from %zu parts.",
+          range_len,
+          range_offset,
+          sliced_len,
+          parts.size()));
+  return sliced_parts;
+}
+
+inline void AppendBlockPartsTail(std::vector<BlockPart>* dst,
+                                 std::vector<BlockPart>* src) {
+  if (src->empty()) return;
+  dst->reserve(dst->size() + src->size());
+  auto begin = src->begin();
+  if (!dst->empty() && TryConcatAdjacentBlockPart(&dst->back(), src->front())) {
+    ++begin;
+  }
+  dst->insert(dst->end(),
+              std::make_move_iterator(begin),
+              std::make_move_iterator(src->end()));
+}
+
+#pragma pack(push, 1)
+struct VmmIpcHeader {
+  uint8_t version;
+  uint16_t flags;
+  uint32_t pid;
+  uint32_t num_entries;
+  uint64_t alloc_size;
+  uint64_t offset;
+  uint64_t reserved_size;
+};
+
+struct VmmIpcEntry {
+  uint8_t handle_type;
+  uint8_t reserved[7];
+  uint64_t rel_offset;
+  uint64_t chunk_size;
+  uint64_t chunk_rel_off;
+};
+#pragma pack(pop)
+
+static_assert(sizeof(VmmIpcHeader) == 35, "VmmIpcHeader size changed");
+static_assert(sizeof(VmmIpcEntry) == 32, "VmmIpcEntry size changed");
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/xpu_allocator.cc

@@ -0,0 +1,57 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/xpu_allocator.h"
+
+#include <string>
+
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/platform/device/xpu/xpu_info.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+bool XPUAllocator::IsAllocThreadSafe() const { return true; }
+void XPUAllocator::FreeImpl(phi::Allocation* allocation) {
+  PADDLE_ENFORCE_EQ(
+      allocation->place(),
+      place_,
+      common::errors::PermissionDenied(
+          "XPU memory is freed in incorrect device. This may be a bug"));
+  platform::RecordedXPUFree(
+      allocation->ptr(), allocation->size(), place_.device);
+  delete allocation;
+}
+
+phi::Allocation* XPUAllocator::AllocateImpl(size_t size) {
+  std::call_once(once_flag_,
+                 [this] { platform::SetXPUDeviceId(place_.device); });
+
+  void* ptr;
+  auto result = platform::RecordedXPUMalloc(&ptr, size, place_.device);
+  if (LIKELY(result == XPU_SUCCESS)) {
+    return new Allocation(ptr, size, Place(place_));
+  }
+
+  PADDLE_THROW_BAD_ALLOC(common::errors::ResourceExhausted(
+      "\n\nOut of memory error on XPU %d. "
+      "Cannot allocate %s memory on XPU %d.\n\n",
+      place_.device,
+      string::HumanReadableSize(size),
+      place_.device));
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/xpu_allocator.h

@@ -0,0 +1,42 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include <mutex>  // NOLINT
+
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+class XPUAllocator : public Allocator {
+ public:
+  explicit XPUAllocator(const XPUPlace& place) : place_(place) {}
+
+  bool IsAllocThreadSafe() const override;
+
+ protected:
+  void FreeImpl(phi::Allocation* allocation) override;
+  phi::Allocation* AllocateImpl(size_t size) override;
+
+ private:
+  XPUPlace place_;
+  std::once_flag once_flag_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/phi/core/memory/allocation/xpu_ipc_allocator.cc

@@ -0,0 +1,109 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef _WIN32
+
+#include "paddle/phi/core/memory/allocation/xpu_ipc_allocator.h"
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <fcntl.h>
+#include <sys/mman.h>
+#include <cstdlib>
+#include <random>
+#include <string>
+
+#include "glog/logging.h"
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/core/platform/device/xpu/xpu_info.h"
+
+namespace paddle::memory::allocation {
+
+namespace {
+// Mutex to protect IPC handle cache.
+std::mutex ipc_mutex_;
+// Cache mapping from handle string to a weak pointer of the opened IPC memory.
+std::unordered_map<std::string, std::weak_ptr<void>> ipc_handle_to_baseptr_;
+}  // namespace
+
+std::shared_ptr<void> GetIpcBasePtr(std::string handle) {
+  std::lock_guard<std::mutex> lock(ipc_mutex_);
+
+  // Get the current device ID.
+  int device_id = platform::GetXPUCurrentDeviceId();
+  paddle::platform::SetXPUDeviceId(device_id);
+
+  auto iter = ipc_handle_to_baseptr_.find(handle);
+  if (iter != ipc_handle_to_baseptr_.end()) {
+    if (auto baseptr = iter->second.lock()) {
+      return baseptr;
+    }
+  }
+  // The IPC memory handle can only be opened once for the same handle,
+  // so we cache the opened pointer.
+  void *baseptr = nullptr;
+  // Interpret the provided handle string as an XPU IPC memory handle.
+  auto ipc_handle =
+      reinterpret_cast<const cudaIpcMemHandle_t *>(handle.c_str());
+  // PADDLE_ENFORCE_XPU_SUCCESS(cudaIpcOpenMemHandle(&baseptr, *ipc_handle,
+  // cudaIpcMemLazyEnablePeerAccess));
+  int ret = cudaIpcOpenMemHandle(
+      &baseptr, *ipc_handle, cudaIpcMemLazyEnablePeerAccess);
+  PADDLE_ENFORCE_XPU_SUCCESS(ret);
+
+  // Create a shared_ptr with a custom deleter that will close the IPC handle.
+  auto sp = std::shared_ptr<void>(baseptr, [handle, device_id](void *ptr) {
+    platform::XPUDeviceGuard guard(device_id);
+    std::lock_guard<std::mutex> lock(ipc_mutex_);
+    PADDLE_ENFORCE_XPU_SUCCESS(cudaIpcCloseMemHandle(ptr));
+    ipc_handle_to_baseptr_.erase(handle);
+    VLOG(6) << "cudaIpcCloseMemHandle for ptr:"
+            << "\t" << ptr;
+  });
+  std::weak_ptr<void> wp = sp;
+  ipc_handle_to_baseptr_.insert({handle, wp});
+
+  return sp;
+}
+
+void IpcCollect() {
+  std::lock_guard<std::mutex> lock(ipc_mutex_);
+  size_t before = ipc_handle_to_baseptr_.size();
+  VLOG(6) << "The number of IPC handles before collection:" << before;
+
+  for (auto it = ipc_handle_to_baseptr_.begin();
+       it != ipc_handle_to_baseptr_.end();) {
+    if (it->second.expired()) {
+      it = ipc_handle_to_baseptr_.erase(it);
+    } else {
+      VLOG(6) << " Valid ipc handle is not expired";
+      ++it;
+    }
+  }
+
+  size_t after = ipc_handle_to_baseptr_.size();
+  size_t collected = before - after;
+  VLOG(1) << "IpcCollect: collected " << collected << " expired IPC handles"
+          << "out of " << before << " total handles";
+}
+
+XpuIpcAllocation::~XpuIpcAllocation() {
+  // Release the underlying IPC resource.
+  shared_ptr_.reset();
+  VLOG(6) << "tensor deleted cudaIpcCloseMemHandle for ptr:"
+          << "\t" << this->ptr();
+}
+}  // namespace paddle::memory::allocation
+
+#endif  // _WIN32

paddle/phi/core/memory/allocation/xpu_ipc_allocator.h

@@ -0,0 +1,59 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef _WIN32
+#pragma once
+
+#include <memory>
+#include <mutex>  // NOLINT
+#include <string>
+#include <unordered_set>
+#include <utility>
+
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/platform/device/xpu/xpu_info.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// Returns a shared pointer that holds the IPC base pointer for the given
+// handle.
+std::shared_ptr<void> GetIpcBasePtr(std::string handle);
+void IpcCollect();
+class XpuIpcAllocation : public Allocation {
+ public:
+  explicit XpuIpcAllocation(void *ptr,
+                            size_t size,
+                            int device_id,
+                            std::shared_ptr<void> shared_ptr)
+      : Allocation(ptr, size, XPUPlace(device_id)),
+        device_id_(device_id),
+        shared_ptr_(std::move(shared_ptr)) {}
+
+  inline const int &device_id() const { return device_id_; }
+
+  ~XpuIpcAllocation() override;
+
+ private:
+  int device_id_;
+  std::shared_ptr<void> shared_ptr_;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle
+
+#endif  // _WIN32

paddle/phi/core/memory/allocation/xpu_pinned_allocator.cc

@@ -0,0 +1,66 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/memory/allocation/xpu_pinned_allocator.h"
+
+#include "paddle/phi/core/memory/stats.h"
+#include "paddle/phi/core/platform/profiler/mem_tracing.h"
+
+#if defined(PADDLE_WITH_XPU)
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include "paddle/phi/backends/xpu/enforce_xpu.h"
+#endif
+
+namespace paddle::memory::allocation {
+
+// Define the destructor so the vtable gets emitted.
+XPUPinnedAllocator::~XPUPinnedAllocator() = default;
+
+bool XPUPinnedAllocator::IsAllocThreadSafe() const { return true; }
+
+void XPUPinnedAllocator::FreeImpl(phi::Allocation* allocation) {
+#if defined(PADDLE_WITH_XPU)
+  PADDLE_ENFORCE_XPU_SUCCESS(cudaFreeHost(allocation->ptr()));
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'XPUPinnedPlace' is not supported. Please re-compile with WITH_XPU."));
+#endif
+  VLOG(10) << "cudaFreeHost " << allocation->ptr();
+  HOST_MEMORY_STAT_UPDATE(Reserved, 0, -allocation->size());
+  platform::RecordMemEvent(allocation->ptr(),
+                           allocation->place(),
+                           allocation->size(),
+                           phi::TracerMemEventType::ReservedFree);
+  delete allocation;
+}
+
+phi::Allocation* XPUPinnedAllocator::AllocateImpl(size_t size) {
+  void* ptr;
+
+#if defined(PADDLE_WITH_XPU)
+  PADDLE_ENFORCE_XPU_SUCCESS(cudaHostAlloc(&ptr, size, cudaHostAllocPortable));
+#else
+  PADDLE_THROW(common::errors::PermissionDenied(
+      "'XPUPinnedPlace' is not supported. Please re-compile with WITH_XPU."));
+#endif
+
+  VLOG(10) << "cudaHostAlloc " << size << " " << ptr;
+  HOST_MEMORY_STAT_UPDATE(Reserved, 0, size);
+  platform::RecordMemEvent(
+      ptr, XPUPinnedPlace(), size, phi::TracerMemEventType::ReservedAllocate);
+  return new Allocation(ptr, size, XPUPinnedPlace());
+}
+
+}  // namespace paddle::memory::allocation

paddle/phi/core/memory/allocation/xpu_pinned_allocator.h

@@ -0,0 +1,37 @@
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+#include "paddle/phi/core/memory/allocation/allocator.h"
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+// Allocator uses `cudaHostAlloc`
+class XPUPinnedAllocator : public Allocator {
+ public:
+  // Add an explicit virtual destructor
+  virtual ~XPUPinnedAllocator();
+
+  bool IsAllocThreadSafe() const override;
+
+ protected:
+  void FreeImpl(phi::Allocation* allocation) override;
+  phi::Allocation* AllocateImpl(size_t size) override;
+};
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle