Arm backend: Add support for building selective prim ops. (#21714)

Prim ops are rare, but take up 12KiB in our example runner. This is due
to the executorch library consisting of executorch_core + prim ops, and
it is linked with
--whole-archive.

Add an option to arm_runner_create_selected_ops_lib, PRIM_OPS, to
generate prim ops. They are built differently than other operator
libraries, so treat it as a special case. Outwards, the api should look
the same.

Use this in the example runner so that the executorch lib no longer
needs --whole-archive for prim ops.

Use this in the example executor_runner. For other builds, link to
executorch_core instead of executorch.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

---------

Signed-off-by: Erik Lundell <erik.lundell@arm.com>

Author: Erik-Lundell

backends/arm/cmake/ArmRunnerUtils.cmake

@@ -31,6 +31,77 @@ function(arm_runner_require_baremetal_targets)
 
 endfunction()
 
+function(verify_targets_exist)
+  cmake_parse_arguments(ARG "" "CONTEXT" "TARGETS" ${ARGN})
+
+  set(_targets ${ARG_TARGETS} ${ARG_UNPARSED_ARGUMENTS})
+  if(NOT ARG_CONTEXT)
+    set(ARG_CONTEXT "verify_targets_exist")
+  endif()
+
+  foreach(_target IN LISTS _targets)
+    if(NOT TARGET ${_target})
+      message(FATAL_ERROR "${ARG_CONTEXT} requires missing target ${_target}.")
+    endif()
+  endforeach()
+endfunction()
+
+# Private function to isolate logic for building prim ops library.
+function(_arm_runner_create_prim_ops_lib)
+  # Parse arguments.
+  set(options INCLUDE_ALL_OPS)
+  set(one_value_args LIB_NAME SELECTED_OPS_YAML)
+  set(multi_value_args DEPS)
+  cmake_parse_arguments(
+    ARG "${options}" "${one_value_args}" "${multi_value_args}" ${ARGN}
+  )
+
+  set(_out_dir "${CMAKE_CURRENT_BINARY_DIR}/${ARG_LIB_NAME}")
+  set(_sources ${EXECUTORCH_ROOT}/kernels/prim_ops/register_prim_ops.cpp)
+
+  # Generate selected operator list if needed.
+  if(NOT ARG_INCLUDE_ALL_OPS)
+    set(_selected_prim_ops_header "${_out_dir}/selected_prim_ops.h")
+    set(_gen_selected_prim_ops_script
+        "${EXECUTORCH_ROOT}/codegen/tools/gen_selected_prim_ops.py"
+    )
+    set(_gen_selected_prim_ops_command
+        "${PYTHON_EXECUTABLE}" -m codegen.tools.gen_selected_prim_ops
+        --op-selection-yaml-path=${ARG_SELECTED_OPS_YAML}
+        --output-dir=${_out_dir}
+    )
+    add_custom_command(
+      COMMENT "Generating selected_prim_ops.h for ${ARG_LIB_NAME}"
+      OUTPUT ${_selected_prim_ops_header}
+      COMMAND ${_gen_selected_prim_ops_command}
+      DEPENDS ${ARG_SELECTED_OPS_YAML} ${_gen_selected_prim_ops_script}
+      WORKING_DIRECTORY ${EXECUTORCH_ROOT}
+    )
+    list(APPEND _sources ${_selected_prim_ops_header})
+  endif()
+
+  # Add prim ops registration library.
+  add_library(${ARG_LIB_NAME} ${_sources})
+  target_include_directories(
+    ${ARG_LIB_NAME} PRIVATE ${EXECUTORCH_ROOT}/kernels/prim_ops ${_out_dir}
+  )
+  if(NOT ARG_INCLUDE_ALL_OPS)
+    target_compile_definitions(
+      ${ARG_LIB_NAME} PRIVATE ET_PRIM_OPS_SELECTIVE_BUILD
+                              EXECUTORCH_ENABLE_PRIM_OPS_SELECTIVE_BUILD
+    )
+  endif()
+  target_link_libraries(${ARG_LIB_NAME} PRIVATE ${ARG_DEPS})
+  executorch_target_link_options_shared_lib(${ARG_LIB_NAME})
+
+  # Add prim ops kernel library.
+  add_library(
+    ${ARG_LIB_NAME}_impl ${EXECUTORCH_ROOT}/kernels/prim_ops/et_copy_index.cpp
+                         ${EXECUTORCH_ROOT}/kernels/prim_ops/et_view.cpp
+  )
+  target_link_libraries(${ARG_LIB_NAME}_impl PRIVATE ${ARG_DEPS})
+endfunction()
+
 #[[
 Create a selected operator registration library for one operator family.
 
@@ -62,12 +133,16 @@ Arguments:
   INCLUDE_ALL_OPS: Generate a non-selective registration library for the
 	selected operator family.
 
+  PRIM_OPS: Build a selective prim ops library instead of a codegen operators lib.
+
 The lib will always be created, even when no operators are selected.
 OP_LIST and OPS_FROM_MODEL can be used additively.
+If neither OP_LIST nor OPS_FROM_MODEL is set, the generated registration
+library is empty. INCLUDE_ALL_OPS explicitly builds a full registration library.
 ]]
 function(arm_runner_create_selected_ops_lib)
   # Parse arguments
