MOM: Matrix Operations in MLIR

TL;DR

This paper introduces MOM, a new MLIR dialect for expressing and optimizing dense linear algebra operations, enabling unified and efficient compilation of matrix computations within the MLIR framework.

Contribution

The paper presents MOM, a novel MLIR dialect that captures matrix properties and facilitates end-to-end compilation of dense linear algebra computations.

Findings

Enables high-level algorithmic transformations of matrix operations.

Integrates seamlessly with existing MLIR dialects for optimized code generation.

Supports end-to-end compilation of matrix computations.

Abstract

Modern research in code generators for dense linear algebra computations has shown the ability to produce optimized code with a performance which compares and often exceeds the one of state-of-the-art implementations by domain experts. However, the underlying infrastructure is often developed in isolation making the interconnection of logically combinable systems complicated if not impossible. In this paper, we propose to leverage MLIR as a unifying compiler infrastructure for the optimization of dense linear algebra operations. We propose a new MLIR dialect for expressing linear algebraic computations including matrix properties to enable high-level algorithmic transformations. The integration of this new dialect in MLIR enables end-to-end compilation of matrix computations via conversion to existing lower-level dialects already provided by the framework.