[Relay][TIR] Add utility to lower Relay func to TIR prim func

[masahi]

This is very useful for creating a TensorIR program that is otherwise hard to so via manual TVMScript or te compute + create_prim_func.

A good example is a TensorIR program having AllocateConst, since it is not possible to create such modules with embedded constants via TVMScript parsing. AllocateConst has been the source of many issues and I've sent many fixes for them, but it has always been difficult to create test cases for such issues. But if we start from Relay, we can easily create such modules by lowering Relay func to TIR with link-params = True.

This also makes it possible to, say, author a complicated module in PyTorch and directly lower it to TensorIR.

As an example of its use, I added a regression test for the issue described in #13605.

@Hzfengsy @junrushao @csullivan

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @Hzfengsy, @junrushao _{See #10317 for details}

_{Generated by tvm-bot}

[csullivan]

Awesome, I love it. One small suggestion if you agree.

[echuraev]

Awesome and very useful improvement! Thanks.

[junrushao]

Looks great!

[Mousius]

👋 Hello @masahi / @csullivan / @echuraev / @junrushao,

Whilst I appreciate this is useful with the use-case demonstrated, I believe it might've needed a bit more time for others to formulate a review as it was up for less than 24 hours before merge?

One of my concerns is that this seems to be creating new public interfaces to the internals of the TECompiler which in the past we've attempted to remove:

• https://discuss.tvm.apache.org/t/rfc-relay-tecompiler-rewrite-existing-compile-engine-to-match-updated-compiler-flow/9233
• Remove compile_engine.h for real this time.  #8775
• [CORE][Relay] Swap and remove compile_engine with te_compiler followup of #8775 #9282

This change means we have to maintain an interface which generates a single tir::PrimFunc from a relay::Function where-as it could potentially lead to multiple tir::PrimFunc with a single IRModule as a result. The next line down:

IRModule tir_mod = PrimFuncToIRModule(f.value());

Implies we wanted it in IRModule form anyway, so creating an IRModule with a Function and running LowerTE(mod) seems the best result here?

[masahi]

@Mousius The new function simply composes the existing interface to the TECompiler, LowerTECompute and CreatePrimFunc that converts lowered TE compute to TensorIR. LowerTECompute was introduced after the compile engine refactor you mentioned completed, and made only possible thanks to that refactoring work. So I don't get your concern here.

What I wanted to add is a simple interface that converts one Relay primitive function to the corresponding TIR prim func. Then I realized that backend/task_extraction.cc already has code that does exactly that, so I simply extracted into a reusable function and exposed it to Python. We require the input to be a primitive function, so we never generates multiple TIR prim funcs.

Implies we wanted it in IRModule form anyway, so creating an IRModule with a Function and running LowerTE(mod) seems the best result here?

This is just a convention of MetaSchedule, storing an IRModule with one prim func as a task. I didn't change any behavior about task extraction in this PR. What we want is a relay::Function -> tir::PrimFunc interface, IRModule -> IRModule is valid but it wouldn't convey what this new API is meant for. LowerTE is also overkill since I don't need to lower schedules and I'm not concerned about external codegen etc.

[Mousius]

@Mousius The new function simply composes the existing interface to the TECompiler, LowerTECompute and CreatePrimFunc that converts lowered TE compute to TensorIR. LowerTECompute was introduced after the compile engine refactor you mentioned completed, and made only possible thanks to that refactoring work. So I don't get your concern here.

I find this equally concerning, given this is also exposing more of the internals of the TE Compiler rather than having clear boundaries within the codebase.

What I wanted to add is a simple interface that converts one Relay primitive function to the corresponding TIR prim func. Then I realized that backend/task_extraction.cc already has code that does exactly that, so I simply extracted into a reusable function and exposed it to Python. We require the input to be a primitive function, so we never generates multiple TIR prim funcs.

Exposing it python and as part of the public API means others can access it, which means that we have to support it as a way of using the compiler rather than as an aspect of meta-scheduler. This is similar to the previous methods that we were unable to deprecate in #13606.

This is just a convention of MetaSchedule, storing an IRModule with one prim func as a task. I didn't change any behavior about task extraction in this PR. What we want is a relay::Function -> tir::PrimFunc interface, IRModule -> IRModule is valid but it wouldn't convey what this new API is meant for. LowerTE is also overkill since I don't need to lower schedules and I'm not concerned about external codegen etc.

I don't think LowerTE is overkill in practice, if it doesn't find any external codegen or other things to do it won't take any action - it's a standardised interface between two modules in the codebase, you've already indicated meta schedule is breaking that boundary and that leads us back to the original issues with CompileEngine.