-
Notifications
You must be signed in to change notification settings - Fork 3.5k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[QoL][Relax] Use SeqExpr in IR types when SeqExpr is required #16859
[QoL][Relax] Use SeqExpr in IR types when SeqExpr is required #16859
Conversation
This PR is currently marked as a draft, to allow time for community discussion. A forum thread (link) for discussion has also been started for this change. |
The Relax IR requires the `FunctionNode::body`, `IfNode::true_branch`, and `IfNode::false_branch` to be instances of `relax::SeqExpr`. If these Relax requirements are violated, correctly-implemented transformations may raise exceptsion (e.g. from `Downcast` in `Downcast<SeqExpr>(func->body)->blocks`), or even segfault (e.g. when `.as` returns a nullptr in `func->body.as<SeqExprNode>()->blocks`). Debugging these failure modes is also difficult, as even the TVMScript printer relies on the body of the function being a `SeqExprNode`. This commit updates the C++ type of `FunctionNode::body`, `IfNode::true_branch`, and `IfNode::false_branch` to be `relax::SeqExpr` instead of `relax::Expr`. This does not impact any well-formed Relax IR, and allows this type of ill-formed Relax IR type to be checked at compile-time. A large number of checks applied during TVM runtime can now be removed, as they duplicate the new compile-time check. To maintain backwards compatibility, this commit adds a new constructor to `relax::SeqExpr`, which accepts a single `Expr body` argument. This constructor provides either an additional reference to the same underlying `relax::SeqExprNode`, if `body` already contains a `relax::SeqExprNode`, and otherwise wraps the body in a `relax::SeqExpr`. For implementations that previously produced well-formed Relax IR, this change has no effect. For implementations that previously produced ill-formed Relax IR, this change results in the equivalent well-formed Relax IR. Alternate implementations considered: * Perform the backwards-compatibility wrapping within the `relax::Function` and `relax::If` constructors. While this would provide the intended conversion when these constructors are used, Relax transforms make frequent use of copy-on-write (e.g. `func.CopyOnWrite()->body = new_body`), which does not use the constructor. Maintaining backwards compatibility for this usage requires the implicit conversion constructor that was chosen for this PR. * Remove the Relax IR requirement for these expressions to be `SeqExpr`. While this would make Relax more internally consistent, such a change would break backwards compatibility that relies on `SeqExpr` being present. While the callsites within TVM could be updated to resolve this breakage, callsites outside of TVM (e.g. MLC-LLM) could not. Exposing the special case within the C++ type, as done in this PR, maintains backwards compatibility.
All breakage was the result of callers relying on ill-formed Relax maintaining that specific type form of ill-formed-ness.
fbcf057
to
0abd154
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks a lot for this. Looks good to me apart from a minor comment about possibly modifying IfNode
in place to improve git history maintenance to just highlight the changed lines instead of the whole class.
The Relax IR requires the
FunctionNode::body
,IfNode::true_branch
, andIfNode::false_branch
to be instances ofrelax::SeqExpr
. If these Relax requirements are violated, correctly-implemented transformations may raise an exception(e.g. from
Downcast
inDowncast<SeqExpr>(func->body)->blocks
), or even segfault (e.g. when.as
returns a nullptr infunc->body.as<SeqExprNode>()->blocks
). Debugging these failure modes is also difficult, as even the TVMScript printer relies on the body of the function being aSeqExprNode
.This commit updates the C++ type of
FunctionNode::body
,IfNode::true_branch
, andIfNode::false_branch
to berelax::SeqExpr
instead ofrelax::Expr
. This does not impact any well-formed Relax IR, and allows this type of ill-formed Relax IR type to be checked at compile-time. A large number of checks applied during TVM runtime can now be removed, as they duplicate the new compile-time check.To maintain backwards compatibility, this commit adds a new constructor to
relax::SeqExpr
, which accepts a singleExpr body
argument. This constructor provides either an additional reference to the same underlyingrelax::SeqExprNode
, ifbody
already contains arelax::SeqExprNode
, and otherwise wraps the body in arelax::SeqExpr
. For implementations that previously produced well-formed Relax IR, this change has no effect. For implementations that previously produced ill-formed Relax IR, this change results in the equivalent well-formed Relax IR.Alternate implementations considered:
Perform the backwards-compatibility wrapping within the
relax::Function
andrelax::If
constructors. While this would provide the intended conversion when these constructors are used, Relax transforms make frequent use of copy-on-write (e.g.func.CopyOnWrite()->body = new_body
), which does not use the constructor. Maintaining backwards compatibility for this usage requires the implicit conversion constructor that was chosen for this PR.Remove the Relax IR requirement for these expressions to be
SeqExpr
. While this would make Relax more internally consistent, such a change would break backwards compatibility that relies onSeqExpr
being present. While the callsites within TVM could be updated to resolve this breakage, callsites outside of TVM (e.g. MLC-LLM) could not. Exposing the special case within the C++ type, as done in this PR, maintains backwards compatibility.