diff --git a/cmake/modules/contrib/CMSISNN.cmake b/cmake/modules/contrib/CMSISNN.cmake index 50a3464a2fd4..73ecd5916df3 100644 --- a/cmake/modules/contrib/CMSISNN.cmake +++ b/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}) diff --git a/python/tvm/driver/tvmc/composite_target.py b/python/tvm/driver/tvmc/composite_target.py index 848af1e4ee4e..eee1ba98ccb2 100644 --- a/python/tvm/driver/tvmc/composite_target.py +++ b/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": { diff --git a/src/relay/backend/contrib/cmsisnn/buffer_size.cc b/src/relay/backend/contrib/cmsisnn/buffer_size.cc new file mode 100644 index 000000000000..2502a09e75d6 --- /dev/null +++ b/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 +#include + +#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 diff --git a/src/relay/backend/contrib/cmsisnn/buffer_size.h b/src/relay/backend/contrib/cmsisnn/buffer_size.h new file mode 100644 index 000000000000..dec3c3eafc48 --- /dev/null +++ b/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 + +#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_ diff --git a/src/relay/backend/contrib/cmsisnn/compiler_attrs.cc b/src/relay/backend/contrib/cmsisnn/compiler_attrs.cc new file mode 100644 index 000000000000..c6fb8f32e7a9 --- /dev/null +++ b/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 +#include + +#include + +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 +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("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 diff --git a/src/relay/backend/contrib/cmsisnn/compiler_attrs.h b/src/relay/backend/contrib/cmsisnn/compiler_attrs.h new file mode 100644 index 000000000000..75005d927a1e --- /dev/null +++ b/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 + +namespace tvm { +namespace relay { +namespace contrib { +namespace cmsisnn { + +/*! \brief Attributes to store the compiler options for CMSIS-NN. */ +struct CMSISNNCompilerConfigNode : public tvm::AttrsNode { + 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_ diff --git a/src/relay/backend/contrib/cmsisnn/relay_to_tir.cc b/src/relay/backend/contrib/cmsisnn/relay_to_tir.cc index 7f1582a5dbab..b8744247e9a6 100644 --- a/src/relay/backend/contrib/cmsisnn/relay_to_tir.cc +++ b/src/relay/backend/contrib/cmsisnn/relay_to_tir.cc @@ -29,6 +29,8 @@ #include "../../../qnn/utils.h" #include "../../../transforms/pattern_utils.h" +#include "buffer_size.h" +#include "compiler_attrs.h" namespace tvm { namespace relay { @@ -159,6 +161,9 @@ class RelayToTIRVisitor : public MixedModeMutator { // CMSIS-NN data structure "cmsis_nn_dims" for ifm expects input layout as NHWC // This is the same layout we expect in Relay Array input_shape = conv2d_call->args[0]->type_as()->shape; + int32_t input_n = qnn::get_const_int(input_shape[0]); + int32_t input_h = qnn::get_const_int(input_shape[1]); + int32_t input_c = qnn::get_const_int(input_shape[3]); // CMSIS-NN data structure "cmsis_nn_dims" for weights expects following layouts // OHWI for Conv2D and IHWO for Depthwise convolutions @@ -167,6 +172,9 @@ class RelayToTIRVisitor : public MixedModeMutator { Array bias_shape{1, 1, 1, out_channels}; Array output_shape = conv2d_call->type_as()->shape; + int32_t output_h = qnn::get_const_int(output_shape[1]); + int32_t output_w = qnn::get_const_int(output_shape[2]); + int32_t output_c = qnn::get_const_int(output_shape[3]); int32_t depth_multiplier = -1; int kernel_pos_o = kernel_layout.find("O"); @@ -179,7 +187,8 @@ class RelayToTIRVisitor : public MixedModeMutator { // original filter_layout for depthwise is HWOI std::string cmsisnn_api = "arm_convolve_wrapper_s8"; - if (depth_multiplier != -1) { + bool is_depthwise = depth_multiplier != -1; + if (is_depthwise) { cmsisnn_api = "arm_depthwise_conv_wrapper_s8"; int filter_pos_h = kernel_layout.find("H"); int filter_pos_w = kernel_layout.find("W"); @@ -187,6 +196,8 @@ class RelayToTIRVisitor : public MixedModeMutator { filter_shape[filter_pos_w], out_channels}; filter_shape = depthwise_filter_shape; } + int32_t filter_h = qnn::get_const_int(filter_shape[1]); + int32_t filter_w = qnn::get_const_int(filter_shape[2]); tvm::Array call_ext_args = {tir::StringImm(cmsisnn_api), input, filter, multiplier}; if (bias_add_call) { @@ -195,11 +206,18 @@ class RelayToTIRVisitor : public MixedModeMutator { call_ext_args.push_back(shift); call_ext_args.push_back(output); - // https://github.com/ARM-software/CMSIS_5/blob/d788fd583984388553391de18afd8b4d2a146868/CMSIS/NN/Source/ConvolutionFunctions/arm_convolve_s8.c#L367 std::string context_buffer_name = "NULL"; - size_t context_buffer_size = - (2 * qnn::get_const_int(input_shape[3]) * qnn::get_const_int(filter_shape[2]) * - qnn::get_const_int(filter_shape[1]) * (int32_t)sizeof(int16_t)); + CMSISNNFlags flags = GetCompilerFlags(transform::PassContext::Current()); + size_t context_buffer_size; + if (is_depthwise) { + context_buffer_size = + DepthwiseConv2dBufferSize(flags, input_n, input_c, output_c, filter_w, filter_h); + } else { + 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); + } + if (context_buffer_size) { context_buffer_name = "context_buffer_" + std::to_string(context_buffer_id_++); } @@ -397,8 +415,9 @@ class RelayToTIRVisitor : public MixedModeMutator { int context_buffer_size = 0; std::string context_buffer_name = "NULL"; if (pool_name == "cmsisnn.qnn_avg_pool2d") { - // TODO(@Mousius): Need to move this into buffer_size calculations - context_buffer_size = qnn::get_const_int(input_shape[3]) * sizeof(int32_t); + CMSISNNFlags flags = GetCompilerFlags(transform::PassContext::Current()); + int32_t input_c = qnn::get_const_int(input_shape[3]); + context_buffer_size = AvgPoolBufferSize(flags, input_c); context_buffer_name = "context_buffer_" + std::to_string(context_buffer_id_++); } tvm::Array context_buffer_args = {tir::StringImm(context_buffer_name), diff --git a/tests/cpp/relay/backend/contrib/cmsisnn/buffer_size_test.cc b/tests/cpp/relay/backend/contrib/cmsisnn/buffer_size_test.cc new file mode 100644 index 000000000000..c83daaf10e9a --- /dev/null +++ b/tests/cpp/relay/backend/contrib/cmsisnn/buffer_size_test.cc @@ -0,0 +1,206 @@ +/* + * 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. + */ + +#ifdef TVM_USE_CMSISNN + +#include "../../../../../../src/relay/backend/contrib/cmsisnn/buffer_size.h" + +#include +#include + +#include +#include +#include + +#include "../../../../../../src/relay/backend/contrib/cmsisnn/compiler_attrs.h" + +namespace tvm { +namespace relay { +namespace contrib { +namespace cmsisnn { + +static std::random_device rd; +static std::mt19937 gen(rd()); +static std::uniform_int_distribution<> fake_parameters(1, 100); + +class CMSISNNCalculatedBufferSize : public testing::TestWithParam> {}; + +TEST(CMSISNNConv2dBufferSize, Conv1x1) { + int32_t any = fake_parameters(gen); + auto conv2d_1x1 = [=](CMSISNNFlags flags, int32_t input_c) { + return Conv2dBufferSize(flags, 0, 0, any, any, input_c, any, any, 1, 1, 1, 1); + }; + + ASSERT_EQ(conv2d_1x1(kNoExt, 4), 0); + ASSERT_EQ(conv2d_1x1(kNoExt, 8), 0); + ASSERT_EQ(conv2d_1x1(kNoExt, 12), 0); + ASSERT_EQ(conv2d_1x1(kNoExt, 16), 0); + ASSERT_EQ(conv2d_1x1(kNoExt, 32), 0); + + ASSERT_EQ(conv2d_1x1(kHasDSP, 4), 0); + ASSERT_EQ(conv2d_1x1(kHasDSP, 8), 0); + ASSERT_EQ(conv2d_1x1(kHasDSP, 12), 0); + ASSERT_EQ(conv2d_1x1(kHasDSP, 16), 0); + ASSERT_EQ(conv2d_1x1(kHasDSP, 32), 0); + + ASSERT_EQ(conv2d_1x1(kHasMVE, 4), 0); + ASSERT_EQ(conv2d_1x1(kHasMVE, 8), 0); + ASSERT_EQ(conv2d_1x1(kHasMVE, 12), 0); + ASSERT_EQ(conv2d_1x1(kHasMVE, 16), 0); + ASSERT_EQ(conv2d_1x1(kHasMVE, 32), 0); +} + +TEST(CMSISNNConv2dBufferSize, Conv1xN) { + int32_t any = fake_parameters(gen); + int32_t input_c = fake_parameters(gen); + int32_t filter_w = fake_parameters(gen); + int32_t filter_h = 1; + int32_t calculated_buffer = (2 * input_c * filter_w * filter_h) * (int32_t)sizeof(int16_t); + + auto conv2d_1xn = [=](CMSISNNFlags flags, int32_t output_w) { + return Conv2dBufferSize(flags, any, any, 1, 1, input_c, 1, output_w, any, any, filter_w, + filter_h); + }; + + ASSERT_EQ(conv2d_1xn(kNoExt, 4), 0); + ASSERT_EQ(conv2d_1xn(kNoExt, 8), 0); + ASSERT_EQ(conv2d_1xn(kNoExt, 12), 0); + ASSERT_EQ(conv2d_1xn(kNoExt, 16), 0); + ASSERT_EQ(conv2d_1xn(kNoExt, 32), 0); + + ASSERT_EQ(conv2d_1xn(kHasDSP, 4), calculated_buffer); + ASSERT_EQ(conv2d_1xn(kHasDSP, 8), calculated_buffer); + ASSERT_EQ(conv2d_1xn(kHasDSP, 12), calculated_buffer); + ASSERT_EQ(conv2d_1xn(kHasDSP, 16), calculated_buffer); + ASSERT_EQ(conv2d_1xn(kHasDSP, 32), calculated_buffer); + + ASSERT_EQ(conv2d_1xn(kHasMVE, 4), 0); + ASSERT_EQ(conv2d_1xn(kHasMVE, 8), 0); + ASSERT_EQ(conv2d_1xn(kHasMVE, 12), 0); + ASSERT_EQ(conv2d_1xn(kHasMVE, 16), 0); + ASSERT_EQ(conv2d_1xn(kHasMVE, 32), 0); +} + +TEST(CMSISNNConv2dBufferSize, Default) { + int32_t any = fake_parameters(gen); + + int32_t input_c = fake_parameters(gen); + int32_t filter_w = fake_parameters(gen); + int32_t filter_h = fake_parameters(gen); + int32_t calculated_buffer = (2 * input_c * filter_w * filter_h) * (int32_t)sizeof(int16_t); + + auto conv2d = [=](CMSISNNFlags flags, int32_t output_w) { + return Conv2dBufferSize(flags, any, any, 1, 1, input_c, 1, output_w, any, any, filter_w, + filter_h); + }; + + ASSERT_EQ(conv2d(kNoExt, 4), 0); + ASSERT_EQ(conv2d(kNoExt, 8), 0); + ASSERT_EQ(conv2d(kNoExt, 12), 0); + ASSERT_EQ(conv2d(kNoExt, 16), 0); + ASSERT_EQ(conv2d(kNoExt, 32), 0); + + ASSERT_EQ(conv2d(kHasDSP, 4), calculated_buffer); + ASSERT_EQ(conv2d(kHasDSP, 8), calculated_buffer); + ASSERT_EQ(conv2d(kHasDSP, 12), calculated_buffer); + ASSERT_EQ(conv2d(kHasDSP, 16), calculated_buffer); + ASSERT_EQ(conv2d(kHasDSP, 32), calculated_buffer); + + ASSERT_EQ(conv2d(kHasMVE, 4), calculated_buffer); + ASSERT_EQ(conv2d(kHasMVE, 8), calculated_buffer); + ASSERT_EQ(conv2d(kHasMVE, 12), calculated_buffer); + ASSERT_EQ(conv2d(kHasMVE, 16), calculated_buffer); + ASSERT_EQ(conv2d(kHasMVE, 32), calculated_buffer); +} + +TEST(CMSISNNDepthwiseConv2dBufferSize, UnEvenChannels) { + int32_t filter_w = fake_parameters(gen); + int32_t filter_h = fake_parameters(gen); + int32_t input_n = 1; + + auto depthwise_conv2d_with_channels = [=](CMSISNNFlags flags, int32_t input_c, int32_t output_c) { + return DepthwiseConv2dBufferSize(flags, input_n, input_c, output_c, filter_w, filter_h); + }; + + ASSERT_EQ(depthwise_conv2d_with_channels(kNoExt, 4, 6), 0); + ASSERT_EQ(depthwise_conv2d_with_channels(kNoExt, 8, 7), 0); + ASSERT_EQ(depthwise_conv2d_with_channels(kHasDSP, 4, 6), 0); + ASSERT_EQ(depthwise_conv2d_with_channels(kHasDSP, 8, 7), 0); + ASSERT_EQ(depthwise_conv2d_with_channels(kHasMVE, 4, 6), 0); + ASSERT_EQ(depthwise_conv2d_with_channels(kHasMVE, 8, 7), 0); +} + +TEST(CMSISNNDepthwiseConv2dBufferSize, MultipleBatches) { + int32_t input_output_c = fake_parameters(gen); + int32_t filter_w = fake_parameters(gen); + int32_t filter_h = fake_parameters(gen); + + auto depthwise_conv2d_with_batch = [=](CMSISNNFlags flags, int32_t input_n) { + return DepthwiseConv2dBufferSize(flags, input_n, input_output_c, input_output_c, filter_w, + filter_h); + }; + + ASSERT_EQ(depthwise_conv2d_with_batch(kNoExt, 4), 0); + ASSERT_EQ(depthwise_conv2d_with_batch(kNoExt, 7), 0); + ASSERT_EQ(depthwise_conv2d_with_batch(kHasDSP, 4), 0); + ASSERT_EQ(depthwise_conv2d_with_batch(kHasDSP, 7), 0); + ASSERT_EQ(depthwise_conv2d_with_batch(kHasMVE, 4), 0); + ASSERT_EQ(depthwise_conv2d_with_batch(kHasMVE, 7), 0); +} + +TEST(CMSISNNDepthwiseConv2dBufferSize, Default) { + int32_t input_output_c = fake_parameters(gen); + int32_t filter_w = fake_parameters(gen); + int32_t filter_h = fake_parameters(gen); + int32_t input_n = 1; + + int32_t mve_calculated_buffer = + (2 * input_output_c * filter_w * filter_h) * (int32_t)sizeof(int16_t) + 4; + int32_t dsp_calculated_buffer = (input_output_c * filter_w * filter_h) * (int32_t)sizeof(int16_t); + + auto depthwise_conv2d = [=](CMSISNNFlags flags) { + return DepthwiseConv2dBufferSize(flags, input_n, input_output_c, input_output_c, filter_w, + filter_h); + }; + + ASSERT_EQ(depthwise_conv2d(kNoExt), 0); + ASSERT_EQ(depthwise_conv2d(kNoExt), 0); + ASSERT_EQ(depthwise_conv2d(kHasDSP), dsp_calculated_buffer); + ASSERT_EQ(depthwise_conv2d(kHasDSP), dsp_calculated_buffer); + ASSERT_EQ(depthwise_conv2d(kHasMVE), mve_calculated_buffer); + ASSERT_EQ(depthwise_conv2d(kHasMVE), mve_calculated_buffer); +} + +TEST(CMSISNNAvgPoolBufferSize, Default) { + int32_t input_c = fake_parameters(gen); + int32_t calculated_buffer = (input_c * sizeof(int32_t)); + + auto avg_pool = [=](CMSISNNFlags flags) { return AvgPoolBufferSize(flags, input_c); }; + + ASSERT_EQ(avg_pool(kNoExt), 0); + ASSERT_EQ(avg_pool(kHasDSP), calculated_buffer); + ASSERT_EQ(avg_pool(kHasMVE), 0); +} + +} // namespace cmsisnn +} // namespace contrib +} // namespace relay +} // namespace tvm + +#endif diff --git a/tests/cpp/relay/backend/contrib/cmsisnn/compiler_attrs_test.cc b/tests/cpp/relay/backend/contrib/cmsisnn/compiler_attrs_test.cc new file mode 100644 index 000000000000..0eb980f14af0 --- /dev/null +++ b/tests/cpp/relay/backend/contrib/cmsisnn/compiler_attrs_test.cc @@ -0,0 +1,157 @@ +/* + * 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. + */ + +#ifdef TVM_USE_CMSISNN + +#include "../../../../../../src/relay/backend/contrib/cmsisnn/compiler_attrs.h" + +#include +#include + +#include +#include + +namespace tvm { +namespace relay { +namespace contrib { +namespace cmsisnn { + +static const char* mveCPUs[] = {"cortex-m55"}; +static const char* dspCPUs[] = {"cortex-m7", "cortex-m33", "cortex-m35p"}; +static const char* noExtensions[] = {"cortex-m0", "cortex-m3", "cortex-m4"}; + +class CMSISNNFlagsMVECPUs : public testing::TestWithParam {}; +class CMSISNNFlagsDSPCPUs : public testing::TestWithParam {}; +class CMSISNNFlagsNoExtensions : public testing::TestWithParam {}; + +static CMSISNNFlags GetFlagsWithCompilerAttrs(String mcpu, String mattr) { + auto context_node = make_object(); + auto cmsisnn_config_node = make_object(); + cmsisnn_config_node->InitBySeq("mcpu", mcpu, "mattr", mattr); + + context_node->config = { + {"relay.ext.cmsisnn.options", CMSISNNCompilerConfig(cmsisnn_config_node)}}; + + tvm::transform::PassContext context = tvm::transform::PassContext(context_node); + return GetCompilerFlags(context); +} + +TEST(CMSISNNFlags, CreateFromUndefined) { + auto context_node = make_object(); + tvm::transform::PassContext context = tvm::transform::PassContext(context_node); + CMSISNNFlags flags = GetCompilerFlags(context); + ASSERT_EQ(flags.mve, false); + ASSERT_EQ(flags.dsp, false); +} + +TEST_P(CMSISNNFlagsMVECPUs, CheckMVESet) { + CMSISNNFlags flags = GetFlagsWithCompilerAttrs(GetParam(), ""); + ASSERT_EQ(flags.dsp, true); + ASSERT_EQ(flags.mve, true); +} + +TEST_P(CMSISNNFlagsMVECPUs, CheckMVEOverrideCPU) { + std::string mcpu = GetParam(); + CMSISNNFlags flags = GetFlagsWithCompilerAttrs(mcpu + "+nomve", ""); + ASSERT_EQ(flags.dsp, false); + ASSERT_EQ(flags.mve, false); +} + +TEST_P(CMSISNNFlagsMVECPUs, CheckDSPOverrideCPU) { + std::string mcpu = GetParam(); + CMSISNNFlags flags = GetFlagsWithCompilerAttrs(mcpu + "+nodsp", ""); + ASSERT_EQ(flags.dsp, false); + ASSERT_EQ(flags.mve, false); +} + +TEST_P(CMSISNNFlagsMVECPUs, CheckCombinedOverrideCPU) { + std::string mcpu = GetParam(); + CMSISNNFlags flags = GetFlagsWithCompilerAttrs(mcpu + "+nodsp+nomve", ""); + ASSERT_EQ(flags.dsp, false); + ASSERT_EQ(flags.mve, false); + flags = GetFlagsWithCompilerAttrs(mcpu + "+nomve+nodsp", ""); + ASSERT_EQ(flags.dsp, false); + ASSERT_EQ(flags.mve, false); +} + +TEST_P(CMSISNNFlagsMVECPUs, CheckMVEOverrideMAttr) { + CMSISNNFlags flags = GetFlagsWithCompilerAttrs(GetParam(), "+nomve"); + ASSERT_EQ(flags.dsp, false); + ASSERT_EQ(flags.mve, false); +} + +TEST_P(CMSISNNFlagsMVECPUs, CheckDSPOverrideMattr) { + CMSISNNFlags flags = GetFlagsWithCompilerAttrs(GetParam(), "+nodsp"); + ASSERT_EQ(flags.dsp, false); + ASSERT_EQ(flags.mve, false); +} + +TEST_P(CMSISNNFlagsMVECPUs, CheckCombinedOverrideMattr) { + CMSISNNFlags flags = GetFlagsWithCompilerAttrs(GetParam(), "+nodsp+nomve"); + ASSERT_EQ(flags.dsp, false); + ASSERT_EQ(flags.mve, false); + flags = GetFlagsWithCompilerAttrs(GetParam(), "+nomve+nodsp"); + ASSERT_EQ(flags.dsp, false); + ASSERT_EQ(flags.mve, false); + flags = GetFlagsWithCompilerAttrs(GetParam(), "+woofles+nomve+nodsp"); + ASSERT_EQ(flags.dsp, false); + ASSERT_EQ(flags.mve, false); +} + +TEST_P(CMSISNNFlagsDSPCPUs, CheckDSPSet) { + CMSISNNFlags flags = GetFlagsWithCompilerAttrs(GetParam(), ""); + ASSERT_EQ(flags.dsp, true); + ASSERT_EQ(flags.mve, false); +} + +TEST_P(CMSISNNFlagsDSPCPUs, CheckDSPOverrideCPU) { + std::string mcpu = GetParam(); + CMSISNNFlags flags = GetFlagsWithCompilerAttrs(mcpu + "+nodsp", ""); + ASSERT_EQ(flags.dsp, false); + ASSERT_EQ(flags.mve, false); + flags = GetFlagsWithCompilerAttrs(mcpu + "+nodsp+woofles", ""); + ASSERT_EQ(flags.dsp, false); + ASSERT_EQ(flags.mve, false); +} + +TEST_P(CMSISNNFlagsDSPCPUs, CheckDSPOverrideMattr) { + CMSISNNFlags flags = GetFlagsWithCompilerAttrs(GetParam(), "+nodsp"); + ASSERT_EQ(flags.dsp, false); + ASSERT_EQ(flags.mve, false); + flags = GetFlagsWithCompilerAttrs(GetParam(), "+nodsp+woofles"); + ASSERT_EQ(flags.dsp, false); + ASSERT_EQ(flags.mve, false); +} + +TEST_P(CMSISNNFlagsNoExtensions, CheckNoFlags) { + CMSISNNFlags flags = GetFlagsWithCompilerAttrs(GetParam(), ""); + ASSERT_EQ(flags.dsp, false); + ASSERT_EQ(flags.mve, false); +} + +INSTANTIATE_TEST_CASE_P(CMSISNNFlags, CMSISNNFlagsMVECPUs, ::testing::ValuesIn(mveCPUs)); +INSTANTIATE_TEST_CASE_P(CMSISNNFlags, CMSISNNFlagsDSPCPUs, ::testing::ValuesIn(dspCPUs)); +INSTANTIATE_TEST_CASE_P(CMSISNNFlags, CMSISNNFlagsNoExtensions, ::testing::ValuesIn(noExtensions)); + +} // namespace cmsisnn +} // namespace contrib +} // namespace relay +} // namespace tvm + +#endif diff --git a/tests/python/relay/aot/aot_test_utils.py b/tests/python/relay/aot/aot_test_utils.py index d335528914b0..e7ca5c649667 100644 --- a/tests/python/relay/aot/aot_test_utils.py +++ b/tests/python/relay/aot/aot_test_utils.py @@ -145,6 +145,11 @@ class AOTTestRunner(NamedTuple): """, includes=["uart.h"], parameters={"NPU_VARIANT": "256"}, + pass_config={ + "relay.ext.cmsisnn.options": { + "mcpu": "cortex-m55", + } + }, )