# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """TVM Script Parser For TIR We use [synr](https://synr.readthedocs.io) to get an AST that is stable over different python versions. Synr also provides an error handling context that we use for error reporting. """ # pylint: disable=invalid-name, inconsistent-return-statements, no-else-return, broad-except import types import json import operator import inspect from typing import Any, Callable, Dict, List, Optional, Union from synr import ast, Transformer, to_ast import tvm from tvm import IRModule from tvm._ffi.base import TVMError from tvm.ir import GlobalVar from tvm.ir.function import BaseFunc from tvm.tir import buffer from tvm.tir.function import PrimFunc from . import _ffi_api from . import tir from .context_maintainer import ContextMaintainer from .meta_unparser import MetaUnparser from .registry import Registry from .diagnostics import TVMDiagnosticCtx from .utils import tvm_span_from_synr, synr_span_from_tvm, call_with_error_reporting from .tir.intrin import Intrin from .tir.node import Slice, BufferSlice from .tir.scope_handler import ScopeHandler, WithScopeHandler, ForScopeHandler from .tir.special_stmt import SpecialStmt from .tir import ty class CallArgumentReader(object): """Helper class to read required arguments from passed arguments. When parsing a function call, we need to match the arguments provided in the AST to the required arguments of the function. This class makes sure all the positional arguments are filled and also fill keyword arguments with thier default value if a different value was not provided. """ def __init__(self, func_name, args, kwargs, parser, node): self.func_name = func_name self.args = args self.kwargs = kwargs self.parser = parser self.node = node def get_pos_only_arg(self, pos, name): """Get corresponding position only function argument from argument list""" if len(self.args) >= pos: arg = self.args[pos - 1] elif name not in self.kwargs: # If no positional argument was found in the AST, we see if it was # defined by name instead. # TODO(tkonolige): this error message is not quite correct. The # number of required arguments is >= pos self.parser.report_error( f"{self.func_name} requires {pos} arguments, but only {len(self.args)} were given.", self.node.span, ) else: arg = self.kwargs[name] return arg def get_kwarg(self, pos, name, default): """Get corresponding keyword function argument from argument list. If the user hasn't provided the argument, set it to the default value. """ if len(self.args) >= pos: arg = self.args[pos - 1] elif name in self.kwargs: arg = self.kwargs[name] else: return default return arg def get_varargs(self, pos): """Get corresponding variable argument from argument list""" if len(self.args) >= pos and len(self.kwargs) == 0: return self.args[pos - 1 :] return [] class TVMScriptParser(Transformer): """Synr AST visitor pass which finally lowers to TIR. Notes for Extension ------------------- 1. To support a new type of AST node, add a function transform_xxx(). 2. To support new functions, add the function to the appropriate registry: We divide allowed function calls in TVM script into 3 categories, intrin, scope_handler and special_stmt. 1. intrin functions are low level functions like mod, load, and constants. They correspond to a tir `IRNode`. They must have a return value. The user can register intrin functions for the parser to use. 2. scope_handler functions have no return value. They take two arguments: the parser and the AST node. scope_handler functions are used in with and for statements. 3. special_stmt functions handle cases that do not have a corresponding tir `IRNode`. These functions take the parser and the AST node as arguments and may return a value. When visiting a Call node, we check the special_stmt registry first. If no registered function is found, we then check the intrin registry. When visiting With node, we check the with_scope registry. When visiting For node, we check the for_scope registry. """ _binop_maker = { ast.BuiltinOp.Add: tvm.tir.Add, ast.BuiltinOp.Sub: tvm.tir.Sub, ast.BuiltinOp.Mul: tvm.tir.Mul, ast.BuiltinOp.Div: tvm.tir.Div, ast.BuiltinOp.FloorDiv: tvm.tir.FloorDiv, ast.BuiltinOp.Mod: tvm.tir.FloorMod, ast.BuiltinOp.BitOr: lambda lhs, rhs, span: operator.or_(lhs, rhs), ast.BuiltinOp.BitAnd: lambda lhs, rhs, span: operator.and_(lhs, rhs), ast.BuiltinOp.BitXor: lambda lhs, rhs, span: operator.xor(lhs, rhs), ast.BuiltinOp.GT: tvm.tir.GT, ast.BuiltinOp.GE: tvm.tir.GE, ast.BuiltinOp.LT: tvm.tir.LT, ast.BuiltinOp.LE: tvm.tir.LE, ast.BuiltinOp.Eq: tvm.tir.EQ, ast.BuiltinOp.NotEq: tvm.tir.NE, ast.BuiltinOp.And: tvm.tir.And, ast.BuiltinOp.Or: tvm.tir.Or, } _unaryop_maker = { ast.BuiltinOp.USub: lambda rhs, span: operator.neg(rhs), ast.BuiltinOp.Invert: lambda rhs, span: operator.invert(rhs), ast.BuiltinOp.Not: tvm.tir.Not, } # pylint gets confused here with synr.Transformer which doesn't have a # custom init, so just disable it def __init__( self, base_lineno, tir_namespace, closure_vars ): # pylint: disable=super-init-not-called self.context = None self.base_lineno = base_lineno self.current_lineno = 0 self.current_col_offset = 0 self.tir_namespace = tir_namespace self.closure_vars = closure_vars self.meta = None self._inside_buffer_sugar = False def init_function_parsing_env(self): """Initialize function parsing environment""" self.context = ContextMaintainer(self.report_error, self.closure_vars) # scope emitter def init_meta(self, meta_dict): if meta_dict is not None: self.meta = tvm.ir.load_json(json.dumps(meta_dict)) def transform(self, node): """Generic transformation for visiting the AST. Dispatches to `transform_ClassName` for the appropriate ClassName.""" old_lineno, old_col_offset = self.current_lineno, self.current_col_offset if hasattr(node, "lineno"): self.current_lineno = self.base_lineno + node.lineno - 1 if hasattr(node, "col_offset"): self.current_col_offset = node.col_offset method = "transform_" + node.__class__.__name__ visitor = getattr(self, method, self.generic_visit) transform_res = visitor(node) self.current_lineno, self.current_col_offset = old_lineno, old_col_offset return transform_res def match_tir_namespace(self, identifier: str) -> bool: """Check if the namespace is equal to tvm.script.tir""" return identifier in self.tir_namespace def report_error(self, message: str, span: Union[ast.Span, tvm.ir.Span]): """Report an error occuring at a location. This just dispatches to synr's DiagnosticContext. Parameters ---------- message : str Error message span : Union[synr.ast.Span, tvm.ir.Span] Location of the error """ if isinstance(span, tvm.ir.Span): span = synr_span_from_tvm(span) self.error(message, span) def parse_body(self, parent): """Parse remaining statements in this scope. Parameters ---------- parent : synr.ast.Node Parent node of this scope. Errors will be reported here. """ body = [] spans = [] stmt = parent while len(self.context.node_stack[-1]) > 0: stmt = self.context.node_stack[-1].pop() spans.append(stmt.span) res = self.transform(stmt) if res is not None: body.append(res) if len(body) == 0: self.report_error( "Expected another statement at the end of this block. Perhaps you " "used a concise statement and forgot to include a body afterwards.", stmt.span, ) else: return ( tvm.tir.SeqStmt(body, tvm_span_from_synr(ast.Span.union(spans))) if len(body) > 1 else body[0] ) def parse_arg_list(self, func, node_call): """Match the arguments of a function call in the AST to the required arguments of the function. This handles positional arguments, positional arguments specified by name, keyword arguments, and varargs. Parameters ---------- func : Function The function that provides the signature node_call: Union[ast.Call, ast.TypeApply, ast.TypeCall] The AST call node that calls into the function. Returns ------- arg_list : list The parsed positional argument. """ assert isinstance(node_call, (ast.Call, ast.TypeApply, ast.TypeCall)) # collect arguments args = [self.transform(arg) for arg in node_call.params] if isinstance(node_call, ast.TypeApply): kw_args = {} # TypeApply (e.g. foo[bar]) doesn't have kwargs defined in synr else: kw_args = { self.transform(k): self.transform(v) for k, v in node_call.keyword_params.items() } # get the name and parameter list of func if isinstance(func, (Intrin, ScopeHandler, SpecialStmt)): func_name, param_list = func.signature() else: self.report_error( "Internal Error: function must be of type Intrin, ScopeHandler or SpecialStmt, " f"but it is {type(func).__name__}", node_call.span, ) # check arguments and parameter list and get a list of arguments reader = CallArgumentReader(func_name, args, kw_args, self, node_call) pos_only, kwargs, varargs = param_list internal_args = list() for i, arg_name in enumerate(pos_only): internal_args.append(reader.get_pos_only_arg(i + 1, arg_name)) for i, arg_info in enumerate(kwargs): arg_name, default = arg_info internal_args.append(reader.get_kwarg(i + 1 + len(pos_only), arg_name, default=default)) if varargs is not None: internal_args.extend(reader.get_varargs(len(pos_only) + len(kwargs) + 1)) elif len(args) + len(kw_args) > len(pos_only) + len(kwargs): self.report_error( "Arguments mismatched. " + f"Expected {len(pos_only) + len(kwargs)} args but got " + f"{len(args) + len(kw_args)}", node_call.span, ) return internal_args def parse_type(self, type_node, parent): """Parse a type annotation. We require the parent object to the type so that we have a place to report the error message if the type does not exist. """ if type_node is None: self.report_error("A type annotation is required", parent.span) res_type = self.transform(type_node) return tvm.ir.TupleType([]) if res_type is None else res_type.evaluate() def generic_visit(self, node): """Fallback visitor if node type is not handled. Reports an error.""" self.report_error(type(node).__name__ + " AST node is not supported", node.span) def transform_Module(self, node): """Module visitor Right now, we only support two formats for TVM Script. Example ------- 1. Generate a PrimFunc (If the code is printed, then it may also contain metadata) .. code-block:: python import tvm @tvm.script def A(...): ... # returns a PrimFunc func = A 2. Generate an IRModule .. code-block:: python import tvm @tvm.script.ir_module class MyMod(): @T.prim_func def A(...): ... @T.prim_func def B(...): ... __tvm_meta__ = ... # returns an IRModule mod = MyMod """ if len(node.funcs) == 1: return self.transform(next(iter(node.funcs.values()))) elif len(node.funcs) == 0: self.report_error( "You must supply at least one class or function definition", node.span ) else: self.report_error( "Only one-function, one-class or function-with-meta source code is allowed", ast.Span.union([x.span for x in list(node.funcs.values())[1:]]), ) def transform_Class(self, node): """Class definition visitor. A class can have multiple function definitions and a single :code:`__tvm_meta__` statement. Each class corresponds to a single :code:`IRModule`. Example ------- .. code-block:: python @tvm.script.ir_module class MyClass: __tvm_meta__ = {} def A(): T.evaluate(0) """ if len(node.assignments) == 1: if not ( len(node.assignments[0].lhs) == 1 and isinstance(node.assignments[0].lhs[0], ast.Var) and node.assignments[0].lhs[0].id.name == "__tvm_meta__" ): self.report_error( "The only top level assignments allowed are `__tvm_meta__ = ...`", node.assignments[0].span, ) self.init_meta( MetaUnparser().do_transform(node.assignments[0].rhs, self._diagnostic_context) ) elif len(node.assignments) > 1: self.report_error( "Only a single top level `__tvm_meta__` is allowed", ast.Span.union([x.span for x in node.assignments[1:]]), ) return IRModule( {GlobalVar(name): self.transform(func) for name, func in node.funcs.items()} ) def transform_Function(self, node): """Function definition visitor. Each function definition is translated to a single :code:`PrimFunc`. There are a couple restrictions on TVM Script functions: 1. Function arguments must have their types specified. 2. The body of the function can contain :code:`func_attr` to specify attributes of the function (like it's name). 3. The body of the function can also contain multiple :code:`buffer_bind`s, which give shape and dtype information to arguments. 4. Return statements are implicit. Example ------- .. code-block:: python @T.prim_func def my_function(x: T.handle): # 1. Argument types T.func_attr({"global_symbol": "mmult"}) # 2. Function attributes X_1 = tir.buffer_bind(x, [1024, 1024]) # 3. Buffer binding T.evaluate(0) # 4. This function returns 0 """ def check_as_torch_decorator(decorator: Union[ast.Call, ast.Var]): if isinstance(decorator, ast.Call): if len(decorator.params) != 1: return False func_name = decorator.func_name else: func_name = decorator if isinstance(func_name, ast.Var): return func_name.id.name == "as_torch" def check_decorator(decorators: List[ast.Expr]) -> bool: """Check the decorator is `T.prim_func""" if len(decorators) > 2 or len(decorators) == 0: return False if len(decorators) == 2 and not check_as_torch_decorator(decorators[0]): return False d: ast.Expr = decorators[-1] return ( isinstance(d, ast.Attr) and isinstance(d.object, ast.Var) and self.match_tir_namespace(d.object.id.name) and d.field.name == "prim_func" ) self.init_function_parsing_env() self.context.enter_scope(nodes=node.body.stmts) # add parameters of function for arg in node.params: # Note that this case is for T.match_buffer syntax sugar if isinstance(arg.ty, (ast.TypeCall, ast.TypeApply)) and isinstance( self.transform(arg.ty.func_name), ty.GenericBufferType ): result = self.handle_match_buffer_type(arg.ty, arg.name) if not isinstance(result, buffer.Buffer): self.report_error( "The result type of evaluating TypeCall and TypeApply stmt" f" is wrong: {type(result)}. It should be a Buffer", node.span, ) arg_name_with_handle = arg.name + "_handle" arg_var = tvm.te.var(arg_name_with_handle, tvm.ir.PrimType("handle")) self.context.func_buffer_map[arg_var] = result self.context.update_symbol(arg.name, result, node) else: arg_var = tvm.te.var(arg.name, self.parse_type(arg.ty, arg)) self.context.update_symbol(arg.name, arg_var, node) self.context.func_params.append(arg_var) if not check_decorator(node.decorators): self.report_error( "All functions should be decorated by `T.prim_func`", node.span, ) # fetch the body of root block body = self.parse_body(node.body) # return a tir.PrimFunc dict_attr = self.context.func_dict_attr ret_type = self.parse_type(node.ret_type, node) if node.ret_type is not None else None func = tvm.tir.PrimFunc( self.context.func_params, body, ret_type, buffer_map=self.context.func_buffer_map, preflattened_buffer_map=self.context.func_preflattened_buffer_map, attrs=tvm.ir.make_node("DictAttrs", **dict_attr) if dict_attr else None, span=tvm_span_from_synr(node.span), ) # New Scope : Implicit root block # Each function contains an implicit root block in TensorIR, # so here we need a block scope for it. # If the PrimFunc is not a TensorIR func (e.g. TE scheduled func or low-level func), # the root block will not be added. The logic to add root block is in `_ffi_api.Complete` # Fix the PrimFunc # 1. generate root block if necessary # 2. generate surrounding loops for blocks if necessary func = call_with_error_reporting( self.report_error, node.span, _ffi_api.Complete, func, self.context.root_alloc_buffers, ) self.context.exit_scope() return func def transform_Lambda(self, node): """Lambda visitor Return an array of input parameters and the transformed lambda body. """ self.context.enter_scope(nodes=[node.body]) # add parameters of the lambda arg_vars = [] for arg in node.params: # Use "void" for dtype here. The actual type is not yet known and will be # determined later. Using void type will allow IRSubstitute to do the # replacement without flagging a type-mismatch error. arg_var = tvm.te.var(arg.name, dtype="") arg_vars.append(arg_var) self.context.update_symbol(arg.name, arg_var, node) # the body of a lambda must be an expr if not isinstance(node.body, ast.Expr): self.report_error("The body of a lambda must be an expression", node.span) # transform the body of the lambda body = self.transform(node.body) self.context.exit_scope() return arg_vars, body def transform_Assign(self, node): """Assign visitor AST abstract grammar: Assign(expr* targets, expr value, string? type_comment) By now 5 patterns of Assign is supported: 1. special stmts with return value 1.1 Buffer = T.match_buffer()/T.buffer_decl() 1.2 Var = T.var() 1.3 Var = T.env_thread() 2. (BufferStore) Buffer[PrimExpr, PrimExpr, ..., PrimExpr] = PrimExpr 3. (Store) Var[PrimExpr] = PrimExpr 4. with scope handlers with concise scoping and var def 4.1 var = T.allocate() 5. A call to a pure python function, consuming and producing TVMScript values. The outputs are inlined into the following body (no variable is created). x, y = f(...) """ if isinstance(node.rhs, ast.Call): # Pattern 1 & Pattern 4 if isinstance(node.rhs.func_name, ast.Op): func = None else: func = self.transform(node.rhs.func_name) if isinstance(func, WithScopeHandler): if not func.concise_scope or not func.def_symbol: self.report_error( "with scope handler " + func.signature()[0] + " is not suitable here", node.rhs.span, ) # Pattern 4 arg_list = self.parse_arg_list(func, node.rhs) func.enter_scope(node, self.context, arg_list, node.rhs.func_name.span) func.body = self.parse_body(node) return func.exit_scope(node, self.context, arg_list, node.rhs.func_name.span) elif isinstance(func, SpecialStmt): # Pattern 1 arg_list = self.parse_arg_list(func, node.rhs) func.handle(node, self.context, arg_list, node.rhs.func_name.span) return self.parse_body(node) elif isinstance(func, types.FunctionType): # Pattern 5 args = [self.transform(arg) for arg in node.rhs.params] try: out = func(*args) except Exception as e: self.report_error( "Error occurred when invoking the function " + func.__name__ + ": \n" + str(e), node.rhs.span, ) if len(node.lhs) == 1 and not isinstance(out, list): out = [out] assert len(out) == len(node.lhs) for var, value in zip(node.lhs, out): self.context.update_symbol(var.id.name, value, node) body = self.parse_body(node) for var, value in zip(node.lhs, out): self.context.remove_symbol(var.id.name) return body if isinstance(node.rhs, (ast.Call, ast.Constant)): # Pattern 4 of let binding value = self.transform(node.rhs) if len(node.lhs) == 1 and not isinstance(node.lhs[0], ast.Var): # This is a little confusing because it only is true when # we have taken this branch. We might need to clarify what # exectly is allowed in Assignments in tvmscript. self.report_error( "Left hand side of assignment must be an unqualified variable", node.span, ) ast_var = node.lhs[0] if node.ty is None and hasattr(value, "dtype"): var_ty = value.dtype else: var_ty = self.parse_type(node.ty, ast_var) var = tvm.te.var( ast_var.id.name, var_ty, span=tvm_span_from_synr(ast_var.span), ) self.context.update_symbol(var.name, var, node) body = self.parse_body(node) self.context.remove_symbol(var.name) return tvm.tir.LetStmt(var, value, body, span=tvm_span_from_synr(node.span)) self.report_error( """Assignments should be one of: 1. A "special statement" with return value 1.1 Buffer = T.match_buffer()/T.buffer_decl() 1.2 Var = T.var() 1.3 Var = T.env_thread() 2. A store into a buffer: Buffer[PrimExpr, PrimExpr, ..., PrimExpr] = PrimExpr 3. A store into a variable: Var[PrimExpr] = PrimExpr 4. A with scope handler with concise scoping and var def 4.1 var = T.allocate() 5. The right-hand side being a call to a pure python function, consuming and producing TVMScript values. x, y = f(...)""", node.span, ) def transform_SubscriptAssign(self, node): """Visitor for statements of the form :code:`x[1] = 2`.""" symbol = self.transform(node.params[0]) indexes = self.transform(node.params[1]) rhs = self.transform(node.params[2]) rhs_span = tvm_span_from_synr(node.params[2].span) if isinstance(symbol, tvm.tir.Buffer): if len(indexes) != len(symbol.shape): self.report_error( f"Buffer {symbol.name} is {len(symbol.shape)}-dimensional, " f"cannot be indexed by {len(indexes)}-dimensional indices.", node.params[1].span, ) def __convert_index(x): if isinstance(x, Slice): return x.as_index_expr(self.report_error) return x # BufferStore indexes = [__convert_index(x) for x in indexes] return tvm.tir.BufferStore( symbol, tvm.runtime.convert(rhs, span=rhs_span), indexes, span=tvm_span_from_synr(node.span), ) else: if symbol.dtype == "handle" and len(indexes) != 1: self.report_error( "Handles only support one-dimensional indexing. Use `T.match_buffer` to " "construct a multidimensional buffer from a handle.", node.params[0].span, ) if len(indexes) != 1: self.report_error( f"Store is only allowed with one index, but {len(indexes)} were provided.", node.params[1].span, ) self.report_error( "Use of tir.Store has been deprecated in favor of tir.BufferStore.", node.span ) def transform_AttrAssign(self, node): """Visitor for statements of the form :code:`x.y = 2`.""" obj = self.transform(node.params[0]) field = node.params[1] value = self.transform(node.params[2]) if not hasattr(obj, field.name): self.error(f"Field {field.name} does not exist", field.span) var = getattr(obj, field.name) if not isinstance(var, tvm.tir.Var): self.error( f"Can only assign to tir.Var attributes, not {type(var).__name__}", node.span ) body = self.parse_body(node) return tvm.tir.LetStmt(var, value, body, span=tvm_span_from_synr(node.span)) def transform_Assert(self, node): """Assert visitor Pattern corresponds to concise mode of :code:`with T.Assert()`. """ condition = self.transform(node.condition) if node.msg is None: self.report_error("Assert statements must have an error message.", node.span) message = self.transform(node.msg) body = self.parse_body(node) return tvm.tir.AssertStmt( condition, tvm.runtime.convert(message), body, span=tvm_span_from_synr(node.span) ) def transform_For(self, node): """For visitor AST abstract grammar: For(expr target, expr iter, stmt* body, stmt* orelse, string? type_comment) By now 1 pattern of For is supported: 1. for scope handler for name in T.serial()/T.parallel()/T.vectorized()/T.unroll()/range()/ T.grid()/T.thread_binding() """ if not isinstance(node.rhs, ast.Call): self.report_error("The loop iterator should be a function call.", node.rhs.span) func = self.transform(node.rhs.func_name) if not isinstance(func, ForScopeHandler): self.report_error( "Only For scope handlers can be used in a for statement.", node.rhs.func_name.span ) # prepare for new for scope old_lineno, old_col_offset = self.current_lineno, self.current_col_offset self.current_lineno = node.span.start_line self.current_col_offset = node.span.start_column self.context.enter_scope(nodes=node.body.stmts) # for scope handler process the scope arg_list = [ tvm.runtime.convert(arg, span=tvm_span_from_synr(node.rhs.span)) for arg in self.parse_arg_list(func, node.rhs) ] func.enter_scope(node, self.context, arg_list, node.rhs.func_name.span) func.body = self.parse_body(node) res = func.exit_scope(node, self.context, arg_list, node.rhs.func_name.span) # exit the scope self.context.exit_scope() self.current_lineno, self.current_col_offset = old_lineno, old_col_offset return res def transform_While(self, node): """While visitor AST abstract grammar: While(expr condition, stmt* body) """ condition = self.transform(node.condition) # body self.context.enter_scope(nodes=node.body.stmts) body = self.parse_body(node) self.context.exit_scope() return tvm.tir.While(condition, body, span=tvm_span_from_synr(node.span)) def transform_With(self, node): """With visitor AST abstract grammar: With(withitem* items, stmt* body, string? type_comment) withitem = (expr context_expr, expr? optional_vars) By now 2 patterns of With is supported: 1. with scope handler with symbol def with T.allocate() as targets: 2. with scope handler without symbol def with T.block(*axes)/T.let()/T.Assert()/T.attr()/T.realize() """ if not isinstance(node.rhs, ast.Call): self.report_error( "The context expression of a `with` statement should be a function call.", node.rhs.span, ) func = self.transform(node.rhs.func_name) if not isinstance(func, WithScopeHandler): self.report_error( f"Function {func} cannot be used in a `with` statement.", node.rhs.func_name.span ) # prepare for new block scope old_lineno, old_col_offset = self.current_lineno, self.current_col_offset self.current_lineno = node.body.span.start_line self.current_col_offset = node.body.span.start_column self.context.enter_block_scope(nodes=node.body.stmts) # with scope handler process the scope arg_list = self.parse_arg_list(func, node.rhs) func.enter_scope(node, self.context, arg_list, node.rhs.func_name.span) func.body = self.parse_body(node) res = func.exit_scope(node, self.context, arg_list, node.rhs.func_name.span) # exit the scope self.context.exit_block_scope() self.current_lineno, self.current_col_offset = old_lineno, old_col_offset return res def transform_If(self, node): """If visitor AST abstract grammar: If(expr test, stmt* body, stmt* orelse) """ condition = self.transform(node.condition) # then body self.context.enter_scope(nodes=node.true.stmts) then_body = self.parse_body(node) self.context.exit_scope() # else body if len(node.false.stmts) > 0: self.context.enter_scope(nodes=node.false.stmts) else_body = self.parse_body(node) self.context.exit_scope() else: else_body = None return tvm.tir.IfThenElse( condition, then_body, else_body, span=tvm_span_from_synr(node.span) ) def transform_Call(self, node): """Call visitor 3 different Call patterns are allowed: 1. Intrin representing a PrimExpr/IterVar 1.1 tir.int/uint/float8/16/32/64/floormod/floordiv/load/cast/ramp/broadcast/max 1.2 tir.range/reduce_axis/scan_axis/opaque_axis 2. tir.Op(dtype, ...) 3. other callable functions """ if isinstance(node.func_name, ast.Op): if node.func_name.name == ast.BuiltinOp.Subscript: return self.transform_Subscript(node) if node.func_name.name in self._binop_maker: lhs = self.transform(node.params[0]) # There is no supertype for everything that can appear in # an expression, so we manually add what we might get here. if not isinstance(lhs, (tvm.tir.PrimExpr, BufferSlice)): # We would really like to report a more specific # error here, but this parser contains no distinction # between parsing statements and parsing expressions. All # rules just call `transform`. self.report_error( f"Left hand side of binary op must be a PrimExpr, " "but it is a {type(lhs).__name__}", node.params[0].span, ) rhs = self.transform(node.params[1]) if not isinstance(rhs, (tvm.tir.PrimExpr, BufferSlice)): self.report_error( f"Right hand side of binary op must be a PrimExpr, " "but it is a {type(rhs).__name__}", node.params[1].span, ) return call_with_error_reporting( self.report_error, node.span, lambda node, lhs, rhs, span: self._binop_maker[node.func_name.name]( lhs, rhs, span=span ), node, lhs, rhs, tvm_span_from_synr(node.span), ) if node.func_name.name in self._unaryop_maker: rhs = self.transform(node.params[0]) return self._unaryop_maker[node.func_name.name]( rhs, span=tvm_span_from_synr(node.span) ) self.report_error(f"Unsupported operator {node.func_name.name}.", node.func_name.span) else: func = self.transform(node.func_name) if isinstance(func, Intrin) and not func.stmt: # pattern 1 arg_list = self.parse_arg_list(func, node) return call_with_error_reporting( self.report_error, node.func_name.span, func.handle, arg_list, node.func_name.span, ) else: args = [self.transform(arg) for arg in node.params] kw_args = { self.transform(k): self.transform(v) for k, v in node.keyword_params.items() } if isinstance(func, tvm.tir.op.Op): if not "dtype" in kw_args.keys(): self.report_error(f"{func} requires a dtype keyword argument.", node.span) # pattern 2 return tvm.tir.Call( kw_args["dtype"], func, args, span=tvm_span_from_synr(node.span) ) elif callable(func): # pattern 3 return func(*args, **kw_args) else: self.report_error( f"Function is neither callable nor a tvm.tir.op.Op (it is a {type(func)}).", node.func_name.span, ) def transform_UnassignedCall(self, node): """Visitor for statements that are function calls. This handles function calls that appear on thier own line like `tir.realize`. Examples -------- .. code-block:: python @T.prim_func def f(): A = T.buffer_decl([10, 10]) T.realize(A[1:2, 1:2], "") # This is an UnassignedCall A[1, 1] = 2 # This is also an UnassignedCall """ # Only allowed builtin operator that can be a statement is x[1] = 3 i.e. subscript assign. if isinstance(node.call.func_name, ast.Op): if node.call.func_name.name == ast.BuiltinOp.SubscriptAssign: return self.transform_SubscriptAssign(node.call) if node.call.func_name.name == ast.BuiltinOp.AttrAssign: return self.transform_AttrAssign(node.call) self.report_error( "Binary and unary operators are not allowed as a statement", node.span ) # handle a regular function call func = self.transform(node.call.func_name) arg_list = self.parse_arg_list(func, node.call) if isinstance(func, tir.scope_handler.AssertHandler): self.report_error( "A standalone `T.Assert` is not allowed. Use `assert condition, message` " "instead.", node.call.func_name.span, ) if isinstance(func, Intrin): if func.stmt: return call_with_error_reporting( self.report_error, node.call.func_name.span, func.handle, arg_list, node.call.func_name.span, ) else: self.report_error(f"This intrinsic cannot be used as a statement.", node.call.span) elif isinstance(func, WithScopeHandler) and func.concise_scope and not func.def_symbol: func.enter_scope(node, self.context, arg_list, node.call.func_name.span) func.body = self.parse_body(node) return func.exit_scope(node, self.context, arg_list, node.call.func_name.span) elif isinstance(func, SpecialStmt) and not func.def_symbol: func.handle(node, self.context, arg_list, node.call.func_name.span) return self.report_error( "Unexpected statement. Expected an assert, an intrinsic, a with statement, or a " f"special statement, but got {type(func).__name__}.", node.call.func_name.span, ) def transform_Slice(self, node): """Index slice visitor.""" start = self.transform(node.start) end = self.transform(node.end) if not ( isinstance(node.step, ast.Constant) and isinstance(node.step.value, int) and node.step.value > 0 ): self.report_error( "Only positive integer step size is supported for slices.", node.step.span ) return Slice(start, end, node.step.value, tvm_span_from_synr(node.span)) def transform_Subscript(self, node): """Array access visitor. By now only 3 types of Subscript are supported: 1. Buffer[index, index, ...], Buffer element access(BufferLoad & BufferStore) Var[index] Buffer element access() 2. Buffer[start: stop, start: stop, ...], BufferRealize(realize(buffer[...])) 3. Array[index], Buffer element access """ symbol = self.transform(node.params[0]) if symbol is None: self.report_error( f"Variable {node.params[0].id.name} is not defined.", node.params[0].span ) indexes = [self.transform(x) for x in node.params[1].values] if isinstance(symbol, tvm.tir.expr.Var): if symbol.dtype == "handle": self.report_error( "Cannot read directly from a handle, use `T.match_buffer` " "to create a buffer to read from.", node.params[0].span, ) if len(indexes) > 1: self.report_error( "Only a single index can be provided when indexing into a `var`.", node.params[1].span, ) index = indexes[0] if not isinstance(index, (tvm.tir.PrimExpr, int)): self.report_error( "Var load index should be an int or PrimExpr, but it is a" + type(index), node.span, ) self.report_error( "Use of tir.Load has been deprecated in favor of tir.BufferLoad", node.span ) elif isinstance(symbol, tvm.tir.Buffer): return BufferSlice( symbol, indexes, self.report_error, span=tvm_span_from_synr(node.span) ) elif isinstance(symbol, tvm.container.Array): if len(indexes) > 1: self.report_error( "Array access should be one-dimension access, but the indices are " + str(indexes), node.span, ) index = indexes[0] if not isinstance(index, (int, tvm.tir.expr.IntImm)): self.report_error( "Array access index expected int or IntImm, but got " + type(index), node.span, ) if int(index) >= len(symbol): self.report_error( f"Array access out of bound, size: {len(symbol)}, got index {index}.", node.span, ) return symbol[int(index)] else: self.report_error( f"Cannot subscript from a {type(symbol).__name__}. Only variables and " "buffers are supported.", node.params[0].span, ) def transform_Attr(self, node): """Visitor for field access of the form `x.y`. This visitor is used to lookup function and symbol names. We have two cases to handle here: 1. If we have a statement of the form `tir.something`, then we lookup `tir.something` in the `Registry`. If the function is not in the registry, then we try to find a `tvm.ir.op.Op` with the same name. 2. All other names `tvm.something` are lookup up in this current python namespace. """ def get_full_attr_name(node: ast.Attr) -> str: reverse_field_names = [node.field.name] while isinstance(node.object, ast.Attr): node = node.object reverse_field_names.append(node.field.name) if isinstance(node.object, ast.Var): reverse_field_names.append(node.object.id.name) return ".".join(reversed(reverse_field_names)) if isinstance(node.object, (ast.Var, ast.Attr)): full_attr_name = get_full_attr_name(node) attr_object, fields = full_attr_name.split(".", maxsplit=1) if self.match_tir_namespace(attr_object): func_name = "tir." + fields res = Registry.lookup(func_name) if res is not None: return res try: return tvm.ir.op.Op.get(func_name) except TVMError as e: # Check if we got an attribute error if e.args[0].find("AttributeError"): self.report_error(f"Unregistered function `tir.{fields}`.", node.span) else: raise e symbol = self.transform(node.object) if symbol is None: self.report_error("Unsupported Attribute expression.", node.object.span) if not hasattr(symbol, node.field.name): self.report_error( f"Type {type(symbol)} does not have a field called `{node.field.name}`.", node.span ) res = getattr(symbol, node.field.name) return res def transform_TypeAttr(self, node): """Visitor for field access of the form `x.y` for types. We have two cases here: 1. If the type is of the form `T.something`, we look up the type in the `tir` namespace in this module. 2. If the type is of the form `tvm.x.something` then we look up `tvm.x.something` in this modules namespace. """ if isinstance(node.object, ast.TypeVar): if self.match_tir_namespace(node.object.id.name): if not hasattr(tir, node.field.name): self.report_error( f"Invalid type annotation `tir.{node.field.name}`.", node.span ) return getattr(tir, node.field.name) symbol = self.transform(node.object) if symbol is None: self.report_error("Unsupported Attribute expression", node.object.span) if not hasattr(symbol, node.field): self.report_error( f"Type {type(symbol)} does not have a field called `{node.field}`.", node.span ) res = getattr(symbol, node.field) return res def transform_DictLiteral(self, node): """Dictionary literal visitor. Handles dictionary literals of the form `{x:y, z:2}`. """ keys = [self.transform(key) for key in node.keys] values = [self.transform(value) for value in node.values] return dict(zip(keys, values)) def transform_Tuple(self, node): """Tuple visitor. Handles tuples of the form `(x, y, 2)`. """ return tuple(self.transform(element) for element in node.values) def transform_ArrayLiteral(self, node): """List literal visitor. Handles lists of the form `[x, 2, 3]`. """ return [self.transform(element) for element in node.values] def transform_Var(self, node): """Variable visitor Handles variables like `x` in `x = 2`. """ name = node.id.name if name == "meta": return self.meta symbol = Registry.lookup(name) if symbol is not None: return symbol symbol = self.context.lookup_symbol(name) if symbol is not None: return symbol self.report_error(f"Unknown identifier {name}.", node.span) def transform_TypeVar(self, node): """Type variable visitor. Equivalent to `transform_Var` but for types. """ name = node.id.name symbol = Registry.lookup(name) or self.context.lookup_symbol(name) if symbol is not None: return symbol self.report_error(f"Unknown identifier {name}.", node.span) def transform_Constant(self, node): """Constant value visitor. Constant values include `None`, `"strings"`, `2` (integers), `4.2` (floats), and `true` (booleans). """ return tvm.runtime.convert(node.value, span=tvm_span_from_synr(node.span)) def transform_TypeConstant(self, node): """Constant value visitor for types. See `transform_Constant`. """ if self._inside_buffer_sugar: return self.transform_Constant(node) return node.value def transform_TypeTuple(self, node): """Tuple value visitor for types. Mostly used in `transform_TypeCall` and `transform_TypeApply`. """ return [self.transform(value) for value in node.values] def transform_TypeCall(self, node): """TypeCall visitor This occurs when an expression is used inside a T.Buffer parameter annotation. """ # ast.Call has the BuiltinOp as node.func_name.name, where # ast.TypeCall has the BuiltinOp as node.func_name. So we can # delegate to self.transform_Call, but the error messages for # unsupported operations will highlight the entire expression # and not just the function itself. op = ast.Op(node.span, node.func_name) call = ast.Call(node.span, op, node.params, node.keyword_params) return self.transform_Call(call) def transform_TypeApply(self, node): """Visitor for Type[Type] expressions. Mostly used for ``T.Ptr`` expressions. """ func = self.transform(node.func_name) if not isinstance(func, ty.TypeGeneric) or not hasattr(func, "__getitem__"): self.report_error( f"Use of type arguments requires a type that accepts type arguments (e.g. T.Ptr), " f"but found {type(func).__name__} instead.", node.span, ) param_types = [] for idx, param in enumerate(node.params): param_type = self.transform(param) if not isinstance(param_type, ty.TypeGeneric) and func.require_type_generic_at(idx): self.report_error( f"Expected a type but found {type(param).__name__} " f"at {idx}th type argument", param.span, ) param_types.append(param_type) if len(param_types) == 1: return func[param_types[0]] else: return func[param_types] def handle_match_buffer_type(self, node, buffer_name): """special function to handle syntax sugar for match buffer. This method is for buffer declarations in the function parameters. """ func = self.transform(node.func_name) assert isinstance(func, SpecialStmt) # parse args and kwargs for TypeCall and TypeApply self._inside_buffer_sugar = True try: arg_list = self.parse_arg_list(func, node) finally: self._inside_buffer_sugar = False # Note that the third element in arg_list would always be the 'name' # TODO: This index is hardcoded as a workaround. Better to make it programmatic if arg_list[2] is None: arg_list[2] = buffer_name buf = func.handle(node, self.context, arg_list, node.func_name.span) return buf def transform_Return(self, node): self.report_error( "TVM script does not support return statements. Instead the last statement in any " "block is implicitly returned.", node.span, ) def get_tir_namespace(script: Union[Callable, type]) -> List[str]: assert inspect.isfunction(script) or inspect.isclass(script) env: Dict[str, Any] = script.__globals__ return [key for key in env.keys() if env[key] == tir] def from_source( input_func: Union[str, Callable], tir_prefix: Optional[List[str]] = None ) -> Union[PrimFunc, IRModule]: """Parse function or string into PrimFunc or IRModule. If possible, pass the TVM script in as a function so that line numbers and filename will be accurate. Parameters ---------- input_module : Union[str, Callable] The python function to be parsed. tir_prefix : Optional[List[str]] The tir prefix list. Only works for str input, default by "tir" and "T". Returns ------- output : Union[Function, Module] The Function or Module in IR. """ if isinstance(input_func, str): tir_prefix = ["T", "tir"] if tir_prefix is None else tir_prefix return to_ast(input_func, TVMDiagnosticCtx(), TVMScriptParser(0, tir_prefix, {})) elif inspect.isfunction(input_func): _, start_line = inspect.getsourcelines(input_func) env: Dict[str, Any] = input_func.__globals__ namespace = [key for key in env.keys() if env[key] is tir] _closure_vars = inspect.getclosurevars(input_func) closure_vars = {**_closure_vars.nonlocals, **_closure_vars.globals} parser = TVMScriptParser(start_line, namespace, closure_vars) result = to_ast(input_func, TVMDiagnosticCtx(), parser) return result else: raise TypeError("Only function definitions are supported.") def ir_module(input_module: type) -> IRModule: """Decorate a python class as tvm IRModule. Parameters ---------- input_module : type The python class to be parsed. Returns ------- output : IRModule The result IRModule. """ if inspect.isclass(input_module): func_dict = { name: f for name, f in input_module.__dict__.items() if isinstance(f, BaseFunc) } return IRModule(func_dict) raise TypeError("Only class definitions are supported.")