From bbc8a7d28a5434b3bdedd0d7e3fdffa4b61378a2 Mon Sep 17 00:00:00 2001
From: zhuhong61 <hong.zhu@intel.com>
Date: Fri, 14 Apr 2023 09:15:06 +0000
Subject: [PATCH] Add ccl event wait to adapt to default oneCCL option

---
 src/gpu/dpcpp_ccl.cpp | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/gpu/dpcpp_ccl.cpp b/src/gpu/dpcpp_ccl.cpp
index c87c27f..acd41c6 100644
--- a/src/gpu/dpcpp_ccl.cpp
+++ b/src/gpu/dpcpp_ccl.cpp
@@ -432,6 +432,7 @@ c10::intrusive_ptr<ProcessGroupCCL::AsyncWorkCCL> XPUCCLStubs::allreduce_(std::v
                                              stream,
                                              attr), stream.get_native());
       });
+      ret_evt.wait();
       return ret_evt;
   },
   c10d::OpType::ALLREDUCE);
@@ -470,6 +471,7 @@ c10::intrusive_ptr<ProcessGroupCCL::AsyncWorkCCL> XPUCCLStubs::reduce_(std::vect
                                 comm,
                                 stream), stream.get_native());
       });
+      ret_evt.wait();
       return ret_evt;
 
   },
@@ -520,6 +522,7 @@ c10::intrusive_ptr<ProcessGroupCCL::AsyncWorkCCL> XPUCCLStubs::_reduce_scatter_b
                                                       comm,
                                                       stream), stream.get_native());
             });
+            ret_evt.wait();
             return ret_evt;
 
           },
@@ -558,6 +561,7 @@ c10::intrusive_ptr<ProcessGroupCCL::AsyncWorkCCL> XPUCCLStubs::broadcast_(std::v
                                              stream,
                                              attr), stream.get_native());
       });
+      ret_evt.wait();
       return ret_evt;
     },
     c10d::OpType::BROADCAST);
@@ -610,7 +614,7 @@ c10::intrusive_ptr<ProcessGroupCCL::AsyncWorkCCL> XPUCCLStubs::allgather_(std::v
                                   comm,
                                   stream), stream.get_native());
       });
-
+      ret_evt.wait();
       return ret_evt;
     },
     c10d::OpType::ALLGATHER);
@@ -658,6 +662,7 @@ c10::intrusive_ptr<ProcessGroupCCL::AsyncWorkCCL> XPUCCLStubs::_allgather_base_(
                                          comm,
                                          stream), stream.get_native());
             });
+            ret_evt.wait();
             return ret_evt;
           },
           c10d::OpType::_ALLGATHER_BASE);
@@ -752,7 +757,7 @@ c10::intrusive_ptr<ProcessGroupCCL::AsyncWorkCCL> XPUCCLStubs::gather_(std::vect
                   }
                 }
               }
-
+              ret_evt.wait();
               return ret_evt;
           },
           c10d::OpType::GATHER);
@@ -803,7 +808,7 @@ c10::intrusive_ptr<ProcessGroupCCL::AsyncWorkCCL> XPUCCLStubs::alltoall_base_(at
                                                       stream,
                                                       attr), stream.get_native());
                 });
-
+                ret_evt.wait();
                 return ret_evt;
             },
             c10d::OpType::ALLTOALL_BASE);
@@ -849,6 +854,7 @@ c10::intrusive_ptr<ProcessGroupCCL::AsyncWorkCCL> XPUCCLStubs::alltoall_base_(at
                                                        stream,
                                                        attr), stream.get_native());
                 });
+                ret_evt.wait();
                 return ret_evt;
             },
             c10d::OpType::ALLTOALL_BASE);
@@ -931,6 +937,7 @@ c10::intrusive_ptr<ProcessGroupCCL::AsyncWorkCCL> XPUCCLStubs::alltoall_(std::ve
                   outputs[i].view({-1}).copy_(flatOutputSplits[i]);
                 }
               }
+              ret_evt.wait();
               return ret_evt;
           },
           c10d::OpType::ALLTOALL);