[Bug]: arm_lstm_unidirectional_s8: gate FC accumulation misses int16 saturation, diverges from TFLM

[Brivost]

Is there an existing issue for this?

• I have searched the existing issues

Problem statement

It seems like there is a bit-exactness gap between arm_lstm_unidirectional_s8 and the TFLM reference LSTM kernel.

arm_nn_lstm_calculate_gate_s8_s16 zeroes the gate buffer, then calls arm_nn_vec_mat_mul_result_acc_s8_s16 twice into it: input FC, then recurrent FC (L60-L87). The kernel accumulates before clamping:

res00 = arm_nn_requantize(res00, dst_multiplier, dst_shift);
res00 += (int32_t)*dst;
res00 = CLAMP(res00, NN_Q15_MAX, NN_Q15_MIN);

On the first call *dst is 0, so that clamp lands on the input FC alone. On the second call *dst already holds the clamped input FC, so the clamp lands on the sum.

TFLM runs the two FCs into separate buffers, clamps each to int16 (the gate tensor is int16, so quantized_activation_min/max are the int16 limits), then adds with saturation: lstm_eval.h L275-L297, per-FC clamp, saturating add.

So:

CMSIS-NN: clamp16(clamp16(fc_input) + fc_recurrent)
TFLM: clamp16(clamp16(fc_input) + clamp16(fc_recurrent))

This is an example scenario of how this could cause a divergence:

fc_input	fc_recurrent	CMSIS-NN	TFLM
10000	40000	clampf16(50000) = 32767	clampf16(10000 + 32767) = 32767	same
-20000	40000	clampf16(20000) = 20000	clampf16(-20000 + 32767) = 12767	different

It seems a potential fix could be clamping the requantized result before adding it to *dst.

Steps To Reproduce

I've attached a minimal C example comparing TFLM and CMSIS-NN for the mismatching case. I built it with this command:

gcc -Wall \
  -I/CMSIS-NN/Include \
  -I/CMSIS-CORE/CMSIS/Core/Include \
  -o repro lstm_gate_saturation_repro.c \
  /libcmsis-nn.a

and am getting this output:

input FC     : -20000  (fits int16)
recurrent FC :  40000  (overflows int16, max 32767)

gate pre-activation  TFLM  12767   CMSIS-NN  20000   diff 7233
gate after sigmoid   TFLM  31378   CMSIS-NN  32521   diff 1143

MISMATCH

lstm_gate_saturation_repro.c