1 | /* |
2 | * Licensed to the Apache Software Foundation (ASF) under one |
3 | * or more contributor license agreements. See the NOTICE file |
4 | * distributed with this work for additional information |
5 | * regarding copyright ownership. The ASF licenses this file |
6 | * to you under the Apache License, Version 2.0 (the |
7 | * "License"); you may not use this file except in compliance |
8 | * with the License. You may obtain a copy of the License at |
9 | * |
10 | * http://www.apache.org/licenses/LICENSE-2.0 |
11 | * |
12 | * Unless required by applicable law or agreed to in writing, |
13 | * software distributed under the License is distributed on an |
14 | * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
15 | * KIND, either express or implied. See the License for the |
16 | * specific language governing permissions and limitations |
17 | * under the License. |
18 | */ |
19 | |
20 | /*! |
21 | * \file src/contrib/ethosu/cascader/tensor_config.h |
22 | * \brief TensorConfig object for the NPU cascader |
23 | */ |
24 | #ifndef TVM_CONTRIB_ETHOSU_CASCADER_TENSOR_CONFIG_H_ |
25 | #define TVM_CONTRIB_ETHOSU_CASCADER_TENSOR_CONFIG_H_ |
26 | |
27 | #include <tvm/node/reflection.h> |
28 | #include <tvm/runtime/object.h> |
29 | |
30 | #include <functional> |
31 | #include <string> |
32 | #include <utility> |
33 | #include <vector> |
34 | |
35 | #include "graph.h" |
36 | #include "stripe_config.h" |
37 | |
38 | namespace tvm { |
39 | namespace contrib { |
40 | namespace ethosu { |
41 | namespace cascader { |
42 | |
43 | class MemoryRegionNode : public Object { |
44 | public: |
45 | void VisitAttrs(AttrVisitor* v); |
46 | |
47 | /*! \brief The name of the region */ |
48 | std::string name; |
49 | /*! \brief The size of the region */ |
50 | int size; |
51 | /*! \brief The read bandwidth of the region in bytes per cycle */ |
52 | int read_bandwidth; |
53 | /*! \brief The write bandwidth of the region in bytes per cycle */ |
54 | int write_bandwidth; |
55 | /*! \brief The read bandwidth of the region in bytes per cycle */ |
56 | int read_latency; |
57 | /*! \brief The write bandwidth of the region in bytes per cycle */ |
58 | int write_latency; |
59 | /*! \brief Length of memory burst */ |
60 | int burst_length; |
61 | |
62 | static constexpr const char* _type_key = "contrib.ethosu.cascader.MemoryRegion" ; |
63 | TVM_DECLARE_FINAL_OBJECT_INFO(MemoryRegionNode, Object) |
64 | }; |
65 | |
66 | class MemoryRegion : public ObjectRef { |
67 | public: |
68 | MemoryRegion(std::string name, int size, int read_bandwidth, int write_bandwidth, |
69 | int read_latency, int write_latency, int burst_length) { |
70 | auto n = make_object<MemoryRegionNode>(); |
71 | n->name = name; |
72 | n->size = size; |
73 | n->read_bandwidth = read_bandwidth; |
74 | n->write_bandwidth = write_bandwidth; |
75 | n->read_latency = read_latency; |
76 | n->write_latency = write_latency; |
77 | n->burst_length = burst_length; |
78 | data_ = std::move(n); |
79 | } |
80 | |
81 | TVM_DEFINE_OBJECT_REF_METHODS(MemoryRegion, ObjectRef, MemoryRegionNode); |
82 | }; |
83 | |
84 | /*! \brief The 'state' of a TensorConfig as used in the Plan generation algorithm. |
85 | * BOUNDARY - Should describe a Plan input/output Tensor. |
86 | * INTERIOR - Should describe an intermediate Tensor in a 'closed' Plan. |
87 | */ |
88 | enum TensorConfigState { BOUNDARY, INTERIOR }; |
89 | |
90 | /*! \brief Node to represent a TensorConfig */ |
91 | class TensorConfigNode : public Object { |
92 | public: |
93 | void VisitAttrs(AttrVisitor* v); |
94 | |
95 | /*! \return The Tensor the config applies to */ |
96 | const Tensor GetTensor() const { return tensor_; } |
97 | /*! \return The region where the tensor is allocated */ |
98 | MemoryRegion GetHomeRegion() const { return home_region_; } |
99 | /*! |
100 | * \return The state of the TensorConfig. |
101 | * \note The TensorConfigState is only used as part of the Plan generation algorithm. For a Plan |
102 | * to be 'closed' (and therefore not subject to any further merging), all the TensorConfigs that |
103 | * describe Plan input or output Tensors must be in the 'BOUNDARY' state with the rest being |
104 | * 'INTERIOR'. If any of the input or output tensors are described by an 'INTERIOR' TensorConfig, |
105 | * then the Plan is 'open' and should be merged with other 'open' Plans until the result becomes |
106 | * 'closed'. |
107 | */ |
108 | TensorConfigState GetState() const { return state_; } |
109 | /*! |
110 | * \return The mode in which the buffer should be realized |
111 | * \note There are multiple buffering strategies by which a tensor may be realized (computed). |
112 | * These affect the amount of recomputation necessary as well as the size of buffer required to |
113 | * store the tensor. See 'BufferMode' for a description of the allowable buffering modes. |
114 | */ |
115 | BufferMode GetBufferMode() const { return buffer_mode_; } |
116 | /*! |
117 | * \return Whether to copy the tensor. |
118 | * \note While a tensor will originally reside in its home region, the TensorConfig may optionally |
119 | * specify that the tensor should be copied (according to the StripeConfigs) into another |
120 | * MemoryRegion. As an example for where this may be used, if a weights tensor initially resides |
121 | * in slow Flash memory then necessarily the home region will be Flash. However, if the weights |
122 | * values are used multiple times by a Part, it may be more performant to choose to copy the |
123 | * weights into a faster memory like SRAM. |
124 | */ |
125 | bool DoCopy() const { return copy_tensor_; } |
126 | /*! \return The region to copy the tensor to */ |
127 | MemoryRegion GetCopyRegion() const { |
128 | if (!copy_tensor_) { |
129 | return home_region_; |
130 | } |
131 | return copy_region_; |
132 | } |
133 | /*! |
134 | * \return The StripeConfigs with which to compute the tensor. |
135 | * \note The StripeConfigs determine the order in which the elements of the tensor should be |
136 | * computed, including potentially computing them multiple times (recompute). Multiple |
137 | * StripeConfigs are used over just a single StripeConfig for the case where the tensor is |
138 | * consumed by two different Parts executing themselves with different StripeConfigs. In this |
139 | * case, there is a StripeConfig per consumer of the tensor. |
140 | */ |
141 | const std::vector<StripeConfig> GetStripeConfigs() const { return stripe_configs_; } |
142 | /*! |
143 | * \return The size of the buffer needed for the TensorConfig. |
144 | * \note The size of buffer necessary to store a tensor being produced using the TensorConfig is |
145 | * not necessarily just the size of the tensor. In Plans, a tensor may be being produced and |
146 | * consumed in 'stripes' which are smaller than the full tensor. Therefore, the buffer necessary |
147 | * to store the tensor may only need to be as large as the stripe. The precise size of the buffer |
148 | * will depend both on the BufferMode and StripeConfigs (as well as, of course, the Tensor). |
149 | */ |
150 | int GetBufferSize() const; |
151 | /*! \return The hash of the TensorConfigNode */ |
152 | size_t GetHash() const { return hash_; } |
153 | |
154 | static constexpr const char* _type_key = "contrib.ethosu.cascader.TensorConfig" ; |
155 | TVM_DECLARE_FINAL_OBJECT_INFO(TensorConfigNode, Object); |
156 | |
157 | protected: |
158 | friend class TensorConfig; |
159 | |
160 | /*! \brief Compute the hash of the TensorConfigNode */ |
161 | void ComputeHash_(); |
162 | |
163 | /*! \return The size of the recompute buffer needed*/ |
164 | int GetRecomputeBufferSize_() const; |
165 | /*! \return The size of the rolling buffer needed*/ |
166 | int GetRollingBufferSize_() const; |
167 | |
168 | /*! \brief The Tensor the config applies to */ |
169 | Tensor tensor_; |
170 | /*! \brief The region where the tensor is allocated */ |
171 | MemoryRegion home_region_; |
172 | /*! \return The state of the TensorConfig */ |
173 | TensorConfigState state_; |
174 | /*! \brief The mode in which the buffer should be realized */ |
175 | BufferMode buffer_mode_; |
176 | /*! \return The StripeConfigs with which to compute the tensor */ |
177 | std::vector<StripeConfig> stripe_configs_; |
178 | /*! \brief Whether to copy the tensor */ |
179 | bool copy_tensor_; |
180 | /*! \brief The region to copy the tensor to */ |
181 | MemoryRegion copy_region_; |
182 | /*! \brief The hash of the TensorConfigNode */ |
183 | size_t hash_{0}; |
184 | }; |
185 | |
186 | /*! |
187 | * \brief A class which describes how to realize a Tensor. |
188 | * \note The TensorConfig describes both how a Tensor is scheduled (the order in which it's |
189 | * produced/consumed) and how its allocated in memory (which region it should reside in and whether |
190 | * it should be copied). For further detail on how TensorConfig stores this information, consult the |
191 | * documentation of TensorConfigNode. |
192 | */ |
193 | class TensorConfig : public ObjectRef { |
194 | public: |
195 | TensorConfig(const Tensor& tensor, const MemoryRegion& home_region, TensorConfigState state, |
196 | BufferMode buffer_mode, const std::vector<StripeConfig>& stripe_configs, |
197 | bool copy_tensor, const MemoryRegion& copy_region); |
198 | /*! |
199 | * \brief Check if two TensorConfigs are equal to each other. |
200 | * \param other TensorConfig to be checked. |
201 | * \return Whether the two TensorConfigs equal each other. |
202 | */ |
203 | bool operator==(const TensorConfig& other) const; |
204 | |
205 | TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(TensorConfig, ObjectRef, TensorConfigNode); |
206 | }; |
207 | |
208 | } // namespace cascader |
209 | } // namespace ethosu |
210 | } // namespace contrib |
211 | } // namespace tvm |
212 | |
213 | // Hash and equal function for TensorConfig |
214 | namespace std { |
215 | |
216 | /*! \brief The equal_to function for tvm::contrib::ethosu::cascader::TensorConfig */ |
217 | template <> |
218 | struct equal_to<::tvm::contrib::ethosu::cascader::TensorConfig> { |
219 | bool operator()(const ::tvm::contrib::ethosu::cascader::TensorConfig& lhs, |
220 | const ::tvm::contrib::ethosu::cascader::TensorConfig& rhs) const { |
221 | return lhs == rhs; |
222 | } |
223 | }; |
224 | |
225 | /*! \brief The hash function for tvm::contrib::ethosu::cascader::TensorConfig */ |
226 | template <> |
227 | struct hash<::tvm::contrib::ethosu::cascader::TensorConfig> { |
228 | std::size_t operator()( |
229 | const ::tvm::contrib::ethosu::cascader::TensorConfig& tensor_config) const { |
230 | return tensor_config->GetHash(); |
231 | } |
232 | }; |
233 | |
234 | } // namespace std |
235 | |
236 | #endif // TVM_CONTRIB_ETHOSU_CASCADER_TENSOR_CONFIG_H_ |
237 | |