MMStencil: Optimizing High-order Stencils on Multicore CPU using Matrix Unit

TL;DR

MMStencil leverages matrix units, SIMD, and memory optimizations to significantly accelerate high-order 3D stencil computations on multicore CPUs, outperforming existing libraries and benefiting real-world HPC applications.

Contribution

This paper introduces a novel matrix-based acceleration approach for 3D high-order stencils on multicore CPUs, including algorithmic, memory, and parallelism optimizations.

Findings

Achieves up to 2.1x speedup over state-of-the-art libraries on Nvidia A100.

Enables 1.8x speedup in real-world HPC applications compared to optimized GPU versions.

Maintains high hardware utilization across diverse stencil shapes and dimensions.

Abstract

Matrix-accelerated stencil computation is a hot research topic, yet its application to three-dimensional (3D) high-order stencils and HPC remains underexplored. With the emergence of matrix units on multicore CPUs, we analyze matrix-based acceleration strategies and tailor an optimal approach for 3D high-order stencils. We introduce algorithmic optimizations based on SIMD and matrix units to address strided memory accesses, alignment conflicts, and redundant accesses. We propose memory optimizations to boost on-package memory efficiency, and a novel multi-thread parallelism paradigm to overcome data-sharing challenges caused by the absence of shared data caches. MMStencil sustains consistently high hardware utilization across diverse stencil shapes and dimensions. Our DMA-based inter-NUMA communication further mitigates NUMA effects and MPI limitations in hybrid parallelism. Combining all the innovations, MMStencil outperforms state-of-the-art libraries on Nvidia A100 GPGPU by up to 2.1x. Moreover, the performance improvements translate directly to real-world HPC applications and enable RTM applications to yield 1.8x speedup versus a highly optimized industrial Nvidia A100 GPGPU version.