1/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#include "tensorflow/core/framework/common_shape_fns.h"
17#include "tensorflow/core/framework/op.h"
18
19namespace tensorflow {
20
21REGISTER_OP("_Send")
22 .Input("tensor: T")
23 .Attr("T: type")
24 .Attr("tensor_name: string")
25 .Attr("send_device: string")
26 .Attr("send_device_incarnation: int")
27 .Attr("recv_device: string")
28 .Attr("client_terminated: bool = false")
29 .SetIsStateful()
30 .SetShapeFn(shape_inference::UnknownShape)
31 .Doc(R"doc(
32Sends the named tensor from send_device to recv_device.
33
34tensor: The tensor to send.
35tensor_name: The name of the tensor to send.
36send_device: The name of the device sending the tensor.
37send_device_incarnation: The current incarnation of send_device.
38recv_device: The name of the device receiving the tensor.
39client_terminated: If set to true, this indicates that the node was added
40 to the graph as a result of a client-side feed or fetch of Tensor data,
41 in which case the corresponding send or recv is expected to be managed
42 locally by the caller.
43)doc");
44
45REGISTER_OP("Send")
46 .Input("tensor: T")
47 .Attr("T: type")
48 .Attr("tensor_name: string")
49 .Attr("send_device: string")
50 .Attr("send_device_incarnation: int")
51 .Attr("recv_device: string")
52 .Attr("client_terminated: bool = false")
53 .SetIsStateful()
54 .SetShapeFn(shape_inference::UnknownShape);
55
56REGISTER_OP("_Recv")
57 .Output("tensor: tensor_type")
58 .Attr("tensor_type: type")
59 .Attr("tensor_name: string")
60 .Attr("send_device: string")
61 .Attr("send_device_incarnation: int")
62 .Attr("recv_device: string")
63 .Attr("client_terminated: bool = false")
64 .SetIsStateful()
65 .SetIsDistributedCommunication()
66 .SetShapeFn(shape_inference::UnknownShape)
67 .Doc(R"doc(
68Receives the named tensor from send_device on recv_device.
69
70tensor: The tensor to receive.
71tensor_name: The name of the tensor to receive.
72send_device: The name of the device sending the tensor.
73send_device_incarnation: The current incarnation of send_device.
74recv_device: The name of the device receiving the tensor.
75client_terminated: If set to true, this indicates that the node was added
76 to the graph as a result of a client-side feed or fetch of Tensor data,
77 in which case the corresponding send or recv is expected to be managed
78 locally by the caller.
79)doc");
80
81REGISTER_OP("Recv")
82 .Output("tensor: tensor_type")
83 .Attr("tensor_type: type")
84 .Attr("tensor_name: string")
85 .Attr("send_device: string")
86 .Attr("send_device_incarnation: int")
87 .Attr("recv_device: string")
88 .Attr("client_terminated: bool = false")
89 .SetIsStateful()
90 .SetIsDistributedCommunication()
91 .SetShapeFn(shape_inference::UnknownShape);
92
93REGISTER_OP("_HostSend")
94 .Input("tensor: T")
95 .Attr("T: type")
96 .Attr("tensor_name: string")
97 .Attr("send_device: string")
98 .Attr("send_device_incarnation: int")
99 .Attr("recv_device: string")
100 .Attr("client_terminated: bool = false")
101 .SetIsStateful()
102 .SetShapeFn(shape_inference::UnknownShape)
103 .Doc(R"doc(
104Sends the named tensor from send_device to recv_device.
105
106_HostSend requires its input on host memory whereas _Send requires its
107input on device memory.
108
109tensor: The tensor to send.
110tensor_name: The name of the tensor to send.
111send_device: The name of the device sending the tensor.
112send_device_incarnation: The current incarnation of send_device.
113recv_device: The name of the device receiving the tensor.
114client_terminated: If set to true, this indicates that the node was added
115 to the graph as a result of a client-side feed or fetch of Tensor data,
116 in which case the corresponding send or recv is expected to be managed
117 locally by the caller.
118)doc");
119
120REGISTER_OP("_HostRecv")
121 .Output("tensor: tensor_type")
122 .Attr("tensor_type: type")
123 .Attr("tensor_name: string")
124 .Attr("send_device: string")
125 .Attr("send_device_incarnation: int")
126 .Attr("recv_device: string")
127 .Attr("client_terminated: bool = false")
128 .SetIsStateful()
129 .SetIsDistributedCommunication()
130 .SetShapeFn(shape_inference::UnknownShape)
131 .Doc(R"doc(
132Receives the named tensor from send_device on recv_device.
133
134_HostRecv produces its output on host memory whereas _Recv produces its
135output on device memory.
136
137tensor: The tensor to receive.
138tensor_name: The name of the tensor to receive.
139send_device: The name of the device sending the tensor.
140send_device_incarnation: The current incarnation of send_device.
141recv_device: The name of the device receiving the tensor.
142client_terminated: If set to true, this indicates that the node was added
143 to the graph as a result of a client-side feed or fetch of Tensor data,
144 in which case the corresponding send or recv is expected to be managed
145 locally by the caller.
146)doc");
147
148} // end namespace tensorflow
149