[RFC] Incorporate existing Relay AoT compiler in mainline

In reference to @tqchen’s previous RFC on ahead-of-time (AoT) compilation, I believe we should consider incorporating the already-existing Relay AoT compiler, mostly written by @MarisaKirisame, into TVM’s codebase as a potential step towards achieving the goals outlined in the previous RFC.

The tool was originally kept separate from the TVM codebase because it was built for rapidly changing research goals, but it has been stable in its current form since April. Having the tool in the official TVM codebase and incorporated in the CI will help ensure that it does not break and make it more easily available for TVM users. Additionally, having more members of the community give feedback about the AoT compiler’s interfaces and implementation would help resolve design issues with it and potentially address low-level bugs arising due to interactions with changing APIs. I think that simply having the existing (quite compact) AoT compiler in the codebase would be useful as a concrete starting point for iterating on the goals in the previous RFC.

Note that the Python debug converter was based on the AoT compiler, so we could potentially restructure the AoT compiler to deduplicate code between these and allow for further extension.

5 Likes

Thanks @slyubomirsky for the proposal

The current relay AoT is a transpiler to c/C++. As per our previous discussion on AoT. In our ultimate solution we want to push for a solution that embeds on the TIR instead, and make it natural part of the compilation pipeline. It would be great if we can start thinking about that goal now to minimize the potential technical debts.

If we decided that it is helpful to incorporating the current implementation, we should improve the implementation to make sure that we re-structure the target generation so that i aligns with other backends, such as those in VM and graph runtime. We also need to make sure that the runtime interface directly make use the unified mechanism for module packaging etc and only depends on the tvm’s runtime API.

In principle, the portion of the current AoT compiler that outputs to C++ is just the back-end, so we can definitely structure the compiler to output to different back-ends (e.g., output Python like in testing/py_converter.py). I will look at how it can be made compatible with TVM’s existing back-ends.

1 Like

As a side note, the new design architecture overview is now up https://tvm.apache.org/docs/dev/index.html

The areas related to the AoT would be the runtime interface part as well as the target translation.

1 Like

So in the spirit of the architecture you posted (specifically this chart), @tqchen, is it correct to say that an integrated AoT compiler should take an IRModule, compile it into some external representation (e.g., a C++ program), and produce a runtime module that would allow for invoking the compiled program?

If my understanding is correct, the primary change from the existing AoT compiler would be to ensure that the output is in the form of TVM PackedFuncs rather than pure C++ functions. Apparently the existing AoT compiler already produces a PackedFunc as the entry point so it should not be very difficult to wrap a runtime module around that.