[MLIR][CodeGen] Native, JIT focused (x86) Code Generation for MLIR without the `llvm` dialect

Hi there. For my final thesis, I’m currently working on a native code generator for MLIR that does not need to go through the llvm dialect.

TL;DR: Does anyone know another project that tries to lower MLIR, or another extensible IR to machine code directly?

There is currently not much to see in the repo, so let me tell you more about it: It defines target specific dialects containing ISA instructions (currently only for x86-64) to which users can lower their instructions, just like they can with the llvm dialect. These instructions are fully qualified with regard to their opcode, so that the encoder can later take an op and know the exact opcode required to lower it. Some in-tree dialects have pre-defined patterns to do the lowering. This is in essence the instruction selection step of the code generation process. The MoNaCo back-end then takes care of SSA-destruction, register allocation and encoding using the fadec encoder, all in one second pass over the IR. As long as all instructions have been lowered to the target specific dialect, the back-end is completely agnostic to what happened during lowering.
Due to the time constraints of a thesis, the entire process is very basic, the instruction selection does not support many dialects yet (ironically the best supported one is probably llvm at the moment, as its easiest to generate from C) and amounts to nothing more than macro expansion, the register allocator is basically purely stack based without any liveness, and at the moment, the only way to use external symbols or call external functions is in the JIT mode, which is the main focus of the code generator, i.e. relocations or object file writing isn’t implemented yet.

This of course leads to worse code quality compared to LLVM, the SPEC CPU 2017 mcf benchmark takes about 4 times as long if compiled with monaco vs LLVM as the baseline(-O0, with fast isel enabled); Compile times are slightly better, monaco JIT compiles mcf in about 75% of the time that LLVM needs (fast isel again).

Does anyone know of a similar project that is ongoing (classic “Related Work” question, I know)? The things that I would consider most similar are the armsve, amx, and x86vector (and maybe nvgpu?) dialects, as well as the intrinsics from llvm, as they are closest to specifying an instruction/opcode. I also found this: https://arxiv.org/pdf/2101.11365.pdf, although I am very unfamiliar with quantum computing and quantum assembly languages, so I don’t know how close this really comes to monaco.
Also, all the dialects and representations I could find still rely on LLVM based code generation in the end, making them less direct than monaco.
So if anyone knows another attempt at lowering MLIR to machine code directly, or knows of another extensible intermediate representation employing direct code generation (‘direct’ meaning not through LLVM), I would greatly appreciate a link :).

2 Likes

xdsl’s python dialect, and MIR have since been suggested on the LLVM discord (Discord), which I will definitely address, thank you Markus!

WASM as a dialect may be an interesting and natural strategy for this, though it may require control-flow restructuring.

Sorry, I can’t quite follow. Do you mean emitting actual wasm from your proposed wasm dialect? An interesting idea, although I don’t think it is in the scope of this project, as MoNaCo is more focused on middle-to-end JIT compilation so to speak. Compiling to wasm is not really ‘the end’, and afaik has little in common with the existing target specific dialect for x86-64. I suppose you were refering to the stack-machine nature of wasm?

Do you mean emitting actual wasm from your proposed wasm dialect?

Yes.

Compiling to wasm is not really ‘the end’, and afaik has little in common with the existing target specific dialect for x86-64.

More what I mean is that wasm is an IR for which a few JITs exist, targeting a variety of machine code languages. If your goal is to avoid the LLVM “jit to a dylib” then a wasm MLIR dialect, or llvm-to-wasm conversion is another way to get yourself to machine code.

WASM is also interpretable, and persistently cacheable. So you could have tiers, where you have a persistent code cache.

An interesting idea, although I don’t think it is in the scope of this project

That’s fair! You asked for a solution, I presented ~3 more problems on the way to a variant :stuck_out_tongue:

A wild idea but I wonder if it’s possible to generate this kind of “ISA dialects” from existing TableGen descriptions written for backends.

1 Like

A wild idea but I wonder if it’s possible to generate this kind of “ISA dialects” from existing TableGen descriptions written for backends.

I had thought about that; in familiarizing myself with TableGen to specify the x86-64 dialect, I read a bit of the target description for the x86 back-end. For moanco at least, it would not be possible to automate this completely, but I wager it could get you a good 80% of the way there, for most ‘boring’ instructions. Doing this without any manual intervention would be quite a challenge.

There is a talk at the LLVM Dev Summit in October on this topic by @AntonLydike and @superlopuh : “Compiler backend design with MLIR” ; looking forward to this :slight_smile:

4 Likes

I’ve not read your implementation yet, but we were thinking of trying out different backends also, including ARM and x86. I’d love to chat about the various challenges and solutions about backends in MLIR. I’m currently planning out the talk for the dev meeting, will you be attending?

1 Like

I would love to attend, but unfortunately I’m not stationed in the US, and can’t afford to travel on my own. In addition the non-student prices (as this is my final thesis, I’m not a student anymore in october) are too high for me.
That said I’m fully open to talking online about it, I’ll be giving a small talk about the subject at my university on sep 28th and be happy to repeat it online (or practice it online beforehand ^^), if there is interest in the matter. I’d also be up for a casual chat about it, after I’ve handed in my thesis (sep 15).

1 Like

@J-MR-T We did exactly the approach mentioned here 2 years back. It was presented in
“Fifth LLVM Performance Workshop at CGO”.

Abstract: The LLVM Compiler Infrastructure Project
Slides : https://llvm.org/devmtg/2021-02-28/slides/Vinay-MLIR-codegen.pdf
Talk : https://www.youtube.com/watch?v=4YZyRJ1ifGQ&t=1681s

We prototyped a generic infrastructure for porting the LLVM Targets to the MLIR framework making use of the existing TableGen representation. We picked X86 as our first Target…

3 Likes

Very interesting work. In your talk you mentioned that you were open sourcing it “soon”, is the code available somewhere?

1 Like

@vinaymadhusudan That indeed looks similar to what I’m doing, and seems very promising for including in the related work section, thank you! I would also really like to have a look at the source though, to understand it more in-depth.

I would not say it is the exact same approach, as monaco only uses a single intermediate step, basically: any dialect → monaco target specific dialect → machine code, so 2 passes over the IR. Your approach looks more general with many conversions: std dialect (now cf, func, arith, etc.) → llvm dialect → generic MIR dialect → x86 based MIR dialect → x86 assembly → machine code.

As I understood your talk, it was more about porting the LLVM code generator to MLIR, than native JIT code generation in MLIR with focus on user specified coversions to closely model domain-specific semantics. Not to downplay it, I just feel like the end-goal is different. Do you agree?