add device context getter (#45790)

paddle/fluid/distributed/collective/ProcessGroup.h

@@ -89,6 +89,11 @@ class ProcessGroup {
   int GetSize() const { return size_; }
 
   virtual const std::string GetBackendName() const = 0;
+  virtual phi::DeviceContext* GetDeviceContext(const Place& place) const {
+    PADDLE_THROW(platform::errors::InvalidArgument(
+        "Does not support to get device_context from ProcessGroup%s.",
+        GetBackendName()));
+  }
 
   // TODO(liyurui): This API will be moved later
   virtual std::shared_ptr<ProcessGroup::Task> AllReduce(

paddle/fluid/distributed/collective/ProcessGroupNCCL.cc

@@ -20,6 +20,7 @@
 #include "paddle/fluid/platform/place.h"
 #include "paddle/phi/api/lib/utils/allocator.h"
 #include "paddle/phi/common/place.h"
+#include "paddle/phi/core/device_context.h"
 
 DECLARE_bool(nccl_blocking_wait);
 DECLARE_bool(use_stream_safe_cuda_allocator);
@@ -1041,5 +1042,16 @@ ncclComm_t ProcessGroupNCCL::NCCLComm(const Place& place) const {
   return iter->second[0]->GetNcclComm();
 }
 
+phi::DeviceContext* ProcessGroupNCCL::GetDeviceContext(
+    const Place& place) const {
+  std::vector<Place> places = {place};
+  const auto& iter = places_to_ctx_.find(GetKeyFromPlaces(places));
+  PADDLE_ENFORCE_NE(iter,
+                    places_to_ctx_.end(),
+                    platform::errors::InvalidArgument(
+                        "Cannot find device context in process group."));
+  return iter->second[0].get();
+}
+
 }  //  namespace distributed
 }  //  namespace paddle

paddle/fluid/distributed/collective/ProcessGroupNCCL.h

@@ -96,6 +96,8 @@ class ProcessGroupNCCL : public ProcessGroupStream {
     return std::string(NCCL_BACKEND_NAME);
   }
 
+  phi::DeviceContext* GetDeviceContext(const Place& place) const override;
+
   std::shared_ptr<ProcessGroup::Task> AllReduce(
       std::vector<phi::DenseTensor>& in_tensors,   // NOLINT
       std::vector<phi::DenseTensor>& out_tensors,  // NOLINT