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#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
21namespace 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.
27class 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.
39class 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.
54class 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.
67class 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.
80class 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.
93class 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.
105class 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.
119class 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