1 | /* Copyright 2017 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 | #include "tensorflow/core/framework/shape_inference.h" |
19 | |
20 | namespace tensorflow { |
21 | |
22 | REGISTER_OP("_XlaSendFromHost" ) |
23 | .Input("inputs: Tinputs" ) |
24 | .Input("dynamic_key: string" ) |
25 | .Attr("Tinputs: list(type) >= 0" ) |
26 | .Attr("key: string" ) |
27 | .Attr("device_ordinal: int" ) |
28 | .SetIsStateful() |
29 | .SetShapeFn(::tensorflow::shape_inference::NoOutputs) |
30 | .Doc(R"doc( |
31 | A placeholder op to send values to a running XLA computation. |
32 | |
33 | inputs: A list of tensors that will be sent to the XLA computation. |
34 | dynamic_key: The key sent at runtime by the compile node to identify which |
35 | execution the transfer corresponds to. |
36 | Tinputs: The element types of each element in `inputs`. |
37 | key: A key that is unique in the computation and associates the send with the consumer in |
38 | the XLA computation. |
39 | device_ordinal: The device id relative to the associated host device. |
40 | )doc" ); |
41 | |
42 | REGISTER_OP("_XlaSendFromHostV2" ) |
43 | .Input("inputs: Tinputs" ) |
44 | .Input("dynamic_key: string" ) |
45 | .Input("device_ordinal: int64" ) |
46 | .Attr("Tinputs: list(type) >= 0" ) |
47 | .Attr("key: string" ) |
48 | .SetIsStateful() |
49 | .SetShapeFn(::tensorflow::shape_inference::NoOutputs) |
50 | .Doc(R"doc( |
51 | A placeholder op to send values to a running XLA computation with support for a runtime device ordinal. |
52 | |
53 | inputs: A list of tensors that will be sent to the XLA computation. |
54 | dynamic_key: The key sent at runtime by the compile node to identify which |
55 | execution the transfer corresponds to. |
56 | device_ordinal: The device id relative to the associated host device. |
57 | Tinputs: The element types of each element in `inputs`. |
58 | key: A key that is unique in the computation and associates the send with the consumer in |
59 | the XLA computation. |
60 | )doc" ); |
61 | |
62 | REGISTER_OP("_XlaRecvAtHost" ) |
63 | .Input("dynamic_key: string" ) |
64 | .Output("outputs: Toutputs" ) |
65 | .Attr("Toutputs: list(type) >= 0" ) |
66 | .Attr("key: string" ) |
67 | .Attr("device_ordinal: int" ) |
68 | .SetIsStateful() |
69 | .SetShapeFn(::tensorflow::shape_inference::UnknownShape) |
70 | .Doc(R"doc( |
71 | A placeholder op to receive values from a running XLA computation. |
72 | |
73 | dynamic_key: The key sent at runtime by the compile node to identify which |
74 | execution the transfer corresponds to. |
75 | outputs: A list of tensors that will be received from the XLA computation. |
76 | Toutputs: The element types of each element in `outputs`. |
77 | key: A key that is unique in the computation and associates the send with the consumer in |
78 | the XLA computation. |
79 | device_ordinal: The device id relative to the associated host device. |
80 | )doc" ); |
81 | |
82 | REGISTER_OP("_XlaRecvAtHostV2" ) |
83 | .Input("dynamic_key: string" ) |
84 | .Input("device_ordinal: int64" ) |
85 | .Output("outputs: Toutputs" ) |
86 | .Attr("Toutputs: list(type) >= 0" ) |
87 | .Attr("key: string" ) |
88 | .SetIsStateful() |
89 | .SetShapeFn(::tensorflow::shape_inference::UnknownShape) |
90 | .Doc(R"doc( |
91 | A placeholder op to receive values from a running XLA computation with support for a runtime device ordinal. |
92 | |
93 | dynamic_key: The key sent at runtime by the compile node to identify which |
94 | execution the transfer corresponds to. |
95 | device_ordinal: The device id relative to the associated host device. |
96 | outputs: A list of tensors that will be received from the XLA computation. |
97 | Toutputs: The element types of each element in `outputs`. |
98 | key: A key that is unique in the computation and associates the send with the consumer in |
99 | the XLA computation. |
100 | )doc" ); |
101 | |
102 | } // namespace tensorflow |
103 | |