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/stripe_config.h |
22 | * \brief StripeConfig object for the NPU cascader |
23 | */ |
24 | #ifndef TVM_CONTRIB_ETHOSU_CASCADER_STRIPE_CONFIG_H_ |
25 | #define TVM_CONTRIB_ETHOSU_CASCADER_STRIPE_CONFIG_H_ |
26 | |
27 | #include <tvm/node/reflection.h> |
28 | #include <tvm/runtime/object.h> |
29 | |
30 | #include <functional> |
31 | #include <map> |
32 | #include <vector> |
33 | |
34 | namespace tvm { |
35 | namespace contrib { |
36 | namespace ethosu { |
37 | namespace cascader { |
38 | |
39 | class StripeConfig; |
40 | class PropagatorNode; |
41 | |
42 | /*! \brief Node to represent a StripeConfig */ |
43 | class StripeConfigNode : public Object { |
44 | public: |
45 | void VisitAttrs(AttrVisitor* v); |
46 | |
47 | /*! |
48 | * \brief Get the shape of the stripe config. |
49 | * \return The shape of the stripe config. |
50 | * \note The shape refers to the size of the stripes in each dimension. |
51 | */ |
52 | inline std::vector<int> GetShape() const { return shape_; } |
53 | /*! |
54 | * \brief Get the extent of the stripe config. |
55 | * \return The extent of the stripe config. |
56 | * \note The extent refers to the extent over which a StripeConfig operates. |
57 | * Specifically, it is the extent in each axis between the lowest value read |
58 | * by a stripe and the highest value read by a stripe. |
59 | */ |
60 | inline std::vector<int> GetExtent() const { return extent_; } |
61 | /*! |
62 | * \brief Get the strides of the stripe config. |
63 | * \return The strides of the stripe config. |
64 | * \note The strides refer to the stride between stripes in each axis. |
65 | * The strides are represented as a float rather than an int to account for |
66 | * cases of 'fractional striding'. The stride should therefore be interpreted |
67 | * as the average striding in each axis. |
68 | * |
69 | * The starting offset of the i-th stripe in axis 'ax' is given by: |
70 | * |
71 | * stripe_offset_i[ax] = offset[ax] + floor(strides[ax]*i) |
72 | * |
73 | * As a concrete example, consider a 2x2 upscaling operation. If an output |
74 | * stripe config with a stride of (3, 3) is chosen, then when this is |
75 | * propagated to the input it will be reduced by a factor of two to become |
76 | * (1.5, 1.5). |
77 | * |
78 | * This means the first stripe in axis 0 should begin at (floor(1.5*0), 0) = (0, 0), |
79 | * the second at (floor(1.5*1), 0) = (1, 0), and the third at (floor(1.5*2), 0) = |
80 | * (3, 0). This results in irregular striding where 'strides' is the average |
81 | * striding value. |
82 | */ |
83 | inline std::vector<float> GetStrides() const { return strides_; } |
84 | /*! |
85 | * \brief Get the order of the stripe config. |
86 | * \return The order of the stripe config. |
87 | * \note The order refers to order in which the axes are iterated over. |
88 | * The first (outermost) axis is labelled as 1 with the rest increasing |
89 | * according to the axis' position. Any axis labelled with 0 isn't iterated over. |
90 | * For example, [1, 3, 2] would mean axis 0 is the outermost iteration axis, |
91 | * then axis 2, then finally axis 1. |
92 | */ |
93 | inline std::vector<int> GetOrder() const { return order_; } |
94 | /*! |
95 | * \brief Get the stripes of the stripe config. |
96 | * \return The stripes of the stripe config. |
97 | * \note The stripes refer to the number of stripes in each axis. |
98 | * There must be at least one stripe in any given axis. |
99 | */ |
100 | inline std::vector<int> GetStripes() const { return stripes_; } |
101 | /*! |
102 | * \brief Get the offset of the stripe config. |
103 | * \return The offset of the stripe config. |
104 | * \note The offset refers to the offset of the first stripe |
105 | * from the first element of the tensor. For example, in a slice operation |
106 | * which only returns the second (4, 8) half of a (8, 8) tensor, the offset |
107 | * would need to be [4, 0]. |
108 | */ |
109 | inline std::vector<int> GetOffset() const { return offset_; } |
110 | /*! \return The hash of the StripeConfigNode */ |
111 | size_t GetHash() const { return hash_; } |
112 | |
113 | static constexpr const char* _type_key = "contrib.ethosu.cascader.StripeConfig" ; |
114 | TVM_DECLARE_FINAL_OBJECT_INFO(StripeConfigNode, Object); |
115 | |
116 | protected: |
117 | friend class StripeConfig; |
118 | friend class PropagatorNode; |
119 | |
120 | /*! \brief Compute the hash of the StripeConfigNode */ |
121 | void ComputeHash_(); |
122 | |
123 | /*! \brief The shape of the stripes */ |
124 | std::vector<int> shape_; |
125 | /*! \brief The extent of region to stripe over */ |
126 | std::vector<int> extent_; |
127 | /*! \brief The strides of the stripes */ |
128 | std::vector<float> strides_; |
129 | /*! \brief The order of the striping axes */ |
130 | std::vector<int> order_; |
131 | /*! \brief The number of stripes in each axis */ |
132 | std::vector<int> stripes_; |
133 | /*! \brief The offset of the first stripe */ |
134 | std::vector<int> offset_; |
135 | /*! \brief The hash of the StripeConfigNode */ |
136 | std::size_t hash_{0}; |
137 | }; |
138 | |
139 | /*! |
140 | * \brief An object to describe how a tensor should be computed as a series |
141 | of n-dimensional tiles, or 'stripes'. |
142 | * \note The StripeConfig is a verbose way of specifying how to tile a tensor. |
143 | * We can imagine taking a 2D tensor of size (12, 12) and wanting to compute |
144 | * it in tiles of (4, 4). The tile is referred to as a stripe here to generalize |
145 | * this to n-dimensional tiles. |
146 | * |
147 | * The size of that stripe in each axis is the 'shape'. The strides is how far |
148 | * you should move between stripes, so also (4, 4) for a simple non-overlappping |
149 | * tiling. However, we explore some overlapping scheduling options so shape != strides |
150 | * in general. Note that the striding may be fractional, for instance (1.5, 1.5). |
151 | * This means the first stripe should begin at (floor(1.5*0), 0) = (0, 0), the second |
152 | * at (floor(1.5*1), 0) = (1, 0), and the third at (floor(1.5*2), 0) = (3, 0). This results |
153 | * in slightly irregular striding where 'strides' should be interpreted as the average |
154 | * striding value. |
155 | * |
156 | * The 'extent' is simply (12, 12), the region over which we're conducting our tiling. |
157 | * |
158 | * The 'order' tells us which axis to iterate over first and which second and the |
159 | * 'stripes' tells us how many stripes we need to compute in each of those axes. |
160 | * |
161 | * Finally, the 'offset' tells us where to start the first stripe. In this simple |
162 | * case the offset is just (0, 0), but in something like a slice operation we |
163 | * may want to start part way through a tensor. |
164 | */ |
165 | class StripeConfig : public ObjectRef { |
166 | public: |
167 | StripeConfig(const std::vector<int>& shape, const std::vector<int>& extent, |
168 | const std::vector<float>& strides, const std::vector<int>& order, |
169 | const std::vector<int>& stripes, const std::vector<int>& offset); |
170 | /*! |
171 | * \brief Check if two StripeConfigs are equals to each other. |
172 | * \param other StripeConfig to be checked. |
173 | * \return Whether the two StripeConfigs equal each other. |
174 | */ |
175 | bool operator==(const StripeConfig& other) const; |
176 | |
177 | TVM_DEFINE_OBJECT_REF_METHODS(StripeConfig, ObjectRef, StripeConfigNode); |
178 | }; |
179 | |
180 | /*! |
181 | * \brief Count the number of stripes of each shape that are executed for a given |
182 | StripeConfig. |
183 | * \param stripe_config The StripeConfig to count the stripes for. |
184 | * \param enable_sliding_window Whether to assume the sliding window optimization. |
185 | * \return A map between stripe shapes and the number of stripes of that shape that need |
186 | * executing. |
187 | * \note If the StripeConfig were to split an (8, 8) tensor into (4, 4) stripes with |
188 | * (4, 4) striding, then this function will return {(4, 4): 4} indicating that 4 (4, 4) |
189 | * stripes will be executed. If instead an (8, 8) were striped using (5, 5) stripes |
190 | * with (5, 5) striding, this function would return: |
191 | * |
192 | * { |
193 | * (5, 5): 1, |
194 | * (3, 5): 1, |
195 | * (5, 3): 1, |
196 | * (3, 3): 1, |
197 | * } |
198 | * |
199 | * This is because some of the stripes will exceed the extent of the tensor and so only part |
200 | * of them will need executing. Therefore, CountStripes will return the exact number of each |
201 | * shape of stripe that is executed, accounting for edge and overlap behaviour which is not |
202 | * explicit in the StripeConfig alone. |
203 | */ |
204 | std::map<std::vector<int>, int> CountStripes(const StripeConfig& stripe_config, |
205 | bool enable_sliding_window); |
206 | |
207 | } // namespace cascader |
208 | } // namespace ethosu |
209 | } // namespace contrib |
210 | } // namespace tvm |
211 | |
212 | // Hash and equal function for StripeConfig |
213 | namespace std { |
214 | |
215 | /*! \brief The equal_to function for tvm::contrib::ethosu::cascader::StripeConfig */ |
216 | template <> |
217 | struct equal_to<::tvm::contrib::ethosu::cascader::StripeConfig> { |
218 | bool operator()(const ::tvm::contrib::ethosu::cascader::StripeConfig& lhs, |
219 | const ::tvm::contrib::ethosu::cascader::StripeConfig& rhs) const { |
220 | return lhs == rhs; |
221 | } |
222 | }; |
223 | |
224 | /*! \brief The hash function for tvm::contrib::ethosu::cascader::StripeConfig */ |
225 | template <> |
226 | struct hash<::tvm::contrib::ethosu::cascader::StripeConfig> { |
227 | std::size_t operator()( |
228 | const ::tvm::contrib::ethosu::cascader::StripeConfig& stripe_config) const { |
229 | return stripe_config->GetHash(); |
230 | } |
231 | }; |
232 | |
233 | } // namespace std |
234 | |
235 | #endif // TVM_CONTRIB_ETHOSU_CASCADER_STRIPE_CONFIG_H_ |
236 | |