-  set(options INCLUDE_ALL_OPS)
+  set(options INCLUDE_ALL_OPS PRIM_OPS)
   set(one_value_args LIB_NAME FUNCTIONS_YAML CUSTOM_OPS_YAML OP_LIST
                      OPS_FROM_MODEL DTYPE_SELECTIVE_BUILD
   )
@@ -101,10 +176,13 @@ function(arm_runner_create_selected_ops_lib)
     )
   endif()
 
-  if(NOT ARG_FUNCTIONS_YAML AND NOT ARG_CUSTOM_OPS_YAML)
+  if(NOT ARG_PRIM_OPS
+     AND NOT ARG_FUNCTIONS_YAML
+     AND NOT ARG_CUSTOM_OPS_YAML
+  )
     message(
       FATAL_ERROR
-        "${_arm_runner_create_selected_ops_lib_context} requires FUNCTIONS_YAML or CUSTOM_OPS_YAML."
+        "${_arm_runner_create_selected_ops_lib_context} requires FUNCTIONS_YAML or CUSTOM_OPS_YAML unless PRIM_OPS is set."
     )
   endif()
 
@@ -128,11 +206,25 @@ function(arm_runner_create_selected_ops_lib)
       DTYPE_SELECTIVE_BUILD
       "${ARG_DTYPE_SELECTIVE_BUILD}"
   )
+  set(_arm_runner_include_all_ops OFF)
   if(ARG_INCLUDE_ALL_OPS)
+    set(_arm_runner_include_all_ops ON)
     list(APPEND _arm_runner_selected_ops_args INCLUDE_ALL_OPS "ON")
   endif()
   gen_selected_ops(${_arm_runner_selected_ops_args})
 
+  if(ARG_PRIM_OPS)
+    set(_arm_runner_prim_ops_args
+        LIB_NAME "${ARG_LIB_NAME}" SELECTED_OPS_YAML
+        "${gen_selected_ops_output_yaml}" DEPS ${ARG_DEPS}
+    )
+    if(_arm_runner_include_all_ops)
+      list(APPEND _arm_runner_prim_ops_args INCLUDE_ALL_OPS)
+    endif()
+    _arm_runner_create_prim_ops_lib(${_arm_runner_prim_ops_args})
+    return()
+  endif()
+
   # Codegen for operator library.
   set(_arm_ops_binding_args LIB_NAME "${ARG_LIB_NAME}")
   if(ARG_FUNCTIONS_YAML)
@@ -198,21 +290,6 @@ function(arm_runner_configure_runtime_output TARGET_NAME FALLBACK_DIR)
   endif()
 endfunction()
 
-function(verify_targets_exist)
-  cmake_parse_arguments(ARG "" "CONTEXT" "TARGETS" ${ARGN})
-
-  set(_targets ${ARG_TARGETS} ${ARG_UNPARSED_ARGUMENTS})
-  if(NOT ARG_CONTEXT)
-    set(ARG_CONTEXT "verify_targets_exist")
-  endif()
-
-  foreach(_target IN LISTS _targets)
-    if(NOT TARGET ${_target})
-      message(FATAL_ERROR "${ARG_CONTEXT} requires missing target ${_target}.")
-    endif()
-  endforeach()
-endfunction()
-
 # Link the provided target with minimal specs. This minimizes code size, but
 # comes with limited support. Notably, printing with %zu is not supported.
 function(arm_runner_link_minimal_specs TARGET_NAME)

