Fast Evaluation of Finite Element Weak Forms Using Python Tensor Contraction Packages

TL;DR

This paper introduces a novel Python-based approach for efficiently evaluating finite element weak forms by translating Einstein summation expressions into optimized tensor contractions, improving performance and automation.

Contribution

It presents a new transpiler that converts Einstein summation expressions into tensor contractions for finite element forms, enabling compatibility with high-performance Python packages.

Findings

Tensor contraction-based evaluation is faster than nested loops.

Memory layout and contraction path optimization significantly impact performance.

Transpiled Python implementations outperform C-based functions in SfePy.

Abstract

In finite element calculations, the integral forms are usually evaluated using nested loops over elements, and over quadrature points. Many such forms (e.g. linear or multi-linear) can be expressed in a compact way, without the explicit loops, using a single tensor contraction expression by employing the Einstein summation convention. To automate this process and leverage existing high performance codes, we first introduce a notation allowing trivial differentiation of multi-linear finite element forms. Based on that we propose and describe a new transpiler from Einstein summation based expressions, augmented to allow defining multi-linear finite element weak forms, to regular tensor contraction expressions. The resulting expressions are compatible with a number of Python scientific computing packages, that implement, optimize and in some cases parallelize the general tensor contractions. We assess the performance of those packages, as well as the influence of operand memory layouts and tensor contraction paths optimizations on the elapsed time and memory requirements of the finite element form evaluations. We also compare the efficiency of the transpiled weak form implementations to the C-based functions available in the finite element package SfePy.