[Relax][Bugfix] Apply FuseOps to nested DataflowBlock

apache / tvm

Open deep learning compiler stack for cpu, gpu and specialized accelerators

https://tvm.apache.org/

Apache License 2.0

11.46k stars 3.41k forks source link

[Relax][Bugfix] Apply FuseOps to nested DataflowBlock #17033

Closed Lunderberg closed 2 months ago

Lunderberg commented 2 months ago

While it is ill-formed for control-flow to occur within a DataflowBlock, it is legal for a DataflowBlock to be contained within a control-flow. Prior to this commit, the FuseOps and FuseOpsByPattern transforms erroneously skipped DataflowBlock instances that were contained within a relax::If node.

This commit updates FuseOps to apply operator fusion to any dataflow block, regardless of whether it is found at the top level of a a Relax function.

Lunderberg commented 2 months ago

This PR is identical to the now-closed PR https://github.com/apache/tvm/pull/17022, as a CI restart prevented the CI from running on the original PR.