chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,952 @@
|
||||
# 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.
|
||||
"""The core parser"""
|
||||
|
||||
import abc
|
||||
import inspect
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable
|
||||
from contextlib import contextmanager
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from tvm.error import DiagnosticError
|
||||
from tvm.ir import GlobalVar
|
||||
|
||||
from . import dispatch, doc
|
||||
from .diagnostics import Diagnostics, Source
|
||||
from .evaluator import eval_assign, eval_expr
|
||||
|
||||
DEFAULT_VISIT = {
|
||||
"Interactive",
|
||||
"Module",
|
||||
"Expression",
|
||||
"Pass",
|
||||
}
|
||||
|
||||
|
||||
def _deferred(exit_f: Callable[[], None]):
|
||||
"""Created context with certain exit function.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
exit_f : Callable[[], None]
|
||||
The function to call when exiting the context.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The created context.
|
||||
"""
|
||||
|
||||
@contextmanager
|
||||
def context():
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
exit_f()
|
||||
|
||||
return context()
|
||||
|
||||
|
||||
def _do_nothing(*args, **kwargs): # pylint: disable=unused-argument
|
||||
pass
|
||||
|
||||
|
||||
class ScriptMacro(abc.ABC):
|
||||
"""Representation of a script macro.
|
||||
|
||||
This is a callable object, intended to be called from the expression evaluator.
|
||||
The evaluator is expected to insert the current parser into the environment
|
||||
undef the name given by "parser_object_name".
|
||||
|
||||
Once called, the ScriptMacro object will locate the current parser, and use it
|
||||
to parse the macro's body and produce the result.
|
||||
|
||||
There were two major considerations for this design:
|
||||
1. Implementing hygienic and non-hygienic macros.
|
||||
2. Implementing macros that return values.
|
||||
|
||||
Macro uses in TIR are only allowed at a statement-level, and they don't produce
|
||||
any values. Parsing of such macros could easily be done by intercepting doc.Call
|
||||
nodes in the TIR parser. If a macro is a value-producing expression, then there
|
||||
may not be a direct way to intercept calls to it if it's embedded in a complex
|
||||
expression. Because macros use function-call syntax, the evaluator will try to
|
||||
call the macro object, which this design relies on to parse and evaluate the macro.
|
||||
"""
|
||||
|
||||
parser_object_name = "__current_script_parser__"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source: Source,
|
||||
closure_vars: dict[str, Any],
|
||||
func: Callable,
|
||||
hygienic: bool,
|
||||
) -> None:
|
||||
self.source = source
|
||||
self.closure_vars = closure_vars
|
||||
self.func = func
|
||||
self.hygienic = hygienic
|
||||
|
||||
def __repr__(self):
|
||||
return self.source.source
|
||||
|
||||
@abc.abstractmethod
|
||||
def parse_macro(self, parser: "Parser") -> Any:
|
||||
"""The main macro parsing function. Different scripts may have different
|
||||
ways to parse a macro, and to return a value to the evaluator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
parser : Parser
|
||||
The parser with the appropriate frame already created and populated depending
|
||||
macro's hygiene settings,
|
||||
|
||||
Returns
|
||||
-------
|
||||
The return value depends on the specifics of the particular script. It can be
|
||||
"None" or any other value or any type.
|
||||
"""
|
||||
|
||||
def _find_parser_def(self):
|
||||
outer_frame_infos = inspect.getouterframes(inspect.currentframe())
|
||||
for finfo in outer_frame_infos:
|
||||
parser = finfo.frame.f_globals.get(ScriptMacro.parser_object_name)
|
||||
if parser is not None:
|
||||
return parser
|
||||
raise RuntimeError(f"{ScriptMacro.parser_object_name} not available")
|
||||
|
||||
def get_macro_def(self):
|
||||
ast_module = self.source.as_ast()
|
||||
for decl in ast_module.body:
|
||||
if isinstance(decl, doc.FunctionDef) and decl.name == self.func.__name__:
|
||||
return decl
|
||||
raise RuntimeError(f"cannot find macro definition for {self.func.__name__}")
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
param_binding = inspect.signature(self.func).bind(*args, **kwargs)
|
||||
param_binding.apply_defaults()
|
||||
local_vars = param_binding.arguments
|
||||
parser = self._find_parser_def()
|
||||
|
||||
with parser.with_diag_source(self.source):
|
||||
if self.hygienic:
|
||||
saved_var_table = parser.var_table
|
||||
parser.var_table = VarTable()
|
||||
|
||||
with parser.var_table.with_frame():
|
||||
for k, v in self.closure_vars.items():
|
||||
parser.var_table.add(k, v)
|
||||
for k, v in local_vars.items():
|
||||
parser.var_table.add(k, v)
|
||||
|
||||
parse_result = self.parse_macro(parser)
|
||||
|
||||
parser.var_table = saved_var_table
|
||||
|
||||
else:
|
||||
with parser.var_table.with_frame():
|
||||
for k, v in local_vars.items():
|
||||
parser.var_table.add(k, v)
|
||||
|
||||
parse_result = self.parse_macro(parser)
|
||||
|
||||
return parse_result
|
||||
|
||||
|
||||
class VarTableFrame:
|
||||
"""The variable table frame.
|
||||
A frame of variable table stores the variables created in one block or scope.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
vars : Set[str]
|
||||
The set of variable names in the variable table frame.
|
||||
"""
|
||||
|
||||
vars: set[str]
|
||||
|
||||
def __init__(self):
|
||||
self.vars = set()
|
||||
|
||||
def add(self, var: str):
|
||||
"""Add a new variable into variable table frame.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
var : str
|
||||
The name of new variable.
|
||||
"""
|
||||
if var in self.vars:
|
||||
raise ValueError(f"Variable {var} already defined in current scope")
|
||||
self.vars.add(var)
|
||||
|
||||
def pop_all(self, fn_pop: Callable[[str], None]):
|
||||
"""Pop out all variable in variable table frame.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
fn_pop : Callable[[str], None]
|
||||
The methods to call when popping each variable.
|
||||
"""
|
||||
for var in self.vars:
|
||||
fn_pop(var)
|
||||
self.vars.clear()
|
||||
|
||||
|
||||
class VarTable:
|
||||
"""The variable table.
|
||||
A variable table stores the all variables when parsing TVMScript.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
frames : List[VarTableFrame]
|
||||
The list or stack of variable table frame.
|
||||
|
||||
name2value : Dict[str, List[Any]]
|
||||
The dictionary for variable table name-based query.
|
||||
"""
|
||||
|
||||
frames: list[VarTableFrame]
|
||||
name2value: dict[str, list[Any]]
|
||||
|
||||
def __init__(self):
|
||||
self.frames = []
|
||||
self.name2value = defaultdict(list)
|
||||
|
||||
def with_frame(self):
|
||||
"""Create a new variable table frame as with statement.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The context with new variable table frame.
|
||||
"""
|
||||
|
||||
def pop_frame():
|
||||
frame = self.frames.pop()
|
||||
frame.pop_all(lambda name: self.name2value[name].pop())
|
||||
|
||||
self.frames.append(VarTableFrame())
|
||||
return _deferred(pop_frame)
|
||||
|
||||
def add(self, var: str, value: Any, allow_shadowing: bool = False):
|
||||
"""Add a new variable to variable table.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
var : str
|
||||
The name of variable.
|
||||
|
||||
value : Any
|
||||
The value of variable.
|
||||
|
||||
allow_shadowing : bool
|
||||
The options of whether variable shadowing allowed for this variable.
|
||||
"""
|
||||
# Skip if the key and value are equal to those in the var_table
|
||||
if self.name2value[var] and isinstance(self.name2value[var][-1], type(value)):
|
||||
if isinstance(value, np.ndarray) and (self.name2value[var][-1] == value).all():
|
||||
return
|
||||
elif self.name2value[var][-1] == value:
|
||||
return
|
||||
if allow_shadowing and var in self.frames[-1].vars:
|
||||
# Shadowing
|
||||
self.name2value[var][-1] = value
|
||||
else:
|
||||
self.frames[-1].add(var)
|
||||
self.name2value[var].append(value)
|
||||
|
||||
def get(self) -> dict[str, Any]:
|
||||
"""Get a variable dictionary of latest variables.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The variable dictionary copy of latest variables.
|
||||
"""
|
||||
return {key: values[-1] for key, values in self.name2value.items() if values}
|
||||
|
||||
def get_at_depth(self, depth: int) -> dict[str, Any]:
|
||||
"""Get variables visible at the given frame depth, using current values.
|
||||
|
||||
For each variable name that appears in frames 0..depth-1, count how many
|
||||
times it was pushed (to handle shadowing), then index into name2value at
|
||||
count-1 to retrieve the latest value visible at that depth.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
depth : int
|
||||
The frame depth (number of frames visible).
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : dict[str, Any]
|
||||
Variable dictionary of values visible at the given depth.
|
||||
"""
|
||||
result: dict[str, Any] = {}
|
||||
name_count: dict[str, int] = defaultdict(int)
|
||||
for frame_idx in range(min(depth, len(self.frames))):
|
||||
for name in self.frames[frame_idx].vars:
|
||||
name_count[name] += 1
|
||||
for name, count in name_count.items():
|
||||
if self.name2value[name]:
|
||||
result[name] = self.name2value[name][count - 1]
|
||||
return result
|
||||
|
||||
def exist(self, value: Any) -> bool:
|
||||
"""Check if any value exists in variable table.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
value : Any
|
||||
The value of variable.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : bool
|
||||
The existence of the value.
|
||||
"""
|
||||
return any(
|
||||
value.same_as(known_value)
|
||||
for known_value_stack in self.name2value.values()
|
||||
for known_value in known_value_stack
|
||||
)
|
||||
|
||||
|
||||
def _dispatch_wrapper(func: dispatch.ParseMethod) -> dispatch.ParseMethod:
|
||||
def _wrapper(self: "Parser", node: doc.AST) -> None:
|
||||
try:
|
||||
return func(self, node)
|
||||
except Exception as err: # pylint: disable=broad-except
|
||||
self.report_error(node, err)
|
||||
raise
|
||||
|
||||
return _wrapper
|
||||
|
||||
|
||||
def _dispatch(self: "Parser", type_name: str) -> dispatch.ParseMethod:
|
||||
for token in [self.dispatch_tokens[-1], "default"]:
|
||||
func = dispatch.get(token=token, type_name=type_name, default=None)
|
||||
if func is not None:
|
||||
return _dispatch_wrapper(func)
|
||||
return _dispatch_wrapper(lambda self, node: self.generic_visit(node))
|
||||
|
||||
|
||||
class Parser(doc.NodeVisitor):
|
||||
"""The TVMScript parser
|
||||
|
||||
Parameters
|
||||
----------
|
||||
diag : Diagnostics
|
||||
The diagnostics for error reporting.
|
||||
|
||||
dispatch_tokens : List[str]
|
||||
The list of dispatching tokens to dispatching parsing method
|
||||
of different IRs and different doc AST structure.
|
||||
|
||||
var_table : VarTable
|
||||
The variable table for parsing.
|
||||
"""
|
||||
|
||||
diag: Diagnostics
|
||||
dispatch_tokens: list[str]
|
||||
function_annotations: dict[str, dict[str, Any]] | None
|
||||
var_table: VarTable
|
||||
inside_function: bool # whether we are within a function
|
||||
current_class: str | None = None # current class being parsed
|
||||
base_py_module_context: bool = False # whether current class inherits from BasePyModule
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source: Source,
|
||||
function_annotations: dict[str, dict[str, Any]],
|
||||
) -> None:
|
||||
self.diag = Diagnostics(source)
|
||||
self.dispatch_tokens = ["default"]
|
||||
self.function_annotations = function_annotations
|
||||
self.var_table = VarTable()
|
||||
self.inside_function = False
|
||||
|
||||
def parse(self, extra_vars: dict[str, Any] | None = None) -> Any:
|
||||
"""The main parse method for parser.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
extra_vars : Optional[Dict[str, Any]]
|
||||
The optional global value table for parsing.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The doc AST node visiting result.
|
||||
"""
|
||||
if extra_vars is None:
|
||||
extra_vars = {}
|
||||
with self.var_table.with_frame():
|
||||
for k, v in extra_vars.items():
|
||||
self.var_table.add(k, v)
|
||||
node = self.diag.source.as_ast()
|
||||
self.visit(node)
|
||||
|
||||
def get_dispatch_token(self, node: doc.FunctionDef) -> str:
|
||||
if not isinstance(node, doc.FunctionDef):
|
||||
self.report_error(node, "Only can get dispatch token for function.")
|
||||
if not node.decorator_list:
|
||||
self.report_error(node, "Function must be decorated")
|
||||
# TODO: only the last decorator is parsed
|
||||
decorator = self.eval_expr(node.decorator_list[-1])
|
||||
if not hasattr(decorator, "dispatch_token"):
|
||||
self.report_error(node, "The parser does not understand the decorator")
|
||||
return decorator.dispatch_token
|
||||
|
||||
def with_dispatch_token(self, token: str):
|
||||
"""Add a new dispatching token as with statement.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
token : str
|
||||
The dispatching token.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The context with new dispatching token.
|
||||
"""
|
||||
|
||||
self.dispatch_tokens.append(token)
|
||||
enter_func = dispatch.get(token=token, type_name="enter_token", default=lambda *args: None)
|
||||
context = enter_func(self)
|
||||
|
||||
def pop_token():
|
||||
exit_func = dispatch.get(
|
||||
token=token, type_name="exit_token", default=lambda *args: None
|
||||
)
|
||||
exit_func(self, context)
|
||||
self.dispatch_tokens.pop()
|
||||
|
||||
return _deferred(pop_token)
|
||||
|
||||
def set_class_context(self, class_name: str, is_base_py_module: bool = False):
|
||||
"""Set the current class context for parsing.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
class_name : str
|
||||
The name of the current class being parsed.
|
||||
is_base_py_module : bool
|
||||
Whether the current class inherits from BasePyModule.
|
||||
"""
|
||||
self.current_class = class_name
|
||||
self.base_py_module_context = is_base_py_module
|
||||
|
||||
def _get_current_class_context(self) -> str | None:
|
||||
"""Get the current class context.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Optional[str]
|
||||
The name of the current class, or None if not in a class context.
|
||||
"""
|
||||
return self.current_class
|
||||
|
||||
def _is_base_py_module_context(self) -> bool:
|
||||
"""Check if the current class context allows Python functions.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bool
|
||||
True if Python functions are allowed in the current context.
|
||||
"""
|
||||
return self.base_py_module_context
|
||||
|
||||
def with_diag_source(self, source: Source):
|
||||
"""Add a new source as with statement.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
source : Source
|
||||
The source for diagnostics.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The context with new source.
|
||||
"""
|
||||
|
||||
last_diag = self.diag
|
||||
self.diag = Diagnostics(source)
|
||||
|
||||
def pop_source():
|
||||
self.diag = last_diag
|
||||
|
||||
return _deferred(pop_source)
|
||||
|
||||
def eval_expr(
|
||||
self,
|
||||
node: doc.Expression | doc.expr,
|
||||
extra_vars: dict[str, Any] | None = None,
|
||||
) -> Any:
|
||||
"""Expression evaluation when parsing.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : Union[doc.expr, doc.Expression]
|
||||
The root node of AST tree node of expression to evaluate.
|
||||
|
||||
extra_vars : Optional[Dict[str, Any]]
|
||||
The optional global value table for expression evaluation.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The evaluation result.
|
||||
"""
|
||||
var_values = self.var_table.get()
|
||||
if extra_vars is not None:
|
||||
for k, v in extra_vars.items():
|
||||
var_values[k] = v
|
||||
var_values[ScriptMacro.parser_object_name] = self
|
||||
return eval_expr(self, node, var_values)
|
||||
|
||||
def _duplicate_lhs_check(self, target: doc.expr) -> bool | set[str]:
|
||||
"""Check whether duplicate lhs exists in assignment.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
target : doc.expr
|
||||
The doc AST expr node for lhs.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Union[bool, Set[str]]
|
||||
The result of true if duplicate lhs exists,
|
||||
or the set of lhs names if no duplicate lhs exists.
|
||||
"""
|
||||
if isinstance(target, doc.Tuple | doc.List):
|
||||
vars: set[str] = set() # pylint: disable=redefined-builtin
|
||||
for i in target.elts:
|
||||
res = self._duplicate_lhs_check(i)
|
||||
if isinstance(res, bool) and res:
|
||||
return True
|
||||
assert isinstance(res, set)
|
||||
if vars & res:
|
||||
return True
|
||||
vars = vars.union(res)
|
||||
return vars
|
||||
elif isinstance(target, doc.Name):
|
||||
return {target.id}
|
||||
elif isinstance(target, doc.Starred):
|
||||
return self._duplicate_lhs_check(target.value)
|
||||
elif isinstance(target, doc.Attribute):
|
||||
# Attribute assignment like packedB.data = ..., treated as rebinding.
|
||||
return {target.attr}
|
||||
else:
|
||||
self.report_error(target, "Invalid type in assign statement")
|
||||
raise NotImplementedError
|
||||
|
||||
def eval_assign(
|
||||
self,
|
||||
target: doc.expr,
|
||||
source: Any,
|
||||
bind_value: Callable[["Parser", doc.expr, str, Any], Any],
|
||||
allow_shadowing: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""Expression assignment evaluation when parsing.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
target : doc.expr
|
||||
The root node of AST tree node of assigned expression to evaluate.
|
||||
|
||||
source : Any
|
||||
The source to be assigned with evaluated expression.
|
||||
|
||||
bind_value : Callable[["Parser", doc.expr, str, Any], Any]
|
||||
The value binding method when assigning the values to variables.
|
||||
|
||||
allow_shadowing : bool
|
||||
The options of whether variable shadowing allowed for assignment.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Dict[str, Any]
|
||||
The dictionary of assignment result.
|
||||
"""
|
||||
if self._duplicate_lhs_check(target) is True:
|
||||
self.report_error(target, "Duplicate vars assigned.")
|
||||
var_values = eval_assign(self, target, source)
|
||||
for k, v in var_values.items():
|
||||
var = bind_value(self, target, k, v)
|
||||
self.var_table.add(k, var, allow_shadowing)
|
||||
return var_values
|
||||
|
||||
def report_error(self, node: doc.AST, err: Exception | str) -> None: # pylint: disable=no-self-use
|
||||
"""The error reporting when parsing.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.AST
|
||||
The doc AST node with errors.
|
||||
|
||||
err: Union[Exception, str]
|
||||
The error to report.
|
||||
"""
|
||||
|
||||
# If the error is already being raised as a DiagnosticError,
|
||||
# re-raise it without wrapping it in a DiagnosticContext.
|
||||
if isinstance(err, DiagnosticError):
|
||||
raise err
|
||||
|
||||
# Only take the last line of the error message
|
||||
if isinstance(err, RuntimeError):
|
||||
lines = list(filter(None, str(err).split("\n")))
|
||||
msg = lines[-1] if lines else (str(err) or type(err).__name__)
|
||||
elif isinstance(err, KeyError):
|
||||
msg = "KeyError: " + str(err)
|
||||
else:
|
||||
msg = str(err)
|
||||
|
||||
try:
|
||||
self.diag.error(node, msg)
|
||||
except Exception as diag_err:
|
||||
# Calling self.diag.error is guaranteed to throw an
|
||||
# exception. When shown to a user, this error should
|
||||
# reference the point of error within the provided
|
||||
# TVMScript. However, when caught in pdb, the full
|
||||
# traceback should be available for debugging.
|
||||
if isinstance(err, Exception):
|
||||
diag_err = diag_err.with_traceback(err.__traceback__)
|
||||
raise diag_err
|
||||
|
||||
def visit(self, node: doc.AST) -> None:
|
||||
"""The general visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.AST
|
||||
The doc AST node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
if isinstance(node, list | tuple):
|
||||
for item in node:
|
||||
self.visit(item)
|
||||
return
|
||||
if not isinstance(node, doc.AST):
|
||||
return
|
||||
name = node.__class__.__name__.split(".")[-1]
|
||||
if name in DEFAULT_VISIT:
|
||||
func = self.generic_visit
|
||||
else:
|
||||
func = getattr(self, "visit_" + name, None)
|
||||
if func is None:
|
||||
raise NotImplementedError(f"Visitor of AST node is not implemented: {name}")
|
||||
try:
|
||||
func(node)
|
||||
except Exception as err: # pylint: disable=broad-except
|
||||
self.report_error(node, err)
|
||||
|
||||
def visit_body(self, node: list[doc.stmt]) -> Any:
|
||||
"""The general body visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : List[doc.stmt]
|
||||
The list of statements in body.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
for stmt in node:
|
||||
self.visit(stmt)
|
||||
|
||||
def visit_tvm_annotation(self, node: doc.expr) -> Any:
|
||||
"""The general TVM annotation visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.expr
|
||||
The doc AST expr node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "tvm_annotation")(self, node)
|
||||
|
||||
def visit_FunctionDef(self, node: doc.FunctionDef) -> None: # pylint: disable=invalid-name
|
||||
"""The general function definition visit method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.FunctionDef
|
||||
The doc FunctionDef node.
|
||||
"""
|
||||
token = self.get_dispatch_token(node)
|
||||
func = dispatch.get(token=token, type_name="FunctionDef", default=None)
|
||||
if func is None:
|
||||
self.report_error(
|
||||
node,
|
||||
"""The parser does not understand the decorator,
|
||||
or visit_FunctionDef is not implemented for the decorator with token: """
|
||||
+ token,
|
||||
)
|
||||
_dispatch(self, "pre_visit_local_function")(self, node)
|
||||
_dispatch_wrapper(func)(self, node)
|
||||
_dispatch(self, "post_visit_local_function")(self, node)
|
||||
|
||||
def visit_tvm_declare_function(self, node: doc.FunctionDef) -> GlobalVar:
|
||||
token = self.get_dispatch_token(node)
|
||||
with self.with_dispatch_token(token):
|
||||
return _dispatch(self, "tvm_declare_function")(self, node)
|
||||
|
||||
def visit_ClassDef(self, node: doc.ClassDef) -> Any: # pylint: disable=invalid-name
|
||||
"""The general class definition visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.ClassDef
|
||||
The doc AST class definition node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
func = dispatch.get(token="ir", type_name="ClassDef", default=None)
|
||||
if func is None:
|
||||
self.report_error(node, "The parser does not understand the decorator")
|
||||
_dispatch_wrapper(func)(self, node)
|
||||
|
||||
def visit_arguments(self, node: doc.arguments) -> Any:
|
||||
"""The general arguments visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.arguments
|
||||
The doc AST arguments node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "arguments")(self, node)
|
||||
|
||||
def visit_For(self, node: doc.For) -> Any: # pylint: disable=invalid-name
|
||||
"""The general for visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.For
|
||||
The doc AST for node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "For")(self, node)
|
||||
|
||||
def visit_While(self, node: doc.While) -> Any: # pylint: disable=invalid-name
|
||||
"""The general while visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.While
|
||||
The doc AST while node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "While")(self, node)
|
||||
|
||||
def visit_With(self, node: doc.With) -> Any: # pylint: disable=invalid-name
|
||||
"""The general with visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.With
|
||||
The doc AST with node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "With")(self, node)
|
||||
|
||||
def visit_Assign(self, node: doc.Assign) -> Any: # pylint: disable=invalid-name
|
||||
"""The general assign visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Assign
|
||||
The doc AST assign node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "Assign")(self, node)
|
||||
|
||||
def visit_AnnAssign(self, node: doc.AnnAssign) -> Any: # pylint: disable=invalid-name
|
||||
"""The general annotated assign visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Assign
|
||||
The doc AST annotated assign node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "AnnAssign")(self, node)
|
||||
|
||||
def visit_Expr(self, node: doc.Expr) -> Any: # pylint: disable=invalid-name
|
||||
"""The general expression visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Expr
|
||||
The doc AST expression node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "Expr")(self, node)
|
||||
|
||||
def visit_If(self, node: doc.If) -> Any: # pylint: disable=invalid-name
|
||||
"""The general if visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.If
|
||||
The doc AST if node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "If")(self, node)
|
||||
|
||||
def visit_AugAssign(self, node: doc.AugAssign) -> Any: # pylint: disable=invalid-name
|
||||
"""The general augmented assignment visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.AugAssign
|
||||
The doc AST augmented assignment node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "AugAssign")(self, node)
|
||||
|
||||
def visit_Assert(self, node: doc.Assert) -> Any: # pylint: disable=invalid-name
|
||||
"""The general assert visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Assert
|
||||
The doc AST assert node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "Assert")(self, node)
|
||||
|
||||
def visit_Return(self, node: doc.Return) -> Any: # pylint: disable=invalid-name
|
||||
"""The general return visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Return
|
||||
The doc AST return node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "Return")(self, node)
|
||||
|
||||
def visit_Continue(self, node: doc.Continue) -> Any: # pylint: disable=invalid-name
|
||||
"""The general continue visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Continue
|
||||
The doc AST continue node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "Continue")(self, node)
|
||||
|
||||
def visit_Break(self, node: doc.Break) -> Any: # pylint: disable=invalid-name
|
||||
"""The general break visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Break
|
||||
The doc AST break node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "Break")(self, node)
|
||||
|
||||
def visit_Nonlocal(self, node: doc.Nonlocal) -> Any: # pylint: disable=invalid-name
|
||||
"""The general nonlocal visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Nonlocal
|
||||
The doc AST nonlocal node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "Nonlocal")(self, node)
|
||||
Reference in New Issue
Block a user