Add fused int8 LSTM support to the Cortex-M (CMSIS-NN) backend#21024
Add fused int8 LSTM support to the Cortex-M (CMSIS-NN) backend#21024rascani wants to merge 1 commit into
Conversation
Lower a single-layer, unidirectional nn.LSTM to a fused cortex_m::quantized_lstm op backed by CMSIS-NN's arm_lstm_unidirectional_s8. Instead of decomposing the LSTM into primitives, aten.lstm.input is preserved through to_edge, boundary-quantized (int8 input/output activations, weights kept in float), and fused in a single pass. The per-gate CMSIS parameters -- Q3.12 gate accumulator, Q0.15 activations, and a power-of-two int16 cell state -- are derived ahead of time from the boundary scales and the float gate weights. Suggested review order: passes/lstm_params.py (the ahead-of-time parameter derivation plus an int8 reference, checked against a float LSTM in test/test_lstm_params.py); then ops/operators.py and ops/op_quantized_lstm.cpp (the op definition and its CMSIS kernel); then quantizer/quantizer.py (LstmBoundaryQuantizer, which annotates only the activation boundaries) and the fusion substitution in passes/aten_to_cortex_m_pass.py. Configurations outside the supported scope (multi-layer, bidirectional, projection, non-zero initial state, or models that consume h_n/c_n) are left unfused and fail lowering with a clear error rather than silently miscompiling. Verified with the dialect and Corstone-300 FVP tests in test/ops/test_lstm.py. This change was authored with Claude Code.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21024
Note: Links to docs will display an error until the docs builds have been completed. ❗ 2 Merge Blocking SEVsThere is 2 active merge blocking SEVs. Please view them below:
If you must merge, use ✅ No FailuresAs of commit ef6ec90 with merge base c2b273e ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
1 participant
Summary
Lower a single-layer, unidirectional nn.LSTM to a fused cortex_m::quantized_lstm op backed by CMSIS-NN's arm_lstm_unidirectional_s8. Instead of decomposing the LSTM into primitives, aten.lstm.input is preserved through to_edge, boundary-quantized (int8 input/output activations, weights kept in float), and fused in a single pass. The per-gate CMSIS parameters -- Q3.12 gate accumulator, Q0.15 activations, and a power-of-two int16 cell state -- are derived ahead of time from the boundary scales and the float gate weights.
Suggested review order: passes/lstm_params.py (the ahead-of-time parameter derivation plus an int8 reference, checked against a float LSTM in test/test_lstm_params.py); then ops/operators.py and ops/op_quantized_lstm.cpp (the op definition and its CMSIS kernel); then quantizer/quantizer.py (LstmBoundaryQuantizer, which annotates only the activation boundaries) and the fusion substitution in passes/aten_to_cortex_m_pass.py. Configurations outside the supported scope (multi-layer, bidirectional, projection, non-zero initial state, or models that consume h_n/c_n) are left unfused and fail lowering with a clear error rather than silently miscompiling.
Test plan
Verified with the dialect and Corstone-300 FVP tests in test/ops/test_lstm.py.
This change was authored with Claude Code.