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