slyubomirsky commented on code in PR #14148: URL: https://github.com/apache/tvm/pull/14148#discussion_r1123788976
########## relax_spec.md: ########## @@ -0,0 +1,803 @@ +# Informal Relax Language Specification + +Note: Text in «double chevrons» indicates features not present in the current prototype. + +In order to develop and test Relax, it is important for compiler developers to agree on what a given program in Relax means and what makes it valid so that test cases can be evaluated independently of any particular Relax implementation. This document is intended to describe Relax's grammar constructs (its [abstract syntax tree](https://en.wikipedia.org/wiki/Abstract_syntax_tree), or AST), the semantics of its grammar (what the different constructs mean), Relax's type system and type-checking rules (what makes a Relax program valid), and its rules for reasoning about structural information (such as tensor shapes) in detailed though still informal terms. If necessary, we may encode these rules more formally to allow for more automated analysis. + +Though this document will use the TVMScript front end for some examples, specifying the mapping from Python's AST to Relax's AST will be deferred until the parser becomes more stable. + +# Table of Contents + +1. [Overview](#overview) +2. [Top-Level Program Organization](#top-level-program-organization-irmodule) +3. [Values in Relax](#values-in-relax) +4. [Variable Scoping](#variable-scoping) +5. [Normal Form](#normal-form) +6. [Well-Formedness Criteria](#well-formedness-criteria) +7. [Structural Information in Relax](#structural-information-in-relax) +8. [Semantics](#detailed-semantics) + +# Overview + +This section will outline the grammar of Relax and give very brief descriptions of the different components, including the semantics and structural information (`StructInfo`) system. The rest of this document will provide more detailed descriptions of these facets of the language, including the validity conditions that the `StructInfo` system upholds. + +## Differences from Relay + +Per the [original workshop paper](https://arxiv.org/abs/1810.00952) and the [later report](https://arxiv.org/abs/1904.08368), Relay was designed to be a high-level functional language for expressing deep learning models at a high level. While Relay is not entirely pure (the `Ref` type is modeled after reference types in SML and similar functional languages), the assumption in Relay is that tensor operators are generally pure, meaning that they do not change the program state other than by producing new values. Additionally, Relay's type system also requires operators to have type relations that infer static tensor types or conclude that a dimension is unknown at compile time (`Any`). The need to register type relations and ensure operators' purity makes it difficult to add new operators to Relay and particularly difficult to call directly into TIR or external libraries, which are often not pure; any such extension requires adding new operators and abstracting over any impurity. + +While Relax aims to be as general and expressive as Relay, Relax is intended to make it much easier to interoperate with external libraries and especially with TIR. In particular, Relax includes a mechanism for calling arbitrary TVM `PackedFunc`s (which can call external libraries) and special support for TIR. The language accordingly does not assume that such operations are pure, though this does require reasoning about aliasing and similar issues. Additionally, tensor shapes are no longer handled during type checking; each expression has associated structural information associated with it, in addition to a type. This structural information supports static reasoning about tensor shapes in many cases, but also facilitates a fallback to dynamic checking when that is not possible. This approach to shapes allows for richer shape constraints and other structural properties to be checked at run time (such as with _symbolic_ shapes, where some dimensions are variables) and allows for mor e quickly integrating calls into TIR or external libraries into Relax code by obviating the need for type relations. + +## Grammar + +Below is a diagram of the various AST constructs in Relax, including types. In code, these are defined on the C++ side in `include/tvm/relax/{expr.h, type.h}` and in Python in `python/tvm/relax/{expr.py, ty.py}`. This diagram will give the names of the AST nodes and the types and names of their members. The semantics will describe what computation each construct represents; an AST is simply data. A Relax program consists of an `IRModule` with global variables bound to Relax functions that implement the computations of interest. + +(On the notation: `[x]` means "a list of `x`," `x?` means "optionally `x`," `{x: y}` means "a map of `x` to `y`," `x | y` means "`x` or `y`," and `#` is used for comments.) + +``` +# PrimExprs are defined in TIR, see include/tvm/tir/expr.h +# They are intended to have the same semantics as in TIR +PrimExpr ::= + Var(name: string) # shape variables + | IntImm(value: int64) + | Add(a: PrimExpr, b: PrimExpr) + | Sub(a: PrimExpr, b: PrimExpr) + | Mul(a: PrimExpr, b: PrimExpr) + | Div(a: PrimExpr, b: PrimExpr) + | Min(a: PrimExpr, b: PrimExpr) + | Max(a: PrimExpr, b: PrimExpr) + | Not(a: PrimExpr) + | And(a: PrimExpr, b: PrimExpr) + | Or(a: PrimExpr, b: PrimExpr) + | Select(condition: PrimExpr, true_value: PrimExpr, false_value: PrimExpr) + # (others may be added later, as deemed necessary) + +# Also from TIR +DataType ::= Int(bits: int, lanes: int) + | UInt(bits: int, lanes: int) + | Float(bits: int, lanes: int) + | Handle(bits: int, lanes: int) + +StructInfo ::= TensorStructInfo(shape: Expr?, dtype: DataType, ndim: int) + | ShapeStructInfo(values: [PrimExpr]?, ndim: int) + | PrimStructInfo(dtype: DataType) + | ObjectStructInfo() + | TupleStructInfo(fields: [StructInfo]) + | FuncStructInfo(params: [StructInfo]?, ret: StructInfo, derive_func: EnvFunc?*) + +# expressions +Expr ::= Constant(data: NDArray) + # scoped to functions or SeqExprs + | Var(name_hint: string, struct_info_annotation: StructInfo?) + # scoped to DataflowBlocks + | DataflowVar(name_hint: string, struct_info_annotation: StructInfo?) + | GlobalVar(name_hint: string) + | Tuple(fields: [Expr]) + | SeqExpr(blocks: [BindingBlock], body: Expr) + | PrimValue(value: PrimExpr) + | StringImm(value: string) + | DataTypeImm(value: DataType) + | Function(params: [Var], body: Expr, ret_struct_info: StructInfo?, attrs: Attrs?) + | If(cond: Expr, true_branch: Expr, false_branch: Expr) + | ExternFunc(global_symbol: string) + | Call(op: Expr, args: [Expr], sinfo_args: [StructInfo], attrs: Attrs?) + | ShapeExpr(values: [PrimExpr]) + | TupleGetItem(tuple_value: Expr, index: int) + | Op(op_name: string) + +# binding blocks (analogous to sequence of statements) +BindingBlock ::= + BindingBlock(bindings: [Binding]) + | DataflowBlock(bindings: [Binding]) + +# bindings (analogous to statements) +Binding ::= + VarBinding(var: Var|DataflowVar, value: Expr) + | MatchCast(var: (Var|DataflowVar)?, struct_info: StructInfo, value: Expr) + +# Relax programs are IRModules. Modules may bind global variables either to +# Relax functions or TIR PrimFuncs (specified separately). +# The Relax compiler may analyze and modify the TIR PrimFUncs as well. +Program ::= IRModule(funcs: {GlobalVar: Function|PrimFunc}) +``` + +### Notes on `derive_func` + +The `derive_func` field of `FuncStructInfo` is a macro in the meta-language: Given a function call and the variable mapping context, return the `StructInfo` of the result. This field is used only at compile time for reasoning about the `StructInfo` of calls to `ExternFunc`s. + +### Notes on `DataType` and Related Terminology + +The representation of datatypes, `DataType`, in the above AST is taken directly from TIR. However, the usage of datatypes in Relax is more restricted than in TIR. +1. The `lanes` field for the `Int`, `UInt`, and `Float` datatypes must always be 1; we do not directly consider vectorized values in Relax. +2. The `lanes` field for the `Handle` datatype must always be 0, indicating that it is `Void` (see below). The `bits` field for `Handle` should always be set to 64 (it will not be used by Relax). + +We also define the following special notation for datatypes, to be used in the rest of the specification: +1. `Bool()`: This is shorthand for `UInt(bits=1, lanes=1)`, since TIR does not have a separate Boolean type. "True" refers to a value of 1 in this datatype and "false" refers to a value of 0. For convenience, we will refer to Boolean values as a separate datatype in the specification, due to their significance in `If` nodes. +2. `Void()`: This is shorthand for `Handle(bits=64, lanes=0)`. TIR uses this datatype to refer to opaque objects; in Relax, it is used to denote an unknown datatype. + +## Expression Survey + +This specification provides a more detailed description of what each expression and `StructInfo` represents and what conditions make them valid. To motivate and provide more context for the full specification later in this document, this section will briefly summarize the purpose of each node. + +1. `Constant` nodes construct tensor constants (n-dimensional arrays of scalars). +2. `Tuple` nodes construct a tuple (immutable fixed-size ordered grouping) of Relax values. +3. `Var`, `DataflowVar`, and `GlobalVar` nodes are all variables, referring to named stored values of different kinds. Variables in Relax must be bound exactly once. `GlobalVar`s are bound in the `IRModule` itself and refer to Relax functions or TIR `PrimFunc`s. `Var` nodes are bound either within functions, where they represent function parameters, or in `VarBinding` or `MatchCast` nodes in `BindingBlock`s, as we will discuss below. `DataflowVar`s are similar to `Var`s and can be bound only within `DataflowBlock`s. +4. `PrimExpr`s are used to represent dimensions of shapes in `ShapeExpr` and `MatchCast` nodes. These represent operations on integers with their own `Var` nodes (`tir::Var`), which we will refer to as "shape variables". Shape variables can only be used in other `PrimExpr`s and are scoped like `Var` nodes (`relax::Var`), which we will call "Relax variables." +5. `ExternFunc` nodes evaluate into `PackedFunc`s; the implementation will look up the registered `PackedFunc` by its global symbol. +6. `PrimValue` nodes construct immutable scalar values from `PrimExpr`s, primarily for interacting with `ExternFunc`s or operators. These scalars are boxed within TVM objects, allowing them to be nested inside TVM's containers. (By contrast, zero-dimensional tensors defined via `Constant` are mutable.) +7. `StringImm` nodes construct strings, intended primarily for interacting with `ExternFunc`s or operators. +8. `DataTypeImm` nodes construct representations of TIR datatypes, intended primarily for interacting with `ExternFunc`s or operators (e.g., for TIR intrinsics that take a datatype as an input). +9. `Call` nodes represent function calls. The callee argument (the `op`) can be an `ExternFunc` node (representing a call to a `PackedFunc`), an `Op` node (representing a call to a Relax operator), or an arbitrary expression. + 1. `Op` nodes refer to built-in Relax operators, which the compiler is free to implement as is deemed appropriate. Certain operators implement important operations, like `call_tir` (allows for calling TIR `PrimFunc`s). + 2. Any other expression must evaluate to a `PackedFunc` or a closure; the result of evaluating `op` will then be called with the given arguments. + + Calls to `ExternFunc`s and operators may perform side effects, hence it is important to reason about whether a function call is permitted inside a `DataflowBlock`. + +10. `If` nodes represent branching control flow. First the condition expression is evaluated, and it must evaluate to a Boolean scalar. If the condition is true, the true branch is evaluated and its result is used; otherwise, the false branch is evaluated and its result is used. +11. `TupleGetItem` nodes represent tuple indexing. The `tuple_value` expression must evaluate to a tuple with at least `index + 1` items and the item with the given index will be returned. +12. `SeqExpr` describes a sequence of binding blocks followed by a return expression. The `SeqExpr` opens a new scope. Its binding blocks are evaluated in order and add new variables to the scope. Binding blocks are either ordinary `BindingBlock`s or `DataflowBlock`s and both consist of a series of bindings. `DataflowBlock`s are the only kind allowed to introduce bindings with `DataflowVar`s and it does not permit any constructs featuring control flow (`If` nodes or recursive calls) or calls to (possibly) impure functions. There are two different kinds of bindings: + 1. `VarBinding`s: The `value` expression (the right-hand side of the binding) of the binding is evaluated first and is bound to the `var` expression, which must be a new `Var` or `DataflowVar` (in a dataflow block). The newly bound variable will have that value for the remainder of the scope (`DataflowVar`s are scoped only to the `DataflowBlock` in which they appear; `Var`s are scoped to the entire `SeqExpr`). Review Comment: Good point. I should rephrase that to mean from the point of its definition in a binding onwards. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
