Skip to content
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

Merged
merged 2 commits into from
Apr 20, 2024

Conversation

Lunderberg
Copy link
Contributor

@Lunderberg Lunderberg commented Apr 9, 2024

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 an exception
(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.

@Lunderberg Lunderberg marked this pull request as draft April 9, 2024 14:00
@Lunderberg
Copy link
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.
@Lunderberg Lunderberg force-pushed the relax_qol_seqexpr_in_cpp_type branch from fbcf057 to 0abd154 Compare April 12, 2024 12:43
@Lunderberg Lunderberg marked this pull request as ready for review April 19, 2024 14:30
Copy link
Contributor

@quic-sanirudh quic-sanirudh left a 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.

include/tvm/relax/expr.h Show resolved Hide resolved
@quic-sanirudh quic-sanirudh merged commit a2511cc into apache:main Apr 20, 2024
22 checks passed
@Lunderberg Lunderberg deleted the relax_qol_seqexpr_in_cpp_type branch April 20, 2024 11:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants