[TE] Light refactoring of TE -> TIR paths.

[Lunderberg]

These refactors were initially part of work to implement apache/tvm-rfcs#39, but ended up being independent of those changes, and do not change existing behavior.

• Added ScheduleToPrimFunc, extracting out common behavior in ScheduleToModule and auto_scheduler's feature extraction.

• Added tvm.driver.build_module.schedule_to_module, to avoid needing the 4-line boilerplate needed to do so in several unit tests. This also makes deviations from the usual path (e.g. debug_keep_trivial_loop) more explicit.

[junrushao]

CC: @Hzfengsy @comaniac

[jroesch]

cc @electriclilies @mikepapadim

[Lunderberg]

Rebased to resolve conflicts with main.

[electriclilies]

I really like the fact that you're removing more duplicated python logic!

However, I'm a bit concerned adding a function that returns a PrimFunc is moving away from a unified lowering world where everything is an IRModule.

It seems like in pretty much all of the cases, you call schedule_to_primfunc, and then wrap that in an IRModule. Could you just use schedule_to_module instead?

[electriclilies]

What does debug_keep_trivial_loop do?

[Lunderberg]

This tells ScheduleOps to keep loops that have an extent of 1, where the default behavior is to replace trivial loops with a Let statement. (impl) The only place that it is used is in lower_ethosu to maintain the previous behavior of keeping the trivial loops. That said, I haven't looked into why the trivial loops are kept in that case.

Edit: And also used in the autotvm feature extraction

[electriclilies]

I like that you've removed this python logic, but why call schedule_to_primfunc and then wrap the func in an IRModule? Why not use schedule_to_module?

[Lunderberg]

Due to silliness and lack of pattern recognition on my part, since I was only looking for cases that could be replaced with schedule_to_primfunc at that point. Changing it to schedule_to_module, and thank you for catching it!

[Hzfengsy]

Thanks @Lunderberg. It's a pretty great PR to simplify lowering path. It overall LGTM

[Hzfengsy]

Why do we need tir.debug_keep_trivial_loop?

[Lunderberg]

This deliberately keeps the loop iterators in-place, even if they have extent=1, rather than the default behavior of replacing trivial iterators with a Let statement. As a result, the itervars can be examined for optimization parameters (e.g. in xgboost).

Longer term, I'd prefer having it always generate the loops with a lowering pass to identify/simplify the trivial loops, but that's a later item.

[Lunderberg]

However, I'm a bit concerned adding a function that returns a PrimFunc is moving away from a unified lowering world where everything is an IRModule.
It seems like in pretty much all of the cases, you call schedule_to_primfunc, and then wrap that in an IRModule. Could you just use schedule_to_module instead?

Good point, and I hadn't considered that aspect, as my initial motivation was just simplifying the call sites. (A different change looked like it would need to pass information from ScheduleOps to SchedulePostProcToPrimFunc, but that ended up not being necessary.) The main reason I didn't use ScheduleToModule directly was because it calls sch.normalize(), whereas auto_scheduler's feature extraction requires sch.normalize_for_feature_extraction().

That said, it looks like calling sch.normalize() after calling sch.normalize_for_feature_extraction() has no effect, so those use cases could still use ScheduleToModule so long as they normalize the schedule first.

[Lunderberg]

@electriclilies I've added a commit to use schedule_to_module directly in all cases, rather than schedule_to_primfunc. I like this version even more, as it removes another line of boilerplate from the usage and as you mentioned keeps the IRModule as the main object to interact with.

[electriclilies]

Thanks Eric, this looks great!
One final thing, would it be possible for you to fold ScheduleToPrimFunc back into ScheduleToModule? Even though it's no longer exposed in python, I don't want people writing C++ code to get the idea that they should be passing bare PrimFuncs around :) Would be OK to do that in a follow up PR as well.
Again, thanks so much for cleaning up the tests and removing the duplicated logic!

[Lunderberg]

Certainly, and I've moved the contents of ScheduleToPrimfunc back inside ScheduleToModule, where they were prior to this PR. That makes sense on avoiding passing around the PrimFuncs where possible, to avoid it appearing like the intended path to a casual reader.