From e4c438f544f6321d79919d7bc5062f404747a6f7 Mon Sep 17 00:00:00 2001
From: HongyuJia <jiahongyu@baidu.com>
Date: Fri, 6 Jan 2023 12:49:32 +0800
Subject: [PATCH] [Custom device stream] Acquire custom_deivce stream, add unit
 test (#49571)

* acquire custom_deivce stream

* regulate file name and unittest
---
 .../fluid/tests/custom_runtime/CMakeLists.txt |  2 +-
 .../{custom_relu_op.cc => custom_op.cc}       | 22 +++++++++++++++++++
 ..._relu_setup.py => test_custom_op_setup.py} | 15 +++++++++++--
 3 files changed, 36 insertions(+), 3 deletions(-)
 rename python/paddle/fluid/tests/custom_runtime/{custom_relu_op.cc => custom_op.cc} (91%)
 rename python/paddle/fluid/tests/custom_runtime/{test_custom_device_relu_setup.py => test_custom_op_setup.py} (96%)

diff --git a/python/paddle/fluid/tests/custom_runtime/CMakeLists.txt b/python/paddle/fluid/tests/custom_runtime/CMakeLists.txt
index 00eef2d5a7..d1885f9f9d 100644
--- a/python/paddle/fluid/tests/custom_runtime/CMakeLists.txt
+++ b/python/paddle/fluid/tests/custom_runtime/CMakeLists.txt
@@ -28,5 +28,5 @@ if(WITH_CUSTOM_DEVICE AND NOT WITH_GPU)
   set_tests_properties(test_custom_cpu_profiler_plugin PROPERTIES TIMEOUT 120)
   set_tests_properties(test_fleet_launch_custom_device PROPERTIES TIMEOUT 120)
   set_tests_properties(test_custom_cpu_to_static PROPERTIES TIMEOUT 120)
-  set_tests_properties(test_custom_device_relu_setup PROPERTIES TIMEOUT 120)
+  set_tests_properties(test_custom_op_setup PROPERTIES TIMEOUT 120)
 endif()
diff --git a/python/paddle/fluid/tests/custom_runtime/custom_relu_op.cc b/python/paddle/fluid/tests/custom_runtime/custom_op.cc
similarity index 91%
rename from python/paddle/fluid/tests/custom_runtime/custom_relu_op.cc
rename to python/paddle/fluid/tests/custom_runtime/custom_op.cc
index da0563ffeb..01dec06604 100644
--- a/python/paddle/fluid/tests/custom_runtime/custom_relu_op.cc
+++ b/python/paddle/fluid/tests/custom_runtime/custom_op.cc
@@ -16,6 +16,7 @@
 #include <vector>
 
 #include "paddle/extension.h"
+#include "paddle/phi/backends/all_context.h"
 
 #define CHECK_CPU_INPUT(x) PD_CHECK(x.is_cpu(), #x " must be a CPU Tensor.")
 #define CHECK_CUSTOM_INPUT(x) \
@@ -191,3 +192,24 @@ PD_BUILD_DOUBLE_GRAD_OP(custom_relu)
     .Outputs({paddle::Grad(paddle::Grad("Out"))})
     .SetKernelFn(PD_KERNEL(ReluDoubleBackward))
     .SetInferShapeFn(PD_INFER_SHAPE(ReluDoubleBackwardInferShape));
+
+std::vector<paddle::Tensor> StreamForward(const paddle::Tensor& x) {
+  CHECK_CUSTOM_INPUT(x);
+
+  auto dev_ctx =
+      paddle::experimental::DeviceContextPool::Instance().Get(x.place());
+  auto custom_ctx = static_cast<const phi::CustomContext*>(dev_ctx);
+  void* stream = custom_ctx->stream();
+
+  PD_CHECK(stream != nullptr);
+  std::cout << "Check stream != nullptr successfully" << std::endl;
+  custom_ctx->Wait();
+  std::cout << "Check Wait successfully" << std::endl;
+
+  return {x};
+}
+
+PD_BUILD_OP(custom_stream)
+    .Inputs({"X"})
+    .Outputs({"Out"})
+    .SetKernelFn(PD_KERNEL(StreamForward));
diff --git a/python/paddle/fluid/tests/custom_runtime/test_custom_device_relu_setup.py b/python/paddle/fluid/tests/custom_runtime/test_custom_op_setup.py
similarity index 96%
rename from python/paddle/fluid/tests/custom_runtime/test_custom_device_relu_setup.py
rename to python/paddle/fluid/tests/custom_runtime/test_custom_op_setup.py
index 760ad56cc3..969fdb2f8a 100644
--- a/python/paddle/fluid/tests/custom_runtime/test_custom_device_relu_setup.py
+++ b/python/paddle/fluid/tests/custom_runtime/test_custom_op_setup.py
@@ -177,14 +177,15 @@ class TestNewCustomOpSetUpInstall(unittest.TestCase):
             )
 
         custom_module = paddle.utils.cpp_extension.load(
-            name='custom_device_relu',
-            sources=['custom_relu_op.cc'],
+            name='custom_device',
+            sources=['custom_op.cc'],
             extra_include_paths=paddle_includes,  # add for Coverage CI
             extra_cxx_cflags=["-w", "-g"],  # test for cc flags
             # build_directory=self.cur_dir,
             verbose=True,
         )
         self.custom_op = custom_module.custom_relu
+        self.custom_stream_op = custom_module.custom_stream
 
         self.dtypes = ["float32", "float64"]
         self.device = "custom_cpu"
@@ -204,6 +205,7 @@ class TestNewCustomOpSetUpInstall(unittest.TestCase):
         self._test_dynamic()
         self._test_double_grad_dynamic()
         self._test_with_dataloader()
+        self._test_stream()
 
     def _test_static(self):
         for dtype in self.dtypes:
@@ -317,6 +319,15 @@ class TestNewCustomOpSetUpInstall(unittest.TestCase):
             if batch_id == 5:
                 break
 
+    def _test_stream(self):
+        import paddle
+
+        paddle.set_device(self.device)
+        x = paddle.ones([2, 2], dtype='float32')
+        out = self.custom_stream_op(x)
+
+        np.testing.assert_array_equal(x.numpy(), out.numpy())
+
 
 if __name__ == "__main__":
     if os.name == 'nt' or sys.platform.startswith('darwin'):
-- 
GitLab