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] Bind symbolic variables in R.match_cast #17034

Closed Lunderberg closed 2 months ago

Lunderberg commented 2 months ago

Prior to this commit, variable replacement by BindSymbolicVars would fail to replace variables that occur within a relax::MatchCast node. This pattern is rare, because the bind_symbolic_vars method can only replace variables that are exposed as part of the function signature, and most uses of relax::MatchCast act as a definition for symbolic variables that are not exposed through the function signature. This pattern is well-formed, though, since the relax::MatchCast node can also act as a user of previously-defined symbolic variables.

The root cause for this bug was in the ExprMutator visitor for relax::MatchCast, which did not visit the struct info field. As a result, the virtual ExprMutator::VisitPrimExpr function was not called for expressions that occur within the StructInfo of a relax::MatchCast. This commit updates ExprMutator to resolve this bug, and applies an analogous fix for ExprVisitor.

Lunderberg commented 2 months ago

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