[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 StructInfo 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.
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]

Dealing with shape variables has been one of the trickiest parts of the specification, in my opinion. Scoping them the same way as Relax variables seemed like a simple-enough rule but it leads to some problematic situations like having the value returned by SeqExpr whose shape expression contains shape variables defined inside the SeqExpr. Per this scoping rule, the shape of the SeqExpr would have to be RuntimeDepShape, which would lose information.

One alternative might be to scope shape variables to the entire function, but there are some potentially problematic cases, like initializing a var in one branch of an If and not the other. Allowing them to scope more broadly may be helpful for checking the shape_ of a SeqExpr, but creates complexity.

[slyubomirsky]

Again, I emphasize that as in the case for getting the shape of a call, it may be advantageous to do these transformations before filling in shapes, though it may be harder for users to reason about when they need a dynamic check in that case. Do we want to consider the possibility of doing canonicalization before shape inference?

[slyubomirsky]

return the result

We may want to make a guarantee about wrapping None results into unit (empty tuples). TVM has a runtime object corresponding to "none" values that is distinct from an empty tuple, but in Relax, we may want to treat these as one case. Alternatively, we can represent empty tuples as None values.

[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.

[sunggg]

@slyubomirsky, thank you for your great efforts to make the language more formal!
Overall, it looks good to me and I left some questions below.
I observed some discrepancy with this doc and our current implementation. Although they deliver the very similar message, I would like to suggest to match them as much as we can for the better readability and ramp-up experience for newcomers.

[sunggg]

Minor comment: it might be more readable if we match the rule with our function declarations by moving optional arguments to the behind.
e.g.,

@tvm._ffi.register_object("relax.expr.MatchShape")
class MatchShape(Binding):
    """Symbolic shape match, binds the variable of the lhs with the rhs."""

    value: Expr
    pattern: List[PrimExpr]
    var: Var

    def __init__(self, value: Expr, pattern: List[PrimExpr], var: Var, span: Span = None) -> None:
        self.__init_handle_by_constructor__(_ffi_api.MatchShape, value, pattern, var, span)

[slyubomirsky]

Maybe. I think it makes more logical sense for the var to come first in a binding

[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]

apache/tvm#14148 Reposted in the unity branch, withdrawing this PR now