[QoL][Relax] Use SeqExpr in IR types when SeqExpr is required#16859
Merged
quic-sanirudh merged 2 commits intoapache:mainfrom Apr 20, 2024
Merged
[QoL][Relax] Use SeqExpr in IR types when SeqExpr is required#16859quic-sanirudh merged 2 commits intoapache:mainfrom
quic-sanirudh merged 2 commits intoapache:mainfrom
Conversation
Contributor
Author
|
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
quic-sanirudh
approved these changes
Apr 19, 2024
Contributor
quic-sanirudh
left a comment
There was a problem hiding this comment.
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The Relax IR requires the
FunctionNode::body,IfNode::true_branch, andIfNode::false_branchto be instances ofrelax::SeqExpr. If these Relax requirements are violated, correctly-implemented transformations may raise an exception(e.g. from
DowncastinDowncast<SeqExpr>(func->body)->blocks), or even segfault (e.g. when.asreturns 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_branchto berelax::SeqExprinstead 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 bodyargument. This constructor provides either an additional reference to the same underlyingrelax::SeqExprNode, ifbodyalready 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::Functionandrelax::Ifconstructors. 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 onSeqExprbeing 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.