apache / tvm

Open deep learning compiler stack for cpu, gpu and specialized accelerators
https://tvm.apache.org/
Apache License 2.0
11.41k stars 3.4k forks source link

[Relax] Apply DefaultGPUSchedule() in default build pipeline #17108

Open Lunderberg opened 2 weeks ago

Lunderberg commented 2 weeks ago

This is a follow-up to https://github.com/apache/tvm/pull/15864, which added LegalizeOps to the default Relax build pipeline. Since legalization may produce additional TIR PrimFuncs that require scheduling, the output of LegalizeOps typically must also be passed through tir.transform.DefaultGPUSchedule(). This PR adds DefaultGPUSchedule() to the relax build pipeline to handle these cases.

Scheduled PrimFunc have the "tir.is_scheduled" attribute set to true, and are ignored by DefaultGPUSchedule(). In addition, the DefaultGPUSchedule transform has no effect on non-GPU targets. Therefore, this change should only impact tvm.relax.build calls that previously resulted in an error due to unscheduled GPU functions, and should not have any impact on existing calls.

yongwww commented 2 weeks ago

Seems we discussed this before in the phase ordering discussion in open dev meeting, can't remember the reason why we didn't add this pass in default pipeline, need to take a look at the meeting note

Lunderberg commented 2 weeks ago

I don't remember that discussion, but I'd be interested in seeing the minutes from it. The main roadblocks to it are ones that I've resolved in previous PRs:

With those fixes applied, I think it's now a pretty big usability benefit to include it in the lowering pipeline.