PSA: Integration with IREE for multi-target and whole model compilation

[nicolasvasilache]

Hi everyone,

Some of us have been doing a recent push to better integrate some of our findings with IREE recently.

The ongoing prototype integration is currently split across 3 branches:

• iree-llvm-sandbox/iree-integration is the branch where the integration with IREE happens. It depends on:
• iree/sandbox branch, which sees the following restructurings:
  • the LinalgExt dialect is folded back into iree/llvm-external-projects where it was originally forked from
  • the transform dialect is transitioning through iree/llvm-external-projects so that IREE can canary such transforms in a larger scaler system while upstreaming to MLIR core is occurring (see Alex's RFC)
• iree-llvm-fork/sandbox is the source of truth for IREE + sandbox

With this setup we are currently able to compile and run IREE with transformations specified in by the transformation dialect. As a consequence, iree-llvm-sandbox/iree-integration compiles with IREE and exposes a simple parallel nevergrad-based search.

In the near-future, we are hoping to extend more of these mechanisms to better iterate with IREE itself (and not just a minimal Numpy API). This should give us a few interesting opportunities:

• compile and run on different devices (mobile, NVIDIA-GPU, mobile GPU) via IREE's AoT (and not just the current JIT CPU)
• extend search to be able to execute and time on those different devices
• start saving and shipping known good transformations with IREE for different devices.

These will open up new exciting work areas.

Until the transition is complete however, some rough edges are expected.

Thanks for reading!

[nicolasvasilache]

cc @ftynse @Mogball @MaheshRavishankar @ThomasRaoux @stellaraccident @matthias-springer @gysit

[stellaraccident]

+1 glad to see this getting bridged

[ftynse]

What's the plan for moving changes between different projects? Let's say I factor out the interface into a separate dialect as proposed in the RFC.

[nicolasvasilache]

@ftynse I'd say strongly favor doing this in core and we'll adopt.

If you need a tmp staging ground, it would be quite better to do it in iree-dialects given the coupling introduced by the integration.

If you must do it directly in the sandbox, we'll either pay the price of manually syncing into IREE or wait until upstreamed to core; depending on perceived cost/benefit.

[stellaraccident]

+1 for things that we think may be "core bound" in some number of months, we make local choices aimed at limiting overhead. This can include a cross project dependency (ie. Sandbox to iree) or copying sub directories (ie. Iree-dialects was made to be "reused by cp" as a staging ground for things that may be core bound at some point).

If it looks like upstreaming something isn't going to happen, then if we still want it, we need to find it a better stable home.

[nicolasvasilache]

The back and forth integration has now landed.

I added a new github action testing IREE integration on the iree/sandbox branch that is meant to be non-blocking for pushes. The main github action remains the thing we should look out for remaining green and updates to iree/sandbox are on a best effort basis for now.

Please let me know if you see something that is cause of concern or experience unnecessary hardships so we can refine the process.