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_STACK_H_
17#define TENSORFLOW_CORE_KERNELS_STACK_H_
18
19// See docs in ../ops/data_flow_ops.cc.
20
21#include "tensorflow/core/framework/op_kernel.h"
22#include "tensorflow/core/framework/types.h"
23#include "tensorflow/core/platform/types.h"
24
25namespace tensorflow {
26
27// A per-run local stack. The stack uses a "per-step" resource manager which
28// ensures that correct garbage collection on error or successful completion.
29class StackOp : public OpKernel {
30 public:
31 explicit StackOp(OpKernelConstruction* context);
32 void Compute(OpKernelContext* ctx) override;
33
34 private:
35 DataType elem_type_;
36 string stack_name_;
37
38 TF_DISALLOW_COPY_AND_ASSIGN(StackOp);
39};
40
41class StackPushOp : public AsyncOpKernel {
42 public:
43 StackPushOp(OpKernelConstruction* context, bool allow_swapping);
44 void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override;
45 bool IsExpensive() override;
46
47 private:
48 bool swap_memory_ = false;
49};
50
51// Templated helper to make it easier to register kernels with or without
52// swapping.
53template <bool allow_swapping>
54class TemplatedStackPushOp : public StackPushOp {
55 public:
56 TemplatedStackPushOp(OpKernelConstruction* context)
57 : StackPushOp(context, allow_swapping) {}
58};
59
60class StackPopOp : public AsyncOpKernel {
61 public:
62 explicit StackPopOp(OpKernelConstruction* context);
63 void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override;
64 bool IsExpensive() override;
65};
66
67class StackCloseOp : public OpKernel {
68 public:
69 explicit StackCloseOp(OpKernelConstruction* context);
70 void Compute(OpKernelContext* ctx) override;
71 bool IsExpensive() override;
72};
73
74} // namespace tensorflow
75
76#endif // TENSORFLOW_CORE_KERNELS_STACK_H_
77