Mass Matrix Assembly on Tensor Cores for Implicit Particle-In-Cell Methods

TL;DR

This paper reformulates mass matrix assembly in implicit particle methods as matrix multiplications suitable for tensor cores, achieving significant performance improvements on modern hardware.

Contribution

It introduces a general reformulation of mass matrix assembly as matrix products for tensor cores, applicable to various interpolation orders and platforms.

Findings

Up to 3x faster than optimized conventional implementations.

Reduced ECSIM runtime by 15%.

Applicable to scalar and tensorial mass matrices in particle-in-cell methods.

Abstract

Matrix-multiply-accumulate (MMA) units, or tensor cores, are now widespread across modern computing architectures. Yet, their use for particle-grid operators remains limited. In implicit particle methods, mass-matrix assembly is a reduction-dominated kernel in which weighted outer products of interpolation weights are accumulated over particle support. We show that this operation can be reformulated exactly, cell by cell, as a sequence of matrix products matched to hardware MMA tiles. The formulation is general with respect to interpolation order and hardware platform, and applies to both scalar mass matrices and the tensorial block mass matrix arising in implicit in the Energy-Conserving Semi-Implicit Method (ECSIM) for Particle-in-Cell simulations. We introduce particle batching and a support-group decomposition for higher-order shape functions whose stencil extends beyond a single cell, specialize the method to first- and second-order B-spline interpolation, and implement it on NVIDIA tensor cores. The resulting kernels achieve up to 3x over optimized conventional implementations and reduce end-to-end ECSIM runtime by 15%.