kparzysz-quic commented on code in PR #15260:
URL: https://github.com/apache/tvm/pull/15260#discussion_r1256638604
##########
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:
Done.
##########
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:
Done.
##########
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:
Done.
--
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]
