Calculate Conv2d 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.

python/tvm/driver/tvmc/composite_target.py

@@ -52,7 +52,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,58 @@
+/*
+ * 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;
+}
+
+}  // namespace cmsisnn
+}  // namespace contrib
+}  // namespace relay
+}  // namespace tvm

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

@@ -0,0 +1,66 @@
+/*
+ * 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);
+
+}  // 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,74 @@
+/*
+ * 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>
+
+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, 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 {false, false};
+  }
+
+  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 {true, true};
+  }
+
+  auto has_dsp = MatchesCpu(mcpu, dspCPUs);
+  if (has_dsp && !nodsp) {
+    return {true, false};
+  }
+
+  return {false, false};
+}
+
+}  // namespace cmsisnn
+}  // namespace contrib
+}  // namespace relay
+}  // namespace tvm

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

@@ -0,0 +1,71 @@
+/*
+ * 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
+};
+
+CMSISNNFlags GetCompilerFlags(const tvm::transform::PassContext& ctx);
+
+}  // namespace cmsisnn
+}  // namespace contrib
+}  // namespace relay
+}  // namespace tvm
+
+#endif  // TVM_RELAY_BACKEND_CONTRIB_CMSISNN_COMPILER_ATTRS_H_

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

@@ -28,6 +28,8 @@
 
 #include "../../../qnn/utils.h"
 #include "../../../transforms/pattern_utils.h"
+#include "buffer_size.h"
+#include "compiler_attrs.h"
 
 namespace tvm {
 namespace relay {
@@ -53,7 +55,6 @@ class RelayToTIRVisitor : public MixedModeVisitor {
         tvm::tir::Call(DataType::Int(8), tir::builtin::call_extern(), call_extern_args));
 
     if (context_buffer_size) {
-      // TODO(@ashutosh-arm) while supporting MVE, we need to move allocation through TVMBAW
       tir::Var buffer_var("context_buffer", PointerType(PrimType(DataType::Int(8)), "global"));
       body = tir::Allocate(buffer_var, DataType::Int(8), {context_buffer_size}, tir::const_true(),
                            body);
@@ -179,8 +180,10 @@ class RelayToTIRVisitor : public MixedModeVisitor {
     func_signature.push_back(const_var5);
     func_signature.push_back(out_var);
 
-    // https://github.com/ARM-software/CMSIS_5/blob/d788fd583984388553391de18afd8b4d2a146868/CMSIS/NN/Source/ConvolutionFunctions/arm_convolve_s8.c#L367
-    size_t context_buffer_size = (2 * input_c * filter_w * filter_h) * (int32_t)sizeof(int16_t);
+    CMSISNNFlags flags = GetCompilerFlags(transform::PassContext::Current());
+    int context_buffer_size =
+        Conv2dBufferSize(flags, padding_w, padding_h, input_n, input_h, input_c, output_h, output_w,
+                         stride_w, stride_h, filter_w, filter_h);
 
     CreatePrimFuncForExtern(func_signature, call_ext_args, context_buffer_size);
   }

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

@@ -208,6 +208,8 @@ class CodeGenCMSISNN : public CodeGenC {
     stream << "printf(\"Failed during execution of " << cmsis_func_name << "().\");\n";
     PrintIndent();
     stream << "}\n";
+
+    ResetBufferContext();
   }
 
   /*!  * \brief Creates a cplusplus guard prefix for extern "C" printing */
@@ -226,8 +228,13 @@ class CodeGenCMSISNN : public CodeGenC {
     ss << "#endif\n";
   }
 
+  void ResetBufferContext() {
+    context_buffer_name_ = "NULL";
+    context_buffer_size_ = 0;
+  }
+
  private:
-  std::string context_buffer_name_ = "Empty";
+  std::string context_buffer_name_ = "NULL";
   int context_buffer_size_ = 0;
 };