Pliron: An extensible compiler IR framework, inspired by MLIR and written in safe Rust

I’ve been working (in my spare time) on a project attempting to write an MLIR like infrastructure in Rust.

GitHub: GitHub - vaivaswatha/pliron: Programming Languages Intermediate Representation

It’s still in a nascent stage, but any feedback is welcome !.

I’ve written an introductory article, and the docs are mostly up-to-date.

Apologies if this wasn’t the right place to post this announcement.

5 Likes

I was thinking of the following (as someone barely aware of Rust / C interactions - so this might be naive): how would Pliron compare with the approach of just using the MLIR CAPI from Rust via an FFI? I didn’t see this covered in your introductory article. Is it feasible, or does it completely destroy the purpose of using Rust? Will stack traces on asserts from the MLIR core infra be visible when used from Rust? How would the compiler development experience be, and will it have memory safety benefits (w/ Rust + MLIR CAPI)? Because if it’s feasible, it’ll still allow one to build compilers in Rust while not duplicating all the core infrastructure that is MLIR.

The rust compiler uses the C-API of LLVM for code-generation. Teaching the rust compiler to go through MLIR and LLVM would be challenging. The rust compiler guys probably refuse to write MLIR passes in C++.

1 Like

I think Uday was more asking about tradeoffs in using actual MLIR in memory and Rust binding to C API vs implementing a system inspired by MLIR in Rust.

2 Likes

Sure, you can use from Rust the C-API to MLIR, but can you write passes in C?

how would Pliron compare with the approach of just using the MLIR CAPI from Rust via an FFI? I didn’t see this covered in your introductory article.

I only briefly mentioned melior in the article, mainly because its use would be unsafe, and I suppose to an extent does destroy the purpose of using Rust.

Will stack traces on asserts from the MLIR core infra be visible when used from Rust? How would the compiler development experience be, and will it have memory safety benefits (w/ Rust + MLIR CAPI)?

I haven’t used MLIR (or even LLVM) through Rust bindings. But I have used LLVM via it’s OCaml bindings. While that was usable, much of the type safety (and I’m guessing in the case of Rust, static memory safety guarantees) is lost when using LLVM through the C bindings. I have seen hard to debug segfaults when using the OCaml bindings.

With bindgen you can create Rust wrappers for C (MLIR/LLVM) APIs. Correct me if I am wrong, you cannot write passes with the C API?

I don’t think you can write passes, create new dialects, create new OpInterfaces (which may be needed for the ops of the new dialects), or create new dialect types and attributes with the C API. I think the CAPI is meant “to use” and not “to create” any of these.

What I had in mind w.r.t my previous comment is that you’d build parallel infrastructure in Rust for all the above: passes, dialects, op interfaces, pattern rewriting, etc., but still use MLIR (the current one) with its CAPI to parse/print and mutate IR, transform/analyze using its core analysis/transformation utilities without having to rewrite all of them. So this would be a much smaller undertaking than rewriting all of MLIR in Rust from the start. It may also provide an incremental path.

Does Rust have all the features to provide a better pass/dialects/op interface etc. creation infrastructure than C++?

1 Like

What I had in mind w.r.t my previous comment is that you’d build parallel infrastructure in Rust for all the above: passes, dialects, op interfaces, pattern rewriting, etc., but still use MLIR (the current one) with its CAPI to parse/print and mutate IR, transform/analyze using its core analysis/transformation utilities without having to rewrite all of them

I’m guessing that this is going to be quite difficult to make it work and maintain. At every point we’ll need a translation from the Rust infrastructure to the C++ infrastructure, or choose to do every Rust<->C++ interaction via textual IR.

Does Rust have all the features to provide a better pass/dialects/op interface etc. creation infrastructure than C++?

I wouldn’t say it’s particularly better, but it’s as good.

My experience with Rust programmers is that they would prefer to have a native Rust infrastructure (where the APIs for example would be in Rust style / idioms and not just a translation of C++ APIs via FFIs).
For example see this comment and that there exists an LLVM like compiler, cranelift, written fully in Rust (they had other reasons too to write that though, such as LLVM being too slow for their purpose).

I think the point would rather be to let MLIR handle the in-memory representation and use bindings to touch on the IR.

The main reason is that a parallel ecosystem is entirely cutting itself from shared infrastructure, and an important part of MLIR is help defragment. Otherwise you have to duplicate everything, tooling but also support libraries (like the dataflow engine for example), but also you’re losing entirely on mixing other dialect (and DialectConversion)
.

I think the point would rather be to let MLIR handle the in-memory representation and use bindings to touch on the IR.

Right, that’s a better option to having Rust representations and then translate them only for analyses and transformations. This already exists today though with the melior project.

The main reason is that a parallel ecosystem is entirely cutting itself from shared infrastructure, and an important part of MLIR is help defragment.

True and I agree with this point. But also, as I pointed out earlier, using an infrastructure via its FFI does have it’s drawbacks. One of the prime example I can think of is from my experience with using LLVM via its OCaml bindings. Even though the bug was in my own OCaml code, to actually debug the memory errors, I had to debug LLVM in a debugger to trace out why something was going wrong. Translating that experience here would mean, for non-trivial cases, a Rust programmer may need to debug MLIR’s C++ code, which he may not even understand.

This is all overlooking the main point that using MLIR via an FFI would just mean not having Rust’s safety guarantees on that anymore. So I can see Rust programmers at some level preferring to have a native Rust implementation (if that’s feasible). So this is ultimately a tradeoff and only time can tell if it was worth attempting this.

