[Draft][WIP] Relax language specification

[slyubomirsky]

Rendered view.

This document is a draft language specification for Relax. The purpose of the language specification is to serve as a technical reference describing the language's behavior in sufficient detail as to clarify the intended behavior of the compiler, hence it is by design a very detailed document rather than an accessible tutorial for the language. Its focus is the "what" and "how" of Relax, but not always the "why," though we can add more sections giving design reasons if that is desired.

Note that «double caret marks» (guillemets) are used to denote parts of the specification discussing functionality that the present prototype doesn't yet support. This notation is somewhat cumbersome, but I wasn't sure how else to proceed because Github Markdown does not support changing the text color (which was how my initial document indicated these areas). The caret marks may look strange in the text, but they have the benefit of being easy to find by text search.

Out of scope, for now, in this document is the subject of parsing: We should eventually document how we intend to parse Python into Relax, but the parser itself is being greatly reworked. We can revisit the issue of documenting its behavior once that work has been completed. Additionally, this specification is intended (for now) to focus on the user-visible behavior of Relax rather than specifying lower-level interfaces or the precise mechanisms of Relax's implementation.

Aspects Requiring Review or Still to be Determined

Since this document is a draft, any part of it is up for review and open to revision, but certain parts of the document have proven particularly challenging to describe and could benefit the most from community discussions.

1. The shape system has been a great challenge to specify and there are many questions as to how it should work. One question is whether "strong shaping" might lead to too many error messages for something that could be checked dynamically. Additionally, many potential shape mismatches could be eliminated using constant propagation or other transformations: If we check shapes without applying transformations, we would force users to add lots of redundant shape checks. On the other hand, if we require these transformations first, that might make the code harder for users to reason about. Some of the more cumbersome rules like those related to function calls or tuple indexing might be possible to eliminate by declaring the results to have RuntimeDepShape, but that forces the user to insert more dynamic checks and decreases the amount of automation in the shape system.
2. The run-time representations of values in the language will be important for determining how PackedFuncs can interact with Relax values. However, embedded targets do not support the TVM object system, so describing values in terms of the TVM object system directly may not work for all settings. Additionally, it should be determined how much detail about the representations should be included in the specification.
3. Operators used for core language functions (like call_tir) should be described in the specification. I am not certain the descriptions presently in the last section are entirely correct, so more review of them should be appreciated. Additionally, are there operators that should be there but are presently missing?
4. Finally, there is the question of process: How will we permit the specification to be revised? Does any change require a fresh RFC? Is there a threshold for changes that can be done as a direct PR? I have not considered this question directly, but the language specification is an important document for the community and changes to the language specification should not be taken lightly.

There are also some more minor TODOs throughout the document.

The Future of This Document

Eventually, we will want this document to be part of the Relax documentation, in which case it will be placed into a different location in the repo and probably be formatted as an rst file. Before that, we will officially RFC the spec into TVM to allow for maximum public discussion as to the design decisions underlying the specification. Hence, I am going to leave this document as a "WIP PR" until it is officially RFC'd.

[slyubomirsky]

See this earlier discussion for thoughts on whether we should use tir::Any to represent a wildcard dimension. For simplicity, the current proposal does not use a per-dimension wildcard and instead requires the entire shape to be relaxed to RuntimeDepShape, but this comes at the cost of partial shape information. It is worth considering whether such partial shape information is worth it for compilation. (As the draft details in the "Possible Extensions to the Shape Expression System" section, we can potentially add such a feature in a later revision without breaking old code.)

[slyubomirsky]

Another issue that has come to my attention thanks to my colleagues @YuchenJin, @sunggg, and @psrivas2 is that while it may be reasonable to leave certain kinds of data structures as opaque "Relax objects" that we can define later because they don't directly interact with many other existing language features, there are some types that we might reasonably want to add into the language that would interact with a lot of existing features.

