[microNPU][1] Add affine analysis structures for the cascader (#9458)

* [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

* Add test guards

Change-Id: I1d7633e20daab33642fa5c4a12e474a4def4d8b8

* Address review comments

Change-Id: Iff5f1effa08e0628de91f5577487d0cecebec824

* Improve docs

Change-Id: I508809d8c1a08d231e3a9b0fd9b3f2639cc2f0e3

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)

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"""

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

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__)

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

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

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 <tvm/ir/expr.h>
+#include <tvm/runtime/container/array.h>
+
+#include <vector>
+
+namespace tvm {
+namespace contrib {
+namespace ethosu {
+namespace cascader {
+
+/*!
+ * \brief Make a tvm::Array<Integer> from an int vector.
+ * \param vec The int vector.
+ * \return The Integer Array.
+ * \note Array<Integer>(std::vector<int>) doesn't work as this implicit
+ * type conversion fails. This is why this helper is required.
+ */
+inline Array<Integer> make_array(const std::vector<int>& vec) {
+  Array<Integer> 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<FloatImm> from an float vector.
+ * \param vec The float vector.
+ * \return The FloatImm Array.
+ */
+inline Array<FloatImm> make_array(const std::vector<float>& vec) {
+  Array<FloatImm> arr;
+  arr.resize(vec.size());
+  for (unsigned int i = 0; i < vec.size(); ++i) {
+    arr.Set(i, FloatImm(DataType::Float(32), static_cast<double>(vec[i])));
+  }
+  return arr;
+}
+
+/*!
+ * \brief Make a vector from a tvm::Array.
+ * \param arr The Array.
+ * \return The vector.
+ */
+template <class T, class tvm_T>
+inline std::vector<T> make_vector(const Array<tvm_T>& arr) {
+  std::vector<T> 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 <class T>
+inline void hash_combine(std::size_t* seed, T const& v) {
+  *seed ^= std::hash<T>()(v) + 0x9e3779b9 + (*seed << 6) + (*seed >> 2);
+}
+
+/*!
+ * \brief Hash a vector.
+ * \param vec The vector to hash.
+ * \return The hash.
+ */
+template <class T>
+inline std::size_t hash_vector(const std::vector<T>& 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_