Calculate CMSIS-NN buffer size with respect to architecture extensions

This correctly calculates the buffer sizes for a variety of targets
based on the `-mcpu` and `-mattr` flags passed to the `cmsis-nn` code
generator.

Added for Conv2d, Depthwise Conv2d and Average Pool.

Author: Mousius

cmake/modules/contrib/CMSISNN.cmake

@@ -16,6 +16,7 @@
 # under the License.
 
 if(USE_CMSISNN)
+  add_definitions(-DTVM_USE_CMSISNN)
   message(STATUS "Build with CMSIS-NN support")
   tvm_file_glob(GLOB RELAY_CONTRIB_CMSISNN_SRCS src/relay/backend/contrib/cmsisnn/*.cc)
   list(APPEND COMPILER_SRCS ${RELAY_CONTRIB_CMSISNN_SRCS})

python/tvm/driver/tvmc/composite_target.py

@@ -53,7 +53,7 @@
         "pass_pipeline": partition_for_arm_compute_lib,
     },
     "cmsis-nn": {
-        "config_key": None,
+        "config_key": "relay.ext.cmsisnn.options",
         "pass_pipeline": partition_for_cmsisnn,
     },
     "ethos-n77": {

src/relay/backend/contrib/cmsisnn/buffer_size.cc

@@ -0,0 +1,78 @@
+/*
+ * 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 <tvm/ir/attrs.h>
+#include <tvm/ir/transform.h>
+
+#include "compiler_attrs.h"
+
+namespace tvm {
+namespace relay {
+namespace contrib {
+namespace cmsisnn {
+
+int Conv2dBufferSize(CMSISNNFlags flags, int32_t padding_w, int32_t padding_h, int32_t input_n,
+                     int32_t input_h, int32_t input_c, int32_t output_h, int32_t output_w,
+                     int32_t stride_w, int32_t stride_h, int32_t filter_w, int32_t filter_h) {
+  bool is1x1 = (padding_w == 0) && (padding_h == 0) && (input_c % 4 == 0) && (stride_w == 1) &&
+               (stride_h == 1) && (filter_w == 1) && (filter_h == 1);
+  bool is1xN =
+      (output_h == 1) && (input_h == 1) && (filter_h == 1) && (output_w % 4 == 0) && (input_n == 1);
+
+  if (is1x1) {
+    return 0;
+  }
+
+  if (is1xN) {
+    if (flags.dsp && !flags.mve) {
+      return (2 * input_c * filter_w * filter_h) * (int32_t)sizeof(int16_t);
+    }
+    return 0;
+  }
+
+  if (flags.dsp) {
+    return (2 * input_c * filter_w * filter_h) * (int32_t)sizeof(int16_t);
+  }
+  return 0;
+}
+
+int DepthwiseConv2dBufferSize(CMSISNNFlags flags, int32_t input_n, int32_t input_c,
+                              int32_t output_c, int32_t filter_w, int32_t filter_h) {
+  if (input_c == output_c && input_n == 1) {
+    if (flags.mve) {
+      return (2 * input_c * filter_w * filter_h) * (int32_t)sizeof(int16_t) + 4;
+    }
+    if (flags.dsp) {
+      return (input_c * filter_w * filter_h) * (int32_t)sizeof(int16_t);
+    }
+  }
+  return 0;
+}
+
+int AvgPoolBufferSize(CMSISNNFlags flags, int32_t input_c) {
+  if (flags.dsp && !flags.mve) {
+    return (input_c * sizeof(int32_t));
+  }
+  return 0;
+}
+
+}  // namespace cmsisnn
+}  // namespace contrib
+}  // namespace relay
+}  // namespace tvm

src/relay/backend/contrib/cmsisnn/buffer_size.h

@@ -0,0 +1,94 @@
+/*
+ * 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/relay/backend/contrib/cmsisnn/buffer_size.h
+ * \brief CMSIS-NN Buffer Size calculation functions
+ */
+
+#ifndef TVM_RELAY_BACKEND_CONTRIB_CMSISNN_BUFFER_SIZE_H_
+#define TVM_RELAY_BACKEND_CONTRIB_CMSISNN_BUFFER_SIZE_H_
+
+#include <tvm/ir/transform.h>
+
+#include "compiler_attrs.h"
+
+namespace tvm {
+namespace relay {
+namespace contrib {
+namespace cmsisnn {
+
+/*!
+ * \brief Calculates the appropriate buffer size for CMSIS-NN Convolutions
+ * See:
+ * https://github.com/ARM-software/CMSIS_5/blob/8c60448c0e1e50e426180b26db9bc31ddf774361/CMSIS/NN/Source/ConvolutionFunctions/arm_convolve_wrapper_s8.c#L108-L127
+ *
+ * \param flags - CMSIS-NN feature flags
+ * \param padding_w - Width padding
+ * \param padding_h - Height padding
+ * \param input_n - Input batch size
+ * \param input_h - Input height
+ * \param input_c - Input channels
+ * \param output_h - Output height
+ * \param output_w - Output width
+ * \param stride_w - Stride width
+ * \param stride_h - Stride height
+ * \param filter_w - Filter width
+ * \param filter_h - Filter height
+ *
+ * \return Size of buffer to allocate for convolution
+ */
+int Conv2dBufferSize(CMSISNNFlags flags, int32_t padding_w, int32_t padding_h, int32_t input_n,
+                     int32_t input_h, int32_t input_c, int32_t output_h, int32_t output_w,
+                     int32_t stride_w, int32_t stride_h, int32_t filter_w, int32_t filter_h);
+
+/*!
+ * \brief Calculates the appropriate buffer size for CMSIS-NN Depthwise Convolutions
+ * See:
+ * https://github.com/ARM-software/CMSIS_5/blob/325443e52637b6c7eedbd160d238a6c462e89c9f/CMSIS/NN/Source/ConvolutionFunctions/arm_depthwise_conv_wrapper_s8.c#L115-L129
+ *
+ * \param flags - CMSIS-NN feature flags
+ * \param input_n - Input batch size
+ * \param input_c - Input channels
+ * \param output_c - Output channels
+ * \param filter_w - Filter width
+ * \param filter_h - Filter height
+ *
+ * \return Size of buffer to allocate for depthwise convolution
+ */
+int DepthwiseConv2dBufferSize(CMSISNNFlags flags, int32_t input_n, int32_t input_c,
+                              int32_t output_c, int32_t filter_w, int32_t filter_h);
+
+/*!
+ * \brief Calculates the appropriate buffer size for CMSIS-NN Average Pooling
+ * See:
+ * https://github.com/ARM-software/CMSIS_5/blob/bff28575f0c96a4ee9008947fea2b018a69b4900/CMSIS/NN/Source/PoolingFunctions/arm_avgpool_s8.c#L388-L398
+ *
+ * \param input_c - Input channels
+ *
+ * \return Size of buffer to allocate for average pooling
+ */
+int AvgPoolBufferSize(CMSISNNFlags flags, int32_t input_c);
+
+}  // namespace cmsisnn
+}  // namespace contrib
+}  // namespace relay
+}  // namespace tvm
+
+#endif  // TVM_RELAY_BACKEND_CONTRIB_CMSISNN_BUFFER_SIZE_H_

src/relay/backend/contrib/cmsisnn/compiler_attrs.cc

@@ -0,0 +1,75 @@
+/*
+ * 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 "compiler_attrs.h"
+
+#include <tvm/ir/attrs.h>
+#include <tvm/ir/transform.h>
+
+#include <string>
+
+namespace tvm {
+namespace relay {
+namespace contrib {
+namespace cmsisnn {
+
+static const char* mveCPUs[] = {"cortex-m55"};
+static const char* dspCPUs[] = {"cortex-m7", "cortex-m33", "cortex-m35p"};
+
+TVM_REGISTER_NODE_TYPE(CMSISNNCompilerConfigNode);
+TVM_REGISTER_PASS_CONFIG_OPTION("relay.ext.cmsisnn.options", CMSISNNCompilerConfig);
+
+template <typename Container>
+static inline bool MatchesCpu(std::string mcpu, const Container& cpus) {
+  auto matches_cpu = [mcpu](const char* cpu) { return mcpu.find(cpu) == 0; };
+  return std::find_if(std::begin(cpus), std::end(cpus), matches_cpu) != std::end(cpus);
+}
+
+static inline bool HasFlag(std::string attr, std::string flag) {
+  return attr.find(flag) != std::string::npos;
+}
+
+CMSISNNFlags GetCompilerFlags(const tvm::transform::PassContext& ctx) {
+  auto cfg = ctx->GetConfig<CMSISNNCompilerConfig>("relay.ext.cmsisnn.options");
+  if (!cfg.defined()) {
+    return kNoExt;
+  }
+
+  std::string mcpu = cfg.value()->mcpu;
+  std::string mattr = cfg.value()->mattr;
+
+  bool nomve = HasFlag(mcpu, "+nomve") || HasFlag(mattr, "+nomve");
+  bool nodsp = HasFlag(mcpu, "+nodsp") || HasFlag(mattr, "+nodsp");
+
+  auto has_mve = MatchesCpu(mcpu, mveCPUs);
+  if (has_mve && !nomve && !nodsp) {
+    return kHasMVE;
+  }
+
+  auto has_dsp = MatchesCpu(mcpu, dspCPUs);
+  if (has_dsp && !nodsp) {
+    return kHasDSP;
+  }
+
+  return kNoExt;
+}
+
+}  // namespace cmsisnn
+}  // namespace contrib
+}  // namespace relay
+}  // namespace tvm

src/relay/backend/contrib/cmsisnn/compiler_attrs.h

@@ -0,0 +1,75 @@
+/*
+ * 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/relay/backend/contrib/cmsisnn/compiler_attrs.h
+ * \brief CMSIS-NN Compiler Attribute functionality
+ */
+
+#ifndef TVM_RELAY_BACKEND_CONTRIB_CMSISNN_COMPILER_ATTRS_H_
+#define TVM_RELAY_BACKEND_CONTRIB_CMSISNN_COMPILER_ATTRS_H_
+
+#include <tvm/ir/transform.h>
+
+namespace tvm {
+namespace relay {
+namespace contrib {
+namespace cmsisnn {
+
+/*! \brief Attributes to store the compiler options for CMSIS-NN. */
+struct CMSISNNCompilerConfigNode : public tvm::AttrsNode<CMSISNNCompilerConfigNode> {
+  String mcpu;
+  String mattr;
+
+  TVM_DECLARE_ATTRS(CMSISNNCompilerConfigNode, "ext.attrs.CMSISNNCompilerConfigNode") {
+    TVM_ATTR_FIELD(mcpu)
+        .describe(
+            "The CPU to configure CMSIS-NN for (i.e. cortex-m55, cortex-m4), can also include "
+            "attributes (i.e. cortex-m55+nomve)")
+        .set_default("");
+    TVM_ATTR_FIELD(mattr)
+        .describe("The attributes to configure CMSIS-NN (i.e. +nodsp, +nomve)")
+        .set_default("");
+  }
+};
+
+class CMSISNNCompilerConfig : public Attrs {
+ public:
+  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(CMSISNNCompilerConfig, Attrs,
+                                            CMSISNNCompilerConfigNode);
+};
+
+/*! \brief Flags to configure the calculations for CMSIS-NN. */
+struct CMSISNNFlags {
+  bool dsp;  // Enable or disable dsp buffers
+  bool mve;  // Enable or disable mve buffers
+};
+
+constexpr CMSISNNFlags kNoExt = {.dsp = false, .mve = false};
+constexpr CMSISNNFlags kHasDSP = {.dsp = true, .mve = false};
+constexpr CMSISNNFlags kHasMVE = {.dsp = true, .mve = true};
+
+CMSISNNFlags GetCompilerFlags(const tvm::transform::PassContext& ctx);
+
+}  // namespace cmsisnn
+}  // namespace contrib
+}  // namespace relay
+}  // namespace tvm
+
+#endif  // TVM_RELAY_BACKEND_CONTRIB_CMSISNN_COMPILER_ATTRS_H_