Lower VM dialect to LLVM IR

[benvanik]

For scenarios where dynamic module loading is not required and entire modules can be compiled into applications we can lower the VM IR to LLVM IR within MLIR's transformation pipeline. Instead of embedding vm.call ops that are dispatched at runtime to things like the HAL we can instead lower to llvm::CallInst to runtime-resolved function pointers. This still enables all of the flexibility of heterogeneous/runtime-determined devices, pluggable diagnostics, and backend composition without any need for flatbuffers or the VM bytecode interpreter.

The VM was designed to make such a lowering easy and the C-style struct-based function pointer registration for runtime modules was designed to make emitting code that used it fairly robust even when linked in dynamically such as when embedded in shared objects.

An extension of this is what we've been calling 'runtimeless mode', where the IREE VM linkage code is statically linked into the binary alongside the generated module LLVM IR. If only a single HAL backend is linked in then (with some build-fu) we should be able to get call devirtualization to reduce code size to precisely the functionality used by the module.

Early versions can be tested by tacking on additional vm-to-llvm conversions on the primary pipeline, and for integration a new iree/compiler/Translation target can be added (or an option can be added here: https://github.com/google/iree/blob/master/iree/compiler/Translation/IREEVM.cpp#L93)

[ergawy]

Hey @benvanik,

I am interested to give Tiny IREE in general a push (recently discovered Tiny ML and that world seems super interesting for an ML novice like myself :D).

Is it fine if I give this a try? Or if #1173 still have some gaps to cover I can also work on that first.

[benvanik]

Good question! Having an iree/compiler/Dialect/VM/Target/LLVMIR/ that spit out a .h + .o (or .bc or .so) is still really interesting, though it may be easier to start with EmitC if interested in making things tiny as it's mostly working and there's other non-VM stuff to fix that would be shared between both approaches. You can see the kind of stuff EmitC produces under iree/compiler/Dialect/VM/Target/C/test/ (with -print-ir-after-all added to the tests).

In the end there should be little difference in the final binaries produced via EmitC and LLVM IR - the major benefit of LLVM IR would be that as a user you'd get a .o you just link against vs needing to setup your own C compiler to produce the .o (and have all of IREE's source). Just like we do on the HAL executable side you could also produce .so's that are directly loadable by some iree/vm/dynamic_module.h so you wouldn't even have to rebuild the runtime. But neither of those directly help IREE in tiny contexts where platform/toolchain details are usually weird (so generating an object file that doesn't use library calls/etc not available on the target platform is trickier and you often want to have to invoke your C compiler so you can polyfill/etc) and you don't want dynamic loading (as that just adds size and the bytecode would suffice). Through that lens, LLVM IR is probably not a great choice for tiny deployments, and why most of the focus for the tiny work has been on EmitC.

But there are still interesting scenarios for it and I think the big unanswered thing for LLVM IR that would be good to explore first is how we avoid needing to manually reproduce all of what we have defined in C without requiring clang to do the parsing. Though we lower through LLVM IR it'd be nice to support non-clang compilers/non-LLVM linkers and that means we may not have clang available and need to emit object files (coff/elf/etc). The EmitC target is nice as it is able to reference all the iree/vm/ types and functions textually while an LLVM IR target would need proper LLVM types for everything. It's not a trivial surface area and where we do have to replicate things today (iree/compiler/Dialect/HAL/Target/LLVM/LibraryBuilder.cpp) it's pretty nasty even though scoped to just a few definitions. The HAL executable side there is also a bit different as it's meant to have a fairly stable API between the runtime and the executable, while the VM output is using less stable types that may change more frequently. That said after https://github.com/google/iree/projects/32 there really shouldn't be many more changes to the VM APIs so it may be fine if we have to replicate things.

There are some other options, though, like something that parsed the headers and dumped them to a bitcode file we could load in the compiler like we do for librt. If it only contains all the type definitions and function declarations it should be relatively small. The compiler code would then load that module, merge it with the one it is generating, and be able to call the inline functions. It still would have a relatively tight coupling to the implementation details of the C, though, so I'm not sure exactly how beneficial it would be besides avoiding all the type/inline function/etc boilerplate implemented in C++ by hand.

So all that said, so far due to the increased complexity/maintenance burden of LLVM IR and smaller scenario-specific benefits to tiny systems we've shied away from it. There's a few big things that would help even if LLVM IR was used instead of EmitC and would be possible to immediately make progress on and see the benefits of. Some are more infrastructure related (a C-compatible benchmark library we could plug in under iree/testing/benchmark.h #5669, proper compiler/linker flags for size-optimized builds #6622, etc), others are fun C trickery (devirtualizing the HAL #6641, finding other places where we can shrink the fixed-cost runtime code size), better use of emitc (#6645), etc. Most of these would be just as useful in cases where we emitted LLVM IR or make it easier, such as instead of vm->emitc we'd be doing vm->llvm and could do the same kind of stuff that #6645 does. But since a majority of the tiny efforts are through EmitC IMO that's a better place to start so everyone benefits :)

[ergawy]

Thanks for the detailed overview. Focusing on the shared bits between the 2 approaches indeed makes sense.

After taking a tour in the issues you linked, I would like to start working on the HAL devirtualization issue (#6641) (and hopefully discovering some other opportunities for code shrinking). Since this is a new corner for me, I will take a few days to explore the HAL API and then dive into the issue.

(Feel free to assign me to #6641 if you like :)).

[benvanik]

Closing for now as we don't have a priority need for this and it's not something that'd be merged yet if it was done (still too much churn).