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