Optimized Convolution Kernels for Edge AI

• 

  ai inference

Inspiration

With the increasing demand for on-device AI inference, optimizing computational performance on edge devices like the Samsung Galaxy S24 Ultra is crucial. Traditional convolution implementations struggle with latency and efficiency, making real-time AI applications impractical. We aimed to develop highly optimized convolution operators that significantly reduce inference time while maintaining numerical accuracy.

What it does

Our project introduces two optimized convolution methods—Direct Block Processing Convolution and Optimized FFT Convolution—tailored for different filter sizes and sequence lengths. By selecting the best approach dynamically, we maximize computational efficiency, enabling faster and more efficient deep learning inference on mobile devices.

How we built it

Comparison of Two Custom C++ Kernels for Convolution Optimization

To maximize efficiency for edge devices, we developed and compared two custom C++ convolution kernels, optimizing them to reduce computational latency and memory footprint.

Kernel 1: Direct Block Processing Convolution

Approach

This kernel uses a block-based overlap-save method combined with distributed arithmetic, making it highly efficient for small filter sizes and short sequence lengths.

Optimizations Implemented

• Bit-level distributed arithmetic: Breaks down multiplications into precomputed bitwise operations for efficiency.
• Sparse Toeplitz matrix multiplication: Reduces redundant computations.
• Recursive partitioning: Divides computation into hierarchical blocks for better memory locality.
• Memory-aware execution: Minimizes cache misses by optimizing memory access patterns.

Kernel 2: Fused FFT-Based Convolution

Approach

This kernel accelerates convolution using the Fast Fourier Transform (FFT), leveraging the frequency domain for efficient multiplication.

Optimizations Implemented

• Precomputed FFT of filters: Avoids redundant FFT computations by storing transformed filters.
• Overlap-save method: Handles boundaries efficiently, reducing memory overhead.
• In-place operations: Minimizes memory allocation and improves cache locality.
• Thread-level parallelism (OpenMP): Utilizes multiple CPU cores for FFT and inverse FFT computation.
• SIMD vectorization: Uses AVX2/FMA instructions for element-wise multiplications in the frequency domain.

Selection Criteria: When to Use Each Kernel

We dynamically select the optimal method based on:

• Filter size and sequence length: Large filters benefit from FFT, while small filters are faster with direct block processing.
• Memory constraints: FFT-based methods require additional memory, making direct methods preferable for lower-end devices.
• Parallel execution capabilities: FFT benefits from multi-core parallelism, while direct methods are more efficient for lower-power CPUs.

Challenges we ran into

• System limitations for Executorch build: My system did not allow for an Executorch build, preventing me from generating .pte files for each use case. Instead, I have uploaded the source code for the optimized operator kernels so they can be integrated and tested on a compatible setup.
  

Accomplishments that we're proud of

• Developed a selection strategy to choose the optimal convolution method based on input characteristics.
  
• Created a robust Executorch-compatible CPP kernels for edge AI inference to reduce through an optimized inference.
  

Submission Drive Link: https://drive.google.com/drive/folders/1OQrYiNcbVAXLIiSEGwfu_WeDb9cZnKy5?usp=sharing

Built With

• bash
• c++
• cmake
• executorch
• ninja
• openmp
• python
• pytorch
• simd