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
34namespace tvm {
35namespace contrib {
36namespace ethosu {
37namespace cascader {
38
39class StripeConfig;
40class PropagatorNode;
41
42/*! \brief Node to represent a StripeConfig */
43class 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 */
165class 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 */
204std::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
213namespace std {
214
215/*! \brief The equal_to function for tvm::contrib::ethosu::cascader::StripeConfig */
216template <>
217struct 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 */
225template <>
226struct 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