[DRAFT][TIR] Prototype of T.macro

[kparzysz-quic]

Initial draft of T.macro as discussed in https://discuss.tvm.apache.org/t/discussion-tir-macros/15247.

The basic idea is to introduce a decorator T.macro that, similarly to T.prim_func would annotate a piece of code as a TIR macro. The macro would not be subject to any parameter restrictions, since the intent is to simply insert the macro's body at the place of use, substituting parameters with the actual arguments. Since no function call is generated, no ABI constraints apply.

A macro is used with a pseudo function call T.include, where the first argument is the macro name followed by actual arguments to the macro.

Functionally, the macro is stored as doc.AST, and no further processing is done on the macro definition (until it's used).

In the TIR parser, when doc.Call is encountered and the call is to T.include, this expression is treated specially:

1. The definition of the macro is located.
2. New variable frame is created.
3. A series of assignments (parameter = argument) is evaluated.
4. The macro's body is further processed by the parser.

Example:

import tvm
from tvm.script import tir as T

@T.macro
def foo1(a, b, n):
    with T.block(n):
        a[0] = b[0]

@T.prim_func
def bar(A: T.Buffer((2,), "int32"), B: T.Buffer((2,), "int32")):
    T.include(foo1, A, B, "block1")

print(bar)
print(foo1)

Output:

# from tvm.script import tir as T

@T.prim_func
def main(A: T.Buffer((2,), "int32"), B: T.Buffer((2,), "int32")):
    with T.block("block1"):
        T.reads(B[0])
        T.writes(A[0])
        A[0] = B[0]
<tvm.script.parser.core.doc_core.Module object at 0x7fa20920de20>

bar has the macro expanded in its body, while foo1 prints as <tvm.script...Module>.

The detailed representation of the `foo1` macro in JSON format

{
  "AST_class": "Module",
  "body": [
    {
      "AST_class": "FunctionDef",
      "name": "foo1",
      "args": {
        "AST_class": "arguments",
        "args": [
          {
            "AST_class": "arg",
            "arg": "a",
            "lineno": 2,
            "col_offset": 9,
            "end_lineno": 2,
            "end_col_offset": 10
          },
          {
            "AST_class": "arg",
            "arg": "b",
            "lineno": 2,
            "col_offset": 12,
            "end_lineno": 2,
            "end_col_offset": 13
          }
        ],
        "kwonlyargs": [],
        "kw_defaults": [],
        "defaults": [],
        "posonlyargs": []
      },
      "body": [
        {
          "AST_class": "With",
          "items": [
            {
              "AST_class": "withitem",
              "context_expr": {
                "AST_class": "Call",
                "func": {
                  "AST_class": "Attribute",
                  "value": {
                    "AST_class": "Name",
                    "id": "T",
                    "ctx": {
                      "AST_class": "Load"
                    },
                    "lineno": 3,
                    "col_offset": 9,
                    "end_lineno": 3,
                    "end_col_offset": 10
                  },
                  "attr": "block",
                  "ctx": {
                    "AST_class": "Load"
                  },
                  "lineno": 3,
                  "col_offset": 9,
                  "end_lineno": 3,
                  "end_col_offset": 16
                },
                "args": [
                  {
                    "AST_class": "Constant",
                    "value": "block1",
                    "s": "block1",
                    "n": "block1",
                    "lineno": 3,
                    "col_offset": 17,
                    "end_lineno": 3,
                    "end_col_offset": 25
                  }
                ],
                "keywords": [],
                "lineno": 3,
                "col_offset": 9,
                "end_lineno": 3,
                "end_col_offset": 26
              }
            }
          ],
          "body": [
            {
              "AST_class": "Assign",
              "targets": [
                {
                  "AST_class": "Subscript",
                  "value": {
                    "AST_class": "Name",
                    "id": "a",
                    "ctx": {
                      "AST_class": "Load"
                    },
                    "lineno": 4,
                    "col_offset": 8,
                    "end_lineno": 4,
                    "end_col_offset": 9
                  },
                  "slice": {
                    "AST_class": "Constant",
                    "value": 0,
                    "s": 0,
                    "n": 0,
                    "lineno": 4,
                    "col_offset": 10,
                    "end_lineno": 4,
                    "end_col_offset": 11
                  },
                  "ctx": {
                    "AST_class": "Store"
                  },
                  "lineno": 4,
                  "col_offset": 8,
                  "end_lineno": 4,
                  "end_col_offset": 12
                }
              ],
              "value": {
                "AST_class": "Subscript",
                "value": {
                  "AST_class": "Name",
                  "id": "b",
                  "ctx": {
                    "AST_class": "Load"
                  },
                  "lineno": 4,
                  "col_offset": 15,
                  "end_lineno": 4,
                  "end_col_offset": 16
                },
                "slice": {
                  "AST_class": "Constant",
                  "value": 0,
                  "s": 0,
                  "n": 0,
                  "lineno": 4,
                  "col_offset": 17,
                  "end_lineno": 4,
                  "end_col_offset": 18
                },
                "ctx": {
                  "AST_class": "Load"
                },
                "lineno": 4,
                "col_offset": 15,
                "end_lineno": 4,
                "end_col_offset": 19
              },
              "lineno": 4,
              "col_offset": 8,
              "end_lineno": 4,
              "end_col_offset": 19
            }
          ],
          "lineno": 3,
          "col_offset": 4,
          "end_lineno": 4,
          "end_col_offset": 19
        }
      ],
      "decorator_list": [
        {
          "AST_class": "Attribute",
          "value": {
            "AST_class": "Name",
            "id": "T",
            "ctx": {
              "AST_class": "Load"
            },
            "lineno": 1,
            "col_offset": 1,
            "end_lineno": 1,
            "end_col_offset": 2
          },
          "attr": "macro",
          "ctx": {
            "AST_class": "Load"
          },
          "lineno": 1,
          "col_offset": 1,
          "end_lineno": 1,
          "end_col_offset": 8
        }
      ],
      "lineno": 2,
      "col_offset": 0,
      "end_lineno": 4,
      "end_col_offset": 19
    }
  ]
}

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

Have a real PR now: #15260