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
38namespace tvm {
39namespace contrib {
40namespace ethosu {
41namespace cascader {
42
43class 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
66class 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 */
88enum TensorConfigState { BOUNDARY, INTERIOR };
89
90/*! \brief Node to represent a TensorConfig */
91class 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 */
193class 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
214namespace std {
215
216/*! \brief The equal_to function for tvm::contrib::ethosu::cascader::TensorConfig */
217template <>
218struct 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 */
226template <>
227struct 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