[TE Compiler] Move header files to includes

With the new refactored TE compiler, we can do more interesting analysis and optimization. For example, one recent work I’m doing is taking a Relay function, lower it to a TE compute and estimate its FLOPS. Ideally, I could use cfunc = PrimFuncFor(func, target) and use a visitor to estimate the FLOPS of cfunc->outputs. However, I realized that this cannot be achieved outside of TVM codebase, because the corresponding APIs and data structures are declared in te_compiler_cache.h and te_compiler.h, and unfortunately they are under src instead of include.

@jroesch @csullivan do you think we should move these header files to include so that they can be used more widely?

2 Likes

+1. I feel the same thing when interfacing with Relay compile engine a year ago.

I think it can be a good idea. On the other hand, our eventual goal is streamline everything to IRModule=>IRModule with no such specialized flow. So we might want to keep that in mind when depending on the particular module

I’m not against the IRModule to IRModule flow. Meanwhile, with this flow as the TVM standard, we could still expose useful APIs to developers for experimenting or building new features.

The thought is that the lowering should remain private with explicit customization for different hardware targets. We want to avoid exposing too much externally as this will make refactoring difficult. I’m definitely supportive of the new analysis and optimization made possible by the TECompiler refactor; perhaps some helpers could be exposed that allow us to maintain the flexibility, while providing this functionality. WDYT @comaniac?

Thanks for the sharing and I agree that it’s a fair concern. In this case, helpers definitely work as well.