Namely, there might be tensor variants that we would want to consider: dense vs sparse tensors, as well as ragged tensors. Adding these into the language design later could complicate how we deal with tensors, so it is likely to be worthwhile to consider these topics at this early stage in the language design process and possibly include these directly in the specification. That would be far preferable to later making changes that could potentially break backwards compatibility or require radically revising compiler passes. I encourage any community members interested in tensor variants like these to comment on those issues as well.

[altanh]

Another issue that has come to my attention thanks to my colleagues @YuchenJin, @sunggg, and @psrivas2 is that while it may be reasonable to leave certain kinds of data structures as opaque "Relax objects" that we can define later because they don't directly interact with many other existing language features, there are some types that we might reasonably want to add into the language that would interact with a lot of existing features.

Namely, there might be tensor variants that we would want to consider: dense vs sparse tensors, as well as ragged tensors. Adding these into the language design later could complicate how we deal with tensors, so it is likely to be worthwhile to consider these topics at this early stage in the language design process and possibly include these directly in the specification. That would be far preferable to later making changes that could potentially break backwards compatibility or require radically revising compiler passes. I encourage any community members interested in tensor variants like these to comment on those issues as well.

I highly recommend taking a look at MLIR's current approach to sparse tensors: https://mlir.llvm.org/docs/Dialects/SparseTensorOps/ (which is inspired by the TACO project). I'm not sure how hard it would be to design something similar for Relax; MLIR has a lot more machinery around "attributes" and just general metaprogramming stuff like tablegen. At a high level though, I think the MLIR approach is to keep high level IR focused on "abstract" tensors with no concern for data layout/representation, and use attributes at some stage to hijack code generation for specialized backend representations.

I think @yzh119 has done a ton of work on the TIR side for sparse stuff, so his thoughts here would definitely be valuable!

[slyubomirsky]

Thanks for the links. Leaving density/sparsity for later would be an option in that case, though ragged tensors would require rules for handling their shapes (perhaps the per-dimension Any would be a viable option there? This is actually not equivalent to using a dummy variable for _, since different tensors could each have a different value for that dimension)

[YuchenJin]

Thanks for the links. Leaving density/sparsity for later would be an option in that case, though ragged tensors would require rules for handling their shapes (perhaps the per-dimension Any would be a viable option there? This is actually not equivalent to using a dummy variable for _, since different tensors could each have a different value for that dimension)

@MasterJH5574 and @VertexC are working on native ragged tensor support in Relax, and they are designing the shape for it.

[slyubomirsky]

In today's community meeting, @tqchen raised the idea that our shape analysis should be better thought of as "best-effort" (making static guarantees when it can but not making claims of completeness), where we treat user shape annotations as assumptions, so I'm wondering if that perspective might allow us to simplify the shape inference areas of the spec. One reason to go with this interpretation is (as this section of the spec discusses) that different passes and transformations might uncover more information that might allow for drawing conclusions in more cases.

Provisionally, I am imagining that (as we discussed at the community meeting), the compiler would be required to insert dynamic checks of shapes at function and operator boundaries unless it can statically prove that shapes match.

From that perspective, would it then make sense not to raise errors or warnings if the compiler cannot conclude whether two shapes are definitely not equal (if it concludes two shapes are definitely not equal, it should certainly raise an error)?

I'm curious for more thoughts on this, since shape inference is certainly the trickiest area of the specification.

[slyubomirsky]

The shape inference issue may be best decided in #293, which could replace shape_ with something a little better-defined.

[slyubomirsky]

I've added Relax's normal form to the specification, since in the Dec. 13 community meeting, we noted that StructInfo (will be specified once we know more about how it should work) will likely rely on it.

[slyubomirsky]

I've overhauled the spec to deal with the newly implemented structural information system, which was a substantial challenge to specify. It is much more powerful than shape_ was previously, but requires quite a bit more machinery in the specification. I invite close review from those working on the structural information system: @tqchen @Hzfengsy @MasterJH5574

[slyubomirsky]

Now updated to account for StructInfo and sinfo_args.

[slyubomirsky]

https://github.com/apache/tvm/pull/14148 Reposted in the unity branch, withdrawing this PR now