Demonstrating a Future for MLIR-native DSL Compilers on a NumPy-like Example

TL;DR

This paper introduces an MLIR-native NumPy-like DSL for tensor kernels, featuring a new dialect-agnostic type checker and a parallel-first lowering scheme, demonstrating effective performance in weather modeling and CFD applications.

Contribution

It presents the first open-source MLIR-native NumPy-like DSL with all frontend actions within MLIR and a novel dialect-agnostic type checker, advancing DSL compiler frameworks.

Findings

Effective performance in weather modeling use cases

Seamless offloading via MLIR dataflow dialect

All frontend actions implemented within MLIR

Abstract

Compilers for general-purpose languages have been shown to be at a disadvantage when it comes to specialized application domains as opposed to their Domain-Specific Language (DSL) counterparts. However, the field of DSL compilers features little consolidation in terms of compiler frameworks and adjacent software ecosystems. As a result, considerable work is duplicated, lost to maintenance issues, or remains undiscovered, and most DSLs are never considered "production-ready". One notable development is the introduction of the Multi-Level Intermediate Representation (MLIR), which promises a similar impact on DSL compilers as LLVM had on general-purpose tooling. In this work, we present a NumPy-like DSL made for offloading numeric tensor kernels that is entirely MLIR-native. In a first for open-source, it implements all frontend actions and semantic analyses directly within MLIR. Most notably, this is made possible by our new dialect-agnostic MLIR type checker, created for the future of DSLs in MLIR. We implement a simple, yet effective, parallel-first lowering scheme that connects our language to another MLIR dataflow dialect for seamless offloading. We show that our approach performs well in real-world use cases from the domain of weather modeling and Computational Fluid Dynamics (CFD) in Fortran.