Build and install cmsis_nn pybinds with executorch

.ci/scripts/wheel/test_base.py

@@ -1,5 +1,6 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -40,6 +41,23 @@ class ModelTest:
     backend: Backend
 
 
+def test_cmsis_nn_install():
+    import executorch.backends.cortex_m.library.cmsis_nn as cmsis_nn
+
+    buf_size = cmsis_nn.convolve_wrapper_buffer_size(
+        cmsis_nn.Backend.MVE,
+        cmsis_nn.DataType.A8W8,
+        input_nhwc=[1, 8, 8, 16],
+        filter_nhwc=[8, 3, 3, 16],
+        output_nhwc=[1, 6, 6, 8],
+        padding_hw=[0, 0],
+        stride_hw=[1, 1],
+        dilation_hw=[1, 1],
+    )
+
+    assert buf_size == 576
+
+
 def run_tests(model_tests: List[ModelTest]) -> None:
     # Test that we can import the portable_lib module - verifies RPATH is correct
     print("Testing portable_lib import...")

.ci/scripts/wheel/test_linux.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -38,6 +39,8 @@
         else:
             print("⚠ VulkanBackend not registered (expected for the default wheel)")
 
+        test_base.test_cmsis_nn_install()
+
     test_base.run_tests(
         model_tests=[
             test_base.ModelTest(

.ci/scripts/wheel/test_macos.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
@@ -9,6 +10,8 @@
 from examples.models import Backend, Model
 
 if __name__ == "__main__":
+    test_base.test_cmsis_nn_install()
+
     test_base.run_tests(
         model_tests=[
             test_base.ModelTest(

.ci/scripts/wheel/test_windows.py

@@ -1,13 +1,16 @@
 #!/usr/bin/env python
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
+# Copyright 2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
 import platform
 from typing import List
 
+import test_base
+
 import torch
 from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
 from executorch.examples.models import Backend, Model, MODEL_NAME_TO_MODEL
@@ -74,6 +77,8 @@ def run_tests(model_tests: List[ModelTest]) -> None:
         else:
             print("⚠ VulkanBackend not registered (expected for the default wheel)")
 
+        test_base.test_cmsis_nn_install()
+
     run_tests(
         model_tests=[
             ModelTest(

CMakeLists.txt

@@ -751,7 +751,7 @@ if(EXECUTORCH_BUILD_XNNPACK)
   list(APPEND _executorch_backends xnnpack_backend)
 endif()
 
-if(EXECUTORCH_BUILD_CORTEX_M)
+if(EXECUTORCH_BUILD_CORTEX_M OR EXECUTORCH_BUILD_CMSIS_NN_PYBINDS)
   add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/backends/cortex_m)
 endif()
 

backends/cortex_m/CMakeLists.txt

@@ -23,7 +23,7 @@ include(FetchContent)
 
 # CMSIS-NN configuration with dynamic path detection
 set(CMSIS_NN_VERSION
-    "d933672e7ca97eec70ef43230baee7b20c2a28ae"
+    "dbf45dbfcc515421dd6099037d3e2637b90748c8"
     CACHE STRING "CMSIS-NN version to download"
 )
 set(CMSIS_NN_LOCAL_PATH
@@ -35,12 +35,19 @@ option(CORTEX_M_ENABLE_RUNTIME_CHECKS
        OFF
 )
 
+set(CMSISNN_BUILD_PYBIND
+    ${EXECUTORCH_BUILD_CMSIS_NN_PYBINDS}
+    CACHE BOOL "Build CMSIS-NN Python bindings" FORCE
+)
+
 # Try to find existing / local CMSIS-NN installation. This is useful for
 # debugging and testing with local changes. This is not common, as the CMSIS-NN
 # library is downloaded via FetchContent in the default/regular case.
 if(CMSIS_NN_LOCAL_PATH AND EXISTS "${CMSIS_NN_LOCAL_PATH}")
   message(STATUS "Using CMSIS-NN from specified path: ${CMSIS_NN_LOCAL_PATH}")
-  add_subdirectory(${CMSIS_NN_LOCAL_PATH} _deps/cmsis_nn-build)
+  add_subdirectory(
+    ${CMSIS_NN_LOCAL_PATH} ${CMAKE_CURRENT_BINARY_DIR}/cmsis_nn-build
+  )
 else()
   # Use FetchContent with automatic fallback
   message(STATUS "Using CMSIS-NN via FetchContent")
@@ -49,6 +56,9 @@ else()
     cmsis_nn
     GIT_REPOSITORY https://github.com/ARM-software/CMSIS-NN.git
     GIT_TAG ${CMSIS_NN_VERSION}
+    SOURCE_DIR ${CMAKE_CURRENT_BINARY_DIR}/cmsis_nn-src BINARY_DIR
+    ${CMAKE_CURRENT_BINARY_DIR}/cmsis_nn-build SUBBUILD_DIR
+    ${CMAKE_CURRENT_BINARY_DIR}/cmsis_nn-subbuild
   )
 
   FetchContent_MakeAvailable(cmsis_nn)
@@ -74,65 +84,72 @@ if(TARGET cmsis-nn)
   endif()
 endif()
 
-# Cortex-M ops kernel sources
-set(_cortex_m_kernels__srcs
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_dequantize_per_tensor.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_maximum.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_minimum.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_pad.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantize_per_tensor.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_activation.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_add.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_avg_pool2d.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_batch_matmul.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_conv2d.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_depthwise_conv2d.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_linear.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_max_pool2d.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_mul.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_transpose_conv2d.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_softmax.cpp
-    ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_transpose.cpp
-)
+if(EXECUTORCH_BUILD_CMSIS_NN_PYBINDS)
+  if(NOT TARGET cmsis_nn)
+    message(
+      FATAL_ERROR
+        "EXECUTORCH_BUILD_CMSIS_NN_PYBINDS is ON, but CMSIS-NN did not define the cmsis_nn pybind target."
+    )
+  endif()
+endif()
+if(EXECUTORCH_BUILD_CORTEX_M)
+  # Cortex-M ops kernel sources
+  set(_cortex_m_kernels__srcs
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_dequantize_per_tensor.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_maximum.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_minimum.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_pad.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantize_per_tensor.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_activation.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_add.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_avg_pool2d.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_batch_matmul.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_conv2d.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_depthwise_conv2d.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_linear.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_max_pool2d.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_mul.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_quantized_transpose_conv2d.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_softmax.cpp
+      ${CMAKE_CURRENT_SOURCE_DIR}/ops/op_transpose.cpp
+  )
 
-# Generate C++ bindings to register kernels into Executorch
-set(_yaml_file ${CMAKE_CURRENT_LIST_DIR}/ops/operators.yaml)
-gen_selected_ops(LIB_NAME "cortex_m_ops_lib" OPS_SCHEMA_YAML "${_yaml_file}")
-generate_bindings_for_kernels(
-  LIB_NAME "cortex_m_ops_lib" CUSTOM_OPS_YAML "${_yaml_file}"
-)
+  # Generate C++ bindings to register kernels into Executorch
+  set(_yaml_file ${CMAKE_CURRENT_LIST_DIR}/ops/operators.yaml)
+  gen_selected_ops(LIB_NAME "cortex_m_ops_lib" OPS_SCHEMA_YAML "${_yaml_file}")
+  generate_bindings_for_kernels(
+    LIB_NAME "cortex_m_ops_lib" CUSTOM_OPS_YAML "${_yaml_file}"
+  )
 
-# Build library for cortex_m_kernels
-add_library(cortex_m_kernels ${_cortex_m_kernels__srcs})
+  # Build library for cortex_m_kernels
+  add_library(cortex_m_kernels ${_cortex_m_kernels__srcs})
 
-# Use PRIVATE for implementation dependencies to avoid INTERFACE pollution
-target_link_libraries(
-  cortex_m_kernels
-  PRIVATE cmsis-nn
-  PRIVATE executorch
-  PRIVATE kernels_util_all_deps
-)
-target_compile_definitions(
-  cortex_m_kernels
-  PRIVATE
-    $<$<BOOL:${CORTEX_M_ENABLE_RUNTIME_CHECKS}>:CORTEX_M_ENABLE_RUNTIME_CHECKS>
-)
+  target_link_libraries(
+    cortex_m_kernels
+    PRIVATE cmsis-nn
+    PRIVATE executorch
+    PRIVATE kernels_util_all_deps
+  )
+  target_compile_definitions(
+    cortex_m_kernels
+    PRIVATE
+      $<$<BOOL:${CORTEX_M_ENABLE_RUNTIME_CHECKS}>:CORTEX_M_ENABLE_RUNTIME_CHECKS>
+  )
 
-# Include directories for cortex_m_kernels
-target_include_directories(
-  cortex_m_kernels PRIVATE ${EXECUTORCH_ROOT}/..
-                           ${EXECUTORCH_ROOT}/runtime/core/portable_type/c10
-)
+  target_include_directories(
+    cortex_m_kernels PRIVATE ${EXECUTORCH_ROOT}/..
+                             ${EXECUTORCH_ROOT}/runtime/core/portable_type/c10
+  )
 
-# cortex_m_ops_lib: Register Cortex-M ops kernels into Executorch runtime
-gen_operators_lib(
-  LIB_NAME "cortex_m_ops_lib" KERNEL_LIBS cortex_m_kernels DEPS executorch
-)
+  gen_operators_lib(
+    LIB_NAME "cortex_m_ops_lib" KERNEL_LIBS cortex_m_kernels DEPS executorch
+  )
 
-install(
-  TARGETS cortex_m_kernels cortex_m_ops_lib cmsis-nn
-  EXPORT ExecuTorchTargets
-  DESTINATION ${CMAKE_INSTALL_LIBDIR}
-  PUBLIC_HEADER
-    DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/backends/cortex_m/ops/
-)
+  install(
+    TARGETS cortex_m_kernels cortex_m_ops_lib cmsis-nn
+    EXPORT ExecuTorchTargets
+    DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    PUBLIC_HEADER
+      DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/backends/cortex_m/ops/
+  )
+endif()

backends/cortex_m/library/cmsis_nn.py

@@ -6,11 +6,15 @@
 
 from __future__ import annotations
 
+import importlib
 from types import ModuleType
 from typing import Any, cast, ClassVar, Sequence, TYPE_CHECKING
 
 _cmsis_nn: ModuleType | None = None
-_cmsis_nn_import_error: ModuleNotFoundError | None = None
+_cmsis_nn_module_candidates = (
+    "executorch.backends.cortex_m.library._cmsis_nn.cmsis_nn",
+    "cmsis_nn",
+)
 
 
 class _EnumValue:
@@ -83,30 +87,28 @@ class DataType:
 
 
 if not TYPE_CHECKING:
-    try:
-        import cmsis_nn as _real_cmsis_nn  # type: ignore[import-not-found, import-untyped]
-    except ModuleNotFoundError as exc:
-        if exc.name != "cmsis_nn":
-            raise
-        _cmsis_nn_import_error = exc
-    else:
-        _cmsis_nn = _real_cmsis_nn
-        Backend = _real_cmsis_nn.Backend
-        CortexM = _real_cmsis_nn.CortexM
-        DataType = _real_cmsis_nn.DataType
-
-
-def _missing_dependencies_error() -> ModuleNotFoundError:
-    return ModuleNotFoundError(
-        "Cortex-M backend dependencies are not installed. "
-        "Install by running `examples/arm/setup.sh --i-agree-to-the-contained-eula`, "
-        "or pip install from the CMSIS-NN repo."
-    )
+    # First try to load cmsis_nn from executorch build, then external.
+    # Load in such a way that we crash only if a cmsis_nn function is required.
+    for module_name in _cmsis_nn_module_candidates:
+        try:
+            _real_cmsis_nn = importlib.import_module(module_name)
+        except ModuleNotFoundError as exc:
+            if exc.name not in module_name:
+                raise exc
+        else:
+            _cmsis_nn = _real_cmsis_nn
+            Backend = _real_cmsis_nn.Backend
+            CortexM = _real_cmsis_nn.CortexM
+            DataType = _real_cmsis_nn.DataType
+            break
 
 
 def _require_cmsis_nn() -> ModuleType:
     if _cmsis_nn is None:
-        raise _missing_dependencies_error() from _cmsis_nn_import_error
+        raise ModuleNotFoundError(
+            f"Cortex-M backend dependencies are not installed (tried {','.join(_cmsis_nn_module_candidates)}). "
+            "Build using install_executorch.sh, or pip install from the CMSIS-NN repo."
+        )
     return _cmsis_nn