MLIR in Mojo
Multi-Level Intermediate Representation (MLIR) is an open-source compiler infrastructure project that is designed to help with the development of optimizing compilers and code transformation tools. MLIR provides a flexible and extensible framework for representing and manipulating various programs and data representations at multiple levels of abstraction.
Now, the question arises: What does MLIR have to do with Mojo? Well, Mojo provides us with low-level MLIR primitives that we can use to create our abstractions. To illustrate how MLIR can be directly interfaced in Mojo, we will implement a boolean type in this Answer using Mojo.
Note: Using our custom types provide us with the advantage of a high-level interface like Python, however, Mojo will optimize those high-level types for each new technology that might appear in the future. This Answer explains the features of Mojo in great detail.
Creating custom abstractions
In Mojo, we can use the structs for our purpose. Let’s start by creating a custom boolean variable:
struct CustomBool:var value: __mlir_type.i1fn __init__(inout self):self.value = __mlir_op.`index.bool.constant`[value = __mlir_attr.`false`,]()
Let’s have a look at the code we wrote above.
The struct has a single member variable, namely
value, that is supposed to store the value of the bool. It will either betrueorfalse. And, we can see that it is represented directly in MLIR by the built-in type,i1.We make use of
__mlir_typeto represent our boolean value using the MLIRi1type.
Note: We can use any MLIR type by prefixing the type name with
__mlir_type.
We also use
__mlir_op, which is a prefix used to invoke the`index.bool.constant`operation. It takes either atrueorfalsecompile-time constant.The compile type constant that the operation needs is passed to it using the
__mlir_attrprefix.
This code allows us to instantiate an object. As such:
let a = CustomBool()
Note: The struct that we created will only create an instance with the
falseMLIR attribute.
Compile-time constants
Let’s modify our struct so that it takes a value for its constructor.
@register_passable("trivial")struct CustomBool:var value: __mlir_type.i1fn __init__(value: __mlir_type.i1) -> Self:return Self{value:value}
Here, we use a register_passable decorator so that the Mojo compiler knows that it can use the only member variable we have in our struct to implement the copy logic. Structs that use this decorator also have to implement the __init__ method. Now, we will make use of the alias keyword to define a compile-time constant that we can use.
alias CustomTrue = CustomBool(__mlir_attr.'true')# oralias CustomFalse: CustomBool = __mlir_attr.'false'
Implementing functionality
Even though we’ve written the code to assign values to our custom boolean abstraction, we still need to make it behave like a boolean value. For example, if we were to use it as shown below, it wouldn’t work because the struct doesn’t implement the method needed to tell the Mojo compiler what it’s supposed to do when the value is either true or false.
let a = CustomTrueif a: # won't work!print("")
This throws an error that we need to implement the __bool__ method for our struct. This makes sense because we didn’t tell our custom variable how it is supposed to act like a boolean variable. We could add a __bool__ method to our struct but that would resultantly be converting the MLIR attribute to Bool. Here’s what the member function for the struct would look like:
fn __bool__(self) -> Bool:return Bool(self.value)
Here’s what the code will look like. It can be seen that we’re able to make use of our custom boolean variable as we would normally use a boolean variable.
@register_passable("trivial")struct CustomBool:var value: __mlir_type.i1fn __init__(value: __mlir_type.i1) -> Self:return Self{value:value}fn __bool__(self) -> Bool:return Bool(self.value)fn main():# define the two instances our abstraction can takealias customTrue = CustomBool(__mlir_attr.`true`)alias customFalse = CustomBool(__mlir_attr.`false`)# define variables with compile-time constantslet a = customTruelet b = customFalse# check functionalityif a:print("a is true")if b:print("b is true")
There is, however, a caveat to this implementation. If we’re eventually converting our MLIR compile-time constant to Mojo’s Bool then there’s no point in creating the custom abstraction. So, let’s implement a member function that will let us use the compile-time constant without having to convert it into Bool.
struct CustomBool:var value: __mlir_type.i1fn __init__(value: __mlir_type.i1) -> Self:return Self{value:value}fn __mlir_i1__(self) -> __mlir_type.i1:return self.value
We’ve implemented the __mlir__i1 method. Our custom bool will act in the same way as Bool does in Mojo because, at the base level, the Mojo compiler converts it into an MLIR i1 value. So now, it should act in the same way as Bool would. Let us check out its implementation in the code below:
@register_passable("trivial")struct CustomBool:var value: __mlir_type.i1fn __init__(value: __mlir_type.i1) -> Self:return Self{value:value}fn __mlir_i1__(self) -> __mlir_type.i1:return self.valuefn main():# define the two instances our abstraction can takealias customTrue = CustomBool(__mlir_attr.`true`)alias customFalse = CustomBool(__mlir_attr.`false`)# define variables with compile-time constantslet a = customTruelet b = customFalse# check functionalityif a:print("a is true")if b:print("b is true")
We’ve successfully implemented our custom boolean variable using low-level MLIR primitives. This, however, was just an example of how MLIR can be leveraged to create more complex data types. Even for our implementation of CustomBool, we still haven’t defined how it will behave outside of conditional statements, for example, if we want to use the inversion operator with our custom type, we will have to implement a __invert__ method in the struct.
Conclusion
In this Answer, we went over the basics of how we can create our own abstractions in Mojo by making use of MLIR. This proves useful for creating custom data types and abstractions that run on custom hardware and not a generalized set of processors. This means that Mojo can be used to write optimized code for processors that use a different architecture as well.
Free Resources
