Right-size operator registry from selective build metadata

[rascani]

Summary

The runtime kernel registry is a fixed-size static array whose default capacity (MAX_KERNEL_NUM=2000) permanently occupies ~24 KiB of BSS on 32-bit machines even when only a handful of kernels are actually registered. When selective build is active the exact set of kernels needed by the model is already known at build time, so the registry can be sized to fit.

A new codegen tool (gen_max_kernel_num.py) counts the (op, kernel_key) tuples in selected_operators.yaml, adds the prim ops registered by register_prim_ops.cpp, and writes the total into a generated header. operator_registry.cpp picks the header up via __has_include. A user-supplied -DMAX_KERNEL_NUM still takes precedence, and builds that don't use selective build keep the 2000 default.

The below example demonstrates examples/selective_build/basic (two selected ops) on a 64-bit machine. The registry's BSS footprint drops from 48000 B (2000 slots) to 840 B (35 slots).

Fixes #18618

Example generated header: executorch/runtime/kernel/selected_max_kernel_num.h

// @generated by executorch/codegen/tools/gen_max_kernel_num.py. Do not edit.
#pragma once
#define EXECUTORCH_SELECTED_MAX_KERNEL_NUM 27

Test plan

# 1. Export a tiny model
python -m examples.portable.scripts.export --model_name="add_mul"

# 2. Auto-sized (no MAX_KERNEL_NUM flag)
cmake -DCMAKE_BUILD_TYPE=Release \
        -DEXECUTORCH_SELECT_OPS_MODEL="./add_mul.pte" \
        -B/tmp/sb_auto examples/selective_build/basic
cmake --build /tmp/sb_auto -j$(nproc)

# 3. User override
cmake -DCMAKE_BUILD_TYPE=Release \
      -DEXECUTORCH_SELECT_OPS_MODEL="./add_mul.pte" \
      -DMAX_KERNEL_NUM=100 \
      -B/tmp/sb_pin examples/selective_build/basic
cmake --build /tmp/sb_pin -j$(nproc)

# 4. Baseline (no selective build)
cmake -DCMAKE_BUILD_TYPE=Release -B/tmp/sb_default .
cmake --build /tmp/sb_default --target executorch_core -j$(nproc)

echo "=== auto-sized ==="
cat /tmp/sb_auto/executorch/executorch_selected_kernels/executorch/runtime/kernel/selected_max_kernel_num.h
nm -S /tmp/sb_auto/executorch/CMakeFiles/executorch_core.dir/runtime/kernel/operator_registry.cpp.o \
  | grep registered_kernels_data
/tmp/sb_auto/selective_build_test --model_path=./add_mul.pte 2>&1 | tail -1

echo "=== user override (MAX_KERNEL_NUM=100) ==="
find /tmp/sb_pin -name selected_max_kernel_num.h 2>/dev/null || echo "(no header generated — expected)"
nm -S /tmp/sb_pin/executorch/CMakeFiles/executorch_core.dir/runtime/kernel/operator_registry.cpp.o \
  | grep registered_kernels_data

echo "=== baseline (no selective build) ==="
nm -S /tmp/sb_default/CMakeFiles/executorch_core.dir/runtime/kernel/operator_registry.cpp.o \
  | grep registered_kernels_data

Expected output (BSS size is the 4-hex column after the address):

=== auto-sized ===
#define EXECUTORCH_SELECTED_MAX_KERNEL_NUM 27
0... 0000000000000288 b ...registered_kernels_data...  # 648 B = 27 × 24
OutputX 0: tensor(sizes=[2, 2], [3., 3., 3., 3.])

=== user override (MAX_KERNEL_NUM=100) ===
(no header generated — expected)
0... 0000000000000960 b ...registered_kernels_data...  # 2400 B = 100 × 24

=== baseline (no selective build) ===
0... 000000000000bb80 b ...registered_kernels_data...  # 48000 B = 2000 × 24

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/19118

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Rolling out OSDC (ARC) runners on pull & trunk workflows in PyTorch main

❌ 1 New Failure, 2 Unrelated Failures

As of commit 97e49a4 with merge base 3ec63f4 ():

NEW FAILURE - The following job has failed:

• pull / unittest-editable / macos / macos-job (gh)
  backends/xnnpack/test/ops/test_conv2d.py::TestConv2d::test_fp16_conv2d

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / unittest / windows / windows-job (gh) (matched win rule in flaky-rules.json)
  ##[error]The operation was canceled.

BROKEN TRUNK - The following job failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest-editable / windows / windows-job (gh) (trunk failure)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[larryliu0820]

Can you attach a sample codegen header in the summary

[rascani]

Can you attach a sample codegen header in the summary

Done