High-Performance FFT Code Generation via MLIR Linalg Dialect and SIMD Micro-Kernels
Yifei He, Stefano Markidis · 2024
Fast Fourier Transform (FFT) libraries are an indispensable and critical component of any High-Performance Computing (HPC) software stack. They are used in many applications, from Partial Differential Equation (PDE) solvers to signal spectral analysis and deep-learning. To design and develop the next generation of HPC FFT libraries, it is essential to leverage modern compiler infrastructures, such as Multi-Level Intermediate Representation (MLIR) and the Low-Level Virtual Machine (LLVM), alongside advanced computer architecture features, including SIMD (Single Instruction, Multiple Data) instructions. In this work, we introduce FFTc 2.0, an MLIR-based domain-specific compilation framework for FFT. We extend the MLIR Linalg dialect with FFT-specific operations to harness high-level tensor-based abstractions for FFT. These abstractions are ideal for formulating various FFT algorithms and facilitating formula rewriting for FFT decomposition and cache-friendly optimization. We employ micro-kernels to increase the performance, particularly in response to the limited support for complex arithmetic in MLIR and the general-purpose compiler LLVM. These micro-kernels are designed to implement small-size FFT kernels for integration into larger FFTs. They utilize SIMD-friendly data layouts for complex arithmetic and feature an optimized memory access pattern, enabling performance enhancements that are not achievable with standard compiler implementations. Our method achieves performance levels comparable to or surpass those of widely used FFT HPC libraries, such as FFTW, FFTE, and Spiral FFT. Additionally, it establishes a robust software infrastructure that facilitates further optimizations and supports additional hardware backends.