diff --git a/paddle/phi/kernels/cpu/graph_send_recv_grad_kernel.cc b/paddle/phi/kernels/cpu/graph_send_recv_grad_kernel.cc index 6a83cee1ae40d1dec6c3d7d2f4499d7eb71390f9..95eeb64afea20943fd16e474173d9a72211888d8 100644 --- a/paddle/phi/kernels/cpu/graph_send_recv_grad_kernel.cc +++ b/paddle/phi/kernels/cpu/graph_send_recv_grad_kernel.cc @@ -118,12 +118,12 @@ void GraphSendRecvGradOpKernelLaunchHelper( template void GraphSendRecvGradKernel(const Context& ctx, - const DenseTensor& out_grad, const DenseTensor& x, - paddle::optional out, const DenseTensor& src_index, const DenseTensor& dst_index, + paddle::optional out, paddle::optional dst_count, + const DenseTensor& out_grad, const std::string& pool_type, DenseTensor* x_grad) { auto index_type = src_index.dtype(); diff --git a/paddle/phi/kernels/gpu/graph_send_recv_grad_kernel.cu b/paddle/phi/kernels/gpu/graph_send_recv_grad_kernel.cu index 8bd3337280d7596ebd7a638b1cc86cf37b70632d..2be0caff79d64b38980c8b06003205db296ea83e 100644 --- a/paddle/phi/kernels/gpu/graph_send_recv_grad_kernel.cu +++ b/paddle/phi/kernels/gpu/graph_send_recv_grad_kernel.cu @@ -102,12 +102,12 @@ void GraphSendRecvGradOpCUDAKernelLaunchHelper( template void GraphSendRecvGradKernel(const Context& ctx, - const DenseTensor& out_grad, const DenseTensor& x, - paddle::optional out, const DenseTensor& src_index, const DenseTensor& dst_index, + paddle::optional out, paddle::optional dst_count, + const DenseTensor& out_grad, const std::string& pool_type, DenseTensor* x_grad) { auto index_type = src_index.dtype(); diff --git a/paddle/phi/kernels/graph_send_recv_grad_kernel.h b/paddle/phi/kernels/graph_send_recv_grad_kernel.h index 3694c8f1e6c9905586ea82d49a87b6b1a12e1e71..c0b1a34d09c008ae41f557d6e7f38fb7c814c7e8 100644 --- a/paddle/phi/kernels/graph_send_recv_grad_kernel.h +++ b/paddle/phi/kernels/graph_send_recv_grad_kernel.h @@ -22,12 +22,12 @@ namespace phi { template void GraphSendRecvGradKernel(const Context& ctx, - const DenseTensor& out_grad, const DenseTensor& x, - paddle::optional out, const DenseTensor& src_index, const DenseTensor& dst_index, + paddle::optional out, paddle::optional dst_count, + const DenseTensor& out_grad, const std::string& pool_type, DenseTensor* x_grad); } // namespace phi diff --git a/paddle/phi/ops/compat/graph_send_recv_sig.cc b/paddle/phi/ops/compat/graph_send_recv_sig.cc index fa4da0704c9871f8b5969fdf79da6dcb8617cae9..cf36b9baa2d03e04ca883fed0ce5f4f80cd8bba7 100644 --- a/paddle/phi/ops/compat/graph_send_recv_sig.cc +++ b/paddle/phi/ops/compat/graph_send_recv_sig.cc @@ -28,7 +28,7 @@ KernelSignature GraphSendRecvGradOpArgumentMapping( const ArgumentMappingContext& ctx) { return KernelSignature( "graph_send_recv_grad", - {GradVarName("Out"), "X", "Out", "Src_index", "Dst_index", "Dst_count"}, + {"X", "Src_index", "Dst_index", "Out", "Dst_count", GradVarName("Out")}, {"pool_type"}, {GradVarName("X")}); } diff --git a/python/paddle/fluid/dygraph/tracer.py b/python/paddle/fluid/dygraph/tracer.py index 05ae17c5e1816702b068b6c2f7c6e357352ad521..6e1ed6b0a1dec80913e27244124754dd5586256d 100644 --- a/python/paddle/fluid/dygraph/tracer.py +++ b/python/paddle/fluid/dygraph/tracer.py @@ -22,6 +22,14 @@ from paddle.fluid import framework from paddle import _C_ops final_state_name_mapping = { + "graph_send_recv": { + "final_op_name": "final_state_graph_send_recv", + "x": "X", + "src_index": "Src_index", + "dst_index": "Dst_index", + "out": "Out", + "dst_count": "Dst_count" + }, "matmul_v2": { "final_op_name": "final_state_matmul", "transpose_x": "trans_x", diff --git a/python/paddle/fluid/tests/unittests/test_graph_send_recv_op.py b/python/paddle/fluid/tests/unittests/test_graph_send_recv_op.py index 30f943e3248e90820abbef26214e81cf691f50a4..c233606c053d8ee44a51b888ff4ef03f766353e4 100644 --- a/python/paddle/fluid/tests/unittests/test_graph_send_recv_op.py +++ b/python/paddle/fluid/tests/unittests/test_graph_send_recv_op.py @@ -17,13 +17,26 @@ import unittest import numpy as np import paddle import paddle.fluid as fluid +from paddle.fluid.framework import _test_eager_guard from op_test import OpTest +def graph_send_recv_wrapper(x, + src_index, + dst_index, + pool_type="sum", + out_size=None, + name=None): + return paddle.incubate.graph_send_recv(x, src_index, dst_index, + pool_type.lower(), out_size, name) + + class TestGraphSendRecvMaxOp(OpTest): def setUp(self): paddle.enable_static() + self.python_api = graph_send_recv_wrapper + self.python_out_sig = ["Out"] self.op_type = "graph_send_recv" x = np.random.random((10, 20)).astype("float64") index = np.random.randint(0, 10, (15, 2)).astype(np.int64) @@ -39,15 +52,18 @@ class TestGraphSendRecvMaxOp(OpTest): self.outputs = {'Out': out} def test_check_output(self): - self.check_output() + self.check_output(check_eager=True) def test_check_grad(self): - self.check_grad(['X'], 'Out', user_defined_grads=[self.gradient]) + self.check_grad( + ['X'], 'Out', user_defined_grads=[self.gradient], check_eager=True) class TestGraphSendRecvMinOp(OpTest): def setUp(self): paddle.enable_static() + self.python_api = graph_send_recv_wrapper + self.python_out_sig = ["Out"] self.op_type = "graph_send_recv" x = np.random.random((10, 20)).astype("float64") index = np.random.randint(0, 10, (15, 2)).astype(np.int64) @@ -64,15 +80,18 @@ class TestGraphSendRecvMinOp(OpTest): self.outputs = {'Out': out} def test_check_output(self): - self.check_output() + self.check_output(check_eager=True) def test_check_grad(self): - self.check_grad(['X'], 'Out', user_defined_grads=[self.gradient]) + self.check_grad( + ['X'], 'Out', user_defined_grads=[self.gradient], check_eager=True) class TestGraphSendRecvSumOp(OpTest): def setUp(self): paddle.enable_static() + self.python_api = graph_send_recv_wrapper + self.python_out_sig = ["Out"] self.op_type = "graph_send_recv" x = np.random.random((10, 20)).astype("float64") index = np.random.randint(0, 10, (15, 2)).astype(np.int64) @@ -88,15 +107,17 @@ class TestGraphSendRecvSumOp(OpTest): self.outputs = {'Out': out} def test_check_output(self): - self.check_output() + self.check_output(check_eager=True) def test_check_grad(self): - self.check_grad(['X'], 'Out') + self.check_grad(['X'], 'Out', check_eager=True) class TestGraphSendRecvMeanOp(OpTest): def setUp(self): paddle.enable_static() + self.python_api = graph_send_recv_wrapper + self.python_out_sig = ["Out"] self.op_type = "graph_send_recv" x = np.random.random((10, 20)).astype("float64") index = np.random.randint(0, 10, (15, 2)).astype(np.int64) @@ -113,10 +134,10 @@ class TestGraphSendRecvMeanOp(OpTest): self.outputs = {'Out': out, 'Dst_count': dst_count} def test_check_output(self): - self.check_output() + self.check_output(check_eager=True) def test_check_grad(self): - self.check_grad(['X'], 'Out') + self.check_grad(['X'], 'Out', check_eager=True) def compute_graph_send_recv_for_sum_mean(inputs, attributes): @@ -333,6 +354,12 @@ class API_GraphSendRecvOpTest(unittest.TestCase): {}\n{}, check diff!" .format(np_res_set_outsize, res_set_outsize)) + def test_api_eager_dygraph(self): + with _test_eager_guard(): + self.test_dygraph() + self.test_int32_input() + self.test_set_outsize_gpu() + if __name__ == '__main__': unittest.main() diff --git a/python/paddle/incubate/operators/graph_send_recv.py b/python/paddle/incubate/operators/graph_send_recv.py index 3972180d86564d3de2160b914a7a172d58fd65ab..80a21aec6cf5ac7f15fa77b2770daf5af4903ac5 100644 --- a/python/paddle/incubate/operators/graph_send_recv.py +++ b/python/paddle/incubate/operators/graph_send_recv.py @@ -13,7 +13,7 @@ # limitations under the License. from paddle.fluid.layer_helper import LayerHelper -from paddle.fluid.framework import _non_static_mode +from paddle.fluid.framework import _non_static_mode, _in_legacy_dygraph, in_dygraph_mode from paddle.fluid.data_feeder import check_variable_and_dtype from paddle.fluid import core from paddle import _C_ops @@ -109,15 +109,30 @@ def graph_send_recv(x, # TODO(daisiming): Should we add judgement for out_size: max(dst_index) + 1. - if _non_static_mode(): - if out_size is None or out_size <= 0: + if out_size is None or out_size <= 0: + if _in_legacy_dygraph(): out, tmp = _C_ops.graph_send_recv(x, src_index, dst_index, 'pool_type', pool_type.upper()) - else: + return out + if in_dygraph_mode(): + return _C_ops.final_state_graph_send_recv(x, src_index, dst_index, + pool_type.upper(), 0) + else: + if _in_legacy_dygraph(): out, tmp = _C_ops.graph_send_recv( x, src_index, dst_index, 'pool_type', pool_type.upper(), 'out_size', out_size) - return out + return out + if in_dygraph_mode(): + if isinstance(out_size, core.eager.Tensor): + if (out_size.size < 1): + raise ValueError( + "out_size should be long type, but received Tensor type." + ) + out_size = out_size.numpy()[0] + return _C_ops.final_state_graph_send_recv(x, src_index, dst_index, + pool_type.upper(), + out_size) check_variable_and_dtype(x, "X", ("float32", "float64", "int32", "int64"), "graph_send_recv") diff --git a/python/paddle/utils/code_gen/api.yaml b/python/paddle/utils/code_gen/api.yaml index 23d4f54d0383ad9a66728b0ac159dbfe46cda3e3..5865d07845fb4ab4964ea5628c143ecc590c62e8 100644 --- a/python/paddle/utils/code_gen/api.yaml +++ b/python/paddle/utils/code_gen/api.yaml @@ -756,6 +756,17 @@ func : gelu backward : gelu_grad +- api : graph_send_recv + args : (Tensor x, Tensor src_index, Tensor dst_index, str pool_type = "SUM", int64_t out_size = 0) + output : Tensor(out), Tensor(dst_count) + infer_meta : + func : GraphSendRecvInferMeta + kernel : + func : graph_send_recv + data_type : x + intermediate : dst_count + backward : graph_send_recv_grad + - api : greater_equal args : (Tensor x, Tensor y, int axis = -1) output : Tensor @@ -1162,7 +1173,7 @@ kernel : func : mean_all backward : mean_all_grad - + - api : meshgrid args : (Tensor[] inputs) output : Tensor[] diff --git a/python/paddle/utils/code_gen/backward.yaml b/python/paddle/utils/code_gen/backward.yaml index 8745e9d038108d84c51bdafcc2c965298c677f12..adfe9c2b99860809b03fd9776eddf41ed9de4ae9 100644 --- a/python/paddle/utils/code_gen/backward.yaml +++ b/python/paddle/utils/code_gen/backward.yaml @@ -537,6 +537,17 @@ kernel : func : gelu_grad +- backward_api : graph_send_recv_grad + forward : graph_send_recv (Tensor x, Tensor src_index, Tensor dst_index, str pool_type = "SUM", int64_t out_size = 0) -> Tensor(out), Tensor(dst_count) + args : (Tensor x, Tensor src_index, Tensor dst_index, Tensor out, Tensor dst_count, Tensor out_grad, str pool_type = "SUM") + output : Tensor(x_grad) + infer_meta : + func : GeneralUnaryGradInferMeta + param : [x] + kernel : + func : graph_send_recv_grad + optional: out, dst_count + - backward_api : hard_shrink_grad forward : hard_shrink (Tensor x, float threshold) -> Tensor(out) args : (Tensor x, Tensor out_grad, float threshold)