apache--tvm/python/tvm/script/parser.py

# 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
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):  # 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.meta = None

    def init_function_parsing_env(self):
        """Initialize function parsing environment"""
        self.context = ContextMaintainer(self.report_error)  # 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.func) == 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_decorator(decorators: List[ast.Expr]) -> bool:
            """Check the decorator is `T.prim_func"""
            if len(decorators) != 1:
                return False
            d: ast.Expr = decorators[0]
            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:
            arg_var = tvm.te.var(arg.name)
            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 3 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()
        """

        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)
            else:
                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]
                var = tvm.te.var(
                    ast_var.id.name,
                    self.parse_type(node.ty, ast_var),
                    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 either
            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()""",
            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,
                )
            # BufferStore
            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=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.block(*axes)/T.allocate() as targets:
            2. with scope handler without symbol def
                with 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):
        start = self.transform(node.start)
        end = self.transform(node.end)
        if not (isinstance(node.step, ast.Constant) and node.step.value == 1):
            self.report_error("Only step size 1 is supported for slices.", node.step.span)
        return Slice(start, end)

    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`.
        """
        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_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 param in node.params:
            param_type = self.transform(param)
            if not isinstance(param_type, ty.TypeGeneric):
                self.report_error(f"Expected a type but found {type(param).__name__}", 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
        arg_list = self.parse_arg_list(func, node)
        # 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]
        parser = TVMScriptParser(start_line, namespace)
        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.")