remove use_calc_stream param [fluid_ops] partial_recv (#68953)

co63oc · web-flow · commit 2359f6d0b863 · 2024-10-31T11:50:07.000+08:00
* Fix

* Fix
diff --git a/paddle/fluid/operators/collective/partial_recv_op.cc b/paddle/fluid/operators/collective/partial_recv_op.cc
@@ -100,10 +100,6 @@ class PartialRecvOpMaker : public framework::OpProtoAndCheckerMaker {
 
     AddAttr<std::vector<int>>("out_shape", "shape of the output tensor.")
         .SetDefault(std::vector<int>());
-    AddAttr<bool>(
-        "use_calc_stream",
-        "(bool default false) eject CUDA operations to calculation stream.")
-        .SetDefault(false);
     AddAttr<int>("num", "(int default 1) The number of Output to be cut.")
         .SetDefault(1);
     AddAttr<int>("id",
diff --git a/paddle/fluid/operators/ops_signature/partial_recv_sig.cc b/paddle/fluid/operators/ops_signature/partial_recv_sig.cc
@@ -0,0 +1,27 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/compat/op_utils.h"
+
+namespace phi {
+
+KernelSignature PartialRecvOpArgumentMapping(
+    const ArgumentMappingContext& ctx) {
+  return KernelSignature(
+      "partial_recv", {}, {"peer", "dtype", "out_shape", "num", "id"}, {"Out"});
+}
+
+}  // namespace phi
+
+PD_REGISTER_ARG_MAPPING_FN(partial_recv, phi::PartialRecvOpArgumentMapping);
diff --git a/paddle/phi/infermeta/nullary.cc b/paddle/phi/infermeta/nullary.cc
@@ -128,8 +128,7 @@ void GaussianInferMeta(const IntArray& shape,
   out->set_layout(DataLayout::NCHW);
 }
 
-void PartialRecvInferMeta(int ring_id,
-                          int peer,
+void PartialRecvInferMeta(int peer,
                           DataType dtype,
                           const std::vector<int>& out_shape,
                           int num,
@@ -140,12 +139,6 @@ void PartialRecvInferMeta(int ring_id,
       0,
       common::errors::InvalidArgument(
           "The peer (%d) for partial_recv op must be non-negative.", peer));
-  PADDLE_ENFORCE_GE(
-      ring_id,
-      0,
-      common::errors::InvalidArgument(
-          "The ring_id (%d) for partial_recv op must be non-negative.",
-          ring_id));
   PADDLE_ENFORCE_GE(num,
                     1,
                     common::errors::InvalidArgument(
diff --git a/paddle/phi/infermeta/nullary.h b/paddle/phi/infermeta/nullary.h
@@ -85,8 +85,7 @@ void RandpermInferMeta(int n, DataType dtype, MetaTensor* out);
 void RandintInferMeta(
     int low, int high, const IntArray& shape, DataType dtype, MetaTensor* out);
 
-void PartialRecvInferMeta(int ring_id,
-                          int peer,
+void PartialRecvInferMeta(int peer,
                           DataType dtype,
                           const std::vector<int>& out_shape,
                           int num,
diff --git a/paddle/phi/kernels/cpu/partial_recv_kernel.cc b/paddle/phi/kernels/cpu/partial_recv_kernel.cc
@@ -18,7 +18,6 @@ namespace phi {
 
 template <typename T, typename Context>
 void PartialRecvKernel(const Context& dev_ctx UNUSED,
-                       int ring_id UNUSED,
                        int peer UNUSED,
                        DataType dtype UNUSED,
                        const std::vector<int>& out_shape UNUSED,
diff --git a/paddle/phi/kernels/gpu/partial_recv_kernel.cu b/paddle/phi/kernels/gpu/partial_recv_kernel.cu
@@ -23,7 +23,6 @@ namespace phi {
 
 template <typename T, typename Context>
 void PartialRecvKernel(const Context& dev_ctx,
-                       int ring_id UNUSED,
                        int peer,
                        DataType type,
                        const std::vector<int>& out_shape,
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -679,9 +679,11 @@
   output : Tensor(out)
   infer_meta :
     func: PartialRecvInferMeta
+    param: [peer, dtype, out_shape, num, id]
   kernel :
     func : partial_recv
     data_type : dtype
+    param: [peer, dtype, out_shape, num, id]
 
 - op : print
   args : (Tensor in, int first_n, str message, int summarize, bool print_tensor_name = true, bool print_tensor_type = true, bool print_tensor_shape = true, bool print_tensor_layout = true, bool print_tensor_lod = true, str print_phase = "BOTH", bool is_forward = true)
