1/* Copyright 2017 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#include "tensorflow/core/framework/shape_inference.h"
19
20namespace tensorflow {
21
22REGISTER_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(
31A placeholder op to send values to a running XLA computation.
32
33inputs: A list of tensors that will be sent to the XLA computation.
34dynamic_key: The key sent at runtime by the compile node to identify which
35execution the transfer corresponds to.
36Tinputs: The element types of each element in `inputs`.
37key: A key that is unique in the computation and associates the send with the consumer in
38the XLA computation.
39device_ordinal: The device id relative to the associated host device.
40)doc");
41
42REGISTER_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(
51A placeholder op to send values to a running XLA computation with support for a runtime device ordinal.
52
53inputs: A list of tensors that will be sent to the XLA computation.
54dynamic_key: The key sent at runtime by the compile node to identify which
55execution the transfer corresponds to.
56device_ordinal: The device id relative to the associated host device.
57Tinputs: The element types of each element in `inputs`.
58key: A key that is unique in the computation and associates the send with the consumer in
59the XLA computation.
60)doc");
61
62REGISTER_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(
71A placeholder op to receive values from a running XLA computation.
72
73dynamic_key: The key sent at runtime by the compile node to identify which
74execution the transfer corresponds to.
75outputs: A list of tensors that will be received from the XLA computation.
76Toutputs: The element types of each element in `outputs`.
77key: A key that is unique in the computation and associates the send with the consumer in
78the XLA computation.
79device_ordinal: The device id relative to the associated host device.
80)doc");
81
82REGISTER_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(
91A placeholder op to receive values from a running XLA computation with support for a runtime device ordinal.
92
93dynamic_key: The key sent at runtime by the compile node to identify which
94execution the transfer corresponds to.
95device_ordinal: The device id relative to the associated host device.
96outputs: A list of tensors that will be received from the XLA computation.
97Toutputs: The element types of each element in `outputs`.
98key: A key that is unique in the computation and associates the send with the consumer in
99the XLA computation.
100)doc");
101
102} // namespace tensorflow
103