[XPU] add stream & event unittests

Author: runzhech

paddle/fluid/pybind/xpu_streams_py.cc

@@ -152,7 +152,11 @@ void BindXpuStream(py::module *m_ptr) {
            [](phi::XPUStreamHandle &self, phi::XPUEventHandle &other) {
              self.wait_event(other.get_event());
            })
-      .def("query", [](phi::XPUStreamHandle &self) { return self.query(); })
+      .def("query",
+           [](phi::XPUStreamHandle &self) {
+             PADDLE_THROW(common::errors::Unavailable(
+                 "Query function for XPUStream is not supported now"));
+           })
       .def("record_event",
            [](phi::XPUStreamHandle &self, phi::XPUEventHandle *event) {
              if (event == nullptr) {
@@ -170,9 +174,9 @@ void BindXpuStream(py::module *m_ptr) {
           Examples:
               .. code-block:: python
 
-                  >>> # doctest: +REQUIRES(env:GPU)
+                  >>> # doctest: +REQUIRES(env:XPU)
                   >>> import paddle
-                  >>> s = paddle.device.cuda.Stream(paddle.CUDAPlace(0), 1)
+                  >>> s = paddle.device.xpu.Stream(paddle.XPUPlace(0), 1)
                   >>> s.synchronize()
 
           )DOC")

paddle/phi/backends/xpu/xpu_context.cc

@@ -645,21 +645,14 @@ void XPUStreamHandle::synchronize() const {
 
 void XPUStreamHandle::set_stream(XPUStream stream_) { stream = stream_; }
 
-bool XPUStreamHandle::query() const {
-  XPUEvent event = XPUEventPool::Instance().CreateEventFromPool();
+void XPUStreamHandle::record_event(XPUEvent event) const {
   int r = xpu_event_record(event, stream);
   PADDLE_ENFORCE_XRE_SUCCESS(r);
-  r = xpu_event_query(event);
-  if (r == XPU_SUCCESS) {
-    return true;
-  } else {
-    return false;
-  }
 }
 
-void XPUStreamHandle::record_event(XPUEvent event) const {
-  int r = xpu_event_record(event, stream);
-  PADDLE_ENFORCE_XRE_SUCCESS(r);
+XPUStreamHandle get_current_stream_handle(int device_id) {
+  auto* dev_ctx = get_xpu_context(device_id);
+  return *dev_ctx->get_current_stream_handle();
 }
 
 XPUStreamHandle get_stream_handle(int device_id) {
@@ -736,7 +729,6 @@ XPUEventHandle::XPUEventHandle(XPUStream stream) {
 
 void XPUEventHandle::record(XPUStream stream) {
   int r = xpu_event_query(event_);
-  printf("====r: %d\n", r);
   PADDLE_ENFORCE_XRE_SUCCESS(xpu_event_record(event_, stream));
 }
 

paddle/phi/backends/xpu/xpu_context.h

@@ -52,7 +52,6 @@ class XPUStreamHandle {
   XPUStream raw_stream() const { return stream; }
   void wait_event(XPUEvent event) const;
   void synchronize() const;
-  bool query() const;
   void record_event(XPUEvent event) const;
   void set_stream(XPUStream stream);
 
@@ -186,6 +185,7 @@ class XPUEventHandle {
   XPUEvent event_;
 };
 
+XPUStreamHandle get_current_stream_handle(int device_id = -1);
 XPUStreamHandle get_stream_handle(int device_id = -1);
 void set_current_stream(XPUStreamHandle* s);
 

test/xpu/test_xpu_stream_event.py

@@ -12,8 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import ctypes
 import unittest
 
+import numpy as np
+
 import paddle
 from paddle.device import xpu
 
@@ -56,7 +59,6 @@ def test_xpu_stream_synchronize(self):
             e2 = paddle.device.xpu.Event()
 
             e1.record(s)
-            print("1111")
             e1.query()
             tensor1 = paddle.to_tensor(paddle.rand([1000, 1000]))
             tensor2 = paddle.matmul(tensor1, tensor1)
@@ -78,7 +80,85 @@ def test_xpu_stream_wait_event_and_record_event(self):
             s2.wait_event(e1)
             s2.synchronize()
 
-            self.assertTrue(e1.query() and s1.query() and s2.query())
+            self.assertTrue(e1.query())
+
+
+class TestXPUEvent(unittest.TestCase):
+    def test_xpu_event(self):
+        if paddle.is_compiled_with_xpu():
+            e = paddle.device.xpu.Event()
+            self.assertIsNotNone(e)
+            s = paddle.device.xpu.current_stream()
+
+    def test_xpu_event_methods(self):
+        if paddle.is_compiled_with_xpu():
+            e = paddle.device.xpu.Event()
+            s = paddle.device.xpu.current_stream()
+            event_query_1 = e.query()
+            tensor1 = paddle.to_tensor(paddle.rand([1000, 1000]))
+            tensor2 = paddle.matmul(tensor1, tensor1)
+            s.record_event(e)
+            e.synchronize()
+            event_query_2 = e.query()
+
+            self.assertTrue(event_query_1)
+            self.assertTrue(event_query_2)
+
+
+class TestStreamGuard(unittest.TestCase):
+    '''
+    Note:
+        The asynchronous execution property of XPU Stream can only be tested offline.
+    '''
+
+    def test_stream_guard_normal(self):
+        if paddle.is_compiled_with_xpu():
+            s = paddle.device.Stream()
+            a = paddle.to_tensor(np.array([0, 2, 4], dtype="int32"))
+            b = paddle.to_tensor(np.array([1, 3, 5], dtype="int32"))
+            c = a + b
+            with paddle.device.stream_guard(s):
+                d = a + b
+                s.synchronize()
+
+            np.testing.assert_array_equal(np.array(c), np.array(d))
+
+    def test_stream_guard_default_stream(self):
+        if paddle.is_compiled_with_xpu():
+            s1 = paddle.device.current_stream()
+            with paddle.device.stream_guard(s1):
+                pass
+            s2 = paddle.device.current_stream()
+
+            self.assertTrue(id(s1.stream_base) == id(s2.stream_base))
+
+    def test_set_current_stream_default_stream(self):
+        if paddle.is_compiled_with_xpu():
+            cur_stream = paddle.device.current_stream()
+            new_stream = paddle.device.set_stream(cur_stream)
+
+            self.assertTrue(
+                id(cur_stream.stream_base) == id(new_stream.stream_base)
+            )
+
+    def test_stream_guard_raise_error(self):
+        if paddle.is_compiled_with_xpu():
+
+            def test_not_correct_stream_guard_input():
+                tmp = np.zeros(5)
+                with paddle.device.stream_guard(tmp):
+                    pass
+
+            self.assertRaises(TypeError, test_not_correct_stream_guard_input)
+
+
+class TestRawStream(unittest.TestCase):
+    def test_xpu_stream(self):
+        if paddle.is_compiled_with_xpu():
+            xpu_stream = paddle.device.xpu.current_stream().xpu_stream
+            print(xpu_stream)
+            self.assertTrue(type(xpu_stream) is int)
+            ptr = ctypes.c_void_p(xpu_stream)
 
 
 if __name__ == "__main__":