PrismStencil: Transforming Stencil Computations into Matrix Multiplications

Abstract

Stencil computations are a fundamental class of workloads in scientific and engineering applications. Although Tensor Cores offer new opportunities for stencil acceleration, existing tensorization methods suffer from substantial redundant computation and insufficient data reuse. We propose PrismStencil, a stencil tensorization algorithm that introduces a novel mapping scheme between matrices and stencil grid points. PrismStencil employs multidimensional vector mapping to accommodate the multidimensional data dependencies of stencil computations and a unified gather--scatter tensorization scheme to exploit the relative sizes of the matrix dimensions in Tensor Core operations, thereby reducing the redundant computation introduced by tensorization. Moreover, it introduces a cross-product data-reuse mechanism that exploits the overlap among neighboring stencil updates to further enhance data reuse. Across representative stencil kernels on an NVIDIA A100 GPU, PrismStencil achieves average speedups of 2.13x and 2.17x over state-of-the-art Tensor Core stencil methods with FP64 and TF32 precision, respectively.