How AI on Microcontrollers Works: Operators and Kernels

The buzz around “edge AI”, which means something slightly different to almost everyone you talk to, is well past reaching a fever pitch. Regardless of what edge AI means to you, the one commonality is typically that the hardware on which inference is being performed is constrained in one or more dimensions, whether it be compute, memory, or network bandwidth. Perhaps the most constrained of these platforms are microcontrollers.

I have found that, while there is much discourse around “running AI” (i.e. performing inference) on microcontrollers, there is a general lack of information about what these systems are actually capable of, and how new hardware advancements impact that equation. It is my hope with this series to peel back some of the layers of terminology and explore what actually happens between supplying inputs to a model and receiving outputs. Along the way, we’ll ground our exploration in performing inference with real models on real constrained hardware.

While “weights” get the majority of the attention with AI models, they alone are not sufficient for performing inference. Depending on how a model is distributed and what runtime is used, additional data or metadata may be supplied alongside the model, or may be defined explicitly in software that interacts with the weights. The most popular runtime for microcontrollers is Tensorflow Lite for Microcontrollers (tflite-micro), which is an optimized version of Tensorflow Lite.

Note: Google recently rebranded Tensorflow Lite to LiteRT, and tflite-micro to LiteRT for Microcontrollers.

tflite-micro uses the .tflite file format, which encodes data using FlatBuffers. Unlike some other model file formats, .tflite files include not only the tensors that encapuslate model weights, but also the computation graph, which informs the runtime of what operations to use when performing inference. In order to do so, there needs to be a defined set of operators. This is somewhat analagous to instructions defined in an instruction set architecture (ISA) for a processor. With an ISA, a compiler will take a higher level programming language and map the behavior onto instructions available in the ISA. Tensorflow supports an extensive set of built-in operators, while Tensorflow Lite, and thus tflite-micro, supports only a subset.

Continuing the analaogy, many processors implement specific versions of the ARM architecture, but that doesn’t mean that processors implementing the same ISA are equivalent. Every instruction that is supported has to be implemented in hardware, and decisions about how the processor is designed can impact performance on multiple dimensions. Similarly, while Tensorflow Lite defines a set of operators, the implementation of those operators, which are referred to as kernels, may vary. Kernels are implemented in software, but depending on the underlying hardware, a kernel might take many instructions to execute, or may be able to optimized to leverage dedicated hardware support.

A simple example is the addition operator (TFL::AddOp). We’ll cover how operators and kernels are registered and invoked in a future post, but let’s start by taking a look at the default tflite-micro addition operator logic.

tensorflow/lite/micro/kernels/add.cc

As can be observed in AddEval(), the type of output we are expecting influences the implementation of the operator. To illustrate how the underlying hardware impacts performance, let’s focus on the case in which we expect kTfLiteInt8 (signed 8-bit integer) or kTfLiteInt16 (signed 16-bit integer) output, meaning that we’ll call EvalAddQuantized().

tensorflow/lite/micro/kernels/add.cc

For kTfLiteInt8 output when broadcast is not required, we make a call to reference_integer_ops::Add().

“Broadcasting” is the process of making arrays to have compatible shapes for arithmetic operations. For example, matrix addition requires that the two input matrices have the same dimensions.

tensorflow/lite/kernels/internal/reference/add.h

As you might expect, this implementation effectively boils down to iterating through the two input tensors and calling input1_data[i] + input2_data[i]. This can be thought of as a lowest common denominator implementation in that it doesn’t leverage any hardware-specific functionality; any processor can perform sequential addition. However, as evidenced by the effectively unlimited demand in the Graphics Processing Unit (GPU) market, there are significant performance gains to be had by parallelizing operations in hardware. Fortunately, many of the operations that are necessary for performing inference are “embarassingly parallel”. For example, rather than iterating through tensors to perform sequential addition, which may take many processor cycles, we could point a processor to the two inputs and, if supported by the hardware, the entire matrix addition operation could be completed in a “single” cycle.

It is unlikely that the operation would literally take one cycle given the complexity of modern processors, but the point is that the runtime of the entire operation could be reduced to the same order of magnitude of cycles as one step of the sequential addition implementation.

Obviously microcontrollers don’t have massive GPUs like those installed in cloud provider datacenters. However, many do implement architecture extensions that enable these common operations to be accelerated. Because Tensorflow Lite enables different kernel implementations, this hardware acceleration can be leveraged when supported.

Many microcontrollers implement Arm Cortex-M cores. For example, chips like the Raspberry Pi RP2350 and Nordic Semiconductor nRF54H20 implement multiple Arm Cortex-M33 cores. The former implements the Armv8-M Digital Signal Processing (DSP) Extension, which adds support for Single Instruction Multiple Data (SIMD) instructions. More capable chips, like the Alif Ensemble E3 implement Cortex-M55 cores, with support for the Armv8-M Vector Extension (MVE), also referred to as Arm Helium. The E3 also includes dedicated accelerators in the form of Arm’s Ethos-U Neural Processing Units (NPU).

Arm provides software that allows for hardware that supports one or more of these extensions to accelerate Tensorflow Lite kernel implementations. For example, the CMSIS-NN library offers kernel implementations that do not leverage optimization (i.e. pure C), leverage just the DSP extension, or leverage the MVE extension (which requires the implementation of the DSP extension). Tensorflow Lite has kernel “ports” that integrate the CMSIS-NN functionality. Let’s take a look at how the add operation differs when using CMSIS-NN kernels.

The setup looks largely the same as the first add operation kernel we observed. However, when we reach EvalAddQuantizedInt8(), we can start to see where hardware acceleration is leveraged.

tensorflow/lite/micro/kernels/cmsis_nn/add.cc

The arm_elementwise_add_s8() function is provided by CMSIS-NN, and the implementation leverages different hardware functionality depending on what extensions are available.

Source/BasicMathFunctions/arm_elementwise_add_s8.c

For example, if the DSP extension is present, the parallel signed 16-bit addition (SAAD16) instruction provided by the extension is used to reduce the number of loops necessary by packing 8-bit signed integers into 16-bit arguments, then calculating 4 outputs on a single iteration. If the MVE extension is present, vector addition instructions (VADD) can be used directly, making the calculation even more efficient.

These optimizations are made available via configuration when compiling tflite-micro. They can be applied to any models that utilize the operations, without the need to modify the model when moving from one architecture to another. Some optimizations do require modifying a model. For example, when using microcontrollers, such as the previously mentioned Alif Ensemble E3, that include Arm’s Ethos-U NPUs you can run .tflite models through the Vela compiler. Converted models replace sequences of built-in operators with a custom ETHOSU operator and a command stream. The application processor notifies the NPU of the address of the command stream and other relevant data, then triggers it to perform inference.

Unlike the addition operator, there is no fallback kernel implementation; models converted via the Vela compiler cannot run on microcontrollers that do not have Ethos-U NPUs. For those that do, we can see the previously described logic in the Eval() implementation for the ETHOSU custom operator.

tensorflow/lite/micro/kernels/ethos_u/ethosu.cc

We’ve now seen the full spectrum of operator optimization, from kernels that are implemented purely in C, to those that leverage hardware instructions provided in architecture extensions, and finally to those that offload inference to a wholly separate processor. In future posts, we’ll explore how operators are encoded in .tflite files, and how the runtime ultimately invokes the underlying kernels.