Pinch: Implementing a borrow-checked language with MLIR

Pinch is a borrow checked language inspired by Rust which is built using LLVM's Multi-level Intermediate Representation infrastructure (MLIR). It is a very simple language, and is mostly just a proof of concept showing the capabilities of MLIR. It was built as a course project for NCSU's CSC 766: compiler optimizations.

// Example program written in Pinch
fn main() {
    let a = 4;
    let b = &mut a;
    print(*b);
}

Borrow checking has become a bit of a buzzword thanks to Rust. It refers to a type of static analysis which ensures no invalid operations occur during any part of a program. Any violations result in a compiler error. MLIR is a framework for creating highly specific IR dialects which can be shared between projects. These dialects can share optimizations, and be lowered into LLVM IR for code generation.

Relevant Links

• The paper
• Source code
• MLIR Project

For the full set of implementation details see the paper. In this post I want to cover some high level ideas, how MLIR helps us achieve them, and my thoughts on MLIR development.

Multi-Level Intermediate Representation (MLIR)

MLIR is one of the sub-projects which is a part of LLVM. Many projects (such as Rust, Tensorflow, Swift, etc) all create a higher level IR which is lowered to LLVM IR. These projects independently implement common optimizations, such as inlining, for their respective IR. MLIR provides a framework for declaring arbitrary IR, and reusing optimizations built for that IR. It is designed to foster code sharing and provide an extra layer on top of LLVM.

If you learn nothing else from this post, just remember that MLIR is absurdly flexible. It can encode almost any type of high level funcionality, and allows for arbitrary analysis of that logic. By lowering your IR into other dialects, you can easily provide huge features with relatively little work. i.e. you can lower to the LLVM IR for traditional code generation, or chose a GPU dialect for generating kernels. The possibilities are endless.

Borrow Checking

Pinch's borrow checker is designed to be very similar to Rust. This decision was purely based on my desire to get a feel for what Rust must be doing under the hood. I did try to see how Rust implements borrow checking, but it was sufficiently different enough from the MLIR model that I basically just went off in my own direction.

Here's a quick and dirty explanation of the type of borrow checking performed by Pinch:

• All data is owned by one named variable (let a = 2)
• Ownership can be transferred by moving a value to a new named variable (let b = a)
  • The value is now only accessible through the destination of the move
• Temporary access to a value can be loaned out as a reference
  • One mutable reference may exist (let c = &mut b)
  • OR multiple shared references may exist (let c = &b)

The atom of MLIR is the operation, which can conceivably wrap any type of functionality. MLIR operations use SSA values, and belong to a dialect which defines them. The dialect specifies optimizations and lowering passes which can be used by any other MLIR project. MLIR operations can also have named attributes attached to them, which is particularly helpful for analysis.

Pinch creates its own dialect which encapsulates the types of data movement and reference mentioned in our ruleset above. It tags all operations with a source and destination, which represent the owner of the data being used or produced. These tags are MLIR attributes which construct a sort of "ownership graph" modeling the flow of values throughout a lexical scope.

The ownership graph attributes look like this:

{ dst = "destination_name", src = "source_name" }

The field names are self-explanatory. We are going to track the names in these fields and compare their usage to our borrow checking rules. These names are based on the lexical scope their corresponding variables were declared in. If the rules are violated, then we can generate an error.

MLIR Generation

Here we can go a little bit off the deep end and show the MLIR operations themselves. Pinch is best explained with an example:

let a = 2;
let b = a;

And here is the Pinch MLIR generated:

%0 = pinch.constant {dst = "a"} 2
%1 = "pinch.move"(%0) {dst = "b", src = "a"} : (ui32) -> ui32

First a new SSA value is initialized from a constant. The destination of this constant operation is marked as ’a’ in the ’dst’ attribute. This signifies that we are creating a new owner. With the "pinch.move" operation we are once again creating a new owner based on the ’dst’ attribute, but we also specify where the moved data originated from with the ’src’ attribute. When the borrow checker is validating this operation it can mark owner ’a’ as non-resident, and any further access of ’a’ should trigger an error.

We can also look at a more complicated scenario where the lifetime of a variable is not tied to a lexical scope:

fn consume_box(a: Box) -> u32 { return *a;
    return *a;
}
fn main() {
    let a = box(2);
    let b = consume_box(a); let c = *a;
}

This program introduces us to the Box type. The box is a heap-allocated integer, and Pinch needs to free the box when its lifetime ends. This example demonstrates how sometimes the lifetime of owned data is not the lexical scope that data was created in. We create a box, and move it into the ’consume box’ function when we pass it by value. We then try to access the box after it was moved, which is not allowed. Let's take a look at the generated MLIR and see where the "pinch.drop" instruction was generated to free this variable:

module {
    func @consume_box(%arg0: !pinch.box) -> ui32 attributes {src = ["a"], sym_visibility = "private"} {
        "pinch.drop"(%arg0) : (!pinch.box) -> () %0 = "pinch.deref"(%arg0) {dst = "return", src = "a"} : (!pinch.box) -> ui32
        pinch.return %0 : ui32 {src = ""}
  }
  func @main() attributes {src = []} {
        %0 = "pinch.box"() {dst = "a", value = 2 : ui32} : () -> !pinch.box
        %1 = "pinch.move"(%0) {dst = "", src = "a"} : (!pinch.box) -> !pinch.box
        %2 = pinch.generic_call @consume_box(%1) {dst = "b"} : (!pinch.box) -> ui32
	%3 = "pinch.deref"(%0) {dst = "c", src = "a"} : (!pinch.box) -> ui32
	pinch.return {src = ""}
}

Notice that "pinch.drop" was generated in "consume_box" instead of in the main function. Pinch detected that the box was moved out of main, and that the box was not moved out of "consume_box". Since the box was moved, it is no longer available in this scope, and our attempts to dereference it will be met with the following error generated by the borrow checker:

error: Trying to dereference already moved variable

The concept of an ownership graph makes this example possible. Because we are keeping track of what data belongs to what owner at operation-level granularity we can perform effective static analysis on the program. As long as the programmer obeys a few simple rules mandated by our borrow checker we can generate safe code and automatically generate drop operations to free data.

For the full set of operations and an explanation of more programs please see the paper.

Development Process

Let me start by saying that MLIR is extremely powerful and is an awesome project. It really is the next big thing, and if you are creating a domain specific language you should probably consider using it.

That being said, my experience developing with MLIR was not pleasant. It's pretty obvious that the target user is someone who is very familiar with LLVM, and they have probably done development with or for it. I do not have either of these qualifications, and getting caught up on all the things I wasn't familiar with was rough. This project was done part time over the course of 7 weeks.

I felt that I spent significantly more time than usual just trying to understand what I had to call to get or set the value I wanted when interacting with MLIR's API. MLIR has an example language which serves as a nice implementation reference, but the less your language resembles tensorflow the less help that example is.

Overall I did not enjoy the majority of the time spent writing the code for Pinch. Note that I am very biased because I do not like C++. LLVM developers apparently love C++, because they went absolutely wild with its advanced features. I did enjoy designing the MLIR dialect, and I enjoyed seeing the end result when it was all put together and code generation was working. LLVM is an awesome project, but not enjoyable to hack if you are a new user.

Maybe you'll feel differently. Maybe you have experience in LLVM and your learning curve will be friendly. Either way, MLIR is awesome and you should definitely give it your attention. Here's to hoping that more projects will utilize its infrastructure.