TensorFlow unveils MLIR for faster machine learning

Sublanguage project promises faster compilation and easier hardware optimization for high-performance machine learning models

TensorFlow unveils MLIR for faster machine learning
GrandeDuc / Getty Images

Engineers working on Google’s TensorFlow machine learning framework have revealed a subproject, MLIR, that is intended to be a common intermediate language for machine learning frameworks.

MLIR, short for Multi-Level Intermediate Representation, will allow projects using TensorFlow and other machine learning libraries to be compiled to more efficient code that takes maximum advantage of underlying hardware. What’s more, MLIR could in time be used by compilers generally, extending its optimization benefits beyond machine learning projects.

MLIR isn’t a language like C++ or Python. It represents an intermediate compilation step between those higher-level languages and machine code. The compiler framework LLVM uses an intermediate representation, or IR, of its own. One of LLVM’s originators, Chris Lattner, is a co-creator of MLIR. Making MLIR an LLVM co-project could be a way to spread its adoption.

In a slide presentation at the EuroLLVM conference earlier this month, Lattner and fellow Googler Tatiana Shpeisman explained how TensorFlow already generates multiple IRs internally, but that these disparate IRs don’t benefit from one another. MLIR provides a single, standard IR for all of those TensorFlow subsystems. TensorFlow is currently migrating to use MLIR internally.

Another benefit MLIR may provide is parallelized compilation. MLIR is designed to allow a compiler to work on different segments of code in parallel, allowing machine learning models—and other sorts of applications—to be pushed to production more quickly.

MLIR could provide other benefits to languages and frameworks outside machine learning. For example, LLVM-based languages like Swift and Rust have had to develop their own internal IRs, because many optimizations used in those languages can’t be expressed in LLVM. MLIR could provide a standard way to express those optimizations, which could in turn be re-used for other languages.

The MLIR project is open source. An official specification is available for those who want to generate MLIR.