Skip to content

feat: Add Nordic AXON NPU backend for nRF54LM20B#18863

Draft
petriok wants to merge 1 commit intopytorch:mainfrom
ioteai:nordic-axon-backend-pr
Draft

feat: Add Nordic AXON NPU backend for nRF54LM20B#18863
petriok wants to merge 1 commit intopytorch:mainfrom
ioteai:nordic-axon-backend-pr

Conversation

@petriok
Copy link
Copy Markdown

@petriok petriok commented Apr 14, 2026

Summary

ExecuTorch backend for Nordic Semiconductor's AXON NPU on the nRF54LM20B (ARM Cortex-M33 + hardware neural network accelerator).

Follows the same composition pattern as the Ethos-U backend: reuses TOSABackend for TOSA lowering, then compiles to AXON command
buffers via Nordic's compiler library.

Python backend (backends/nordic/):

  • AxonBackend, AxonPartitioner, AxonQuantizer, AxonCompileSpec
  • TOSA → AXON compiler bridge with per-op converters
  • Subgraph naming (content-hash), marker format, header code generation
  • Operator support: FC (2048), Conv2D (16x16), Pool (32x32), Add, Mul, ReLU/ReLU6, sigmoid/tanh/softmax op extensions

C++ runtime (backends/nordic/runtime/):

  • Delegate with marker-based multi-subgraph lookup and profiling API
  • Single-precision sigmoid/tanh op extension callbacks

Zephyr integration:

  • EXECUTORCH_BUILD_NORDIC_AXON Kconfig option
  • Auto-link executorch_delegate_axon in CMake

Examples (examples/nordic/):

Example Model AXON subgraphs
hello_axon sin(x) regression 1
multi_layer Chained FC classifier 1
simple_rnn RNN with recurrent state 2

Each example includes an export script, Zephyr firmware, and step-by-step README.

Nordic's sdk-edge-ai (containing the AXON compiler library) is an external dependency, not redistributed. Discovered via
SDK_EDGE_AI_PATH environment variable — same pattern as Ethos-U's dependency on Vela.

Test plan

  • 48 unit tests pass without Nordic SDK (TOSA lowering, operator constraints, codegen, per-op conversion)
  • 6 SDK-dependent tests skip gracefully with clear message
  • All three examples hardware-verified on nRF54LM20DK:

hello_axon: sin(1.57) = 0.967, 163 us @ 128 MHz, 1 AXON delegate
multi_layer: 4/4 classes correct, ~210 us, 1 delegate (chained layers)
simple_rnn: 4 RNN steps, ~690 us/step, 2 delegates (tanh breaks chain)

Add ExecuTorch backend for Nordic Semiconductor's AXON NPU, targeting
the nRF54LM20B (ARM Cortex-M33 + hardware neural network accelerator).
Follows the same composition pattern as the Ethos-U backend: reuses
TOSABackend for TOSA lowering, then compiles to AXON command buffers
via Nordic's compiler library.

Python backend (backends/nordic/):
- AxonBackend: @Final BackendDetails with TOSA composition
- AxonPartitioner: extends TOSAPartitioner with AXON constraint checks
- AxonQuantizer: wraps TOSAQuantizer with AXON INT8 defaults
- AxonCompileSpec: hardware constraints and SDK path configuration
- TOSA-to-AXON compiler bridge with per-op converters
- Subgraph naming (content-hash), marker format, header code generation
- Operator support checks (FC 2048, Conv 16x16, Pool 32x32, max 2 inputs)

C++ runtime (backends/nordic/runtime/):
- AxonBackend delegate: marker-based multi-subgraph lookup, profiling API
- Op extensions: single-precision sigmoid/tanh CPU callbacks

Zephyr integration:
- Kconfig: EXECUTORCH_BUILD_NORDIC_AXON (depends on NRF_AXON)
- CMakeLists: auto-link executorch_delegate_axon

Examples (examples/nordic/):
- hello_axon: minimal sin(x) regression (1 AXON subgraph)
- multi_layer: chained FC classifier with profiling (1 subgraph)
- simple_rnn: RNN with recurrent state — multi-subgraph delegation (2 subgraphs)
- Each includes export script, Zephyr firmware, and setup instructions

Tests: 48 passed, 6 skipped (require Nordic SDK)

Hardware verified on nRF54LM20DK: all three examples produce correct
inference output with AXON NPU acceleration.

The Nordic sdk-edge-ai (containing the AXON compiler library) is an
external dependency, not redistributed. Discovered via SDK_EDGE_AI_PATH
environment variable — same pattern as Ethos-U's dependency on Vela.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot bot commented Apr 14, 2026

🔗 Helpful Links

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

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:

⚠️ 11 Awaiting Approval

As of commit 99fdc20 with merge base 411ede2 (image):

AWAITING APPROVAL - The following workflows need approval before CI can run:

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

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Apr 14, 2026
@github-actions
Copy link
Copy Markdown

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant