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 | #ifndef TENSORFLOW_CORE_COMMON_RUNTIME_ENTRY_H_ |
16 | #define TENSORFLOW_CORE_COMMON_RUNTIME_ENTRY_H_ |
17 | |
18 | #include "tensorflow/core/framework/allocator.h" |
19 | #include "tensorflow/core/framework/tensor.h" |
20 | #include "tensorflow/core/lib/gtl/inlined_vector.h" |
21 | #include "tensorflow/core/lib/gtl/manual_constructor.h" |
22 | |
23 | namespace tensorflow { |
24 | |
25 | class Tensor; |
26 | |
27 | // An Entry store a single input value for an individual kernel invocation in |
28 | // an executor. |
29 | // |
30 | // Either a tensor pointer (pass-by-reference) or a tensor (pass-by-value). |
31 | struct Entry { |
32 | enum class State { |
33 | NO_VALUE = 0, // The default state for a newly-created Entry. |
34 | HAS_VALUE, // `this->val` is valid. |
35 | HAS_CONST_TENSOR, // `this->const_tensor` is valid. |
36 | HAS_REF_TENSOR, // `this->ref_tensor` is valid. |
37 | }; |
38 | |
39 | Entry() : state(State::NO_VALUE) {} |
40 | Entry(const Entry& other) : state(other.state), alloc_attr(other.alloc_attr) { |
41 | switch (state) { |
42 | case State::NO_VALUE: |
43 | break; |
44 | case State::HAS_VALUE: |
45 | val.Init(*other.val); |
46 | break; |
47 | case State::HAS_CONST_TENSOR: |
48 | const_tensor = other.const_tensor; |
49 | break; |
50 | case State::HAS_REF_TENSOR: |
51 | ref_tensor = other.ref_tensor; |
52 | break; |
53 | } |
54 | } |
55 | |
56 | ~Entry() { |
57 | if (state == State::HAS_VALUE) val.Destroy(); |
58 | } |
59 | |
60 | Entry& operator=(const Entry& other) { |
61 | if (state == State::HAS_VALUE) { |
62 | val.Destroy(); |
63 | } |
64 | state = other.state; |
65 | alloc_attr = other.alloc_attr; |
66 | switch (state) { |
67 | case State::NO_VALUE: |
68 | break; |
69 | case State::HAS_VALUE: |
70 | val.Init(*other.val); |
71 | break; |
72 | case State::HAS_CONST_TENSOR: |
73 | const_tensor = other.const_tensor; |
74 | break; |
75 | case State::HAS_REF_TENSOR: |
76 | ref_tensor = other.ref_tensor; |
77 | break; |
78 | } |
79 | return *this; |
80 | } |
81 | |
82 | Entry& operator=(Entry&& other) { |
83 | if (state == State::HAS_VALUE) { |
84 | val.Destroy(); |
85 | } |
86 | state = other.state; |
87 | alloc_attr = other.alloc_attr; |
88 | switch (state) { |
89 | case State::NO_VALUE: |
90 | break; |
91 | case State::HAS_VALUE: |
92 | val.Init(std::move(*other.val)); |
93 | break; |
94 | case State::HAS_CONST_TENSOR: |
95 | const_tensor = other.const_tensor; |
96 | break; |
97 | case State::HAS_REF_TENSOR: |
98 | ref_tensor = other.ref_tensor; |
99 | break; |
100 | } |
101 | return *this; |
102 | } |
103 | |
104 | // Clears the <val> field, and sets this entry to the `NO_VALUE` state. |
105 | void ClearVal() { |
106 | if (state == State::HAS_VALUE) { |
107 | val.Destroy(); |
108 | } |
109 | state = State::NO_VALUE; |
110 | } |
111 | |
112 | union { |
113 | // A tensor value. Valid iff `state_ == HAS_VALUE`. |
114 | ManualConstructor<Tensor> val; |
115 | |
116 | // A pointer to a constant tensor value. Valid iff `state_ == |
117 | // HAS_CONST_TENSOR`. |
118 | const Tensor* const_tensor; |
119 | |
120 | // A tensor reference and associated mutex. Valid iff `state_ == |
121 | // HAS_REF_TENSOR`. |
122 | struct { |
123 | Tensor* tensor; |
124 | mutex* mu; |
125 | } ref_tensor; |
126 | }; |
127 | |
128 | // The current state of this entry, indicating which member of the above |
129 | // union is active. |
130 | State state; |
131 | |
132 | // The attributes of the allocator that creates the tensor. |
133 | AllocatorAttributes alloc_attr; |
134 | }; |
135 | |
136 | // TODO(b/152925936): Re-evaluate this constant with current usage patterns. |
137 | typedef gtl::InlinedVector<Entry, 4> EntryVector; |
138 | |
139 | } // namespace tensorflow |
140 | |
141 | #endif // TENSORFLOW_CORE_COMMON_RUNTIME_ENTRY_H_ |
142 | |