Examples: Optimizing LiteRT (TFLite) kernels

The Coral NPU GitHub repository contains several examples of programming optimized LiteRT (TensorFlow Lite) kernels. These examples run on the RVV vector execution engine of Coral NPU, not the specialized matrix execution engine. Using the vector execution engine maximizes performance, as compared to traditional scalar processor cores.

2D convolution

The conv.cc file in GitHub shows how to optimize a LiteRT kernel using C intrinsics for RISC-V RVV vector operations and inline assembly code. The compiler's register allocator is used, but inline assembly code performs the multiply-accumulate (MAC).

The code starting at line 115 shows a specialized Conv2D_4x4 kernel with LMUL=4:

Sample code

// Specialized 4x4 Conv2D kernel with LMUL=4 for high channel counts (up to 64)
void Conv2D_4x4(const tflite::ConvParams& params,
            const coralnpu_v2::opt::litert_micro::OpDataConvCustom& data,
            const tflite::RuntimeShape& input_shape,
            const int8_t* input_data,
            const tflite::RuntimeShape& filter_shape,
            const int8_t* filter_data, const int32_t* bias_data,
            const tflite::RuntimeShape& output_shape, int8_t* output_data,
            const int8_t* repacked_weights, TfLiteContext* context) {
   const int stride_width = params.stride_width;
   const int stride_height = params.stride_height;
...

A Python test program for conv.cc is shown in cocotb_conv2d.py.

To compile and run the code on the Coral NPU Verilator simulator, use this Bazel command (from the GitHub repository root):

bazel run //tests/cocotb/tutorial/tfmicro:cocotb_conv2d

To run on the MPACT-CoralNPU behavioral simulator, use this Bazel command:

bazel run //tests/cocotb/tutorial/tfmicro:npusim_conv2d

Note: At a high level, convolution is a specialized, constrained form of matrix multiplication. While general matrix multiplication allows every input element to interact with every output element independently, a convolution program enforces local connectivity and weight sharing. Mathematically, any discrete convolution can be converted into a standard matrix multiplication, but programming them reveals major differences in memory patterns, parameter sizes, and hardware optimization.

Depthwise convolution

The depthwise_conv.cc file in GitHub shows another coding example.

Depthwise convolution is natively supported in LiteRT by the built-in operator DEPTHWISE_CONV_2D. It is a highly optimized operation specifically designed to minimize processing power and memory footprint on mobile and edge devices.

Unlike a standard convolution that processes spatial dimensions and all channels simultaneously, a depthwise convolution applies a single, independent filter to each input channel.

A Python test program for depthwise_conv.cc is shown in cocotb_depthwise_conv.py.

To compile and run the code on the Coral NPU Verilator simulator, use this Bazel command (from the GitHub repository root):

bazel run //tests/cocotb/tutorial/tfmicro:depthwise_conv_test

To run on the MPACT-CoralNPU behavioral simulator, use this Bazel command:

bazel run //tests/cocotb/tutorial/tfmicro:npusim_depthwise_conv

---