[Bug] Match Cast Missing After Legalization

[zxybazh]

When we were trying to get Expand operator supported, we tried to legalize broadcast_to op and found that the operator would fail well formed check after legalization. The cause is match cast went missing during legalization.

The reproducible script is here: https://gist.github.com/zxybazh/5b09bad2cf13ea7d15329a86b03a6abf

CC @MasterJH5574 @tqchen

[MasterJH5574]

LegalizeOps does a step of RemoveAllUnused here.

relax/python/tvm/relax/transform/legalize_ops.py

Line 777 in ce5c7f4

updated_func = remove_all_unused(updated_func)

It looks to me that the RemoveAllUnused regards the MatchCast as an unused binding, and thereby removes it. While in fact, the MatchCast is a binding with side-effect (where the side-effect is “defining x_0 and x_1).

As a workaround, you can probably try removing the step of RemoveAllUnused. In the long term, I think RemoveAllUnused should be updated to skip the bindings with side-effects, or we introducing a better DCE pass.

[tqchen]

This should be a bug of remove unused not considering dep through symbolic variables and we should fix it. As a quick fix, we can avoid removing match cast for now. Then apply deeper analysis later