Flexible Multi-Dimensional FFTs for Plane Wave Density Functional Theory Codes

TL;DR

This paper introduces FFTB, a flexible distributed framework for multi-dimensional Fourier transforms tailored for complex data structures in plane wave Density Functional Theory, improving efficiency and scalability on supercomputers.

Contribution

We present FFTB, a versatile framework supporting regular and irregular data grids, batched transforms, and CPU/GPU platforms, specifically designed for plane wave DFT applications.

Findings

Improved execution time on HP Cray EX supercomputer

Supports both regular and irregular data grids

Scalable performance for large-scale transforms

Abstract

Multi-dimensional Fourier transforms are key mathematical building blocks that appear in a wide range of applications from materials science, physics, chemistry and even machine learning. Over the past years, a multitude of software packages targeting distributed multi-dimensional Fourier transforms have been developed. Most variants attempt to offer efficient implementations for single transforms applied on data mapped onto rectangular grids. However, not all scientific applications conform to this pattern, i.e. plane wave Density Functional Theory codes require multi-dimensional Fourier transforms applied on data represented as batches of spheres. Typically, the implementations for this use case are hand-coded and tailored for the requirements of each application. In this work, we present the Fastest Fourier Transform from Berkeley (FFTB) a distributed framework that offers flexible implementations for both regular/non-regular data grids and batched/non-batched transforms. We provide a flexible implementations with a user-friendly API that captures most of the use cases. Furthermore, we provide implementations for both CPU and GPU platforms, showing that our approach offers improved execution time and scalability on the HP Cray EX supercomputer. In addition, we outline the need for flexible implementations for different use cases of the software package.