backends/arm/test/test_arm_backend.sh

@@ -187,6 +187,7 @@ test_run_ethos_u55() {
     # Cortex-M op tests
     echo "${TEST_SUITE_NAME}: Test target Cortex-M55 (on Ethos-U55)"
     examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=add --bundleio --no_delegate --select_ops_list="aten::add.out"
+    examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=examples/arm/example_modules/prim_ops.py --bundleio --no_delegate --no_quantize --select_ops_list="aten::add.out,aten::select_copy.int_out"
     examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=qadd --bundleio
     examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=qops --bundleio
     examples/arm/run.sh --et_build_root=arm_test/test_run --target=ethos-u55-128 --model_name=qops --bundleio --no_delegate --select_ops_list="aten::sub.out,aten::add.out,aten::mul.out"

backends/cortex_m/CMakeLists.txt

@@ -128,7 +128,7 @@ if(EXECUTORCH_BUILD_CORTEX_M)
   target_link_libraries(
     cortex_m_kernels
     PRIVATE cmsis-nn
-    PRIVATE executorch
+    PRIVATE executorch_core
     PRIVATE kernels_util_all_deps
   )
   target_compile_definitions(
@@ -143,7 +143,8 @@ if(EXECUTORCH_BUILD_CORTEX_M)
   )
 
   gen_operators_lib(
-    LIB_NAME "cortex_m_ops_lib" KERNEL_LIBS cortex_m_kernels DEPS executorch
+    LIB_NAME "cortex_m_ops_lib" KERNEL_LIBS cortex_m_kernels DEPS
+    executorch_core
   )
 
   install(

codegen/tools/gen_selected_prim_ops.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.
@@ -8,9 +9,11 @@
 
 import argparse
 import os
+import re
 import sys
 from typing import Any, List
 
+import yaml
 from torchgen.code_template import CodeTemplate  # type: ignore[import-not-found]
 
 
@@ -34,13 +37,32 @@ def normalize_op_name(op_name: str) -> str:
     normalized = op_name.replace("::", "_")
     # Replace dots with underscores
     normalized = normalized.replace(".", "_")
+    # Collapse cases like aten::_local_scalar_dense to the macro spelling used
+    # by register_prim_ops.cpp.
+    normalized = re.sub("_+", "_", normalized)
     # Convert to uppercase
     normalized = normalized.upper()
     # Add INCLUDE_ prefix
     normalized = f"INCLUDE_{normalized}"
     return normalized
 
 
+def read_prim_op_names_from_yaml(op_selection_yaml_path: str) -> List[str]:
+    with open(op_selection_yaml_path, "r") as f:
+        selected_operators = yaml.safe_load(f) or {}
+
+    prim_op_names = set()
+    operators = selected_operators.get("operators", {})
+    if operators:
+        prim_op_names.update(operators.keys())
+
+    et_kernel_metadata = selected_operators.get("et_kernel_metadata", {})
+    if et_kernel_metadata:
+        prim_op_names.update(et_kernel_metadata.keys())
+
+    return sorted(prim_op_names)
+
+
 def write_selected_prim_ops(prim_op_names: List[str], output_dir: str) -> None:
     """
     Generate selected_prim_ops.h from a list of prim op names.
@@ -73,7 +95,13 @@ def main(argv: List[Any]) -> None:
         "--prim-op-names",
         "--prim_op_names",
         help="Comma-separated list of prim op names to include",
-        required=True,
+        required=False,
+    )
+    parser.add_argument(
+        "--op-selection-yaml-path",
+        "--op_selection_yaml_path",
+        help="Path to selected_operators.yaml containing prim op names to include",
+        required=False,
     )
     parser.add_argument(
         "--output-dir",
@@ -84,12 +112,21 @@ def main(argv: List[Any]) -> None:
 
     options = parser.parse_args(argv)
 
-    # Parse comma-separated prim op names
-    prim_op_names = [
-        name.strip() for name in options.prim_op_names.split(",") if name.strip()
-    ]
-
-    write_selected_prim_ops(prim_op_names, options.output_dir)
+    if options.prim_op_names is None and not options.op_selection_yaml_path:
+        parser.error("one of --prim-op-names or --op-selection-yaml-path is required")
+
+    prim_op_names: List[str] = []
+    if options.prim_op_names is not None:
+        # Parse comma-separated prim op names
+        prim_op_names.extend(
+            name.strip() for name in options.prim_op_names.split(",") if name.strip()
+        )
+    if options.op_selection_yaml_path:
+        prim_op_names.extend(
+            read_prim_op_names_from_yaml(options.op_selection_yaml_path)
+        )
+
+    write_selected_prim_ops(sorted(set(prim_op_names)), options.output_dir)
 
 
 if __name__ == "__main__":

codegen/tools/targets.bzl

@@ -172,6 +172,9 @@ def define_common_targets(is_fbcode = False):
         srcs = ["gen_selected_prim_ops.py"],
         base_module = "executorch.codegen.tools",
         visibility = ["//executorch/..."],
+        deps = [
+            "fbsource//third-party/pypi/pyyaml:pyyaml",
+        ],
         external_deps = ["torchgen"],
     )
 
@@ -188,7 +191,21 @@ def define_common_targets(is_fbcode = False):
         _is_external_target = True,
     )
 
-    
+    runtime.python_test(
+        name = "test_gen_selected_prim_ops",
+        srcs = [
+            "test/test_gen_selected_prim_ops.py",
+        ],
+        package_style = "inplace",
+        visibility = [
+            "PUBLIC",
+        ],
+        deps = [
+            ":gen_selected_prim_ops_lib",
+        ],
+        _is_external_target = True,
+    )
+
     runtime.cxx_python_extension(
         name = "selective_build",
         srcs = [

codegen/tools/test/test_gen_selected_prim_ops.py

@@ -0,0 +1,65 @@
+# 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 tempfile
+import unittest
+from pathlib import Path
+
+import yaml
+
+from executorch.codegen.tools.gen_selected_prim_ops import (
+    normalize_op_name,
+    write_selected_prim_ops,
+)
+
+
+class TestGenSelectedPrimOps(unittest.TestCase):
+    def test_normalizes_aten_op_with_leading_underscore(self) -> None:
+        self.assertEqual(
+            normalize_op_name("aten::_local_scalar_dense"),
+            "INCLUDE_ATEN_LOCAL_SCALAR_DENSE",
+        )
+
+    def test_writes_selected_prim_ops_from_yaml(self) -> None:
+        with tempfile.TemporaryDirectory() as temp_dir:
+            yaml_path = Path(temp_dir) / "selected_operators.yaml"
+            yaml_path.write_text(
+                yaml.safe_dump(
+                    {
+                        "operators": {
+                            "executorch_prim::et_view.default": {},
+                            "aten::_local_scalar_dense": {},
+                        },
+                        "et_kernel_metadata": {
+                            "aten::sym_size.int": ["default"],
+                        },
+                    }
+                )
+            )
+
+            from executorch.codegen.tools.gen_selected_prim_ops import main
+
+            main(
+                [
+                    f"--op-selection-yaml-path={yaml_path}",
+                    f"--output-dir={temp_dir}",
+                ]
+            )
+
+            header = (Path(temp_dir) / "selected_prim_ops.h").read_text()
+            self.assertIn("#define INCLUDE_ATEN_LOCAL_SCALAR_DENSE", header)
+            self.assertIn("#define INCLUDE_ATEN_SYM_SIZE_INT", header)
+            self.assertIn("#define INCLUDE_EXECUTORCH_PRIM_ET_VIEW_DEFAULT", header)
+
+    def test_writes_empty_header_for_empty_op_list(self) -> None:
+        with tempfile.TemporaryDirectory() as temp_dir:
+            write_selected_prim_ops([], temp_dir)
+
+            header = (Path(temp_dir) / "selected_prim_ops.h").read_text()
+            self.assertNotIn("#define INCLUDE_", header)
+
+
+if __name__ == "__main__":
+    unittest.main()