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 | #ifndef TENSORFLOW_CORE_KERNELS_CONTROL_FLOW_OPS_H_ |
17 | #define TENSORFLOW_CORE_KERNELS_CONTROL_FLOW_OPS_H_ |
18 | |
19 | #include "tensorflow/core/framework/op_kernel.h" |
20 | |
21 | namespace tensorflow { |
22 | |
23 | // A ControlTriggerOp is similar to a NoOp. However, it always treats the input |
24 | // control edges as Live edges. Its primary use so far is in the scheduling of |
25 | // recvs, where we add ControlTrigger nodes and use them to trigger recvs. We |
26 | // allow ControlTrigger nodes to be enabled by dead nodes. |
27 | class ControlTriggerOp : public OpKernel { |
28 | public: |
29 | explicit ControlTriggerOp(OpKernelConstruction* context) |
30 | : OpKernel(context) {} |
31 | void Compute(OpKernelContext* context) override {} |
32 | bool IsExpensive() override { return false; } |
33 | }; |
34 | |
35 | // A switch op has two inputs and two outputs. It forwards the value of |
36 | // Input:0 to the output specified by input:1. Input:1 is a boolean tensor. |
37 | // Input:0 is forwarded to output:0 if input:1 is false, otherwise to |
38 | // output:1. |
39 | class SwitchOp : public OpKernel { |
40 | public: |
41 | explicit SwitchOp(OpKernelConstruction* context) : OpKernel(context) {} |
42 | void Compute(OpKernelContext* context) override; |
43 | bool IsExpensive() override { return false; } |
44 | ~SwitchOp() override {} |
45 | |
46 | TF_DISALLOW_COPY_AND_ASSIGN(SwitchOp); |
47 | }; |
48 | |
49 | // An n-way switch op has two inputs and N outputs. It forwards the value of |
50 | // Input:0 to the output specified by Input:1. Input:1 is an integer tensor. |
51 | // Input:0 is forwarded to output:0 if Input:1 is 0, to output:1 if 1, and so |
52 | // forth. If Input:1 is <0 or >=num_outputs(), Input:0 is forwarded to |
53 | // output:num_outputs()-1. |
54 | class SwitchNOp : public OpKernel { |
55 | public: |
56 | explicit SwitchNOp(OpKernelConstruction* context) : OpKernel(context) {} |
57 | void Compute(OpKernelContext* context) override; |
58 | bool IsExpensive() override { return false; } |
59 | ~SwitchNOp() override {} |
60 | |
61 | TF_DISALLOW_COPY_AND_ASSIGN(SwitchNOp); |
62 | }; |
63 | |
64 | // A merge op has n inputs and two outputs. It forwards the value of the |
65 | // first input that becomes available to its first output, and the |
66 | // index of the first input to its second output. |
67 | class MergeOp : public OpKernel { |
68 | public: |
69 | explicit MergeOp(OpKernelConstruction* context); |
70 | void Compute(OpKernelContext* context) override; |
71 | bool IsExpensive() override { return false; } |
72 | ~MergeOp() override {} |
73 | |
74 | TF_DISALLOW_COPY_AND_ASSIGN(MergeOp); |
75 | }; |
76 | |
77 | // An enter op has one input and one output. It creates or finds |
78 | // the child frame that is uniquely identified by the frame_name, |
79 | // and makes its input available to the child frame. |
80 | class EnterOp : public OpKernel { |
81 | public: |
82 | explicit EnterOp(OpKernelConstruction* context) : OpKernel(context) {} |
83 | void Compute(OpKernelContext* context) override; |
84 | bool IsExpensive() override { return false; } |
85 | ~EnterOp() override {} |
86 | |
87 | TF_DISALLOW_COPY_AND_ASSIGN(EnterOp); |
88 | }; |
89 | |
90 | // An exit op has one input and one output. It exits the current |
91 | // frame to its parent frame, and makes its input available to the |
92 | // parent frame. |
93 | class ExitOp : public OpKernel { |
94 | public: |
95 | explicit ExitOp(OpKernelConstruction* context) : OpKernel(context) {} |
96 | void Compute(OpKernelContext* context) override; |
97 | bool IsExpensive() override { return false; } |
98 | ~ExitOp() override {} |
99 | |
100 | TF_DISALLOW_COPY_AND_ASSIGN(ExitOp); |
101 | }; |
102 | |
103 | // A next_iteration op has one input and one output. It makes its input |
104 | // available to the next iteration. |
105 | class NextIterationOp : public OpKernel { |
106 | public: |
107 | explicit NextIterationOp(OpKernelConstruction* context) : OpKernel(context) {} |
108 | void Compute(OpKernelContext* context) override; |
109 | bool IsExpensive() override { return false; } |
110 | ~NextIterationOp() override {} |
111 | |
112 | TF_DISALLOW_COPY_AND_ASSIGN(NextIterationOp); |
113 | }; |
114 | |
115 | // A LoopCond op has one input and one output. The input is a boolean |
116 | // scalar representing the taken branches of the "pivot" Switch that |
117 | // determines loop termination. As a contract, any high-level front-end |
118 | // should always use port '0' of the "pivot" switches for loop exit. |
119 | class LoopCondOp : public OpKernel { |
120 | public: |
121 | explicit LoopCondOp(OpKernelConstruction* context); |
122 | ~LoopCondOp() override; |
123 | |
124 | void Compute(OpKernelContext* context) override; |
125 | |
126 | bool IsExpensive() override; |
127 | |
128 | TF_DISALLOW_COPY_AND_ASSIGN(LoopCondOp); |
129 | }; |
130 | |
131 | } // namespace tensorflow |
132 | |
133 | #endif // TENSORFLOW_CORE_KERNELS_CONTROL_FLOW_OPS_H_ |
134 | |