As a background info---the script parser works by visiting a "statement"
(or top-level expression) at a time. The expression parts of the state-
ment are evaluated, and then the IR corresponding to the statement is
constructed if necessary.
In TIR, macro calls can only occur at the statement level, and they don't
produce any values. This means that the statement visitor (visit_expr_stmt)
can see these calls directly in its node parameter. At this point it could
simply visit the body of the macro instead, which is the basis of the
existing implementation.
In other dialects there may be a need for macros to produce values. This
means that macro calls can occur in the middle of complex expressions.
As a result, these calls will not be present at the statement level, and
the TIR approach by intercepting them in visit_expr_stmt will no longer
work. Instead, these macros delay the visiting of the macro body to the
evaluation time. A macro is represented by an ScriptMacro (TIRMacro in
the current implementation) object (created via macro decorator). When the
evaluator evaluates an expression with a macro call, it will call the
macro object (since macro calls use function call syntax). It is in the
macro object's __call__ function where the macro parsing picks up. The
remaining issue was to pass the Parser object to the __call__ function.
This is done by injecting it into the global dictionary under a reserved
name.
It turns out that the same approach also works for TIR, and the macro
processing can be generalized, leaving only language-specific details to
the language-specific language macro objects.
* [TIR] Implement TIR macros
This patch introduces two new symbols: `T.macro` and `T.insert`.
`T.macro` is a decorator that, when applied to a function, turns the
body of that function into a piece of TIR that can be inserted via
`T.insert` into a PrimFunc.
For example:
```python
@T.macro
def copy_backwards(dst, src, size):
with T.block("backwards"):
for i in T.serial(size):
ai = T.axis.remap("S", [i])
T.reads(src[0:size])
T.writes(dst[0:size])
dst[ai] = src[size - ai - 1]
@T.prim_func
def foo_int32(A: T.Buffer((128,), "int32"), B: T.Buffer((128,), "int32")):
T.insert(copy_backwards, A, B, 128)
@T.prim_func
def foo_int8(A: T.Buffer((128,), "int8"), B: T.Buffer((128,), "int8")):
T.insert(copy_backwards, A, B, 128)
```
The above will generate two PrimFuncs that do the same backwards copy,
but applied to buffers with different data types.
Semantics:
- Function that is decorated with @T.macro can have any parameters that
follow Python syntax, i.e. positional, keyword, etc. Type annotations
are not required, but are allowed.
- The arguments to `T.insert` are macro name followed by the argument
list.
For `T.insert(arg1, arg2, arg3, ...)`, the values are substituted into
the body of the macro as in the call `arg1(arg2, arg3, ...)`.
The body with the substituted values is then inserted at the point
where the `T.insert` is located.
* Fix linter
* Fix linter again
One linter suggested something that the other didn't like...
* Get rid of T.insert, apply macro via function-call syntax
* Store closure vars in TIRMacro
* ast.parse always returns ast.Module, hence doc is doc.Module
* Simplify `expand_macro`, capture environment variables
* Implement macro hygiene
* Fix linter
* Make T.macro work same as T.macro()
The previous commit inadvertently made T.macro (without parentheses)
illegal, only abbreviated form allowed was T.macro(). Restore T.macro
as a valid decorator use.
* Edit comment: insertion -> expansion
* Add import pytest
* One more typo...
* Remove stale testcase
This PR introduces some minor restructuring of the `python/tvm/script`
folder structure to make it more convenient for future upstreaming.
Co-authored-by: Yaxing Cai <caiyaxing666@gmail.com>
Tianqi Chen