I’ve been working on a similar project (though in Python), and I saw that there is quite a large community that is interested in using MLIR, but is not trying it just because of the language (C++) barrier. For us, we chose Python because we wanted to teach MLIR concepts to students, and wanted to work with some people from the HPC community that aren’t familiar with C++, but are very familiar with Python. Rewriting MLIR core also allowed us to have a more Pythonic API, compared to the Python bindings which are quite low-level. (We presented this at EuroLLVM this year).

One thing we are actively working on though, is to try to make these parallel implementations compatible in some way, instead of having them totally disjoint. Optimally, one could still use both MLIR and Pliron at different level in the pipeline, to not completely split the efforts.

I guess the most important thing is to keep the same textual format, which I don’t think Pliron is currently doing? This would allow users to write their high-level dialects in Pliron, and then compile it down to MLIR by passing the textual representation. That would allow you to both stay compatible with MLIR, while keeping a pure Rust implementation when needed. We found out that it was quite easy to stay up to date with MLIR textual syntax, the only major change in the last two years being the addition of properties.

The second major pain point is the porting of dialects from one framework to the other. While ODS allows you to quickly write a dialect definition in MLIR, it is still heavily using C++. I’ve been working on IRDL, which is a more declarative representation for dialect definitions. The main point is that since it is representation as a dialect, you could use this front-end in pliron to generate your dialect definitions, avoiding the manual reimplementation of the MLIR dialects. Currently, we are working on a translation from ODS to IRDL, as well as adding more features in IRDL to express dialects more declaratively. Help is always welcome :wink:

3 Likes

Thank you for the detailed response @math-fehr . I didn’t know about xDSL and it looks cool.

try to make these parallel implementations compatible in some way

I agree

I guess the most important thing is to keep the same textual format, which I don’t think Pliron is currently doing?

Textual format compatibility was an initial goal of mine, but I kind of neglected it as I progressed. But given that pliron is still in an early stage, working towards achieving this shouldn’t be difficult. So I’ll strive towards that. It’s a very useful suggestion, thank you.

I see your point on porting of dialects. I haven’t given this much thought yet. I will soon.

We have folks writing Python “passes” using C API, but it’s not a pass as in C++ one indeed, rather a function that takes an operation and mutates it. No one has added such an interface C side, it should be less complicated than extensible dialects I think (at least for just top-level passes, how to handle parallel execution on random isolated from above operations I haven’t thought about). Just no one has much tried :slight_smile: The former is used for experiments and the latter is used when scaling.

I think it would be very interesting to think how these could really be mixed like this. With haskell bindings at some point we were going to look at that “next” but unfortunately didn’t have time to yet.

I started building out an AST matcher against upstream OpViews a bit ago: https://github.com/makslevental/nelli/blob/main/nelli/mlir/ast/visitors.py#L151. I also added replace_all_uses_with to the bindings.

Using these two my plan was to replicate some of libcst. Personal use value has decreased as I’ve gotten more familiar with the C++ API - ironically that happened as I was trying to enable myself and others to avoid the C++ API, same as @math-fehr :slightly_smiling_face: - so I haven’t done much more than generate the matcher. But I think there’s a pretty clear path to victory there. If there’s grassroots interest I could put together an RFC (for upstreaming) but I believe prior when I’ve tried to add functionality to the python bindings that enables IR creation, @ftynse has cautioned against.

Let me clarify. I’m specifically cautious about the following situation: IR is constructed in language A, then a mutation in language B erases some parts of the IR, language A keeps holding a reference to the erased pieces of the IR and the entire system crashes when it attempts to access it. A specific example is Python holding a reference to an operation that gets erased by a C++ pass (that’s why we have a validation mechanism at that level), but it could also happen in another direction. Debugging this would be a nightmare. So mutations that don’t result in such behavior are okay to be present in bindings. More general mutations need a clear and careful design to avoid this.

1 Like

Actually this situation will crash with the current bindings:

module = builtin.ModuleOp.parse("""
  module {
    arith.constant 10
  }
""")
const_op = module.body.operations[0]
print(const_op)
PassManager.parse('builtin.module(canonicalize)').run(module)
print(const_op)  # Segmentation fault (core dumped)

And I don’t think we have the tools to prevent this today, we would need something like the IR listeners proposal to properly track invalidated references. And even then, this wouldn’t really work in the opposite direction (C++ references invalidated by outside code) since C++ holds raw pointers rather than handles that can be invalidated.

1 Like

I was under the impression that this was guarded against. Specifically, that the invalidation mechanism (https://github.com/llvm/llvm-project/blob/582e1d58bd43af138033b98b736d97cc89c7b228/mlir/lib/Bindings/Python/IRModule.h#L668) is expected to mark as invalid any operations nested under the pass manager root. Maybe that was only discussed and never implemented, but it looks trivial to add. Certainly this will not update the pointers, but will at least give a proper exception in Python instead of segfaulting.

I think it was half-implemented. The plumbing was added to invalidate on “external IR mutations” but I don’t think the Python bindings for the PassManager existed then, and once it did, it is plausible we never went back and made it do this…

This is probably good enough for most use-cases, but it assumes the pass is written properly and doesn’t erase ops above the pass root, so I think in the rust case the API would still be unsafe. Though maybe you could defend against that by detaching the root op from its parent before running the passes.

(Anyways probably a meta-point here is: even if possible, is it worth jumping through all these hoops to try and make memory-safe bindings, vs re-writing the infrastructure with this mind in the first place? As someone already invested in MLIR I think it is, but for people coming from other languages I can see why it’s not so appealing)

1 Like