kparzysz-quic opened a new pull request, #15238: URL: https://github.com/apache/tvm/pull/15238
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): with T.block("block1"): a[0] = b[0] @T.prim_func def bar(A: T.Buffer((2,), "int32"), B: T.Buffer((2,), "int32")): T.include(foo1, A, B) 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 0x7f06c4d63e80> ``` `bar` has the macro expanded in its body, while `foo1` prints as `<tvm.script...Module>`. <details> <summary>The detailed representation of the `foo1` macro in JSON format</summary> ``` { "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 } ] } ``` </details> -- 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]
