Krzysztof Parzyszek
cf15375e20
* Compute common type for shape elements in BroadcastHelper The corresponding dimensions in the input/output tensors in a broadcast operations may have the same value, but different types (e.g. int32 vs int64). When the broadcast helper tries to unify the dimensions it also needs to compute the common type to hold the dimension. * Cast and simplify both members of `Range` Only the `min` member was type-casted, which could lead to ranges with different types for `min` and `extent`. Move the casts to the argument of Simplify, so that they can be eliminated if they aren't needed. * Type-check iv domain ranges, use cast only if needed in MakeLoopNest In some cases the domain ranges had the `min` and the `extent` values be of different types (e.g. [(int64)0, 32)). This is an error, and it can lead to compilation failures later on. Add a check for equal types here to catch this early. Also, only add the cast operation when the desired type differs from the current one to keep the expressions simpler. * Check that variable and substituted expression have same types Add a check to IRSubstitute to detect when the type of a variable and the type of the expression to replace it with have different types. * Add testcase * [TVMScript] Use void for lambda parameters, allow mismatch in Substitute When the script parser deals with lambdas, it creates Var objects for each parameter. Their actual types are not known at the time, and the properly typed variables are subtituted in the body later. Since the default dtype of a Var is "int32", this could lead to a type mismatch in Substitute. To deal with this scenario, use "void" for newly created Vars in the parser, and add an exception to Substitute to allow replacing void Vars with expressions of any type. * Fix type error in test_reduce_combiner_simplify * Restart CI Co-authored-by: Jiawei Liu <jaway.liu@gmail.com>
1374 lines
53 KiB
Python
1374 lines
53 KiB
Python
# 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.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:
|
|
# 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 occured 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.")
|