From ef6ec90c5d99766651554615a383024028c4f42e Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 17 Jul 2026 11:00:50 -0700 Subject: [PATCH] Add fused int8 LSTM support to the Cortex-M (CMSIS-NN) backend 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. --- backends/cortex_m/CMakeLists.txt | 1 + backends/cortex_m/ops/op_quantized_lstm.cpp | 176 ++++++++ backends/cortex_m/ops/operators.py | 131 ++++++ backends/cortex_m/ops/operators.yaml | 6 + .../cortex_m/passes/aten_to_cortex_m_pass.py | 166 ++++++++ backends/cortex_m/passes/lstm_params.py | 389 ++++++++++++++++++ .../cortex_m/passes/scratch_buffer_sizes.py | 20 + backends/cortex_m/quantizer/quantizer.py | 112 ++++- backends/cortex_m/test/ops/test_lstm.py | 131 +++--- backends/cortex_m/test/test_lstm_params.py | 103 +++++ backends/cortex_m/test/tester.py | 1 + 11 files changed, 1183 insertions(+), 53 deletions(-) create mode 100644 backends/cortex_m/ops/op_quantized_lstm.cpp create mode 100644 backends/cortex_m/passes/lstm_params.py create mode 100644 backends/cortex_m/test/test_lstm_params.py diff --git a/backends/cortex_m/CMakeLists.txt b/backends/cortex_m/CMakeLists.txt index 85e52b5782a..78fb991c3e1 100644 --- a/backends/cortex_m/CMakeLists.txt +++ b/backends/cortex_m/CMakeLists.txt @@ -107,6 +107,7 @@ if(EXECUTORCH_BUILD_CORTEX_M) ${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_lstm.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 diff --git a/backends/cortex_m/ops/op_quantized_lstm.cpp b/backends/cortex_m/ops/op_quantized_lstm.cpp new file mode 100644 index 00000000000..7d5deac260b --- /dev/null +++ b/backends/cortex_m/ops/op_quantized_lstm.cpp @@ -0,0 +1,176 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include "cortex_m_ops_common.h" + +namespace cortex_m { +namespace native { +using KernelRuntimeContext = torch::executor::KernelRuntimeContext; + +namespace { + +// PyTorch stacks gate weights in IFGO order (input, forget, cell, output); +// CMSIS-NN addresses gates by name. Only the cell gate uses tanh. +constexpr int kNumGates = 4; +constexpr int kCellGateIdx = 2; + +void fail_arg(KernelRuntimeContext& context, const char* msg) { + ET_LOG(Error, "quantized_lstm_out: %s", msg); + context.fail(Error::InvalidArgument); +} + +} // namespace + +// cppcheck-suppress unusedFunction +Tensor& quantized_lstm_out( + KernelRuntimeContext& context, + const Tensor& input, + const Tensor& input_weights, + const Tensor& hidden_weights, + const Tensor& input_effective_bias, + const Tensor& hidden_effective_bias, + const Int64ArrayRef input_multipliers, + const Int64ArrayRef input_shifts, + const Int64ArrayRef hidden_multipliers, + const Int64ArrayRef hidden_shifts, + const int64_t input_offset, + const int64_t output_offset, + const int64_t forget_to_cell_multiplier, + const int64_t forget_to_cell_shift, + const int64_t input_to_cell_multiplier, + const int64_t input_to_cell_shift, + const int64_t output_multiplier, + const int64_t output_shift, + const int64_t cell_scale_power, + const int64_t cell_clip, + const bool time_major, + const Tensor& temp1, + const Tensor& temp2, + const Tensor& cell_state, + Tensor& out) { + if (input.dim() != 3 || out.dim() != 3 || input_weights.dim() != 2 || + hidden_weights.dim() != 2) { + fail_arg(context, "input/out must be 3-D, weights 2-D"); + return out; + } + if (input.scalar_type() != ScalarType::Char || + out.scalar_type() != ScalarType::Char || + input_weights.scalar_type() != ScalarType::Char || + hidden_weights.scalar_type() != ScalarType::Char) { + fail_arg(context, "input/output/weights must be int8"); + return out; + } + if (input_effective_bias.scalar_type() != ScalarType::Int || + hidden_effective_bias.scalar_type() != ScalarType::Int) { + fail_arg(context, "effective biases must be int32"); + return out; + } + if (input_weights.size(0) % kNumGates != 0) { + fail_arg(context, "input_weights rows must be 4*hidden"); + return out; + } + + const int32_t hidden_size = + static_cast(input_weights.size(0)) / kNumGates; + const int32_t input_size = static_cast(input_weights.size(1)); + const int32_t time_steps = + static_cast(time_major ? input.size(0) : input.size(1)); + const int32_t batch_size = + static_cast(time_major ? input.size(1) : input.size(0)); + + // Shape self-consistency: CMSIS strides these tensors using the dims above. + const int64_t gated_hidden = static_cast(kNumGates) * hidden_size; + if (input.size(2) != input_size || hidden_weights.size(0) != gated_hidden || + hidden_weights.size(1) != hidden_size || + input_effective_bias.numel() != gated_hidden || + hidden_effective_bias.numel() != gated_hidden || + out.numel() != + static_cast(time_steps) * batch_size * hidden_size) { + fail_arg(context, "tensor shapes are inconsistent"); + return out; + } + if (input_multipliers.size() != kNumGates || + input_shifts.size() != kNumGates || + hidden_multipliers.size() != kNumGates || + hidden_shifts.size() != kNumGates) { + fail_arg(context, "per-gate multiplier/shift lists must be length 4"); + return out; + } + + // Each working buffer holds batch*hidden int16 values. + const size_t buffer_bytes = + static_cast(batch_size) * hidden_size * sizeof(int16_t); + if (temp1.nbytes() < buffer_bytes || temp2.nbytes() < buffer_bytes || + cell_state.nbytes() < buffer_bytes) { + fail_arg(context, "scratch buffers too small"); + return out; + } + + const int8_t* input_data = input.const_data_ptr(); + const int8_t* iw = input_weights.const_data_ptr(); + const int8_t* hw = hidden_weights.const_data_ptr(); + const int32_t* ieb = input_effective_bias.const_data_ptr(); + const int32_t* heb = hidden_effective_bias.const_data_ptr(); + int8_t* output_data = out.mutable_data_ptr(); + + cmsis_nn_lstm_gate gates[kNumGates]; + for (int g = 0; g < kNumGates; ++g) { + gates[g].input_multiplier = static_cast(input_multipliers[g]); + gates[g].input_shift = static_cast(input_shifts[g]); + gates[g].input_weights = iw + g * hidden_size * input_size; + gates[g].input_effective_bias = ieb + g * hidden_size; + gates[g].hidden_multiplier = static_cast(hidden_multipliers[g]); + gates[g].hidden_shift = static_cast(hidden_shifts[g]); + gates[g].hidden_weights = hw + g * hidden_size * hidden_size; + gates[g].hidden_effective_bias = heb + g * hidden_size; + gates[g].bias = nullptr; // unused by the s8 gate kernel + gates[g].activation_type = (g == kCellGateIdx) ? ARM_TANH : ARM_SIGMOID; + } + + cmsis_nn_lstm_params params; + params.time_major = time_major ? 1 : 0; + params.batch_size = batch_size; + params.time_steps = time_steps; + params.input_size = input_size; + params.hidden_size = hidden_size; + params.input_offset = static_cast(input_offset); + params.forget_to_cell_multiplier = + static_cast(forget_to_cell_multiplier); + params.forget_to_cell_shift = static_cast(forget_to_cell_shift); + params.input_to_cell_multiplier = + static_cast(input_to_cell_multiplier); + params.input_to_cell_shift = static_cast(input_to_cell_shift); + params.cell_clip = static_cast(cell_clip); + params.cell_scale_power = static_cast(cell_scale_power); + params.output_multiplier = static_cast(output_multiplier); + params.output_shift = static_cast(output_shift); + params.output_offset = static_cast(output_offset); + params.input_gate = gates[0]; + params.forget_gate = gates[1]; + params.cell_gate = gates[kCellGateIdx]; + params.output_gate = gates[3]; + + cmsis_nn_lstm_context buffers; + buffers.temp1 = reinterpret_cast(temp1.mutable_data_ptr()); + buffers.temp2 = reinterpret_cast(temp2.mutable_data_ptr()); + buffers.cell_state = + reinterpret_cast(cell_state.mutable_data_ptr()); + + const arm_cmsis_nn_status status = + arm_lstm_unidirectional_s8(input_data, output_data, ¶ms, &buffers); + if (status != ARM_CMSIS_NN_SUCCESS) { + ET_LOG( + Error, "quantized_lstm_out: CMSIS-NN failed with status [%d]", status); + context.fail(Error::Internal); + return out; + } + return out; +} + +} // namespace native +} // namespace cortex_m diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py index f10802d3695..e33accca346 100644 --- a/backends/cortex_m/ops/operators.py +++ b/backends/cortex_m/ops/operators.py @@ -1436,3 +1436,134 @@ def quantized_max_pool2d_impl( ) result = torch.clamp(result, activation_min, activation_max) return result.to(torch.int8).contiguous(memory_format=torch.channels_last) + + +# =================================================================== +# QUANTIZED LSTM OPERATION DEFINITION +# =================================================================== +# +# Fused single-layer unidirectional int8 LSTM backed by CMSIS-NN's +# arm_nn_lstm_step_s8. Per-gate tensors are stacked in PyTorch IFGO order +# (input, forget, cell, output); the C++ kernel maps them to CMSIS's by-name +# gate structs. Gate/activation Q-formats are fixed by CMSIS (see lstm_params); +# the AOT pass supplies the derived multipliers/shifts and effective biases. + +_LSTM_ARGS = ( + "Tensor input, " + "Tensor input_weights, " + "Tensor hidden_weights, " + "Tensor input_effective_bias, " + "Tensor hidden_effective_bias, " + "int[] input_multipliers, " + "int[] input_shifts, " + "int[] hidden_multipliers, " + "int[] hidden_shifts, " + "int input_offset, " + "int output_offset, " + "int forget_to_cell_multiplier, " + "int forget_to_cell_shift, " + "int input_to_cell_multiplier, " + "int input_to_cell_shift, " + "int output_multiplier, " + "int output_shift, " + "int cell_scale_power, " + "int cell_clip, " + "bool time_major, " + "Tensor temp1, " + "Tensor temp2, " + "Tensor cell_state" +) + +lib.define(f"quantized_lstm({_LSTM_ARGS}) -> Tensor") +lib.define(f"quantized_lstm.out({_LSTM_ARGS}, *, Tensor(a!) out) -> Tensor(a!)") + + +@register_fake("cortex_m::quantized_lstm") # type: ignore[misc] +def quantized_lstm_meta( + input: torch.Tensor, + input_weights: torch.Tensor, + hidden_weights: torch.Tensor, + input_effective_bias: torch.Tensor, + hidden_effective_bias: torch.Tensor, + input_multipliers: Sequence[int], + input_shifts: Sequence[int], + hidden_multipliers: Sequence[int], + hidden_shifts: Sequence[int], + input_offset: int, + output_offset: int, + forget_to_cell_multiplier: int, + forget_to_cell_shift: int, + input_to_cell_multiplier: int, + input_to_cell_shift: int, + output_multiplier: int, + output_shift: int, + cell_scale_power: int, + cell_clip: int, + time_major: bool, + temp1: torch.Tensor, + temp2: torch.Tensor, + cell_state: torch.Tensor, +) -> torch.Tensor: + # Output keeps the input's [dim0, dim1] layout (time-major or batch-first). + hidden_size = input_weights.shape[0] // 4 + shape = (input.shape[0], input.shape[1], hidden_size) + return torch.empty(shape, dtype=torch.int8, device=input.device) + + +@impl(lib, "quantized_lstm", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_lstm_impl( + input: torch.Tensor, + input_weights: torch.Tensor, + hidden_weights: torch.Tensor, + input_effective_bias: torch.Tensor, + hidden_effective_bias: torch.Tensor, + input_multipliers: Sequence[int], + input_shifts: Sequence[int], + hidden_multipliers: Sequence[int], + hidden_shifts: Sequence[int], + input_offset: int, + output_offset: int, + forget_to_cell_multiplier: int, + forget_to_cell_shift: int, + input_to_cell_multiplier: int, + input_to_cell_shift: int, + output_multiplier: int, + output_shift: int, + cell_scale_power: int, + cell_clip: int, + time_major: bool, + temp1: torch.Tensor, + temp2: torch.Tensor, + cell_state: torch.Tensor, +) -> torch.Tensor: + # Imported lazily: at operators.py load time the passes package (via its + # __init__) references edge cortex_m ops that this module has not finished + # registering yet, so a top-level import would deadlock. + from executorch.backends.cortex_m.passes.lstm_params import ( + lstm_params_from_op_args, + run_quantized_lstm, + ) + + params = lstm_params_from_op_args( + input_weights, + hidden_weights, + input_effective_bias, + hidden_effective_bias, + list(input_multipliers), + list(input_shifts), + list(hidden_multipliers), + list(hidden_shifts), + input_offset, + output_offset, + forget_to_cell_multiplier, + forget_to_cell_shift, + input_to_cell_multiplier, + input_to_cell_shift, + output_multiplier, + output_shift, + cell_scale_power, + cell_clip, + ) + x_tm = input if time_major else input.transpose(0, 1) + out_tm = run_quantized_lstm(params, x_tm) + return out_tm if time_major else out_tm.transpose(0, 1).contiguous() diff --git a/backends/cortex_m/ops/operators.yaml b/backends/cortex_m/ops/operators.yaml index b4babe8f4a5..aca910449bc 100644 --- a/backends/cortex_m/ops/operators.yaml +++ b/backends/cortex_m/ops/operators.yaml @@ -106,3 +106,9 @@ kernels: - arg_meta: null kernel_name: cortex_m::quantized_batch_matmul_out + +- func: cortex_m::quantized_lstm.out(Tensor input, Tensor input_weights, Tensor hidden_weights, Tensor input_effective_bias, Tensor hidden_effective_bias, int[] input_multipliers, int[] input_shifts, int[] hidden_multipliers, int[] hidden_shifts, int input_offset, int output_offset, int forget_to_cell_multiplier, int forget_to_cell_shift, int input_to_cell_multiplier, int input_to_cell_shift, int output_multiplier, int output_shift, int cell_scale_power, int cell_clip, bool time_major, Tensor temp1, Tensor temp2, Tensor cell_state, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_lstm_out diff --git a/backends/cortex_m/passes/aten_to_cortex_m_pass.py b/backends/cortex_m/passes/aten_to_cortex_m_pass.py index 4cf3f2c17a0..7ce1b9a1c3f 100644 --- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py +++ b/backends/cortex_m/passes/aten_to_cortex_m_pass.py @@ -73,6 +73,31 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: result.graph_module ) + # The LSTM substitution rewires getitem(0)'s users onto a new + # single-output op, leaving getitem(0) dead and the tuple-returning + # aten.lstm.input with only that dead getitem as a user. DCE drops the + # dead getitem, but fx treats lstm.input as impure and won't remove it + # even once user-less, so erase it explicitly, then DCE its now-dead + # inputs (initial hidden/cell) and recompile. + graph = result.graph_module.graph + graph.eliminate_dead_code() + for n in list(graph.nodes): + if n.target == exir_ops.edge.aten.lstm.input and len(n.users) == 0: + graph.erase_node(n) + graph.eliminate_dead_code() + + # A preserved aten.lstm.input that is still live was not fused (an + # unsupported configuration). There is no portable LSTM kernel, so fail + # here with a clear message rather than later with a missing-kernel error. + for n in graph.nodes: + if n.target == exir_ops.edge.aten.lstm.input: + raise RuntimeError( + "cortex_m: an aten.lstm.input survived lowering unfused. Only " + "single-layer, unidirectional, biased, zero-initial-state, " + "output-only LSTM is supported by the Cortex-M backend." + ) + result.graph_module.recompile() + for node in result.graph_module.graph.nodes: self._initialize_alloc_node_size(node) @@ -186,6 +211,147 @@ def _restore_max_pool2d_with_indices_fallback( _SOFTMAX_INPUT_INTEGER_BITS = 5 +# Power-of-two cell-state scale exponent used when lowering the fused LSTM. +# The cell state is internal (never observed), so its scale cannot be +# calibrated here; -12 gives an int16 range of +/-8, ample for typical cell +# magnitudes. Values beyond the range saturate at cell_clip rather than wrap. +# TODO: derive from calibration statistics (project_cortex_m_lstm R3). +_DEFAULT_CELL_SCALE_POWER = -12 + + +@AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.lstm.input) +def _get_lstm_replacement( + node: Node, dialect_pass: AtenToDialectPass +) -> DialectNodeSpec | None: + """Fuse a preserved, boundary-quantized aten.lstm.input into a single + cortex_m.quantized_lstm. Gate weights are quantized per-gate ahead of time + from the float params; boundary scales come from the input DQ (folded onto + this node) and the output Q (folded onto getitem(0)). + + The op is single-output, so this rewires getitem(0)'s consumers directly and + returns None; the surrounding pass DCEs the now-dead tuple op. + """ + from executorch.backends.cortex_m.passes.lstm_params import ( + derive_lstm_params, + flatten_lstm_params, + ) + + input_qparams = node.meta.get("input_qparams") + if not input_qparams or 0 not in input_qparams: + # Not boundary-annotated (unsupported config); leave it for call() to + # flag as an un-fused, unlowerable LSTM. + return None + + getitem0 = next( + (u for u in node.users if u.target == operator.getitem and u.args[1] == 0), + None, + ) + if getitem0 is None: + return None + output_qparams = getitem0.meta.get("output_qparams") + if not output_qparams or 0 not in output_qparams: + return None + + exported_program = dialect_pass.exported_program + in_qp = input_qparams[0] + out_qp = output_qparams[0] + + params = cast(list[Node], node.args[2]) + weight_ih = cast(torch.Tensor, get_param_tensor(exported_program, params[0])) + weight_hh = cast(torch.Tensor, get_param_tensor(exported_program, params[1])) + bias = cast(torch.Tensor, get_param_tensor(exported_program, params[2])) + cast( + torch.Tensor, get_param_tensor(exported_program, params[3]) + ) + + lstm_params = derive_lstm_params( + weight_ih, + weight_hh, + bias, + float(in_qp.scale), + int(in_qp.zp), + float(out_qp.scale), + int(out_qp.zp), + _DEFAULT_CELL_SCALE_POWER, + ) + iw, hw, ieb, heb, in_mult, in_shift, hid_mult, hid_shift = flatten_lstm_params( + lstm_params + ) + + graph = node.graph + # Const placeholders must live in the placeholder region (before user + # inputs), so insert them after an existing weight placeholder. + with graph.inserting_after(params[0]): + iw_c = create_constant_placeholder( + exported_program, + graph, + node.name + "_input_weights", + InputKind.PARAMETER, + iw, + ) + hw_c = create_constant_placeholder( + exported_program, + graph, + node.name + "_hidden_weights", + InputKind.PARAMETER, + hw, + ) + ieb_c = create_constant_placeholder( + exported_program, + graph, + node.name + "_input_eff_bias", + InputKind.PARAMETER, + ieb, + ) + heb_c = create_constant_placeholder( + exported_program, + graph, + node.name + "_hidden_eff_bias", + InputKind.PARAMETER, + heb, + ) + # CMSIS needs three int16 working buffers (temp1, temp2, cell_state); the + # multi-buffer scratch API sizes each one independently (scratch_buffer_sizes). + temp1 = _create_uninitialized_alloc_node(node, exported_program) + temp2 = _create_uninitialized_alloc_node(node, exported_program) + cell_state = _create_uninitialized_alloc_node(node, exported_program) + + time_major = not bool(node.args[8]) + op_args = ( + node.args[0], + iw_c, + hw_c, + ieb_c, + heb_c, + in_mult, + in_shift, + hid_mult, + hid_shift, + lstm_params.input_offset, + lstm_params.output_offset, + lstm_params.forget_to_cell_multiplier, + lstm_params.forget_to_cell_shift, + lstm_params.input_to_cell_multiplier, + lstm_params.input_to_cell_shift, + lstm_params.output_multiplier, + lstm_params.output_shift, + lstm_params.cell_scale_power, + lstm_params.cell_clip, + time_major, + temp1, + temp2, + cell_state, + ) + with graph.inserting_before(node): + op_node = graph.create_node( + "call_function", + target=exir_ops.edge.cortex_m.quantized_lstm.default, + args=op_args, + ) + op_node.meta = dict(getitem0.meta) + op_node.meta["val"] = getitem0.meta["val"].to(torch.int8) + getitem0.replace_all_uses_with(op_node) + return None + def _to_int_pair( value: Argument, default: Optional[tuple[int, int]] diff --git a/backends/cortex_m/passes/lstm_params.py b/backends/cortex_m/passes/lstm_params.py new file mode 100644 index 00000000000..191f79ce826 --- /dev/null +++ b/backends/cortex_m/passes/lstm_params.py @@ -0,0 +1,389 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Ahead-of-time quantization math for the fused CMSIS-NN int8 LSTM kernel +(``arm_lstm_unidirectional_s8`` / ``arm_nn_lstm_step_s8``). + +The CMSIS kernel does not consume observed gate-activation scales: the gate +pre-activation is fixed at Q3.12 (``2**-12``) and the sigmoid/tanh outputs at +Q0.15 (``2**-15``). Everything the kernel needs is derived here from four +quantities — the input activation scale, the output/hidden activation scale, a +chosen (power-of-two) cell-state scale, and per-gate weight scales — mirroring +CMSIS-NN's own generator (``Tests/UnitTest/RefactoredTestGen/Lib/op_lstm.py``). + +``derive_lstm_params`` produces the CMSIS parameter structs (as a dataclass); +``quantized_lstm_reference`` is a faithful Python execution of the kernel used +both as the op's reference implementation and to validate the derivation +against a float ``torch.nn.LSTM``. PyTorch stores gates in IFGO order; CMSIS +addresses them by name (input/forget/cell/output). +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import torch + +from executorch.backends.cortex_m.passes.passes_utils import ( + quantize_multiplier_aot, + requantize_cmsis, +) + +# CMSIS-NN fixed intermediate Q-formats (see arm_nn_lstm_calculate_gate_s8_s16 / +# arm_nn_lstm_step_s8): gate accumulator is Q3.12, activation outputs are Q0.15. +_GATE_ACC_SCALE = 2.0**-12 +_ACTIVATION_SCALE = 2.0**-15 + +_Q15_MIN = -32768 +_Q15_MAX = 32767 +_INT8_MIN = -128 +_INT8_MAX = 127 + +# PyTorch weight_ih/weight_hh row order within the 4*H stack. +_IFGO = ("input", "forget", "cell", "output") + + +@dataclass +class GateParams: + input_multiplier: int + input_shift: int + hidden_multiplier: int + hidden_shift: int + input_weights: torch.Tensor # int8 [hidden, input] + hidden_weights: torch.Tensor # int8 [hidden, hidden] + input_effective_bias: torch.Tensor # int32 [hidden] + hidden_effective_bias: torch.Tensor # int32 [hidden] + is_tanh: bool # cell gate uses tanh; input/forget/output use sigmoid + + +@dataclass +class LstmParams: + input_size: int + hidden_size: int + input_offset: int + output_offset: int + forget_to_cell_multiplier: int + forget_to_cell_shift: int + input_to_cell_multiplier: int + input_to_cell_shift: int + output_multiplier: int + output_shift: int + cell_scale_power: int + cell_clip: int + gates: dict[str, GateParams] # keyed by "input"/"forget"/"cell"/"output" + + +def choose_cell_scale_power(max_abs_cell: float, headroom_bits: int = 2) -> int: + """Pick the power-of-two cell-state scale exponent so the observed cell + magnitude fits the int16 range with headroom for calibration drift. + + CMSIS stores the cell state as int16 with scale ``2**cell_scale_power`` and + ``cell_clip = int16_max``; the scale must be a power of two (only the + exponent is stored). ``ceil`` alone can leave as little as ~1x headroom, so + inference values slightly above ``max_abs_cell`` would saturate; the + ``headroom_bits`` margin (default 2 -> 4-8x) guards against that. Values that + still exceed the range saturate at ``cell_clip`` rather than wrapping. + """ + if max_abs_cell <= 0.0: + return -15 + return math.ceil(math.log2(max_abs_cell / _Q15_MAX)) + headroom_bits + + +def _quantize_weight_per_tensor(w: torch.Tensor) -> tuple[torch.Tensor, float]: + """Symmetric per-tensor int8 quantization of a weight block.""" + max_abs = float(w.abs().max()) + scale = max_abs / _INT8_MAX if max_abs > 0 else 1.0 + q = torch.clamp(torch.round(w / scale), _INT8_MIN, _INT8_MAX).to(torch.int8) + return q, scale + + +def _effective_bias( + w_q: torch.Tensor, bias_q: torch.Tensor | None, offset: int +) -> torch.Tensor: + """CMSIS effective bias = bias + offset * row_sum(weights) (int32).""" + kernel_sum = w_q.to(torch.int32).sum(dim=1) + eff = offset * kernel_sum + if bias_q is not None: + eff = eff + bias_q.to(torch.int32) + return eff.to(torch.int32) + + +def derive_lstm_params( + weight_ih: torch.Tensor, # float [4H, input], IFGO row order + weight_hh: torch.Tensor, # float [4H, hidden] + bias: torch.Tensor | None, # float [4H] (b_ih + b_hh), IFGO + input_scale: float, + input_zp: int, + output_scale: float, + output_zp: int, + cell_scale_power: int, +) -> LstmParams: + """Derive the CMSIS-NN LSTM parameter set from float weights and the three + boundary scales (input act, output/hidden act, chosen power-of-two cell). + + Single-layer, unidirectional only (``weight_ih`` is ``[4*hidden, input]``); + the quantizer's pattern check rejects other configurations upstream. + """ + four_h, input_size = weight_ih.shape + if four_h % 4 != 0: + raise ValueError(f"weight_ih first dim {four_h} is not 4*hidden_size") + hidden_size = four_h // 4 + cell_scale = 2.0**cell_scale_power + + input_offset = -input_zp + hidden_offset = -output_zp # hidden state reuses the output scale/zp + + gates: dict[str, GateParams] = {} + for gate_idx, name in enumerate(_IFGO): + rows = slice(gate_idx * hidden_size, (gate_idx + 1) * hidden_size) + w_ih_q, w_ih_scale = _quantize_weight_per_tensor(weight_ih[rows]) + w_hh_q, w_hh_scale = _quantize_weight_per_tensor(weight_hh[rows]) + + bias_q = None + if bias is not None: + # Bias lives in the input-path accumulator domain. + bias_acc_scale = input_scale * w_ih_scale + bias_q = torch.round(bias[rows] / bias_acc_scale).to(torch.int32) + + in_mult, in_shift = quantize_multiplier_aot( + input_scale * w_ih_scale / _GATE_ACC_SCALE + ) + hid_mult, hid_shift = quantize_multiplier_aot( + output_scale * w_hh_scale / _GATE_ACC_SCALE + ) + + gates[name] = GateParams( + input_multiplier=in_mult, + input_shift=in_shift, + hidden_multiplier=hid_mult, + hidden_shift=hid_shift, + input_weights=w_ih_q, + hidden_weights=w_hh_q, + input_effective_bias=_effective_bias(w_ih_q, bias_q, input_offset), + hidden_effective_bias=_effective_bias(w_hh_q, None, hidden_offset), + is_tanh=(name == "cell"), + ) + + f2c_mult, f2c_shift = quantize_multiplier_aot(_ACTIVATION_SCALE) + i2c_mult, i2c_shift = quantize_multiplier_aot( + _ACTIVATION_SCALE * _ACTIVATION_SCALE / cell_scale + ) + out_mult, out_shift = quantize_multiplier_aot( + _ACTIVATION_SCALE * _ACTIVATION_SCALE / output_scale + ) + + return LstmParams( + input_size=input_size, + hidden_size=hidden_size, + input_offset=input_offset, + output_offset=output_zp, + forget_to_cell_multiplier=f2c_mult, + forget_to_cell_shift=f2c_shift, + input_to_cell_multiplier=i2c_mult, + input_to_cell_shift=i2c_shift, + output_multiplier=out_mult, + output_shift=out_shift, + cell_scale_power=cell_scale_power, + cell_clip=_Q15_MAX, + gates=gates, + ) + + +def flatten_lstm_params( + params: LstmParams, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + list[int], + list[int], + list[int], + list[int], +]: + """Stack the four per-gate tensors (IFGO order) into the flat form carried by + the op schema: (input_weights[4H,in], hidden_weights[4H,H], + input_effective_bias[4H], hidden_effective_bias[4H], and the four 4-element + per-gate multiplier/shift lists for the input and hidden paths).""" + g = [params.gates[n] for n in _IFGO] + return ( + torch.cat([x.input_weights for x in g], dim=0), + torch.cat([x.hidden_weights for x in g], dim=0), + torch.cat([x.input_effective_bias for x in g], dim=0), + torch.cat([x.hidden_effective_bias for x in g], dim=0), + [x.input_multiplier for x in g], + [x.input_shift for x in g], + [x.hidden_multiplier for x in g], + [x.hidden_shift for x in g], + ) + + +def lstm_params_from_op_args( + input_weights: torch.Tensor, + hidden_weights: torch.Tensor, + input_effective_bias: torch.Tensor, + hidden_effective_bias: torch.Tensor, + input_multipliers: list[int], + input_shifts: list[int], + hidden_multipliers: list[int], + hidden_shifts: list[int], + input_offset: int, + output_offset: int, + forget_to_cell_multiplier: int, + forget_to_cell_shift: int, + input_to_cell_multiplier: int, + input_to_cell_shift: int, + output_multiplier: int, + output_shift: int, + cell_scale_power: int, + cell_clip: int, +) -> LstmParams: + """Rebuild the ``LstmParams`` struct from the op's flat args (inverse of + ``flatten_lstm_params`` plus the scalar fields carried directly).""" + four_h, input_size = input_weights.shape + hidden_size = four_h // 4 + gates = {} + for gate_idx, name in enumerate(_IFGO): + rows = slice(gate_idx * hidden_size, (gate_idx + 1) * hidden_size) + gates[name] = GateParams( + input_multiplier=int(input_multipliers[gate_idx]), + input_shift=int(input_shifts[gate_idx]), + hidden_multiplier=int(hidden_multipliers[gate_idx]), + hidden_shift=int(hidden_shifts[gate_idx]), + input_weights=input_weights[rows], + hidden_weights=hidden_weights[rows], + input_effective_bias=input_effective_bias[rows], + hidden_effective_bias=hidden_effective_bias[rows], + is_tanh=(name == "cell"), + ) + return LstmParams( + input_size=input_size, + hidden_size=hidden_size, + input_offset=input_offset, + output_offset=output_offset, + forget_to_cell_multiplier=forget_to_cell_multiplier, + forget_to_cell_shift=forget_to_cell_shift, + input_to_cell_multiplier=input_to_cell_multiplier, + input_to_cell_shift=input_to_cell_shift, + output_multiplier=output_multiplier, + output_shift=output_shift, + cell_scale_power=cell_scale_power, + cell_clip=cell_clip, + gates=gates, + ) + + +def _activation_q15( + x_q16: torch.Tensor, real_scale: float, is_tanh: bool +) -> torch.Tensor: + """Faithful stand-in for arm_nn_activation_s16: interpret the int16 input at + ``real_scale``, apply sigmoid/tanh in float, emit Q0.15 int16. This is the + exact operation the CMSIS table approximates (to ~1 LSB).""" + x = x_q16.to(torch.float64) * real_scale + y = torch.tanh(x) if is_tanh else torch.sigmoid(x) + q = torch.round(y / _ACTIVATION_SCALE) + return torch.clamp(q, _Q15_MIN, _Q15_MAX).to(torch.int32) + + +def _gate( + x_q: torch.Tensor, # int8 [input] + h_q: torch.Tensor | None, # int8 [hidden] or None at t=0 + gp: GateParams, +) -> torch.Tensor: + """One gate: two accumulating int8xint8->int16 matmuls + activation (Q0.15).""" + acc = gp.input_effective_bias + ( + x_q.to(torch.int32) * gp.input_weights.to(torch.int32) + ).sum(dim=1) + out = torch.clamp( + requantize_cmsis(acc, gp.input_multiplier, gp.input_shift), _Q15_MIN, _Q15_MAX + ) + if h_q is not None: + acc_h = gp.hidden_effective_bias + ( + h_q.to(torch.int32) * gp.hidden_weights.to(torch.int32) + ).sum(dim=1) + out = torch.clamp( + requantize_cmsis(acc_h, gp.hidden_multiplier, gp.hidden_shift) + out, + _Q15_MIN, + _Q15_MAX, + ) + return _activation_q15(out, _GATE_ACC_SCALE, gp.is_tanh) + + +def run_quantized_lstm(params: LstmParams, x_q: torch.Tensor) -> torch.Tensor: + """Core int8 CMSIS LSTM: int8 input [T, B, input] -> int8 output + [T, B, hidden], full sequence, zero initial state, time-major. Shared by the + op's reference implementation and the float-model validation wrapper. + """ + T, B, _ = x_q.shape + H = params.hidden_size + cell_scale = 2.0**params.cell_scale_power + x_q = x_q.to(torch.int32) + + out = torch.empty(T, B, H, dtype=torch.int32) + for b in range(B): + h_q: torch.Tensor | None = None + cell = torch.zeros(H, dtype=torch.int32) + for t in range(T): + xt = x_q[t, b] + forget = _gate(xt, h_q, params.gates["forget"]) + # cell = clamp(requant(forget * cell), Q15) -- overwrite + cell = torch.clamp( + requantize_cmsis( + forget * cell, + params.forget_to_cell_multiplier, + params.forget_to_cell_shift, + ), + _Q15_MIN, + _Q15_MAX, + ) + inp = _gate(xt, h_q, params.gates["input"]) + cellg = _gate(xt, h_q, params.gates["cell"]) + # cell += clamp(requant(inp * cellg), [-cell_clip, cell_clip]) + cell = torch.clamp( + requantize_cmsis( + inp * cellg, + params.input_to_cell_multiplier, + params.input_to_cell_shift, + ) + + cell, + -params.cell_clip, + params.cell_clip, + ) + outp = _gate(xt, h_q, params.gates["output"]) + tanh_cell = _activation_q15(cell, cell_scale, is_tanh=True) + hidden = ( + requantize_cmsis( + outp * tanh_cell, + params.output_multiplier, + params.output_shift, + ) + + params.output_offset + ) + hidden_q = torch.clamp(hidden, _INT8_MIN, _INT8_MAX).to(torch.int32) + out[t, b] = hidden_q + h_q = hidden_q + return out.to(torch.int8) + + +def quantized_lstm_reference( + x: torch.Tensor, # float input, [T, B, input] (time-major) + params: LstmParams, + input_scale: float, + output_scale: float, + dequantize: bool = False, +) -> torch.Tensor: + """Validation wrapper: quantize a float input, run the int8 LSTM, and + optionally dequantize the output for comparison against a float model. + """ + input_zp = -params.input_offset + x_q = torch.clamp(torch.round(x / input_scale) + input_zp, _INT8_MIN, _INT8_MAX) + out8 = run_quantized_lstm(params, x_q) + if dequantize: + return (out8.to(torch.int32) - params.output_offset).to( + torch.float32 + ) * output_scale + return out8 diff --git a/backends/cortex_m/passes/scratch_buffer_sizes.py b/backends/cortex_m/passes/scratch_buffer_sizes.py index b247e2be944..5b21a13dc51 100644 --- a/backends/cortex_m/passes/scratch_buffer_sizes.py +++ b/backends/cortex_m/passes/scratch_buffer_sizes.py @@ -268,8 +268,28 @@ def cmsis_nn_avgpool_buffer_size( ] +def cmsis_nn_lstm_buffer_size( + backend: cmsis_nn.Backend, + lstm_node: torch.fx.Node, +) -> list[int]: + # arm_lstm_unidirectional_s8 needs three int16 working buffers (temp1, + # temp2, cell_state), each batch*hidden elements. CMSIS-NN has no buffer + # sizing helper for LSTM, so compute it directly; the multi-buffer scratch + # API allocates one alloc node per size. + del backend + x = cast(torch.fx.Node, lstm_node.args[0]) + input_weights = cast(torch.fx.Node, lstm_node.args[1]) + time_major = bool(lstm_node.args[19]) + shape = _shape_from_node(x) + batch = int(shape[1] if time_major else shape[0]) + hidden = int(_shape_from_node(input_weights)[0]) // 4 + buffer_bytes = batch * hidden * 2 # int16 + return [buffer_bytes, buffer_bytes, buffer_bytes] + + _target_to_buffer_sizes_registry: dict[Any, BufferSizeFunction] = { exir_ops.edge.cortex_m.quantized_conv2d.default: cmsis_nn_conv_buffer_size, + exir_ops.edge.cortex_m.quantized_lstm.default: cmsis_nn_lstm_buffer_size, exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default: cmsis_nn_depthwise_conv_buffer_size, exir_ops.edge.cortex_m.quantized_batch_matmul.default: cmsis_nn_batch_matmul_buffer_size, exir_ops.edge.cortex_m.quantized_transpose_conv2d.default: cmsis_nn_transpose_conv_buffer_size, diff --git a/backends/cortex_m/quantizer/quantizer.py b/backends/cortex_m/quantizer/quantizer.py index e331c80ed1c..8741118d2f1 100644 --- a/backends/cortex_m/quantizer/quantizer.py +++ b/backends/cortex_m/quantizer/quantizer.py @@ -4,10 +4,13 @@ # LICENSE file in the root directory of this source tree. +import operator from typing import cast, List, Optional +import torch from executorch.backends.arm.quantizer.arm_quantizer_utils import ( _mark_node_as_quantized, + is_annotated, PatternCheck, PatternQuantizer, SharedQspecQuantizer, @@ -19,6 +22,7 @@ ) from executorch.backends.cortex_m.quantizer.pattern_matcher import PatternMatcher from executorch.backends.cortex_m.quantizer.quantization_configs import ( + INT8_ACTIVATION_PER_TENSOR_QSPEC, INT8_PER_CHANNEL_CONFIG, INT8_PER_TENSOR_CONFIG, ) @@ -28,12 +32,117 @@ CONV_TRANSPOSE_OP_PATTERNS, CORTEX_M_QUANTIZER_SUPPORT_DICT, ) -from executorch.backends.cortex_m.quantizer_reporter import QuantizerReporter +from executorch.backends.cortex_m.quantizer_reporter import ( + QuantizerInfo, + QuantizerReporter, + QuantizerReporterUser, +) from torch._ops import OpOverload from torch.fx import GraphModule from torchao.quantization.pt2e.quantizer import ComposableQuantizer, Quantizer +def _is_zero_init(node) -> bool: + """True if an initial hidden/cell state node is a constant zero tensor.""" + if not isinstance(node, torch.fx.Node): + return False + if node.target == torch.ops.aten.zeros.default: + return True + if node.target == torch.ops.aten.full.default: + return bool(len(node.args) > 1 and node.args[1] == 0) + return False + + +class LstmBoundaryQuantizer(Quantizer, QuantizerReporterUser): + """Annotate a single-layer, unidirectional, biased, zero-initial-state, + output-only ``aten.lstm.input`` at its activation boundaries only: int8 + per-tensor on the input activation and on the primary output + (``getitem(0)``). The weight/bias params are left in float so the lowering + pass can quantize each gate per-tensor ahead of time (a single per-tensor + scale over the 4*hidden weight stack would be too coarse for the CMSIS + kernel). + + Configurations outside that scope are left unannotated. Because the flow + preserves ``aten.lstm.input`` through to_edge and there is no portable LSTM + kernel, an unannotated LSTM cannot be lowered; ``AtenToCortexMPass`` raises a + clear error if one survives rather than failing later with a missing kernel. + """ + + _TARGET = torch.ops.aten.lstm.input + + def __init__(self) -> None: + super().__init__() + QuantizerReporterUser.__init__(self) + + def get_quantizer_info(self) -> QuantizerInfo: + return QuantizerInfo( + self.__class__.__name__, + str(self._TARGET), + "int8 per-tensor activations (weights kept float)", + __name__ + ".LstmBoundaryQuantizer", + ) + + @staticmethod + def _rejection_reason(node) -> Optional[str]: + # args: (input, hx, params, has_biases, num_layers, dropout, train, + # bidirectional, batch_first) + hx, params = node.args[1], node.args[2] + has_biases, num_layers, bidirectional = ( + node.args[3], + node.args[4], + node.args[7], + ) + if not has_biases: + return "bias is required" + if num_layers != 1: + return "only single-layer is supported" + if bidirectional: + return "only unidirectional is supported" + # A projection LSTM appends weight_hr, so params has 5 (biased) entries. + if len(params) != 4: + return "projection LSTM is not supported" + if not all(_is_zero_init(h) for h in hx): + return "only a zero initial state is supported" + # The fused op emits only the sequence output; h_n/c_n (getitem 1/2) + # would leave the original LSTM live and unlowerable. + for user in node.users: + if ( + user.target == operator.getitem + and user.args[1] != 0 + and len(user.users) > 0 + ): + return "h_n / c_n outputs are not supported" + return None + + def annotate(self, model: GraphModule) -> None: # type: ignore[override] + for node in model.graph.nodes: + if node.target != self._TARGET or is_annotated(node): + continue + reason = self._rejection_reason(node) + if reason is not None: + self.report_reject([node], f"cortex_m LSTM: {reason}.") + continue + _mark_node_as_quantized( + node, + {node.args[0]: INT8_ACTIVATION_PER_TENSOR_QSPEC}, + None, + is_quantized=True, + ) + for user in node.users: + if ( + user.target == operator.getitem + and user.args[1] == 0 + and not is_annotated(user) + ): + _mark_node_as_quantized( + user, {}, INT8_ACTIVATION_PER_TENSOR_QSPEC, is_quantized=True + ) + self.report_accept([node]) + + def validate(self, model: GraphModule) -> bool: # type: ignore[override] + return True + + def mark_node_as_annotated( node, input_qspec_map, @@ -71,6 +180,7 @@ def __init__(self) -> None: node_finder=GlobalNodeFinder(), pattern_matcher=pattern_matcher, ), + LstmBoundaryQuantizer(), SharedQspecQuantizer(), ] super().__init__(quantizers) diff --git a/backends/cortex_m/test/ops/test_lstm.py b/backends/cortex_m/test/ops/test_lstm.py index a7328496018..0e937e8dec3 100644 --- a/backends/cortex_m/test/ops/test_lstm.py +++ b/backends/cortex_m/test/ops/test_lstm.py @@ -1,4 +1,5 @@ -# Copyright 2025-2026 Arm Limited and/or its affiliates. +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -6,6 +7,7 @@ import pytest import torch +from executorch.backends.arm.test.common import parametrize from executorch.backends.cortex_m.test.tester import ( CortexMTester, McuTestCase, @@ -14,57 +16,27 @@ class CortexMLSTM(torch.nn.Module): + # A single-layer unidirectional nn.LSTM is preserved as one aten.lstm.input + # node, boundary-quantized, and fused into cortex_m.quantized_lstm. ops_before_transforms = { + "executorch_exir_dialects_edge__ops_aten_lstm_input": 1, + "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 2, + "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 2, "executorch_exir_dialects_edge__ops_aten_full_default": 2, - "executorch_exir_dialects_edge__ops_aten_squeeze_copy_dims": 4, - "executorch_exir_dialects_edge__ops_aten_unsqueeze_copy_default": 2, - "executorch_exir_dialects_edge__ops_aten_view_copy_default": 6, - "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 3, - "executorch_exir_dialects_edge__ops_aten_addmm_default": 3, - "executorch_exir_dialects_edge__ops_aten_slice_copy_Tensor": 2, - "executorch_exir_dialects_edge__ops_aten_add_Tensor": 4, - "executorch_exir_dialects_edge__ops_aten_split_with_sizes_copy_default": 2, - "executorch_exir_dialects_edge__ops_aten_sigmoid_default": 6, - "executorch_exir_dialects_edge__ops_aten_tanh_default": 4, - "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 6, - "executorch_exir_dialects_edge__ops_aten_cat_default": 1, } - ops_after_transforms: dict[str, int] = {} - - def __init__(self, input_size: int = 4, hidden_size: int = 3) -> None: - super().__init__() - self.lstm = torch.nn.LSTM(input_size=input_size, hidden_size=hidden_size) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - y, _ = self.lstm(x) - return y - - -class CortexMQuantizableLSTM(torch.nn.Module): - ops_before_transforms = { - "executorch_exir_dialects_edge__ops_aten_add_Tensor": 4, - "executorch_exir_dialects_edge__ops_aten_addmm_default": 4, - "executorch_exir_dialects_edge__ops_aten_cat_default": 1, - "executorch_exir_dialects_edge__ops_aten_full_default": 1, - "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 6, - "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 4, - "executorch_exir_dialects_edge__ops_aten_select_copy_int": 2, - "executorch_exir_dialects_edge__ops_aten_sigmoid_default": 6, - "executorch_exir_dialects_edge__ops_aten_split_with_sizes_copy_default": 2, - "executorch_exir_dialects_edge__ops_aten_squeeze_copy_dims": 1, - "executorch_exir_dialects_edge__ops_aten_tanh_default": 4, - "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, - "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 34, - "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 27, + ops_after_transforms = { + "executorch_exir_dialects_edge__ops_cortex_m_quantized_lstm_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1, } - ops_after_transforms: dict[str, int] = {} - - def __init__(self, input_size: int = 4, hidden_size: int = 3) -> None: + def __init__( + self, input_size: int = 4, hidden_size: int = 3, batch_first: bool = False + ) -> None: super().__init__() - self.lstm = torch.ao.nn.quantizable.LSTM( - input_size=input_size, hidden_size=hidden_size + self.lstm = torch.nn.LSTM( + input_size=input_size, hidden_size=hidden_size, batch_first=batch_first ) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -77,26 +49,81 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: model=CortexMLSTM(), example_inputs=(ramp_tensor(-1, 1, (2, 1, 4)),), ), - "lstm_quantizable": McuTestCase( - model=CortexMQuantizableLSTM(), - example_inputs=(ramp_tensor(-1, 1, (2, 1, 4)),), + # Wider hidden size exercises the CMSIS MVE main loop (hidden % 4 == 0). + "lstm_fp32_wide": McuTestCase( + model=CortexMLSTM(input_size=8, hidden_size=16), + example_inputs=(ramp_tensor(-1, 1, (3, 1, 8)),), + ), + # batch_first -> CMSIS batch-major (time_major=0) with the [B, T, F] layout. + "lstm_batch_first": McuTestCase( + model=CortexMLSTM(batch_first=True), + example_inputs=(ramp_tensor(-1, 1, (1, 3, 4)),), + ), + # batch > 1 exercises the per-batch state/scratch handling. + "lstm_batched": McuTestCase( + model=CortexMLSTM(), + example_inputs=(ramp_tensor(-1, 1, (2, 2, 4)),), ), } -@pytest.mark.skip("Not implemented yet.") +@parametrize("test_case", test_cases) def test_dialect_lstm(test_case: McuTestCase, cortex_m_target) -> None: tester = CortexMTester( test_case.model, test_case.example_inputs, target_config=cortex_m_target ) tester.test_dialect( - test_case.model.ops_before_transforms, test_case.model.ops_after_transforms + test_case.model.ops_before_transforms, + test_case.model.ops_after_transforms, + atol=0.05, + qtol=1, ) -@pytest.mark.skip("Not implemented yet.") +@parametrize("test_case", test_cases) def test_implementation_lstm(test_case: McuTestCase, cortex_m_target) -> None: tester = CortexMTester( test_case.model, test_case.example_inputs, target_config=cortex_m_target ) - tester.test_implementation() + # The reference impl matches CMSIS-NN's arm_lstm_unidirectional_s8 within a + # LSB (its int16 activation LUT vs the float reference), so allow one int8 + # step of slack on top of the quantization tolerance. + tester.test_implementation(atol=0.1, qtol=1) + + +class MultiLayerLSTM(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.lstm = torch.nn.LSTM(input_size=4, hidden_size=3, num_layers=2) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + y, _ = self.lstm(x) + return y + + +class StateReturningLSTM(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.lstm = torch.nn.LSTM(input_size=4, hidden_size=3) + + def forward(self, x: torch.Tensor): + y, (h_n, _) = self.lstm(x) + return y, h_n + + +unsupported_cases = { + "multi_layer": MultiLayerLSTM, + "uses_h_n": StateReturningLSTM, +} + + +@parametrize("model_cls", unsupported_cases) +def test_unsupported_lstm_fails_clearly(model_cls, cortex_m_target) -> None: + # Unsupported configurations are not fused; because aten.lstm.input is + # preserved and has no portable kernel, lowering must raise a clear error + # rather than silently emitting a broken graph. + tester = CortexMTester( + model_cls(), (ramp_tensor(-1, 1, (2, 1, 4)),), target_config=cortex_m_target + ) + with pytest.raises(Exception, match="AtenToCortexMPass"): + tester.quantize().export().to_edge().run_passes() diff --git a/backends/cortex_m/test/test_lstm_params.py b/backends/cortex_m/test/test_lstm_params.py new file mode 100644 index 00000000000..d39d893d448 --- /dev/null +++ b/backends/cortex_m/test/test_lstm_params.py @@ -0,0 +1,103 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Unit tests for the CMSIS-NN LSTM AoT math: the derived parameters + the +int8 reference implementation must track a float torch.nn.LSTM within the +quantization noise floor.""" + +from typing import cast + +import torch +import torch.nn as nn +from executorch.backends.arm.test.common import parametrize +from executorch.backends.cortex_m.passes.lstm_params import ( + choose_cell_scale_power, + derive_lstm_params, + quantized_lstm_reference, +) + + +def _affine_int8_qparams(t: torch.Tensor) -> tuple[float, int]: + lo = min(0.0, float(t.min())) + hi = max(0.0, float(t.max())) + scale = (hi - lo) / 255.0 if hi > lo else 1e-6 + zp = max(-128, min(127, int(round(-128 - lo / scale)))) + return scale, zp + + +def _lstm_weights( + lstm: nn.LSTM, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # nn.Module.__getattr__ types these as Tensor | Module; they are tensors. + weight_ih = cast(torch.Tensor, lstm.weight_ih_l0) + weight_hh = cast(torch.Tensor, lstm.weight_hh_l0) + bias = cast(torch.Tensor, lstm.bias_ih_l0) + cast(torch.Tensor, lstm.bias_hh_l0) + return weight_ih, weight_hh, bias + + +def _float_cell_max_abs(x: torch.Tensor, lstm: nn.LSTM) -> float: + h_size = lstm.hidden_size + weight_ih, weight_hh, layer_bias = _lstm_weights(lstm) + w_ih, w_hh, bias = weight_ih.detach(), weight_hh.detach(), layer_bias.detach() + time_steps, batch = x.shape[0], x.shape[1] + max_abs = 0.0 + for b in range(batch): + h = torch.zeros(h_size) + c = torch.zeros(h_size) + for t in range(time_steps): + g = x[t, b] @ w_ih.T + h @ w_hh.T + bias + i, f = torch.sigmoid(g[0:h_size]), torch.sigmoid(g[h_size : 2 * h_size]) + gg, o = torch.tanh(g[2 * h_size : 3 * h_size]), torch.sigmoid( + g[3 * h_size :] + ) + c = f * c + i * gg + h = o * torch.tanh(c) + max_abs = max(max_abs, float(c.abs().max())) + return max_abs + + +def _sqnr_db(ref: torch.Tensor, test: torch.Tensor) -> float: + noise = (ref - test).pow(2).mean() + if noise == 0: + return float("inf") + return float(10 * torch.log10(ref.pow(2).mean() / noise)) + + +shapes = { + "small": (4, 3, 8), + "wide": (8, 16, 12), + "single_step": (4, 5, 1), +} + + +@parametrize("shape", shapes) +def test_reference_tracks_float_lstm(shape: tuple[int, int, int]) -> None: + input_size, hidden_size, time_steps = shape + lstm = nn.LSTM(input_size, hidden_size).eval() + x = torch.randn(time_steps, 1, input_size) + with torch.no_grad(): + y_float, _ = lstm(x) + + input_scale, input_zp = _affine_int8_qparams(x) + output_scale, output_zp = _affine_int8_qparams(y_float) + cell_scale_power = choose_cell_scale_power(_float_cell_max_abs(x, lstm)) + + weight_ih, weight_hh, bias = _lstm_weights(lstm) + params = derive_lstm_params( + weight_ih.detach(), + weight_hh.detach(), + bias.detach(), + input_scale, + input_zp, + output_scale, + output_zp, + cell_scale_power, + ) + y_ref = quantized_lstm_reference( + x, params, input_scale, output_scale, dequantize=True + ) + + assert _sqnr_db(y_float, y_ref) > 30.0 diff --git a/backends/cortex_m/test/tester.py b/backends/cortex_m/test/tester.py index 1f7d7f3059d..5b1201c83be 100644 --- a/backends/cortex_m/test/tester.py +++ b/backends/cortex_m/test/tester.py @@ -37,6 +37,7 @@ class CortexMToEdge(ToEdge): def __init__(self): config = EdgeCompileConfig( preserve_ops=[ + torch.ops.aten.lstm.input, torch.ops.aten.linear.default, torch.ops.aten.hardsigmoid.default, torch.ops.aten.hardsigmoid_.default,