From d05b76b9755b5552e2943f4108f63d98ae2286c6 Mon Sep 17 00:00:00 2001 From: Matthew Barrett Date: Mon, 7 Jun 2021 15:26:17 +0100 Subject: [PATCH 1/4] [ETHOSU][1] Add affine analysis structures for the cascader The cascader relies heavily on being able to determine data dependencies between operators. This is so that it can calculate how stripes should be propagated through a cascade. To do this, two data structures are defined: StripeConfig and Propagator. StripeConfig stores information for how a tensor should be broken up into stripes and executed. Propagator transforms a StripeConfig using an affine transform matrix, allowing an input StripeConfig for an operator to be determined by 'propagating' the output StripeConfig. By chaining together Propagators, we can analyse how data dependencies vary throughout a cascade and therefore calculate the memory requirements (and approximate the performance). Change-Id: If7176fea961c631be4a6c195303da536030d957b --- cmake/modules/contrib/EthosU.cmake | 6 +- python/tvm/contrib/ethosu/__init__.py | 17 ++ .../tvm/contrib/ethosu/cascader/__init__.py | 23 ++ .../tvm/contrib/ethosu/cascader/_ffi_api.py | 21 ++ .../tvm/contrib/ethosu/cascader/propagator.py | 60 +++++ .../contrib/ethosu/cascader/stripe_config.py | 86 +++++++ src/contrib/ethosu/cascader/common.h | 111 +++++++++ src/contrib/ethosu/cascader/propagator.cc | 132 +++++++++++ src/contrib/ethosu/cascader/propagator.h | 137 +++++++++++ src/contrib/ethosu/cascader/stripe_config.cc | 181 +++++++++++++++ src/contrib/ethosu/cascader/stripe_config.h | 202 +++++++++++++++++ .../contrib/test_ethosu/cascader/__init__.py | 17 ++ .../test_ethosu/cascader/test_propagator.py | 134 +++++++++++ .../cascader/test_stripe_config.py | 213 ++++++++++++++++++ 14 files changed, 1338 insertions(+), 2 deletions(-) create mode 100644 python/tvm/contrib/ethosu/__init__.py create mode 100644 python/tvm/contrib/ethosu/cascader/__init__.py create mode 100644 python/tvm/contrib/ethosu/cascader/_ffi_api.py create mode 100644 python/tvm/contrib/ethosu/cascader/propagator.py create mode 100644 python/tvm/contrib/ethosu/cascader/stripe_config.py create mode 100644 src/contrib/ethosu/cascader/common.h create mode 100644 src/contrib/ethosu/cascader/propagator.cc create mode 100644 src/contrib/ethosu/cascader/propagator.h create mode 100644 src/contrib/ethosu/cascader/stripe_config.cc create mode 100644 src/contrib/ethosu/cascader/stripe_config.h create mode 100644 tests/python/contrib/test_ethosu/cascader/__init__.py create mode 100644 tests/python/contrib/test_ethosu/cascader/test_propagator.py create mode 100644 tests/python/contrib/test_ethosu/cascader/test_stripe_config.py 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..2f6415f89a6a --- /dev/null +++ b/python/tvm/contrib/ethosu/cascader/propagator.py @@ -0,0 +1,60 @@ +# 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) + + def benchmark_propagate(self, stripe_config, repeats): + return _ffi_api.PropagatorPropagateBenchmark(self, stripe_config, repeats) + + @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..e3d2b3e58573 --- /dev/null +++ b/src/contrib/ethosu/cascader/propagator.cc @@ -0,0 +1,132 @@ +/* + * 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_GLOBAL("contrib.ethosu.cascader.PropagatorPropagateBenchmark") + .set_body_typed([](Propagator propagator, StripeConfig stripe_config, int repeats) { + for (int i = 0; i < repeats; i++) { + propagator->propagate(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..ce6cb901007b --- /dev/null +++ b/src/contrib/ethosu/cascader/stripe_config.h @@ -0,0 +1,202 @@ +/* + * 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'. This may happen, for instance, with an + * upscaling operation where elements of the affine transformation matrix + * are not integers. In this case we can't simply round the strides as the + * error will compound when we need to multiply the strides by the number of + * stripes along a given axis. + */ + 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 2D padding operation + * that is padding by 1 in every dimension, the offset would be [-1, -1]. + */ + 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. 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 padding operation we + * may want to start from a negative index, which is captured by the offset. + */ +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. + */ +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..5506db1b0396 --- /dev/null +++ b/tests/python/contrib/test_ethosu/cascader/test_propagator.py @@ -0,0 +1,134 @@ +# 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. +from tvm.contrib.ethosu.cascader import StripeConfig, Propagator + +import pytest +from math import isclose + + +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..1ea3bf8e7082 --- /dev/null +++ b/tests/python/contrib/test_ethosu/cascader/test_stripe_config.py @@ -0,0 +1,213 @@ +# 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. +from tvm.contrib.ethosu.cascader.stripe_config import StripeConfig, count_stripes + +import pytest + + +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__]) From 6c93b3a2cfabc3bc58aa80da04d4c52ca0a38843 Mon Sep 17 00:00:00 2001 From: Matthew Barrett Date: Fri, 5 Nov 2021 15:56:38 +0000 Subject: [PATCH 2/4] Add test guards Change-Id: I1d7633e20daab33642fa5c4a12e474a4def4d8b8 --- .../python/contrib/test_ethosu/cascader/test_propagator.py | 6 ++++-- .../contrib/test_ethosu/cascader/test_stripe_config.py | 6 ++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/python/contrib/test_ethosu/cascader/test_propagator.py b/tests/python/contrib/test_ethosu/cascader/test_propagator.py index 5506db1b0396..2a6f442f1221 100644 --- a/tests/python/contrib/test_ethosu/cascader/test_propagator.py +++ b/tests/python/contrib/test_ethosu/cascader/test_propagator.py @@ -14,10 +14,12 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from tvm.contrib.ethosu.cascader import StripeConfig, Propagator - import pytest + +pytest.importorskip("ethosu.vela") + from math import isclose +from tvm.contrib.ethosu.cascader import StripeConfig, Propagator def test_propagator(): diff --git a/tests/python/contrib/test_ethosu/cascader/test_stripe_config.py b/tests/python/contrib/test_ethosu/cascader/test_stripe_config.py index 1ea3bf8e7082..2ca1838b7f34 100644 --- a/tests/python/contrib/test_ethosu/cascader/test_stripe_config.py +++ b/tests/python/contrib/test_ethosu/cascader/test_stripe_config.py @@ -14,10 +14,12 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from tvm.contrib.ethosu.cascader.stripe_config import StripeConfig, count_stripes - 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] From 912114128b1e8de525122a5522b69c2e1e058899 Mon Sep 17 00:00:00 2001 From: Matthew Barrett Date: Thu, 18 Nov 2021 13:19:16 +0000 Subject: [PATCH 3/4] Address review comments Change-Id: Iff5f1effa08e0628de91f5577487d0cecebec824 --- python/tvm/contrib/ethosu/cascader/propagator.py | 3 --- src/contrib/ethosu/cascader/propagator.cc | 8 -------- src/contrib/ethosu/cascader/stripe_config.h | 16 ++++++++++++++++ 3 files changed, 16 insertions(+), 11 deletions(-) diff --git a/python/tvm/contrib/ethosu/cascader/propagator.py b/python/tvm/contrib/ethosu/cascader/propagator.py index 2f6415f89a6a..636c265923cc 100644 --- a/python/tvm/contrib/ethosu/cascader/propagator.py +++ b/python/tvm/contrib/ethosu/cascader/propagator.py @@ -34,9 +34,6 @@ def __init__(self, transform, offset): def propagate(self, stripe_config): return _ffi_api.PropagatorPropagate(self, stripe_config) - def benchmark_propagate(self, stripe_config, repeats): - return _ffi_api.PropagatorPropagateBenchmark(self, stripe_config, repeats) - @property def transform(self): """Get the transform matrix""" diff --git a/src/contrib/ethosu/cascader/propagator.cc b/src/contrib/ethosu/cascader/propagator.cc index e3d2b3e58573..25b711a53d05 100644 --- a/src/contrib/ethosu/cascader/propagator.cc +++ b/src/contrib/ethosu/cascader/propagator.cc @@ -116,14 +116,6 @@ TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.PropagatorPropagate") return propagator->propagate(stripe_config); }); -TVM_REGISTER_GLOBAL("contrib.ethosu.cascader.PropagatorPropagateBenchmark") - .set_body_typed([](Propagator propagator, StripeConfig stripe_config, int repeats) { - for (int i = 0; i < repeats; i++) { - propagator->propagate(stripe_config); - } - return propagator->propagate(stripe_config); - }); - TVM_REGISTER_NODE_TYPE(PropagatorNode); } // namespace cascader diff --git a/src/contrib/ethosu/cascader/stripe_config.h b/src/contrib/ethosu/cascader/stripe_config.h index ce6cb901007b..c2f129800a69 100644 --- a/src/contrib/ethosu/cascader/stripe_config.h +++ b/src/contrib/ethosu/cascader/stripe_config.h @@ -167,6 +167,22 @@ class StripeConfig : public ObjectRef { * \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); From b6dd634a9b8d4e13c6077d3c9747765d088f058f Mon Sep 17 00:00:00 2001 From: Matthew Barrett Date: Thu, 25 Nov 2021 11:36:18 +0000 Subject: [PATCH 4/4] Improve docs Change-Id: I508809d8c1a08d231e3a9b0fd9b3f2639cc2f0e3 --- src/contrib/ethosu/cascader/stripe_config.h | 39 +++++++++++++++------ 1 file changed, 28 insertions(+), 11 deletions(-) diff --git a/src/contrib/ethosu/cascader/stripe_config.h b/src/contrib/ethosu/cascader/stripe_config.h index c2f129800a69..95759c7e4f03 100644 --- a/src/contrib/ethosu/cascader/stripe_config.h +++ b/src/contrib/ethosu/cascader/stripe_config.h @@ -63,11 +63,22 @@ class StripeConfigNode : public Object { * \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'. This may happen, for instance, with an - * upscaling operation where elements of the affine transformation matrix - * are not integers. In this case we can't simply round the strides as the - * error will compound when we need to multiply the strides by the number of - * stripes along a given axis. + * 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_; } /*! @@ -91,8 +102,9 @@ class StripeConfigNode : public Object { * \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 2D padding operation - * that is padding by 1 in every dimension, the offset would be [-1, -1]. + * 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 */ @@ -135,15 +147,20 @@ class StripeConfigNode : public Object { * 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. The 'extent' is simply (12, 12), the region over which we're conducting - * our tiling. + * 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 padding operation we - * may want to start from a negative index, which is captured by the offset. + * 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: