[Relax] Macro recfactor, R.macro prototype

[kparzysz-quic]

If this is acceptable, I'll break it up into the TIR part and the Relax part, and create PRs in main and unity respectively.

Edit: added description below

---

As a background info---the parser (for either dialect, TIR or Relax) works by visiting a "statement" (or top-level expression) at a time. The expression parts of the statement 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 Relax, on the other hand, there was 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, Relax 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.

---

Relax macros generate a SeqExpr containing the evaluated body of the macro. A Relax macro should end with a return statement with the expression that will become the value of the macro. A return is normally not allowed in a SeqExpr, so macro parsing intercepts it to get its argument, after which the return statement itself is discarded. The argument is then set as the "body" of the SeqExpr. Everything preceding the return becomes a binding block in the SeqExpr.

Example:

import tvm
from tvm import relax
from tvm.script import relax as R

@R.macro
def alloc_and_shape(dtype: str):
    alloc = R.builtin.alloc_tensor(R.shape([4, 4]), runtime_device_index=0, dtype=dtype)
    shape = R.shape_of(alloc)
    return shape

@R.function
def foo(x: R.Tensor((4, 4), "float32")):
    shape = alloc_and_shape(dtype="float32")
    return shape


print(alloc_and_shape)
print()
print(foo)

Output:

@R.macro
def alloc_and_shape(dtype: str):
    alloc = R.builtin.alloc_tensor(R.shape([4, 4]), runtime_device_index=0, dtype=dtype)
    shape = R.shape_of(alloc)
    return shape


# from tvm.script import relax as R

@R.function
def foo(x: R.Tensor((4, 4), dtype="float32")) -> R.Shape([4, 4]):
    alloc: R.Tensor((4, 4), dtype="float32") = R.builtin.alloc_tensor(R.shape([4, 4]), R.dtype("float32"), R.prim_value(0))
    shape: R.Shape([4, 4]) = R.shape_of(alloc)
    shape_1: R.Shape([4, 4]) = shape
    return shape_1

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

_{Generated by tvm-bot}

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

_{Generated by tvm-bot}

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

_{Generated by tvm-bot}

[kparzysz-quic]

cc: @Hzfengsy @yzh119 @cyx-6

[tqchen]

also cc @slyubomirsky

[kparzysz-quic]

Ping.

[Hzfengsy]

I'm sorry for being late. I thought it was WIP as it's a draft PR.

The current implementation looks good to me and please add testcases, then we can get it in

[kparzysz-quic]

Thanks! Created a PR for the non-Relax part here: #15432

[kparzysz-quic]

The Relax part is in #15455.