Code generation for in-place stencils

Mohamed Essadki (1) , Bertrand Michel (1) , Bruno Maugars (1) , Oleksandr Zinenko (2) , Nicolas Vasilache (3) , Albert Cohen (2)

1 DAAA, ONERA, Université Paris-Saclay [Châtillon]
2 Google Research [Paris]
3 Google Switzerland

Numerical simulation often resorts to iterative in-place stencils such as the Gauss-Seidel or Successive Overrelaxation (SOR) methods. Writing high performance implementations of such stencils requires significant effort and time; it also involves non-local transformations beyond the stencil kernel itself. While automated code generation is a mature technology for image processing stencils, convolutions and out-of-place iterative stencils (such as the Jacobi method), the optimization of in-place stencils requires manual craftsmanship. Building on recent advances in tensor compiler construction, we propose the first domain-specific code generator for iterative in-place stencils. Starting from a generic tensor compiler implemented in the MLIR framework, tensor abstractions are incrementally refined and lowered down to parallel, tiled, fused and vectorized code. We used our generator to implement a realistic, implicit solver for structured meshes, and demonstrate results competitive with an industrial computational fluid dynamics framework. We also compare with stand-alone stencil kernels for dense tensors.

Mots clés: Computational Fluid Dynamics; Implicit Meth- ods; Gauss-Seidel; SOR; Iterative In-place Stencil; Domain-Specific; MLIR; Vectorization; Tiling; Code Generation

Mots clés: GENERATION CODE; STENCIL; CFD; METHODE GAUSS-SEIDEL; PARALLELISME; VECTORISATION; BLOCAGE CACHE

Domaines: Physique [physics]; Sciences de l'ingénieur [physics]

Mohamed Essadki, Bertrand Michel, Bruno Maugars, Oleksandr Zinenko, Nicolas Vasilache, et al.. Code generation for in-place stencils. CGO '23: 21st ACM/IEEE International Symposium on Code Generation and Optimization, Feb 2023, Montreal, Canada. pp.2-13, ⟨10.1145/3579990.3580006⟩. ⟨hal-04116904⟩