Skip to content

Commit ad95bdf

Browse files
authored
Add pybindings for host buffer size functions (#213)
The purpose is to make the host-side buffer size functions available in a Python environment. Also adds pytest-based checks for the new python glue. Check README files for more details.
1 parent 437bfad commit ad95bdf

23 files changed

+2014
-6
lines changed

CMakeLists.txt

Lines changed: 86 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
#
2-
# SPDX-FileCopyrightText: Copyright 2010-2023 Arm Limited and/or its affiliates <open-source-office@arm.com>
2+
# SPDX-FileCopyrightText: Copyright 2010-2023, 2026 Arm Limited and/or its affiliates <open-source-office@arm.com>
33
#
44
# SPDX-License-Identifier: Apache-2.0
55
#
@@ -19,7 +19,7 @@
1919
cmake_minimum_required(VERSION 3.15.6)
2020
set(CMSIS_OPTIMIZATION_LEVEL "-Ofast" CACHE STRING "Compiler optimization level.")
2121

22-
project(CMSISNN)
22+
project(CMSISNN C CXX)
2323

2424
add_library(cmsis-nn STATIC)
2525

@@ -28,3 +28,87 @@ target_compile_options(cmsis-nn PRIVATE ${CMSIS_OPTIMIZATION_LEVEL})
2828
target_include_directories(cmsis-nn PUBLIC "Include")
2929

3030
add_subdirectory(Source)
31+
32+
option(CMSISNN_BUILD_PYBIND "Build pybind11 Python module for CMSIS-NN helpers" OFF)
33+
34+
if(CMSISNN_BUILD_PYBIND)
35+
36+
# Add shared library for testing bindings with ctest
37+
add_library(cmsis-nn-shared SHARED
38+
$<TARGET_PROPERTY:cmsis-nn,SOURCES>
39+
)
40+
set_target_properties(cmsis-nn-shared PROPERTIES OUTPUT_NAME cmsis-nn)
41+
target_include_directories(cmsis-nn-shared PRIVATE
42+
$<TARGET_PROPERTY:cmsis-nn,INCLUDE_DIRECTORIES>
43+
)
44+
target_compile_options(cmsis-nn-shared PRIVATE ${CMSIS_OPTIMIZATION_LEVEL})
45+
46+
find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
47+
48+
find_package(pybind11 CONFIG QUIET)
49+
50+
if(pybind11_FOUND)
51+
message(STATUS "Found pybind11 via CMake package")
52+
else()
53+
execute_process(
54+
COMMAND ${Python3_EXECUTABLE} -c "import pybind11; print(pybind11.get_include())"
55+
OUTPUT_VARIABLE PYBIND11_INCLUDE_DIR
56+
OUTPUT_STRIP_TRAILING_WHITESPACE
57+
)
58+
endif()
59+
60+
if(PYBIND11_INCLUDE_DIR)
61+
message(STATUS "pybind11 include: ${PYBIND11_INCLUDE_DIR}")
62+
elseif(NOT pybind11_FOUND)
63+
message(FATAL_ERROR "pybind11 not found. Install pybind11 or set CMSISNN_BUILD_PYBIND=OFF.")
64+
endif()
65+
66+
add_library(cmsis_nn MODULE
67+
Source/Bindings/arm_py_module.cpp
68+
Source/Bindings/arm_py_backend.cpp
69+
Source/Bindings/arm_py_avgpool_buffer_size.cpp
70+
Source/Bindings/arm_py_conv_buffer_size.cpp
71+
Source/Bindings/arm_py_depthwise_conv_buffer_size.cpp
72+
Source/Bindings/arm_py_fully_connected_buffer_size.cpp
73+
Source/Bindings/arm_py_svdf_buffer_size.cpp
74+
Source/Bindings/arm_py_transpose_conv_buffer_size.cpp
75+
)
76+
set_target_properties(cmsis_nn PROPERTIES PREFIX "")
77+
78+
if(Python3_INCLUDE_DIRS)
79+
target_include_directories(cmsis_nn PRIVATE
80+
${Python3_INCLUDE_DIRS}
81+
)
82+
else()
83+
execute_process(
84+
COMMAND ${Python3_EXECUTABLE} -c "import sysconfig; print(sysconfig.get_paths()['include'])"
85+
OUTPUT_VARIABLE PYTHON_INCLUDE_DIR
86+
OUTPUT_STRIP_TRAILING_WHITESPACE
87+
)
88+
target_include_directories(cmsis_nn PRIVATE
89+
${PYTHON_INCLUDE_DIR}
90+
)
91+
endif()
92+
93+
target_link_libraries(cmsis_nn PRIVATE cmsis-nn)
94+
95+
target_include_directories(cmsis_nn PRIVATE
96+
${CMAKE_CURRENT_SOURCE_DIR}/Include
97+
)
98+
if(pybind11_FOUND)
99+
target_link_libraries(cmsis_nn PRIVATE pybind11::headers)
100+
else()
101+
target_include_directories(cmsis_nn PRIVATE
102+
${PYBIND11_INCLUDE_DIR}
103+
)
104+
endif()
105+
106+
target_compile_options(cmsis_nn PRIVATE ${CMSIS_OPTIMIZATION_LEVEL})
107+
108+
# Put the built .so into the python package directory inside the wheel
109+
install(TARGETS cmsis_nn
110+
LIBRARY DESTINATION cmsis_nn
111+
RUNTIME DESTINATION cmsis_nn
112+
)
113+
114+
endif()

README.md

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,47 @@ cmake .. -DCMAKE_TOOLCHAIN_FILE=</path/to/ethos-u-core-platform>/cmake/toolchain
9595
cmake .. -DCMAKE_TOOLCHAIN_FILE=</path/to/ethos-u-core-platform>/cmake/toolchain/armclang.cmake -DTARGET_CPU=cortex-m3
9696
```
9797

98+
## Python bindings (optional)
99+
The Python helpers are built as a `cmsis_nn` extension module using pybind11.
100+
101+
The purpose is to expose the CMSIS-NN host buffer size getter functions so they can be accessed and used from Python.
102+
103+
Build the extension with CMake:
104+
```
105+
cmake -S . -B build -DCMSISNN_BUILD_PYBIND=ON
106+
cmake --build build
107+
```
108+
109+
This produces a `cmsis_nn` shared library in the build tree. For a pip-installable wheel, use:
110+
```
111+
pip wheel . -w dist
112+
pip install dist/cmsis_nn-*.whl
113+
```
114+
115+
Example usage:
116+
```
117+
import cmsis_nn
118+
119+
backend = cmsis_nn.Backend.MVE
120+
121+
buf_size = cmsis_nn.convolve_wrapper_buffer_size(
122+
backend,
123+
cmsis_nn.DataType.A8W8,
124+
input_nhwc=[1, 8, 8, 16],
125+
filter_nhwc=[8, 3, 3, 16],
126+
output_nhwc=[1, 6, 6, 8],
127+
padding_hw=[0, 0],
128+
stride_hw=[1, 1],
129+
dilation_hw=[1, 1],
130+
)
131+
```
132+
133+
Optionally backend can be derived, e.g.:
134+
```
135+
backend = cmsis_nn.resolve_backend(cmsis_nn.CortexM.M55)
136+
```
137+
138+
98139
### Compiler Options
99140
Default optimization level is set at Ofast. This can be overwritten with CMake on command line by using <nobr>*"-DCMSIS_OPTIMIZATION_LEVEL"*</nobr>. Please change according to project needs.
100141
Just bear in mind this can impact performance. With only optimization level -O0, *ARM_MATH_AUTOVECTORIZE* needs to be defined for processors with Helium
Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
/*
2+
* SPDX-FileCopyrightText: Copyright 2026 Arm Limited and/or its affiliates <open-source-office@arm.com>
3+
*
4+
* SPDX-License-Identifier: Apache-2.0
5+
*
6+
* Licensed under the Apache License, Version 2.0 (the License); you may
7+
* not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an AS IS BASIS, WITHOUT
14+
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
/* ----------------------------------------------------------------------
20+
* Project: CMSIS NN Library
21+
* Title: arm_py_avgpool_buffer_size.cpp
22+
* Description: Average pooling buffer size pybinds (optional Python module)
23+
*
24+
* $Date: 23 Feb 2026
25+
* $Revision: V.1.0.0
26+
*
27+
* Target : Host/Python
28+
* -------------------------------------------------------------------- */
29+
30+
#include <sstream>
31+
32+
#include "arm_py_common.hpp"
33+
34+
extern "C" {
35+
#include "arm_nnfunctions.h"
36+
}
37+
38+
namespace py = pybind11;
39+
40+
void avgpool_buffer_size(py::module_ &m)
41+
{
42+
m.def(
43+
"avgpool_buffer_size",
44+
[](Backend backend, DataType data_type, int32_t dim_dst_width, int32_t ch_src) -> int32_t {
45+
switch (data_type)
46+
{
47+
case DataType::A8W8:
48+
switch (backend)
49+
{
50+
case Backend::MVE:
51+
return arm_avgpool_s8_get_buffer_size_mve(dim_dst_width, ch_src);
52+
case Backend::DSP:
53+
return arm_avgpool_s8_get_buffer_size_dsp(dim_dst_width, ch_src);
54+
case Backend::SCALAR:
55+
return arm_avgpool_s8_get_buffer_size(dim_dst_width, ch_src);
56+
}
57+
break;
58+
case DataType::A16W8:
59+
switch (backend)
60+
{
61+
case Backend::MVE:
62+
return arm_avgpool_s16_get_buffer_size_mve(dim_dst_width, ch_src);
63+
case Backend::DSP:
64+
return arm_avgpool_s16_get_buffer_size_dsp(dim_dst_width, ch_src);
65+
case Backend::SCALAR:
66+
return arm_avgpool_s16_get_buffer_size(dim_dst_width, ch_src);
67+
}
68+
break;
69+
}
70+
std::ostringstream msg;
71+
msg << "invalid Backend/DataType combination: backend=" << static_cast<int>(backend)
72+
<< " data_type=" << static_cast<int>(data_type);
73+
throw py::value_error(msg.str());
74+
},
75+
py::arg("backend"),
76+
py::arg("data_type"),
77+
py::arg("dim_dst_width"),
78+
py::arg("ch_src"));
79+
}

Source/Bindings/arm_py_backend.cpp

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
/*
2+
* SPDX-FileCopyrightText: Copyright 2026 Arm Limited and/or its affiliates <open-source-office@arm.com>
3+
*
4+
* SPDX-License-Identifier: Apache-2.0
5+
*
6+
* Licensed under the Apache License, Version 2.0 (the License); you may
7+
* not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an AS IS BASIS, WITHOUT
14+
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
/* ----------------------------------------------------------------------
20+
* Project: CMSIS NN Library
21+
* Title: arm_py_backend.cpp
22+
* Description: Backend resolution helpers for optional Python module
23+
*
24+
* $Date: 23 Feb 2026
25+
* $Revision: V.1.0.0
26+
*
27+
* Target : Host/Python
28+
* -------------------------------------------------------------------- */
29+
30+
#include "arm_py_common.hpp"
31+
32+
namespace py = pybind11;
33+
34+
void backend_helpers(py::module_ &m)
35+
{
36+
py::enum_<CortexM>(m, "CortexM")
37+
.value("M0", CortexM::M0)
38+
.value("M0PLUS", CortexM::M0PLUS)
39+
.value("M3", CortexM::M3)
40+
.value("M4", CortexM::M4)
41+
.value("M7", CortexM::M7)
42+
.value("M23", CortexM::M23)
43+
.value("M33", CortexM::M33)
44+
.value("M35P", CortexM::M35P)
45+
.value("M55", CortexM::M55)
46+
.value("M85", CortexM::M85);
47+
48+
m.def(
49+
"resolve_backend",
50+
[](CortexM core) -> Backend {
51+
switch (core)
52+
{
53+
case CortexM::M55:
54+
case CortexM::M85:
55+
return Backend::MVE;
56+
case CortexM::M4:
57+
case CortexM::M7:
58+
case CortexM::M33:
59+
case CortexM::M35P:
60+
return Backend::DSP;
61+
case CortexM::M0:
62+
case CortexM::M0PLUS:
63+
case CortexM::M3:
64+
case CortexM::M23:
65+
return Backend::SCALAR;
66+
}
67+
throw py::value_error("unsupported Cortex-M core");
68+
},
69+
py::arg("core"));
70+
}

Source/Bindings/arm_py_common.hpp

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
/*
2+
* SPDX-FileCopyrightText: Copyright 2026 Arm Limited and/or its affiliates <open-source-office@arm.com>
3+
*
4+
* SPDX-License-Identifier: Apache-2.0
5+
*
6+
* Licensed under the Apache License, Version 2.0 (the License); you may
7+
* not use this file except in compliance with the License.
8+
* You may obtain a copy of the License at
9+
*
10+
* www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an AS IS BASIS, WITHOUT
14+
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
19+
/* ----------------------------------------------------------------------
20+
* Project: CMSIS NN Library
21+
* Title: arm_py_common.hpp
22+
* Description: Common helpers and types for optional Python module
23+
*
24+
* $Date: 23 Feb 2026
25+
* $Revision: V.1.0.0
26+
*
27+
* Target : Host/Python
28+
* -------------------------------------------------------------------- */
29+
30+
#ifndef ARM_PY_BINDINGS_COMMON_HPP
31+
#define ARM_PY_BINDINGS_COMMON_HPP
32+
33+
#include <array>
34+
35+
#include <pybind11/pybind11.h>
36+
#include <pybind11/stl.h>
37+
38+
extern "C" {
39+
#include "arm_nn_types.h"
40+
}
41+
42+
namespace py = pybind11;
43+
44+
enum class Backend
45+
{
46+
MVE,
47+
DSP,
48+
SCALAR,
49+
};
50+
51+
enum class DataType
52+
{
53+
A8W4,
54+
A8W8,
55+
A16W8,
56+
};
57+
58+
enum class CortexM
59+
{
60+
M0,
61+
M0PLUS,
62+
M3,
63+
M4,
64+
M7,
65+
M23,
66+
M33,
67+
M35P,
68+
M55,
69+
M85,
70+
};
71+
72+
static inline cmsis_nn_dims make_dims(const std::array<int32_t, 4> &nhwc)
73+
{
74+
cmsis_nn_dims d;
75+
d.n = nhwc[0];
76+
d.h = nhwc[1];
77+
d.w = nhwc[2];
78+
d.c = nhwc[3];
79+
return d;
80+
}
81+
82+
#endif // ARM_PY_BINDINGS_COMMON_HPP

0 commit comments

Comments
 (0)