Replace platform::DeviceContext [fluid_ops] part2 (#66082)

co63oc · web-flow · commit 1fc5e50cb99c · 2024-07-18T11:24:27.000+08:00
diff --git a/paddle/fluid/imperative/bkcl_context.cc b/paddle/fluid/imperative/bkcl_context.cc
@@ -202,12 +202,10 @@ void BKCLParallelContext::Broadcast(framework::Variable *src, int ring_id) {
                     platform::errors::Unavailable("bkcl_broadcast failed"));
 }
 
-paddle::platform::DeviceContext *BKCLParallelContext::GetDeviceContext(
-    int ring_id) {
-  return static_cast<platform::DeviceContext *>(
-      platform::BKCLCommContext::Instance()
-          .Get(ring_id, place_)
-          ->dev_context());
+phi::DeviceContext *BKCLParallelContext::GetDeviceContext(int ring_id) {
+  return static_cast<phi::DeviceContext *>(platform::BKCLCommContext::Instance()
+                                               .Get(ring_id, place_)
+                                               ->dev_context());
 }
 
 void BKCLParallelContext::WaitCompute(int ring_id) {
diff --git a/paddle/fluid/imperative/bkcl_context.h b/paddle/fluid/imperative/bkcl_context.h
@@ -46,7 +46,7 @@ class BKCLParallelContext : public ParallelContext {
 
   void Broadcast(framework::Variable* src, int ring_id) override;
 
-  paddle::platform::DeviceContext* GetDeviceContext(int ring_id) override;
+  phi::DeviceContext* GetDeviceContext(int ring_id) override;
 
   void WaitCompute(int ring_id) override;
 
diff --git a/paddle/fluid/imperative/execution_context.h b/paddle/fluid/imperative/execution_context.h
@@ -33,7 +33,7 @@ class DygraphExecutionContext : public framework::ExecutionContext {
  public:
   DygraphExecutionContext(const framework::OperatorBase& op,
                           const framework::Scope& scope,
-                          const platform::DeviceContext& device_context,
+                          const phi::DeviceContext& device_context,
                           const framework::RuntimeContext& ctx,
                           const NameVarMap<VarType>& var_map_in,
                           const NameVarMap<VarType>& var_map_out,
diff --git a/paddle/fluid/imperative/gloo_context.cc b/paddle/fluid/imperative/gloo_context.cc
@@ -196,8 +196,7 @@ void GLOOParallelContext::Broadcast(framework::Variable *src, int ring_id) {
       "Unimplemented inter-broadcast for CPU now."));
 }
 
-paddle::platform::DeviceContext *GLOOParallelContext::GetDeviceContext(
-    int ring_id) {
+phi::DeviceContext *GLOOParallelContext::GetDeviceContext(int ring_id) {
   // return the CPUContext
   return device_.get();
 }
diff --git a/paddle/fluid/imperative/gloo_context.h b/paddle/fluid/imperative/gloo_context.h
@@ -51,7 +51,7 @@ class GLOOParallelContext : public ParallelContext {
 
   void Broadcast(framework::Variable* src, int ring_id) override;
 
-  paddle::platform::DeviceContext* GetDeviceContext(int ring_id) override;
+  phi::DeviceContext* GetDeviceContext(int ring_id) override;
 
   void WaitCompute(int ring_id) override;
 
diff --git a/paddle/fluid/imperative/gradient_accumulator.cc b/paddle/fluid/imperative/gradient_accumulator.cc
@@ -317,7 +317,7 @@ void SelectedRowsAddToTensor(const VarType& src, VarType* dst) {
 
 #define PADDLE_SELECTED_ROWS_ADD_TO_TENSOR(dev_ctx_type, cpp_type)       \
   if (data_type == framework::DataTypeTrait<cpp_type>::DataType()) {     \
-    paddle::platform::DeviceContext* dev_ctx = pool.Get(place);          \
+    phi::DeviceContext* dev_ctx = pool.Get(place);                       \
     phi::funcs::SelectedRowsAddToTensor<dev_ctx_type, cpp_type> functor; \
     functor(*(dynamic_cast<dev_ctx_type*>(dev_ctx)),                     \
             src_selected_rows,                                           \
@@ -438,7 +438,7 @@ std::shared_ptr<ReturnVarType> SelectedRowsMerge(const VarType& src1,
 
 #define PADDLE_SELECTED_ROWS_ADD(dev_ctx_type, cpp_type)             \
   if (data_type == framework::DataTypeTrait<cpp_type>::DataType()) { \
-    paddle::platform::DeviceContext* dev_ctx = pool.Get(place);      \
+    phi::DeviceContext* dev_ctx = pool.Get(place);                   \
     phi::funcs::scatter::MergeAdd<dev_ctx_type, cpp_type> merge_add; \
     merge_add(*(dynamic_cast<dev_ctx_type*>(dev_ctx)),               \
               src_selected_rows,                                     \
diff --git a/paddle/fluid/imperative/heter_ccl_context.cc b/paddle/fluid/imperative/heter_ccl_context.cc
@@ -175,8 +175,7 @@ void HeterParallelContext::Broadcast(framework::Variable *src, int ring_id) {
   PADDLE_THROW(platform::errors::Unimplemented("Unimplemented function."));
 }
 
-paddle::platform::DeviceContext *HeterParallelContext::GetDeviceContext(
-    int ring_id) {
+phi::DeviceContext *HeterParallelContext::GetDeviceContext(int ring_id) {
   // directly call the implementation of target parallel ctx.
   return node_parallel_ctx_->GetDeviceContext(ring_id);
 }
diff --git a/paddle/fluid/imperative/heter_ccl_context.h b/paddle/fluid/imperative/heter_ccl_context.h
@@ -54,7 +54,7 @@ class HeterParallelContext : public ParallelContext {
 
   void Broadcast(framework::Variable* src, int ring_id) override;
 
-  paddle::platform::DeviceContext* GetDeviceContext(int ring_id) override;
+  phi::DeviceContext* GetDeviceContext(int ring_id) override;
 
   void WaitCompute(int ring_id) override;
 
diff --git a/paddle/fluid/imperative/nccl_context.cc b/paddle/fluid/imperative/nccl_context.cc
@@ -153,12 +153,10 @@ void NCCLParallelContext::Broadcast(framework::Variable *src, int ring_id) {
       src_ptr, src_tensor->numel(), nccl_dtype, 0, comm->comm(), stream));
 }
 
-paddle::platform::DeviceContext *NCCLParallelContext::GetDeviceContext(
-    int ring_id) {
-  return static_cast<platform::DeviceContext *>(
-      platform::NCCLCommContext::Instance()
-          .Get(ring_id, place_)
-          ->dev_context());
+phi::DeviceContext *NCCLParallelContext::GetDeviceContext(int ring_id) {
+  return static_cast<phi::DeviceContext *>(platform::NCCLCommContext::Instance()
+                                               .Get(ring_id, place_)
+                                               ->dev_context());
 }
 
 void NCCLParallelContext::WaitCompute(int ring_id) {
diff --git a/paddle/fluid/imperative/nccl_context.h b/paddle/fluid/imperative/nccl_context.h
@@ -64,7 +64,7 @@ class NCCLParallelContext : public ParallelContext {
 
   void Broadcast(framework::Variable* src, int ring_id) override;
 
-  paddle::platform::DeviceContext* GetDeviceContext(int ring_id) override;
+  phi::DeviceContext* GetDeviceContext(int ring_id) override;
 
   void WaitCompute(int ring_id) override;
 
diff --git a/paddle/fluid/imperative/parallel_context.h b/paddle/fluid/imperative/parallel_context.h
@@ -55,7 +55,7 @@ class ParallelContext {
 
   virtual void Broadcast(framework::Variable* src, int ring_id) = 0;
 
-  virtual paddle::platform::DeviceContext* GetDeviceContext(int ring_id) = 0;
+  virtual phi::DeviceContext* GetDeviceContext(int ring_id) = 0;
 
   // comm_stream[ring_id] wait compute_stream.
   // if CPU, should do nothing.
diff --git a/paddle/fluid/imperative/prepared_operator.cc b/paddle/fluid/imperative/prepared_operator.cc
@@ -121,7 +121,7 @@ PreparedOp::PreparedOp(const framework::OperatorBase& op,
                        const framework::OperatorWithKernel::OpKernelFunc& func,
                        const phi::ArgumentMappingFn* arg_map_fn,
                        const phi::KernelSignature* default_kernel_signature,
-                       platform::DeviceContext* dev_ctx)
+                       phi::DeviceContext* dev_ctx)
     : op_(op),
       ctx_(ctx),
       kernel_key_(kernel_key),
@@ -138,7 +138,7 @@ PreparedOp::PreparedOp(const framework::OperatorBase& op,
                        const phi::KernelSignature* default_kernel_signature,
                        phi::KernelSignature&& kernel_signature,
                        const phi::Kernel& phi_kernel,
-                       platform::DeviceContext* dev_ctx)
+                       phi::DeviceContext* dev_ctx)
     : op_(op),
       ctx_(ctx),
       kernel_key_(kernel_key),
@@ -511,7 +511,7 @@ static void PreparedOpRunImpl(
     const framework::OperatorWithKernel::OpKernelFunc& func,
     const phi::ArgumentMappingFn* arg_map_fn,
     const phi::KernelSignature* default_kernel_signature,
-    platform::DeviceContext* dev_ctx,
+    phi::DeviceContext* dev_ctx,
     const NameVarMap<VarType>& ins,
     const NameVarMap<VarType>& outs,
     const framework::AttributeMap& attrs,
@@ -586,7 +586,7 @@ static void PreparedOpRunPtImpl(
     const phi::KernelSignature& kernel_signature,
     const phi::Kernel& phi_kernel,
     const framework::RuntimeContext& ctx,
-    platform::DeviceContext* dev_ctx,
+    phi::DeviceContext* dev_ctx,
     const NameVarMap<VarType>& ins,
     const NameVarMap<VarType>& outs,
     const framework::AttributeMap& attrs,
diff --git a/paddle/fluid/imperative/prepared_operator.h b/paddle/fluid/imperative/prepared_operator.h
@@ -157,7 +157,7 @@ class PreparedOp {
              const framework::OperatorWithKernel::OpKernelFunc& func,
              const phi::ArgumentMappingFn* arg_map_fn,
              const phi::KernelSignature* default_kernel_signature,
-             platform::DeviceContext* dev_ctx);
+             phi::DeviceContext* dev_ctx);
 
   PreparedOp(const framework::OperatorBase& op,
              const framework::RuntimeContext& ctx,
@@ -166,7 +166,7 @@ class PreparedOp {
              const phi::KernelSignature* default_kernel_signature,
              phi::KernelSignature&& kernel_signature,
              const phi::Kernel& phi_kernel,
-             platform::DeviceContext* dev_ctx);
+             phi::DeviceContext* dev_ctx);
 
   static PreparedOp Prepare(const NameVarMap<VarBase>& ins,
                             const NameVarMap<VarBase>& outs,
@@ -213,7 +213,7 @@ class PreparedOp {
   const framework::RuntimeContext& ctx_;
   phi::KernelKey kernel_key_;
   framework::OperatorWithKernel::OpKernelFunc func_;
-  platform::DeviceContext* dev_ctx_;
+  phi::DeviceContext* dev_ctx_;
   // NOTE(chenweihang): Similar op members are used to adapt to
   // new phi kernel, if there is a better design in the future,
   // we may polish the implementation here
@@ -252,7 +252,7 @@ void BuildDygraphPhiKernelContext(const phi::KernelSignature& kernel_signature,
                                   const NameVarMap<VarType>& outs,
                                   const framework::AttributeMap& attrs,
                                   const framework::AttributeMap& default_attrs,
-                                  platform::DeviceContext* dev_ctx,
+                                  phi::DeviceContext* dev_ctx,
                                   phi::KernelContext* kernel_ctx) {
   kernel_ctx->SetDeviceContext(dev_ctx);
 
diff --git a/paddle/fluid/imperative/reducer.cc b/paddle/fluid/imperative/reducer.cc
@@ -33,7 +33,7 @@ namespace imperative {
     defined(PADDLE_WITH_XPU_BKCL) || defined(PADDLE_WITH_GLOO) || \
     defined(PADDLE_WITH_CUSTOM_DEVICE)
 // div the nranks
-void Group::DivNRanks(const platform::DeviceContext &context, int64_t nranks) {
+void Group::DivNRanks(const phi::DeviceContext &context, int64_t nranks) {
   phi::DenseTensor *tensor =
       is_sparse_
           ? sparse_contents_->GetMutable<phi::SelectedRows>()->mutable_value()
@@ -224,7 +224,7 @@ void SplitTensorsWithType<platform::XPUDeviceContext>(
 }
 #endif
 
-void Group::ConcatTensors(const platform::DeviceContext &context) {
+void Group::ConcatTensors(const phi::DeviceContext &context) {
   auto place = context.GetPlace();
   if (phi::is_gpu_place(place)) {  // NOLINT
 #if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
@@ -260,7 +260,7 @@ void Group::ConcatTensors(const platform::DeviceContext &context) {
   }
 }
 
-void Group::SplitTensors(const platform::DeviceContext &context) {
+void Group::SplitTensors(const phi::DeviceContext &context) {
   auto place = context.GetPlace();
   if (phi::is_gpu_place(place)) {  // NOLINT
 #if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
diff --git a/paddle/fluid/imperative/reducer.cu b/paddle/fluid/imperative/reducer.cu
@@ -20,7 +20,7 @@ namespace imperative {
 #if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
 void Group::DivNRanks(phi::DenseTensor *tensor,
                       int64_t nranks,
-                      const platform::DeviceContext &context) {
+                      const phi::DeviceContext &context) {
 #ifdef PADDLE_WITH_HIP
   if (dtype_ == paddle::framework::proto::VarType_Type_BF16) {
     PADDLE_THROW(phi::errors::Fatal("Unsupport BF16 in DataParallel for now"));
diff --git a/paddle/fluid/imperative/reducer.h b/paddle/fluid/imperative/reducer.h
@@ -63,10 +63,10 @@ template <typename Dex>
 struct DivNRanksForAllReduce {
   phi::DenseTensor* in_;
   int64_t nranks_;
-  const platform::DeviceContext& ctx_;
+  const phi::DeviceContext& ctx_;
   DivNRanksForAllReduce(phi::DenseTensor* in,
                         int64_t nranks,
-                        const platform::DeviceContext& ctx)
+                        const phi::DeviceContext& ctx)
       : in_(in), nranks_(nranks), ctx_(ctx) {}
 
   template <typename T>
@@ -105,17 +105,17 @@ class Group {
   framework::proto::VarType::Type dtype_;
 
   // context is used to select the stream for concat
-  void ConcatTensors(const platform::DeviceContext& context);
+  void ConcatTensors(const phi::DeviceContext& context);
 
   // context is used to select the stream for split
-  void SplitTensors(const platform::DeviceContext& context);
+  void SplitTensors(const phi::DeviceContext& context);
 
   // use it in CUDA
   void DivNRanks(phi::DenseTensor* tensor,
                  int64_t nranks,
-                 const platform::DeviceContext& context);
+                 const phi::DeviceContext& context);
 
-  void DivNRanks(const platform::DeviceContext& context, int64_t nranks);
+  void DivNRanks(const phi::DeviceContext& context, int64_t nranks);
 
   friend std::ostream& operator<<(std::ostream&, const Group&);
 };
diff --git a/paddle/fluid/imperative/xccl_context.cc b/paddle/fluid/imperative/xccl_context.cc
@@ -210,9 +210,8 @@ void XCCLParallelContext::Broadcast(framework::Variable *src, int ring_id) {
                                    *stream);
 }
 
-paddle::platform::DeviceContext *XCCLParallelContext::GetDeviceContext(
-    int ring_id) {
-  return static_cast<platform::DeviceContext *>(
+phi::DeviceContext *XCCLParallelContext::GetDeviceContext(int ring_id) {
+  return static_cast<phi::DeviceContext *>(
       platform::XCCLCommContext::Instance(place_.GetDeviceType())
           .Get(ring_id, place_)
           ->dev_context());
diff --git a/paddle/fluid/imperative/xccl_context.h b/paddle/fluid/imperative/xccl_context.h
@@ -51,7 +51,7 @@ class XCCLParallelContext : public ParallelContext {
 
   void Broadcast(framework::Variable* src, int ring_id) override;
 
-  paddle::platform::DeviceContext* GetDeviceContext(int ring_id) override;
+  phi::DeviceContext* GetDeviceContext(int ring_id) override;
 
   void WaitCompute(int ring_id) override;
 
diff --git a/paddle/fluid/inference/tensorrt/convert/test_io_converter.cc b/paddle/fluid/inference/tensorrt/convert/test_io_converter.cc
@@ -21,7 +21,7 @@ namespace paddle {
 namespace inference {
 namespace tensorrt {
 
-void IOConverterTester(const platform::DeviceContext& ctx) {
+void IOConverterTester(const phi::DeviceContext& ctx) {
   cudaStream_t stream;
   ASSERT_EQ(0, cudaStreamCreate(&stream));
 
diff --git a/paddle/fluid/inference/tensorrt/convert/ut_helper.h b/paddle/fluid/inference/tensorrt/convert/ut_helper.h
@@ -47,7 +47,7 @@ float random(float low, float high) {
 
 void RandomizeTensor(phi::DenseTensor* tensor,
                      const phi::Place& place,
-                     const platform::DeviceContext& ctx) {
+                     const phi::DeviceContext& ctx) {
   auto dims = tensor->dims();
   size_t num_elements = analysis::AccuDims(dims, dims.size());
   PADDLE_ENFORCE_GT(
diff --git a/paddle/fluid/pir/dialect/operator/transforms/param_to_variable.cc b/paddle/fluid/pir/dialect/operator/transforms/param_to_variable.cc
@@ -40,7 +40,7 @@ ParameterConvertInterface::ParameterToVariable(pir::Parameter *parameter) {
         parameter->type().dyn_cast<DenseTensorType>().lod(),
         parameter->type().dyn_cast<DenseTensorType>().offset());
     tensor->set_meta(meta);
-    paddle::platform::DeviceContext *dev_ctx =
+    phi::DeviceContext *dev_ctx =
         phi::DeviceContextPool::Instance().Get(phi::CPUPlace());
     dev_ctx->Alloc(tensor,
                    TransToPhiDataType(
diff --git a/paddle/fluid/platform/device_context_test.cu b/paddle/fluid/platform/device_context_test.cu
@@ -20,7 +20,7 @@ limitations under the License. */
 #include "paddle/phi/core/dense_tensor.h"
 
 TEST(Device, Init) {
-  using paddle::platform::DeviceContext;
+  using phi::DeviceContext;
   using phi::GPUContext;
   using phi::GPUPlace;
 
diff --git a/paddle/fluid/platform/device_context_test_cuda_graph.cu b/paddle/fluid/platform/device_context_test_cuda_graph.cu
@@ -20,7 +20,7 @@ limitations under the License. */
 #include "paddle/fluid/platform/device_context.h"
 
 TEST(Device, DeviceContextWithCUDAGraph) {
-  using paddle::platform::DeviceContext;
+  using phi::DeviceContext;
   using phi::DeviceContextPool;
   using phi::GPUContext;
   using phi::GPUPlace;
diff --git a/paddle/fluid/platform/device_context_xpu_test.cc b/paddle/fluid/platform/device_context_xpu_test.cc
@@ -18,8 +18,8 @@ limitations under the License. */
 #include "paddle/fluid/platform/device_context.h"
 
 TEST(Device, Init) {
-  using paddle::platform::DeviceContext;
   using paddle::platform::XPUDeviceContext;
+  using phi::DeviceContext;
   using phi::XPUPlace;
 
   int count = paddle::platform::GetXPUDeviceCount();
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
diff --git a/paddle/fluid/pybind/slice_utils.h b/paddle/fluid/pybind/slice_utils.h
diff --git a/test/cpp/fluid/framework/op_registry_test.cc b/test/cpp/fluid/framework/op_registry_test.cc
diff --git a/test/cpp/fluid/pscore/send_and_recv_op_gpu_test.cc b/test/cpp/fluid/pscore/send_and_recv_op_gpu_test.cc
diff --git a/test/cpp/new_executor/standalone_executor_pir_test.cc b/test/cpp/new_executor/standalone_executor_pir_test.cc
diff --git a/test/cpp/pir/pattern_rewrite/pattern_rewrite_test.cc b/test/cpp/pir/pattern_rewrite/pattern_rewrite_test.cc

Original file line number	Diff line number	Diff line change
`@@ -196,8 +196,7 @@ void GLOOParallelContext::Broadcast(framework::Variable *src, int ring_id) {`
`196`	`196`	`"Unimplemented inter-broadcast for CPU now."));`
`197`	`197`	`}`
`198`	`198`
`199`		`-paddle::platform::DeviceContext *GLOOParallelContext::GetDeviceContext(`
`200`		`- int ring_id) {`
	`199`	`+phi::DeviceContext *GLOOParallelContext::GetDeviceContext(int ring_id) {`
`201`	`200`	`// return the CPUContext`
`202`	`201`	`return device_.get();`
`203`	`202`	`}`
Original file line number	Diff line number	Diff line change
`@@ -175,8 +175,7 @@ void HeterParallelContext::Broadcast(framework::Variable *src, int ring_id) {`
`175`	`175`	`PADDLE_THROW(platform::errors::Unimplemented("Unimplemented function."));`
`176`	`176`	`}`
`177`	`177`
`178`		`-paddle::platform::DeviceContext *HeterParallelContext::GetDeviceContext(`
`179`		`- int ring_id) {`
	`178`	`+phi::DeviceContext *HeterParallelContext::GetDeviceContext(int ring_id) {`
`180`	`179`	`// directly call the implementation of target parallel ctx.`
`181`	`180`	`return node_parallel_ctx_->GetDeviceContext(ring_id);`
`182`	`181`	`}`