diff --git a/cmake/modules/contrib/EthosU.cmake b/cmake/modules/contrib/EthosU.cmake index 8f3e09b8179b..f66e4af55626 100644 --- a/cmake/modules/contrib/EthosU.cmake +++ b/cmake/modules/contrib/EthosU.cmake @@ -16,6 +16,8 @@ # under the License. if(USE_ETHOSU) - file(GLOB ETHOSU_RELAY_CONTRIB_SRC src/relay/backend/contrib/ethosu/*) - list(APPEND COMPILER_SRCS ${ETHOSU_RELAY_CONTRIB_SRC}) + file(GLOB COMPILER_ETHOSU_SRCS + CONFIGURE_DEPENDS src/relay/backend/contrib/ethosu/* + CONFIGURE_DEPENDS src/contrib/ethosu/cascader/*) + list(APPEND COMPILER_SRCS ${COMPILER_ETHOSU_SRCS}) endif(USE_ETHOSU) \ No newline at end of file diff --git a/python/tvm/contrib/ethosu/__init__.py b/python/tvm/contrib/ethosu/__init__.py new file mode 100644 index 000000000000..0ac5badae572 --- /dev/null +++ b/python/tvm/contrib/ethosu/__init__.py @@ -0,0 +1,17 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Namespace for Arm(R) Ethos(TM)-U NPU contrib functionality""" diff --git a/python/tvm/contrib/ethosu/cascader/__init__.py b/python/tvm/contrib/ethosu/cascader/__init__.py new file mode 100644 index 000000000000..009359287fca --- /dev/null +++ b/python/tvm/contrib/ethosu/cascader/__init__.py @@ -0,0 +1,23 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""The NPU cascading planner. + +This component performs inter-operator scheduling to optimize +for both performance and memory usage on Arm(R) Ethos(TM)-U NPUs. +""" +from .stripe_config import StripeConfig +from .propagator import Propagator diff --git a/python/tvm/contrib/ethosu/cascader/_ffi_api.py b/python/tvm/contrib/ethosu/cascader/_ffi_api.py new file mode 100644 index 000000000000..9f098ad3df74 --- /dev/null +++ b/python/tvm/contrib/ethosu/cascader/_ffi_api.py @@ -0,0 +1,21 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""FFI APIs for the NPU cascader.""" +import tvm._ffi + + +tvm._ffi._init_api("contrib.ethosu.cascader", __name__) diff --git a/python/tvm/contrib/ethosu/cascader/propagator.py b/python/tvm/contrib/ethosu/cascader/propagator.py new file mode 100644 index 000000000000..636c265923cc --- /dev/null +++ b/python/tvm/contrib/ethosu/cascader/propagator.py @@ -0,0 +1,57 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Propagator class.""" +# pylint: disable=invalid-name +import tvm._ffi + +from tvm.runtime import Object + +from . import _ffi_api + + +@tvm._ffi.register_object("contrib.ethosu.cascader.Propagator") +class Propagator(Object): + """Propagator class""" + + def __init__(self, transform, offset): + float_transform = list([list(float(v) for v in row) for row in transform]) + self.__init_handle_by_constructor__(_ffi_api.Propagator, float_transform, offset) + + def propagate(self, stripe_config): + return _ffi_api.PropagatorPropagate(self, stripe_config) + + @property + def transform(self): + """Get the transform matrix""" + new_matrix = [] + for row in self._transform: + new_row = [] + for v in row: + new_row.append(v.value) + + new_matrix.append(new_row) + + return new_matrix + + @property + def offset(self): + """Get the offset matrix""" + new_vec = [] + for v in self._offset: + new_vec.append(v.value) + + return new_vec diff --git a/python/tvm/contrib/ethosu/cascader/stripe_config.py b/python/tvm/contrib/ethosu/cascader/stripe_config.py new file mode 100644 index 000000000000..a575e1c20689 --- /dev/null +++ b/python/tvm/contrib/ethosu/cascader/stripe_config.py @@ -0,0 +1,86 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Stripe config class to hold tensor striping information.""" +# pylint: disable=invalid-name +import tvm._ffi + +from tvm.runtime import Object + +from . import _ffi_api + + +@tvm._ffi.register_object("contrib.ethosu.cascader.StripeConfig") +class StripeConfig(Object): + """StripeConfig class""" + + def __init__(self, shape, extent, strides, order, stripes, offset): + strides = list([float(v) for v in strides]) + self.__init_handle_by_constructor__( + _ffi_api.StripeConfig, shape, extent, strides, order, stripes, offset + ) + + @property + def shape(self): + return list(self._shape) + + @property + def extent(self): + return list(self._extent) + + @property + def strides(self): + return list([float(v.value) for v in self._strides]) + + @property + def order(self): + return list(self._order) + + @property + def stripes(self): + return list(self._stripes) + + @property + def offset(self): + return list(self._offset) + + def __hash__(self): + return self._hash + + def __eq__(self, other): + return _ffi_api.StripeConfigEqual(self, other) + + def __repr__(self): + return ( + f"StripeConfig(shape={self.shape}, " + f"extent={self.extent}, " + f"strides={self.strides}, " + f"order={self.order}, " + f"stripes={self.stripes}, " + f"offset={self.offset}" + ) + + +def count_stripes(stripe_config: StripeConfig, enable_sliding_window: bool = False): + stripe_counts = dict(_ffi_api.CountStripes(stripe_config, enable_sliding_window)) + # Some code to 'de-TVM' the data types and make them pure Python + clean_stripe_counts = dict() + for stripe, count in stripe_counts.items(): + clean_stripe = tuple([int(v) for v in stripe]) + clean_count = int(count) + clean_stripe_counts[clean_stripe] = clean_count + + return clean_stripe_counts diff --git a/src/contrib/ethosu/cascader/common.h b/src/contrib/ethosu/cascader/common.h new file mode 100644 index 000000000000..07e1a365f0d3 --- /dev/null +++ b/src/contrib/ethosu/cascader/common.h @@ -0,0 +1,111 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * \file src/contrib/ethosu/cascader/common.h + * \brief Common functions used in the NPU cascader + */ +#ifndef TVM_CONTRIB_ETHOSU_CASCADER_COMMON_H_ +#define TVM_CONTRIB_ETHOSU_CASCADER_COMMON_H_ + +#include +#include + +#include + +namespace tvm { +namespace contrib { +namespace ethosu { +namespace cascader { + +/*! + * \brief Make a tvm::Array from an int vector. + * \param vec The int vector. + * \return The Integer Array. + * \note Array(std::vector) doesn't work as this implicit + * type conversion fails. This is why this helper is required. + */ +inline Array make_array(const std::vector& vec) { + Array arr; + arr.resize(vec.size()); + for (unsigned int i = 0; i < vec.size(); ++i) { + arr.Set(i, Integer(vec[i])); + } + return arr; +} + +/*! + * \brief Make a tvm::Array from an float vector. + * \param vec The float vector. + * \return The FloatImm Array. + */ +inline Array make_array(const std::vector& vec) { + Array arr; + arr.resize(vec.size()); + for (unsigned int i = 0; i < vec.size(); ++i) { + arr.Set(i, FloatImm(DataType::Float(32), static_cast(vec[i]))); + } + return arr; +} + +/*! + * \brief Make a vector from a tvm::Array. + * \param arr The Array. + * \return The vector. + */ +template +inline std::vector make_vector(const Array& arr) { + std::vector vec(arr.size()); + for (unsigned int i = 0; i < arr.size(); ++i) { + vec[i] = arr[i]->value; + } + return vec; +} + +/*! + * \brief Create a combined hash. + * \param seed The current hash value. + * \param v The value to combine into the hash. + * \return The combined hash. + */ +template +inline void hash_combine(std::size_t* seed, T const& v) { + *seed ^= std::hash()(v) + 0x9e3779b9 + (*seed << 6) + (*seed >> 2); +} + +/*! + * \brief Hash a vector. + * \param vec The vector to hash. + * \return The hash. + */ +template +inline std::size_t hash_vector(const std::vector& vec) { + std::size_t seed = vec.size(); + for (const auto& elem : vec) { + hash_combine(&seed, elem); + } + return seed; +} + +} // namespace cascader +} // namespace ethosu +} // namespace contrib +} // namespace tvm + +#endif // TVM_CONTRIB_ETHOSU_CASCADER_COMMON_H_ diff --git a/src/contrib/ethosu/cascader/propagator.cc b/src/contrib/ethosu/cascader/propagator.cc new file mode 100644 index 000000000000..25b711a53d05 --- /dev/null +++ b/src/contrib/ethosu/cascader/propagator.cc @@ -0,0 +1,124 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +#include "propagator.h" + +#include +#include +#include + +#include +#include + +#include "common.h" +#include "stripe_config.h" + +namespace tvm { +namespace contrib { +namespace ethosu { +namespace cascader { + +void PropagatorNode::VisitAttrs(AttrVisitor* v) { + Array > tmp_transform; + for (const auto& vec : transform_) { + tmp_transform.push_back(make_array(vec)); + } + v->Visit("_transform", &tmp_transform); + Array tmp_arr = make_array(offset_); + v->Visit("_offset", &tmp_arr); +} + +Propagator::Propagator(const std::vector >& transform, + const std::vector& offset) { + auto n = make_object(); + size_t rows = transform.size(); + ICHECK_GT(rows, 0) << "The transform matrix must have at least 1 row."; + size_t columns = transform[0].size(); + for (const auto& row : transform) { + ICHECK_EQ(row.size(), columns) + << "All rows of the transform matrix must be of the same length."; + } + ICHECK_EQ(offset.size(), rows - 1) + << "The offset vector length must be equal to the transform matrix rows - 1."; + n->transform_ = std::move(transform); + n->offset_ = std::move(offset); + data_ = std::move(n); +} + +StripeConfig PropagatorNode::propagate(const StripeConfig& stripe_config) const { + size_t input_dimensions = transform_[0].size() - 1; + size_t output_dimensions = transform_.size() - 1; + auto n = make_object(); + n->shape_.resize(output_dimensions); + n->extent_.resize(output_dimensions); + n->strides_.resize(output_dimensions); + n->order_.resize(output_dimensions); + n->stripes_.resize(output_dimensions); + n->offset_.resize(output_dimensions); + for (size_t i = 0; i < output_dimensions; i++) { + float new_shape_acc{}; + float new_extent_acc{}; + const float* row = &transform_[i][0]; + for (size_t j = 0; j < input_dimensions; j++) { + new_shape_acc += row[j] * stripe_config->shape_[j]; + new_extent_acc += row[j] * stripe_config->extent_[j]; + n->strides_[i] += row[j] * stripe_config->strides_[j]; + // Order, stripes and offset should only get re-ordered, so we only + // care about whether or not transform elements are non-zero. + int non_zero = row[j] != 0; + n->order_[i] += non_zero * stripe_config->order_[j]; + n->stripes_[i] += non_zero * stripe_config->stripes_[j]; + n->offset_[i] += non_zero * stripe_config->offset_[j]; + } + // Shape and extent gain an additional constant term + new_shape_acc += row[input_dimensions]; + new_extent_acc += row[input_dimensions]; + // Shape and extent are ceil-rounded back to integers + n->shape_[i] = std::ceil(new_shape_acc); + n->extent_[i] += std::ceil(new_extent_acc); + // Apply the offset + n->offset_[i] += offset_[i]; + // No axis can have '0 stripes', so change all 0 elements to 1 + n->stripes_[i] = n->stripes_[i] == 0 ? 1 : n->stripes_[i]; + } + // Remember to compute the hash + n->ComputeHash_(); + return StripeConfig(n); +} + +TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.Propagator") + .set_body_typed([](Array > transform, Array offset) { + std::vector > vtransform; + for (const auto& vec : transform) { + vtransform.push_back(make_vector(vec)); + } + std::vector voffset = make_vector(offset); + return Propagator(vtransform, voffset); + }); + +TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.PropagatorPropagate") + .set_body_typed([](Propagator propagator, StripeConfig stripe_config) { + return propagator->propagate(stripe_config); + }); + +TVM_REGISTER_NODE_TYPE(PropagatorNode); + +} // namespace cascader +} // namespace ethosu +} // namespace contrib +} // namespace tvm diff --git a/src/contrib/ethosu/cascader/propagator.h b/src/contrib/ethosu/cascader/propagator.h new file mode 100644 index 000000000000..2d4bd0d0154a --- /dev/null +++ b/src/contrib/ethosu/cascader/propagator.h @@ -0,0 +1,137 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * \file src/contrib/ethosu/cascader/propagator.h + * \brief Propagator class for the NPU cascader + */ +#ifndef TVM_CONTRIB_ETHOSU_CASCADER_PROPAGATOR_H_ +#define TVM_CONTRIB_ETHOSU_CASCADER_PROPAGATOR_H_ + +#include +#include + +#include + +namespace tvm { +namespace contrib { +namespace ethosu { +namespace cascader { + +class Propagator; +class StripeConfig; + +/*! \brief Node to represent a Propagator */ +class PropagatorNode : public Object { + public: + void VisitAttrs(AttrVisitor* v); + + /*! \return The transform matrix to apply to the StripeConfigs */ + const std::vector > GetTransform() const { return transform_; } + /*! \return The offset vector to apply to the StripeConfigs */ + const std::vector GetOffset() const { return offset_; } + /*! \return The number of input dimensions */ + size_t GetInputDims() const { return offset_.size(); } + /*! \return The number of output dimensions */ + size_t GetOutputDims() const { return transform_[0].size() - 1; } + /*! + * \brief Propagate a StripeConfig through the transform and offset matrices. + * \param stripe_config The StripeConfig to propagate. + * \return The transformed StripeConfig. + * \note The propagation proceeds as follows: + * + * Both the stripe shape and extent have 1 appended to them (so they pick up + * constant factors from the affine transform) and are then multiplied by the + * transform matrix. The result is then ceil-rounded and has the trailing 1 + * stripped to give the new shape and extent. + * + * The strides has 0 appended to it (so it doesn't pick up constant factors) + * and is then multiplied by the transform matrix. The trailing 0 is stripped. + * + * For the remaining three values we introduce the concept of the 'binarized' + * transform matrix. This is the transform matrix but with every non-zero element + * set to 1. It represents how axes get re-ordered as part of the propagation. + * + * [2, 0, 0, 1] [1, 0, 0, 1] + * [0, 0, 0.4, 2] binarize [0, 0, 1, 1] + * [0, 1.5, 0, 0] ----> [0, 1, 0, 0] + * [0, 0, 0, 1] [0, 0, 0, 1] + * + * The order has 0 appended to it and is multiplied by the 'binarized' transform + * matrix. The trailing 0 is then stripped. + * + * The stripes has 0 appended to it and multiplied by the 'binarized' transform + * matrix. The trailing 0 is then stripped and any remaining 0 elements that + * were introduced by the transform are set instead to 1. + * + * The stripe offset is multiplied by the 'binarized' transform matrix and is + * then summed with the propagator offset. + */ + StripeConfig propagate(const StripeConfig& stripe_config) const; + + static constexpr const char* _type_key = "contrib.ethosu.cascader.Propagator"; + TVM_DECLARE_FINAL_OBJECT_INFO(PropagatorNode, Object); + + protected: + friend class Propagator; + + /*! \brief The transform matrix to apply to the StripeConfigs */ + std::vector > transform_; + /*! \brief The offset vector to apply to the StripeConfigs */ + std::vector offset_; +}; + +/*! + * \brief A class to transform StripeConfigs according to the data dependencies + between Part outputs and inputs. The dependency is represented as an affine + transformation matrix + an offset vector. Using this, an output StripeConfig + can be propagated through a Part to arrive at the input StripeConfigs. + * \note The transform matrix should be a 2D affine transform matrix. + * As an example, consider a (1, 1, 2, 32) output stripe for an NHWC pooling + * operation with a 3x3 pool size: + * + * [1, 0, 0, 0, 0] [ 1] [ 1] + * [0, 1, 0, 0, 2] [ 1] [ 3] + * [0, 0, 1, 0, 2] x [ 2] = [ 4] + * [0, 0, 0, 1, 0] [32] [32] + * [0, 0, 0, 0, 1] [ 1] [ 1] + * + * Using the appropriate affine matrix we see that the required input data to + * produce that output stripe is a (1, 3, 4, 32) stripe. These matrices should + * be derived for the Parts to relate input and output data dependencies. + * + * The offset is a 1D vector representing the first tensor element to read. + * Often this is just the 0 element, but for an operator such as pad it may be + * negative. For instance, a symmetric padding by 1 of a 2D tensor would require + * the offset vector [-1, -1]. Additionally, positive offsets may be required + * for operators like strided_slice where only part of a tensor is read from. + */ +class Propagator : public ObjectRef { + public: + Propagator(const std::vector >& transform, const std::vector& offset); + + TVM_DEFINE_OBJECT_REF_METHODS(Propagator, ObjectRef, PropagatorNode); +}; + +} // namespace cascader +} // namespace ethosu +} // namespace contrib +} // namespace tvm + +#endif // TVM_CONTRIB_ETHOSU_CASCADER_PROPAGATOR_H_ diff --git a/src/contrib/ethosu/cascader/stripe_config.cc b/src/contrib/ethosu/cascader/stripe_config.cc new file mode 100644 index 000000000000..4a75730e5e39 --- /dev/null +++ b/src/contrib/ethosu/cascader/stripe_config.cc @@ -0,0 +1,181 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +#include "stripe_config.h" + +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "common.h" + +namespace tvm { +namespace contrib { +namespace ethosu { +namespace cascader { + +template +std::map, int> MultiplyCombinations(std::vector> values) { + if (values.size() == 1) { + std::map, int> combs; + for (const auto& it : values[0]) { + combs[std::vector(1, it.first)] = it.second; + } + return combs; + } + auto combs = + MultiplyCombinations(std::vector>(values.begin(), values.end() - 1)); + std::map, int> new_combs; + for (const auto& val_it : values.back()) { + for (const auto& comb_it : combs) { + auto new_comb = std::vector(comb_it.first); + new_comb.push_back(val_it.first); + new_combs[new_comb] = val_it.second * comb_it.second; + } + } + return new_combs; +} + +std::map, int> CountStripes(const StripeConfig& stripe_config, + bool enable_sliding_window = false) { + std::vector> per_axis_sizes(stripe_config->GetOrder().size()); + for (size_t axis = 0; axis < stripe_config->GetOrder().size(); axis++) { + int start = stripe_config->GetOffset()[axis]; + size_t stripe_count = static_cast(stripe_config->GetStripes()[axis]); + int stride = stripe_config->GetStrides()[axis]; + int shape = stripe_config->GetShape()[axis]; + int extent = stripe_config->GetExtent()[axis]; + int low; + int high = std::numeric_limits::min(); + for (size_t i = 0; i < stripe_count; i++) { + // Calculate the 'non-edge case' sizes in one go to save effort + if (!enable_sliding_window || i > 0) { + if (start >= 0 && extent - shape - start >= 0 && stride > 0) { + int whole_stripes = + std::min(static_cast(stripe_count - i), (extent - shape - start) / stride + 1); + if (enable_sliding_window) { + per_axis_sizes[axis][stride] += whole_stripes; + } else { + per_axis_sizes[axis][shape] += whole_stripes; + } + i += whole_stripes - 1; + start += whole_stripes * stride; + high = std::min(start - stride + shape, extent); + continue; + } + } + low = std::max(start, 0); + if (enable_sliding_window) { + low = std::max(low, high); + } + high = std::min(start + shape, extent); + int size = high - low; + if (size > 0) { + per_axis_sizes[axis][size]++; + } + start += stride; + } + } + return MultiplyCombinations(per_axis_sizes); +} + +TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.CountStripes") + .set_body_typed([](StripeConfig stripe_config, bool enable_sliding_window) { + Map, Integer> ret; + auto stripe_counts = CountStripes(stripe_config, enable_sliding_window); + for (const auto& it : stripe_counts) { + ret.Set(make_array(it.first), it.second); + } + return ret; + }); + +void StripeConfigNode::VisitAttrs(AttrVisitor* v) { + Array tmp_arr = make_array(shape_); + v->Visit("_shape", &tmp_arr); + tmp_arr = make_array(extent_); + v->Visit("_extent", &tmp_arr); + tmp_arr = make_array(order_); + v->Visit("_order", &tmp_arr); + tmp_arr = make_array(stripes_); + v->Visit("_stripes", &tmp_arr); + tmp_arr = make_array(offset_); + v->Visit("_offset", &tmp_arr); + Array tmp_float_arr = make_array(strides_); + v->Visit("_strides", &tmp_float_arr); + int64_t tmp_hash = static_cast(hash_); + v->Visit("_hash", &tmp_hash); +} + +void StripeConfigNode::ComputeHash_() { + hash_ = hash_vector(shape_); + hash_combine(&hash_, hash_vector(extent_)); + hash_combine(&hash_, hash_vector(strides_)); + hash_combine(&hash_, hash_vector(order_)); + hash_combine(&hash_, hash_vector(stripes_)); + hash_combine(&hash_, hash_vector(offset_)); +} + +StripeConfig::StripeConfig(const std::vector& shape, const std::vector& extent, + const std::vector& strides, const std::vector& order, + const std::vector& stripes, const std::vector& offset) { + auto n = make_object(); + n->shape_ = std::move(shape); + n->extent_ = std::move(extent); + n->strides_ = std::move(strides); + n->order_ = std::move(order); + n->stripes_ = std::move(stripes); + n->offset_ = std::move(offset); + n->ComputeHash_(); + data_ = std::move(n); +} + +inline bool StripeConfig::operator==(const StripeConfig& other) const { + if (get() == other.get()) return true; + if (get() == nullptr || other.get() == nullptr) return false; + return ((*this)->shape_ == other->shape_ && (*this)->extent_ == other->extent_ && + (*this)->strides_ == other->strides_ && (*this)->order_ == other->order_ && + (*this)->stripes_ == other->stripes_ && (*this)->offset_ == other->offset_); +} + +TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.StripeConfig") + .set_body_typed([](Array shape, Array extent, Array strides, + Array order, Array stripes, Array offset) { + std::vector vshape = make_vector(shape); + std::vector vextent = make_vector(extent); + std::vector vstrides = make_vector(strides); + std::vector vorder = make_vector(order); + std::vector vstripes = make_vector(stripes); + std::vector voffset = make_vector(offset); + return StripeConfig(vshape, vextent, vstrides, vorder, vstripes, voffset); + }); + +TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.StripeConfigEqual") + .set_body_method(&StripeConfig::operator==); + +TVM_REGISTER_NODE_TYPE(StripeConfigNode); + +} // namespace cascader +} // namespace ethosu +} // namespace contrib +} // namespace tvm diff --git a/src/contrib/ethosu/cascader/stripe_config.h b/src/contrib/ethosu/cascader/stripe_config.h new file mode 100644 index 000000000000..95759c7e4f03 --- /dev/null +++ b/src/contrib/ethosu/cascader/stripe_config.h @@ -0,0 +1,235 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * \file src/contrib/ethosu/cascader/stripe_config.h + * \brief StripeConfig object for the NPU cascader + */ +#ifndef TVM_CONTRIB_ETHOSU_CASCADER_STRIPE_CONFIG_H_ +#define TVM_CONTRIB_ETHOSU_CASCADER_STRIPE_CONFIG_H_ + +#include +#include + +#include +#include +#include + +namespace tvm { +namespace contrib { +namespace ethosu { +namespace cascader { + +class StripeConfig; +class PropagatorNode; + +/*! \brief Node to represent a StripeConfig */ +class StripeConfigNode : public Object { + public: + void VisitAttrs(AttrVisitor* v); + + /*! + * \brief Get the shape of the stripe config. + * \return The shape of the stripe config. + * \note The shape refers to the size of the stripes in each dimension. + */ + inline std::vector GetShape() const { return shape_; } + /*! + * \brief Get the extent of the stripe config. + * \return The extent of the stripe config. + * \note The extent refers to the extent over which a StripeConfig operates. + * Specifically, it is the extent in each axis between the lowest value read + * by a stripe and the highest value read by a stripe. + */ + inline std::vector GetExtent() const { return extent_; } + /*! + * \brief Get the strides of the stripe config. + * \return The strides of the stripe config. + * \note The strides refer to the stride between stripes in each axis. + * The strides are represented as a float rather than an int to account for + * cases of 'fractional striding'. The stride should therefore be interpreted + * as the average striding in each axis. + * + * The starting offset of the i-th stripe in axis 'ax' is given by: + * + * stripe_offset_i[ax] = offset[ax] + floor(strides[ax]*i) + * + * As a concrete example, consider a 2x2 upscaling operation. If an output + * stripe config with a stride of (3, 3) is chosen, then when this is + * propagated to the input it will be reduced by a factor of two to become + * (1.5, 1.5). + * + * This means the first stripe in axis 0 should begin at (floor(1.5*0), 0) = (0, 0), + * the second at (floor(1.5*1), 0) = (1, 0), and the third at (floor(1.5*2), 0) = + * (3, 0). This results in irregular striding where 'strides' is the average + * striding value. + */ + inline std::vector GetStrides() const { return strides_; } + /*! + * \brief Get the order of the stripe config. + * \return The order of the stripe config. + * \note The order refers to order in which the axes are iterated over. + * The first (outermost) axis is labelled as 1 with the rest increasing + * according to the axis' position. Any axis labelled with 0 isn't iterated over. + * For example, [1, 3, 2] would mean axis 0 is the outermost iteration axis, + * then axis 2, then finally axis 1. + */ + inline std::vector GetOrder() const { return order_; } + /*! + * \brief Get the stripes of the stripe config. + * \return The stripes of the stripe config. + * \note The stripes refer to the number of stripes in each axis. + * There must be at least one stripe in any given axis. + */ + inline std::vector GetStripes() const { return stripes_; } + /*! + * \brief Get the offset of the stripe config. + * \return The offset of the stripe config. + * \note The offset refers to the offset of the first stripe + * from the first element of the tensor. For example, in a slice operation + * which only returns the second (4, 8) half of a (8, 8) tensor, the offset + * would need to be [4, 0]. + */ + inline std::vector GetOffset() const { return offset_; } + /*! \return The hash of the StripeConfigNode */ + size_t GetHash() const { return hash_; } + + static constexpr const char* _type_key = "contrib.ethosu.cascader.StripeConfig"; + TVM_DECLARE_FINAL_OBJECT_INFO(StripeConfigNode, Object); + + protected: + friend class StripeConfig; + friend class PropagatorNode; + + /*! \brief Compute the hash of the StripeConfigNode */ + void ComputeHash_(); + + /*! \brief The shape of the stripes */ + std::vector shape_; + /*! \brief The extent of region to stripe over */ + std::vector extent_; + /*! \brief The strides of the stripes */ + std::vector strides_; + /*! \brief The order of the striping axes */ + std::vector order_; + /*! \brief The number of stripes in each axis */ + std::vector stripes_; + /*! \brief The offset of the first stripe */ + std::vector offset_; + /*! \brief The hash of the StripeConfigNode */ + std::size_t hash_{0}; +}; + +/*! + * \brief An object to describe how a tensor should be computed as a series + of n-dimensional tiles, or 'stripes'. + * \note The StripeConfig is a verbose way of specifying how to tile a tensor. + * We can imagine taking a 2D tensor of size (12, 12) and wanting to compute + * it in tiles of (4, 4). The tile is referred to as a stripe here to generalize + * this to n-dimensional tiles. + * + * The size of that stripe in each axis is the 'shape'. The strides is how far + * you should move between stripes, so also (4, 4) for a simple non-overlappping + * tiling. However, we explore some overlapping scheduling options so shape != strides + * in general. Note that the striding may be fractional, for instance (1.5, 1.5). + * This means the first stripe should begin at (floor(1.5*0), 0) = (0, 0), the second + * at (floor(1.5*1), 0) = (1, 0), and the third at (floor(1.5*2), 0) = (3, 0). This results + * in slightly irregular striding where 'strides' should be interpreted as the average + * striding value. + * + * The 'extent' is simply (12, 12), the region over which we're conducting our tiling. + * + * The 'order' tells us which axis to iterate over first and which second and the + * 'stripes' tells us how many stripes we need to compute in each of those axes. + * + * Finally, the 'offset' tells us where to start the first stripe. In this simple + * case the offset is just (0, 0), but in something like a slice operation we + * may want to start part way through a tensor. + */ +class StripeConfig : public ObjectRef { + public: + StripeConfig(const std::vector& shape, const std::vector& extent, + const std::vector& strides, const std::vector& order, + const std::vector& stripes, const std::vector& offset); + /*! + * \brief Check if two StripeConfigs are equals to each other. + * \param other StripeConfig to be checked. + * \return Whether the two StripeConfigs equal each other. + */ + bool operator==(const StripeConfig& other) const; + + TVM_DEFINE_OBJECT_REF_METHODS(StripeConfig, ObjectRef, StripeConfigNode); +}; + +/*! + * \brief Count the number of stripes of each shape that are executed for a given + StripeConfig. + * \param stripe_config The StripeConfig to count the stripes for. + * \param enable_sliding_window Whether to assume the sliding window optimization. + * \return A map between stripe shapes and the number of stripes of that shape that need + * executing. + * \note If the StripeConfig were to split an (8, 8) tensor into (4, 4) stripes with + * (4, 4) striding, then this function will return {(4, 4): 4} indicating that 4 (4, 4) + * stripes will be executed. If instead an (8, 8) were striped using (5, 5) stripes + * with (5, 5) striding, this function would return: + * + * { + * (5, 5): 1, + * (3, 5): 1, + * (5, 3): 1, + * (3, 3): 1, + * } + * + * This is because some of the stripes will exceed the extent of the tensor and so only part + * of them will need executing. Therefore, CountStripes will return the exact number of each + * shape of stripe that is executed, accounting for edge and overlap behaviour which is not + * explicit in the StripeConfig alone. + */ +std::map, int> CountStripes(const StripeConfig& stripe_config, + bool enable_sliding_window); + +} // namespace cascader +} // namespace ethosu +} // namespace contrib +} // namespace tvm + +// Hash and equal function for StripeConfig +namespace std { + +/*! \brief The equal_to function for tvm::contrib::ethosu::cascader::StripeConfig */ +template <> +struct equal_to<::tvm::contrib::ethosu::cascader::StripeConfig> { + bool operator()(const ::tvm::contrib::ethosu::cascader::StripeConfig& lhs, + const ::tvm::contrib::ethosu::cascader::StripeConfig& rhs) const { + return lhs == rhs; + } +}; + +/*! \brief The hash function for tvm::contrib::ethosu::cascader::StripeConfig */ +template <> +struct hash<::tvm::contrib::ethosu::cascader::StripeConfig> { + std::size_t operator()( + const ::tvm::contrib::ethosu::cascader::StripeConfig& stripe_config) const { + return stripe_config->GetHash(); + } +}; + +} // namespace std + +#endif // TVM_CONTRIB_ETHOSU_CASCADER_STRIPE_CONFIG_H_ diff --git a/tests/python/contrib/test_ethosu/cascader/__init__.py b/tests/python/contrib/test_ethosu/cascader/__init__.py new file mode 100644 index 000000000000..7a08f7d693ed --- /dev/null +++ b/tests/python/contrib/test_ethosu/cascader/__init__.py @@ -0,0 +1,17 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Test infrastructure for the NPU planner""" diff --git a/tests/python/contrib/test_ethosu/cascader/test_propagator.py b/tests/python/contrib/test_ethosu/cascader/test_propagator.py new file mode 100644 index 000000000000..2a6f442f1221 --- /dev/null +++ b/tests/python/contrib/test_ethosu/cascader/test_propagator.py @@ -0,0 +1,136 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +import pytest + +pytest.importorskip("ethosu.vela") + +from math import isclose +from tvm.contrib.ethosu.cascader import StripeConfig, Propagator + + +def test_propagator(): + transform = [ + [1, 0, 0, 0], + [0, 1 / 2, 0, 0], + [0, 0, -1, 0], + [0, 0, 0, 1], + ] + offset = [-1, 1, 2] + propagator = Propagator( + transform=transform, + offset=offset, + ) + assert list(propagator.offset) == offset + for i, row in enumerate(transform): + for j, value in enumerate(row): + assert isclose(propagator.transform[i][j], value) + + +@pytest.mark.parametrize( + ["propagator", "input_stripe_config", "output_stripe_config"], + [ + ( + Propagator( + transform=[ + [1, 0, 0, 0, 0], + [0, 1, 0, 0, 0], + [0, 0, 0, 1 / 16, 0], + [0, 0, 1, 0, 0], + [0, 0, 0, 0, 16], + [0, 0, 0, 0, 1], + ], + offset=[0, 0, 0, 0, 0], + ), + StripeConfig( + shape=[1, 12, 14, 36], + extent=[1, 24, 18, 72], + strides=[1, 12, 14, 36], + order=[1, 2, 3, 4], + stripes=[1, 2, 2, 2], + offset=[0, 0, 0, 0], + ), + StripeConfig( + shape=[1, 12, 3, 14, 16], + extent=[1, 24, 5, 18, 16], + strides=[1, 12, 2.25, 14, 0], + order=[1, 2, 4, 3, 0], + stripes=[1, 2, 2, 2, 1], + offset=[0, 0, 0, 0, 0], + ), + ), + ( + Propagator( + transform=[ + [0.5, 0, 0], + [0, 0.5, 0], + [0, 0, 1], + ], + offset=[0, 0], + ), + StripeConfig( + shape=[3, 5], + extent=[27, 50], + strides=[3, 5], + order=[1, 2], + stripes=[9, 10], + offset=[0, 0], + ), + StripeConfig( + shape=[2, 3], + extent=[14, 25], + strides=[1.5, 2.5], + order=[1, 2], + stripes=[9, 10], + offset=[0, 0], + ), + ), + ( + Propagator( + transform=[ + [2, 0, 0, 4], + [0, 1, 0, 2], + [0, 0, 0, 8], + [0, 0, 0, 1], + ], + offset=[-2, -1, 0], + ), + StripeConfig( + shape=[4, 6, 32], + extent=[48, 60, 64], + strides=[4, 6, 32], + order=[1, 2, 3], + stripes=[12, 10, 2], + offset=[0, 0, 0], + ), + StripeConfig( + shape=[12, 8, 8], + extent=[100, 62, 8], + strides=[8, 6, 0], + order=[1, 2, 0], + stripes=[12, 10, 1], + offset=[-2, -1, 0], + ), + ), + ], +) +def test_propagate(propagator, input_stripe_config, output_stripe_config): + result_stripe_config = propagator.propagate(input_stripe_config) + assert result_stripe_config == output_stripe_config + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/tests/python/contrib/test_ethosu/cascader/test_stripe_config.py b/tests/python/contrib/test_ethosu/cascader/test_stripe_config.py new file mode 100644 index 000000000000..2ca1838b7f34 --- /dev/null +++ b/tests/python/contrib/test_ethosu/cascader/test_stripe_config.py @@ -0,0 +1,215 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +import pytest + +pytest.importorskip("ethosu.vela") + +from tvm.contrib.ethosu.cascader.stripe_config import StripeConfig, count_stripes + + +def test_stripe_config(): + shape = [1, 2, 3] + extent = [2, 3, 4] + strides = [3, 4, 5] + order = [4, 5, 6] + stripes = [5, 6, 7] + offset = [6, 7, 8] + hash_value = 3107995860559090954 + stripe_config = StripeConfig( + shape=shape, + extent=extent, + strides=strides, + order=order, + stripes=stripes, + offset=offset, + ) + assert stripe_config.shape == shape + assert stripe_config.extent == extent + assert stripe_config.strides == strides + assert stripe_config.order == order + assert stripe_config.stripes == stripes + assert stripe_config.offset == offset + assert hash(stripe_config) == hash_value + + +@pytest.mark.parametrize( + "mismatch", [None, "shape", "extent", "strides", "order", "stripes", "offset"] +) +def test_stripe_config_equal(mismatch): + init_dict = { + "shape": [1, 2, 3], + "extent": [2, 3, 4], + "strides": [3, 4, 5], + "order": [4, 5, 6], + "stripes": [5, 6, 7], + "offset": [6, 7, 8], + } + stripe_config_a = StripeConfig(**init_dict) + if mismatch: + init_dict[mismatch] = [1, 1, 1] + stripe_config_b = StripeConfig(**init_dict) + if not mismatch: + assert stripe_config_a == stripe_config_b + else: + assert stripe_config_a != stripe_config_b + + +@pytest.mark.parametrize( + ["stripe_config", "expected_stripe_counts"], + [ + ( + StripeConfig( + shape=[3, 3, 3], + extent=[9, 9, 9], + strides=[3, 3, 3], + order=[1, 2, 3], + stripes=[3, 3, 3], + offset=[0, 0, 0], + ), + { + (3, 3, 3): 27, + }, + ), + ( + StripeConfig( + shape=[3, 3], + extent=[10, 10], + strides=[2, 2], + order=[1, 2], + stripes=[5, 5], + offset=[0, 0], + ), + { + (3, 3): 16, + (2, 3): 4, + (3, 2): 4, + (2, 2): 1, + }, + ), + ( + StripeConfig( + shape=[3, 3, 9], + extent=[9, 9, 9], + strides=[3, 3, 0], + order=[1, 2, 3], + stripes=[3, 3, 1], + offset=[0, 0, 0], + ), + { + (3, 3, 9): 9, + }, + ), + ( + StripeConfig( + shape=[5, 5], + extent=[8, 8], + strides=[5, 5], + order=[1, 2], + stripes=[2, 2], + offset=[0, 0], + ), + { + (5, 5): 1, + (3, 5): 1, + (5, 3): 1, + (3, 3): 1, + }, + ), + ( + StripeConfig( + shape=[5, 5], + extent=[8, 8], + strides=[5, 5], + order=[1, 2], + stripes=[2, 2], + offset=[-1, -2], + ), + { + (4, 3): 2, + (4, 5): 2, + }, + ), + ( + StripeConfig( + shape=[13, 7], + extent=[128, 73], + strides=[13, 7], + order=[1, 2], + stripes=[11, 12], + offset=[-10, -5], + ), + { + (3, 1): 1, + (3, 2): 1, + (8, 7): 10, + (8, 2): 1, + (13, 7): 90, + (13, 1): 9, + (8, 1): 1, + (3, 7): 10, + (13, 2): 9, + }, + ), + ], +) +def test_count_stripes(stripe_config, expected_stripe_counts): + assert count_stripes(stripe_config) == expected_stripe_counts + + +@pytest.mark.parametrize( + ["stripe_config", "expected_stripe_counts"], + [ + ( + StripeConfig( + shape=[4, 4], + extent=[16, 16], + strides=[2, 2], + order=[1, 2], + stripes=[7, 7], + offset=[0, 0], + ), + { + (4, 4): 1, + (2, 4): 6, + (4, 2): 6, + (2, 2): 36, + }, + ), + ( + StripeConfig( + shape=[4, 4], + extent=[8, 8], + strides=[2, 2], + order=[1, 2], + stripes=[6, 3], + offset=[-5, 0], + ), + { + (1, 4): 2, + (2, 4): 3, + (2, 2): 6, + (1, 2): 4, + }, + ), + ], +) +def test_count_stripes_sliding_window(stripe_config, expected_stripe_counts): + assert count_stripes(stripe_config, enable_sliding_window=True) == expected_stripe_counts + + +if __name__ == "__main__": + pytest.main([__file__])