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