Lunderberg commented on code in PR #15260:
URL: https://github.com/apache/tvm/pull/15260#discussion_r1256377949
##########
python/tvm/script/parser/tir/parser.py:
##########
@@ -427,6 +428,20 @@ def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
node : doc.Expr
The doc AST Expr node.
"""
+
+ def is_insert_macro(node: doc.Call) -> bool:
+ if not isinstance(node.func, doc.Attribute):
+ return False
+ attr = node.func
+ if not isinstance(attr.value, doc.Name):
+ return False
+ if attr.value.id != "T" or attr.attr != "insert":
Review Comment:
If we keep a separate `T.insert`, then we should still check for it using
`self._eval_expr`, in case the same `T.insert` object is access through another
namespace (e.g. `tvm.script.tir.insert`, or a redefinition using `T.meta_var`).
```python
callee = self._eval_expr(node.func)
is_macro = callee is tvm.script.insert
```
##########
python/tvm/script/parser/tir/parser.py:
##########
@@ -427,6 +428,20 @@ def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
node : doc.Expr
The doc AST Expr node.
"""
+
+ def is_insert_macro(node: doc.Call) -> bool:
Review Comment:
Is the separate `T.insert(...)` necessary? If the macro is called directly
(e.g. `my_macro(buffer_A, buffer_B)`, then this check could be implemented as
```python
callee = self._eval_expr(node.func)
is_macro = isinstance(callee, TIRMacro)
```
##########
python/tvm/script/parser/tir/entry.py:
##########
@@ -50,6 +50,48 @@ def prim_func(func: Callable) -> Union[PrimFunc, Callable]:
setattr(prim_func, "dispatch_token", "tir")
+# Semantics of TIR macros:
+# - 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 (either as value, or as
+# a string with the 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.
+
+
+class TIRMacro:
+ """Representation of T.macro: consists of the doc.AST and the text of the
source."""
+
+ def __init__(self, node, source):
+ self.doc = node
+ self.source = source
+
+ def __repr__(self):
+ return self.source
+
+
+def macro(func: Callable) -> doc.AST:
+ obj = TIRMacro(*parse_macro(func))
+ setattr(obj, "__name__", func.__name__)
+ # We don't need to explicitly store the return value anywhere.
Review Comment:
It may be useful to also store the original `func`, to be used in
`inspect.signature` later. (A later commment below assumes `obj.__orig_func__
= func`)
##########
python/tvm/script/parser/tir/parser.py:
##########
@@ -528,3 +543,76 @@ def visit_tvm_declare_function(self: Parser, node:
doc.FunctionDef) -> GlobalVar
# Only ret_type is needed for func_signature.
func_signature = tvm.tir.PrimFunc([], None, ret_type=ret_type)
return I.decl_function(node.name, func_signature)
+
+
+def process_insert_macro(self: Parser, call: doc.Call) -> None:
+ """Bind arguments to T.insert to the parameters of the macro, and pass the
macro body
+ for further parsing.
+ """
+
+ def find_macro_def(name: str, decl_list: doc.AST) ->
Union[doc.FunctionDef, Any]:
+ for decl in decl_list:
+ if isinstance(decl, doc.FunctionDef) and decl.name == name:
+ return decl
+ return None
+
+ macro_name = call.args[0]
+
+ if not isinstance(macro_name, doc.Name):
+ self.report_error(call, "Invalid macro name in T.insert")
+ macro_name = macro_name.id
+
+ macro = self.var_table.get().get(macro_name)
+ if macro is None:
+ self.report_error(node, f"Undefined macro '{macro_name}'")
+
+ if isinstance(macro.doc, doc.Module):
+ macro_def = find_macro_def(macro_name, macro.doc.body)
+ elif not isinstance(macro.doc, doc.FunctionDef) or macro.doc.name !=
macro_name:
+ macro_def = None
+
+ if macro_def is None:
+ self.report_error(call, f"Undefined macro {macro_name}")
+
+ # `macro_def` is a FunctionDef of the macro.
+
+ # We have the AST for the macro definition, and the AST for the call. We
need to
+ # substitute the actual arguments from the call for the parameters from the
+ # definition. To allow full flexibility of python, i.e. positional,
unnamed, and
+ # keyword parameters, get the python interpreter to do the work: create
and execute
+ # the following:
+ # ```
+ # def macro_name(...macro parameters...)
+ # return locals()
+ # tmp = macro_name(...arguments from the call...)
+ # ```
+ # Obtain the dictionary `tmp` resulting from the execution, and update the
var_table
+ # with it.
+
+ # Construct the function with the macro's parameters, and returning
locals().
+ macro_ast = doc.from_doc(macro_def)
+ macro_ast.body = [
+ ast.Return(value=ast.Call(func=ast.Name("locals", ctx=ast.Load()),
args=[], keywords=[]))
+ ]
+ macro_ast.decorator_list = []
+
+ # Construct the assignment with the call.
+ call_ast = doc.from_doc(call)
+ call_ast.func = ast.Name(macro_name, ctx=ast.Load())
+ call_ast.args = call_ast.args[1:]
+ tmp_name = "__tmp_param_eval_64e98b523301204b"
+ assign_ast = ast.Assign(targets=[ast.Name(tmp_name, ctx=ast.Store())],
value=call_ast)
+
+ # Finalize and execute the module:
+ module_ast = ast.Module(body=[macro_ast, assign_ast], type_ignores=[])
+ module_ast = ast.fix_missing_locations(module_ast)
+ cmacro = compile(module_ast, filename="<tmp-string>", mode="exec")
+ local_vars = {}
+ exec(cmacro, self.var_table.get(), local_vars) # pylint: disable=exec-used
+ local_vars = local_vars[tmp_name]
+
+ with self.var_table.with_frame():
Review Comment:
As currently implemented, the macros have access to all variables accessible
at the `T.insert` scope, and no variables accessible at the `T.macro` scope.
While this does mimic C-style macros, this could cause some confusion, such as
the (admittedly contrived) example below.
```python
def test_macro_scoping():
x = 'extern_func_name'
@T.macro
def macro():
# By python scoping rules, I'd expect this to refer to
# the .
T.call_extern(x)
@T.prim_func
def func():
for x in range(16):
# Because "x" is in the calling scope of the macro, the
# expanded macro would use the iterator "x" within the
# macro body.
T.insert(macro)
```
To forbid this access and provide hygienic macros, we could add an optional
argument to `with_frame()` that indicates whether the frame should allow access
to the parent frames variables (default), or block access to parent frames
(used for macros). If this argument is passed, `with_frame` would remove all
values from `name2value` on entering the scope, repopulate with the top-most
frame (global variables such as `T` and `tvm`), and restore the values into
`name2value` on exiting the frame.
That said, there are also cases where having access to variables from the
calling scope of the macro can be useful, so we should probably also support
that case (e.g. `@T.macro(hygienic=False)`).
##########
python/tvm/script/parser/tir/entry.py:
##########
@@ -50,6 +50,48 @@ def prim_func(func: Callable) -> Union[PrimFunc, Callable]:
setattr(prim_func, "dispatch_token", "tir")
+# Semantics of TIR macros:
+# - 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 (either as value, or as
+# a string with the 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.
+
+
+class TIRMacro:
+ """Representation of T.macro: consists of the doc.AST and the text of the
source."""
+
+ def __init__(self, node, source):
+ self.doc = node
+ self.source = source
+
+ def __repr__(self):
+ return self.source
+
+
+def macro(func: Callable) -> doc.AST:
+ obj = TIRMacro(*parse_macro(func))
+ setattr(obj, "__name__", func.__name__)
+ # We don't need to explicitly store the return value anywhere.
+ # This function is a decorator, so the return value will replace
+ # the function definition (to which the decorator it is applied)
+ # in that function's name space.
+ return obj
Review Comment:
We should collect any closure variables that are used within the function,
so that they can be later be added to the macro's variables. This would be
similar to how the `@T.prim_func` and `@I.ir_module` decorators capture the
closure variables for use in metaprogramming.
```python
obj.__closure_vars = utils.inspect_function_capture(func)
```
##########
python/tvm/script/parser/tir/parser.py:
##########
@@ -528,3 +543,76 @@ def visit_tvm_declare_function(self: Parser, node:
doc.FunctionDef) -> GlobalVar
# Only ret_type is needed for func_signature.
func_signature = tvm.tir.PrimFunc([], None, ret_type=ret_type)
return I.decl_function(node.name, func_signature)
+
+
+def process_insert_macro(self: Parser, call: doc.Call) -> None:
+ """Bind arguments to T.insert to the parameters of the macro, and pass the
macro body
+ for further parsing.
+ """
+
+ def find_macro_def(name: str, decl_list: doc.AST) ->
Union[doc.FunctionDef, Any]:
+ for decl in decl_list:
+ if isinstance(decl, doc.FunctionDef) and decl.name == name:
+ return decl
+ return None
+
+ macro_name = call.args[0]
+
+ if not isinstance(macro_name, doc.Name):
+ self.report_error(call, "Invalid macro name in T.insert")
+ macro_name = macro_name.id
+
+ macro = self.var_table.get().get(macro_name)
+ if macro is None:
+ self.report_error(node, f"Undefined macro '{macro_name}'")
+
+ if isinstance(macro.doc, doc.Module):
+ macro_def = find_macro_def(macro_name, macro.doc.body)
Review Comment:
What would be the use case for this branch? If I'm reading correctly, it
only applies when the macro is an entire module, but discards all functions
within the module other than the one with a function name matching the module
name. I suppose I'm not seeing the advantage of allowing that construct, as
opposed to requiring that a macro be a function.
##########
python/tvm/script/parser/tir/entry.py:
##########
@@ -50,6 +50,48 @@ def prim_func(func: Callable) -> Union[PrimFunc, Callable]:
setattr(prim_func, "dispatch_token", "tir")
+# Semantics of TIR macros:
+# - 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 (either as value, or as
+# a string with the 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.
+
+
+class TIRMacro:
Review Comment:
Should this be specific to TIR, or should it apply to any dialect supported
by TVMScript? Thinking that this would be quite useful on the unity branch as
well, where a Relax method for an end-to-end model often contains many repeated
elements. Implementing those as a macro would also allow Relax's shape
propagation to resolve differently for each expansion of the macro (e.g. in a
chain of convolutions).
If we want it to be more general, we could move the implementation over to
the `tvm.script.parser.ir` namespace instead.
##########
python/tvm/script/parser/tir/parser.py:
##########
@@ -528,3 +543,76 @@ def visit_tvm_declare_function(self: Parser, node:
doc.FunctionDef) -> GlobalVar
# Only ret_type is needed for func_signature.
func_signature = tvm.tir.PrimFunc([], None, ret_type=ret_type)
return I.decl_function(node.name, func_signature)
+
+
+def process_insert_macro(self: Parser, call: doc.Call) -> None:
+ """Bind arguments to T.insert to the parameters of the macro, and pass the
macro body
+ for further parsing.
+ """
+
+ def find_macro_def(name: str, decl_list: doc.AST) ->
Union[doc.FunctionDef, Any]:
+ for decl in decl_list:
+ if isinstance(decl, doc.FunctionDef) and decl.name == name:
+ return decl
+ return None
+
+ macro_name = call.args[0]
+
+ if not isinstance(macro_name, doc.Name):
+ self.report_error(call, "Invalid macro name in T.insert")
+ macro_name = macro_name.id
+
+ macro = self.var_table.get().get(macro_name)
+ if macro is None:
+ self.report_error(node, f"Undefined macro '{macro_name}'")
+
+ if isinstance(macro.doc, doc.Module):
+ macro_def = find_macro_def(macro_name, macro.doc.body)
+ elif not isinstance(macro.doc, doc.FunctionDef) or macro.doc.name !=
macro_name:
+ macro_def = None
+
+ if macro_def is None:
+ self.report_error(call, f"Undefined macro {macro_name}")
+
+ # `macro_def` is a FunctionDef of the macro.
+
+ # We have the AST for the macro definition, and the AST for the call. We
need to
+ # substitute the actual arguments from the call for the parameters from the
+ # definition. To allow full flexibility of python, i.e. positional,
unnamed, and
+ # keyword parameters, get the python interpreter to do the work: create
and execute
+ # the following:
+ # ```
+ # def macro_name(...macro parameters...)
+ # return locals()
+ # tmp = macro_name(...arguments from the call...)
+ # ```
+ # Obtain the dictionary `tmp` resulting from the execution, and update the
var_table
+ # with it.
+
+ # Construct the function with the macro's parameters, and returning
locals().
+ macro_ast = doc.from_doc(macro_def)
Review Comment:
It looks like the ast construction and execution is to provide the binding
of parameter to arguments for use in the var map. Would it be easier to use
`inspect.signature` to get the parameter binding directly? (Assumes
`macro_def.__orig_func__` is available from earlier comment.)
```python
args = [self._eval_expr(arg) for arg in macro_def.args]
kwargs = {kw.arg:self._eval_expr(kw.value) for kw in macro_def.keywords}
param_binding = inspect.signature(macro_def.__orig_func__).bind(*args,
**kwargs)
param_binding.apply_defaults()
local_vars = param_binding.arguments
```
##########
python/tvm/script/parser/tir/entry.py:
##########
@@ -50,6 +50,48 @@ def prim_func(func: Callable) -> Union[PrimFunc, Callable]:
setattr(prim_func, "dispatch_token", "tir")
+# Semantics of TIR macros:
+# - 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 (either as value, or as
+# a string with the 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.
+
+
+class TIRMacro:
+ """Representation of T.macro: consists of the doc.AST and the text of the
source."""
+
+ def __init__(self, node, source):
+ self.doc = node
+ self.source = source
+
+ def __repr__(self):
+ return self.source
+
+
+def macro(func: Callable) -> doc.AST:
+ obj = TIRMacro(*parse_macro(func))
+ setattr(obj, "__name__", func.__name__)
Review Comment:
Why use `setattr` instead of just `obj.__name__ = func.__name__`? I usually
use `setattr` only if I have a dynamic value for the second argument.
--
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]
