chore: import upstream snapshot with attribution

test/xpu/cpp/overload_xpu_alloc_test.cc

@@ -0,0 +1,104 @@
+// 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/backends/xpu/enforce_xpu.h"
+#include "paddle/phi/backends/xpu/xpu_context.h"
+#include "paddle/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/stats.h"
+
+#include "gtest/gtest.h"
+
+namespace paddle {
+namespace memory {
+TEST(XPUOverloadAllocTest, EnvTest) {
+  setenv("XPUAPI_DEFAULT_SIZE", "4096", 1);
+  // use alloc overload
+  unsetenv("XPU_PADDLE_DISABLE_ALLOC_OVERLOAD");
+  phi::XPUContext dev_ctx_overload(
+      phi::XPUPlace(phi::backends::xpu::GetXPUCurrentDeviceId()));
+  EXPECT_STREQ(dev_ctx_overload.x_context()->get_option("XPUAPI_DEFAULT_SIZE"),
+               "1");
+  EXPECT_NE(dev_ctx_overload.x_context()->overload_alloc_gm, nullptr);
+  // do not use alloc overload
+  setenv("XPU_PADDLE_DISABLE_ALLOC_OVERLOAD", "1", 1);
+  phi::XPUContext dev_ctx_origin(
+      phi::XPUPlace(phi::backends::xpu::GetXPUCurrentDeviceId()));
+  EXPECT_STREQ(dev_ctx_origin.x_context()->get_option("XPUAPI_DEFAULT_SIZE"),
+               "4096");
+  EXPECT_EQ(dev_ctx_origin.x_context()->overload_alloc_gm, nullptr);
+  unsetenv("XPU_PADDLE_DISABLE_ALLOC_OVERLOAD");
+  unsetenv("XPUAPI_DEFAULT_SIZE");
+}
+
+TEST(XPUOverloadAllocTest, BasicTest) {
+  phi::XPUContext dev_ctx(
+      phi::XPUPlace(phi::backends::xpu::GetXPUCurrentDeviceId()));
+  int numel = 64;
+  int alignment = phi::backends::xpu::XPUMinChunkSize();
+  int expected_alloc_size =
+      allocation::AlignedSize(numel * sizeof(int), alignment);
+  xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
+  int pre_alloc_value = DEVICE_MEMORY_STAT_CURRENT_VALUE(
+      Allocated, dev_ctx.GetPlace().GetDeviceId());
+  int* buffer = RAII_GUARD.alloc<int>(numel);
+  int after_alloc_value = DEVICE_MEMORY_STAT_CURRENT_VALUE(
+      Allocated, dev_ctx.GetPlace().GetDeviceId());
+  EXPECT_NE(buffer, nullptr);
+  EXPECT_EQ(after_alloc_value - pre_alloc_value, expected_alloc_size);
+}
+
+TEST(XPUOverloadAllocTest, NestedScopeTest) {
+  phi::XPUContext dev_ctx(
+      phi::XPUPlace(phi::backends::xpu::GetXPUCurrentDeviceId()));
+  xpu::ctx_guard RAII_GUARD1(dev_ctx.x_context());
+  int pre_alloc_value = DEVICE_MEMORY_STAT_CURRENT_VALUE(
+      Allocated, dev_ctx.GetPlace().GetDeviceId());
+  int* buffer_outer = RAII_GUARD1.alloc<int>(64);
+  EXPECT_NE(buffer_outer, nullptr);
+  {
+    // The destruction of inner guard should not free the memory allocated from
+    // outer guard.
+    xpu::ctx_guard RAII_GUARD2(dev_ctx.x_context());
+    int* buffer_inner = RAII_GUARD2.alloc<int>(64);
+    EXPECT_NE(buffer_inner, nullptr);
+  }
+  int post_alloc_value = DEVICE_MEMORY_STAT_CURRENT_VALUE(
+      Allocated, dev_ctx.GetPlace().GetDeviceId());
+  EXPECT_NE(post_alloc_value, pre_alloc_value);
+}
+
+TEST(XPUOverloadAllocTest, MultiStreamTest) {
+  // Test whether stream 1 use the memory poll of stream 0.
+  int size = 64;
+  setenv("XPU_CDNN_CLUSTER_PARALLEL", "1", 1);
+  phi::XPUContext dev_ctx(
+      phi::XPUPlace(phi::backends::xpu::GetXPUCurrentDeviceId()));
+  xpu::ctx_guard RAII_GUARD0(dev_ctx.x_context(0));
+  xpu::ctx_guard RAII_GUARD1(dev_ctx.x_context(1));
+  int pre_alloc_value = DEVICE_MEMORY_STAT_CURRENT_VALUE(
+      Allocated, dev_ctx.GetPlace().GetDeviceId());
+  int* buffer0 = RAII_GUARD1.alloc<int>(size);
+  EXPECT_NE(buffer0, nullptr);
+  {
+    int* buffer1 = RAII_GUARD0.alloc<int>(size);
+    EXPECT_NE(buffer1, nullptr);
+  }
+  int post_alloc_value = DEVICE_MEMORY_STAT_CURRENT_VALUE(
+      Allocated, dev_ctx.GetPlace().GetDeviceId());
+
+  EXPECT_NE(pre_alloc_value, post_alloc_value);
+  unsetenv("XPU_CDNN_CLUSTER_PARALLEL");
+}
+}  // namespace memory
+}  // namespace paddle