[C++ API GetAllocator] Add C++ `GetAllocator` interface (#50813)

* [C++ API GetAllocator] Add C++ GetAllocator interface

* move api to accurate directory

paddle/phi/api/include/context_pool.h

@@ -87,3 +87,10 @@ class PADDLE_API DeviceContextPool {
 
 }  // namespace experimental
 }  // namespace paddle
+
+namespace phi {
+class Allocator;
+
+PADDLE_API Allocator* GetAllocator(const Place& place);
+
+}  // namespace phi

paddle/phi/api/lib/CMakeLists.txt

@@ -289,7 +289,7 @@ cc_library(
 cc_library(
   context_pool
   SRCS context_pool.cc
-  DEPS phi_backends phi_enforce place init)
+  DEPS phi_backends phi_enforce place init phi_device_context)
 
 cc_library(
   kernel_dispatch

paddle/phi/api/lib/context_pool.cc

@@ -15,6 +15,7 @@ limitations under the License. */
 #include "paddle/phi/api/include/context_pool.h"
 
 #include "paddle/phi/backends/all_context.h"
+#include "paddle/phi/core/allocator.h"
 #include "paddle/phi/core/enforce.h"
 
 #include "paddle/fluid/platform/init.h"
@@ -50,3 +51,13 @@ phi::DeviceContext* DeviceContextPool::GetMutable(const Place& place) {
 
 }  // namespace experimental
 }  // namespace paddle
+
+namespace phi {
+
+PADDLE_API Allocator* GetAllocator(const Place& place) {
+  const DeviceContext* dev_ctx =
+      paddle::experimental::DeviceContextPool::Instance().Get(place);
+  return const_cast<phi::Allocator*>(&dev_ctx->GetAllocator());
+}
+
+}  // namespace phi

paddle/phi/core/allocator.cc

@@ -14,4 +14,7 @@ limitations under the License. */
 
 #include "paddle/phi/core/allocator.h"
 
+#include "paddle/phi/api/include/context_pool.h"
+#include "paddle/phi/core/device_context.h"
+
 namespace phi {}  // namespace phi

paddle/phi/core/allocator.h

@@ -18,6 +18,7 @@ limitations under the License. */
 #include <functional>
 #include <memory>
 
+#include "paddle/phi/api/include/dll_decl.h"
 #include "paddle/phi/common/place.h"
 
 namespace phi {

paddle/phi/tests/api/CMakeLists.txt

@@ -5,11 +5,19 @@ if(WITH_GPU)
     test_phi_tensor
     SRCS test_phi_tensor.cc
     DEPS glog selected_rows ${COMMON_API_TEST_DEPS})
+  nv_test(
+    test_allocator
+    SRCS test_allocator.cu
+    DEPS memory place device_context context_pool)
 elseif(WITH_ROCM)
   hip_test(
     test_phi_tensor
     SRCS test_phi_tensor.cc
     DEPS glog selected_rows ${COMMON_API_TEST_DEPS})
+  hip_test(
+    test_allocator
+    SRCS test_allocator.cu
+    DEPS memory place device_context context_pool)
 else()
   cc_test(
     test_phi_tensor

paddle/phi/tests/api/test_allocator.cu

+/* 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 <gtest/gtest.h>
+
+#include "paddle/phi/api/include/context_pool.h"
+
+#include "paddle/fluid/memory/memcpy.h"
+#include "paddle/phi/common/place.h"
+#include "paddle/phi/common/transform.h"
+#include "paddle/phi/core/allocator.h"
+#include "paddle/phi/core/device_context.h"
+
+TEST(Allocator, CPU) {
+  phi::Allocator* allocator = phi::GetAllocator(phi::CPUPlace());
+  auto cpu_allocation = allocator->Allocate(sizeof(float) * 4);
+  float* cpu_buf = static_cast<float*>(cpu_allocation->ptr());
+  ASSERT_NE(cpu_buf, nullptr);
+  cpu_buf[0] = 1.0f;
+  cpu_buf[1] = 2.0f;
+  cpu_buf[2] = 3.0f;
+  cpu_buf[3] = 4.0f;
+  for (size_t i = 0; i < 4; ++i) {
+    cpu_buf[i] = cpu_buf[i] + 1;
+  }
+  for (size_t i = 0; i < 4; ++i) {
+    ASSERT_NEAR(cpu_buf[i], static_cast<float>(2.0 + i), 1e-5);
+  }
+}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+using paddle::memory::Copy;
+
+template <typename T>
+class Scale {
+ public:
+  explicit Scale(const T& scale) : scale_(scale) {}
+  HOSTDEVICE T operator()(const T& a) const { return a * scale_; }
+
+ private:
+  T scale_;
+};
+
+TEST(Allocator, GPU) {
+  phi::GPUPlace gpu0(0);
+  float cpu_buf[4] = {0.1, 0.2, 0.3, 0.4};
+  phi::Allocator* allocator = phi::GetAllocator(gpu0);
+  auto gpu_allocation = allocator->Allocate(sizeof(cpu_buf));
+  float* gpu_buf = static_cast<float*>(gpu_allocation->ptr());
+
+  phi::DeviceContextPool& pool = phi::DeviceContextPool::Instance();
+  auto* ctx = reinterpret_cast<phi::GPUContext*>(pool.Get(gpu0));
+  Copy(gpu0, gpu_buf, phi::CPUPlace(), cpu_buf, sizeof(cpu_buf), ctx->stream());
+  phi::Transform<phi::GPUContext> trans;
+  trans(*ctx, gpu_buf, gpu_buf + 4, gpu_buf, Scale<float>(10));
+  ctx->Wait();
+  Copy(phi::CPUPlace(), cpu_buf, gpu0, gpu_buf, sizeof(cpu_buf), ctx->stream());
+  for (int i = 0; i < 4; ++i) {
+    ASSERT_NEAR(cpu_buf[i], static_cast<float>(i + 1), 1e-5);
+  }
+}
+#endif