chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,210 @@
|
||||
# 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 entry point of TVM parser."""
|
||||
|
||||
import inspect
|
||||
from typing import Any
|
||||
|
||||
import tvm
|
||||
|
||||
from ....ir.module import IRModule
|
||||
from ...ir_builder import IRBuilder
|
||||
from . import doc
|
||||
from .diagnostics import Source
|
||||
from .error import ParserError
|
||||
from .parser import Parser
|
||||
|
||||
WELL_FORMED_ERROR_MESSAGE = (
|
||||
"Program is not well-formed. If this is deliberate, consider "
|
||||
"setting check_well_formed in the top-level decorator to False "
|
||||
"(e.g., @I.ir_module(check_well_formed=False) or "
|
||||
"@R.function(check_well_formed=False))."
|
||||
)
|
||||
|
||||
|
||||
def _default_globals() -> dict[str, Any]:
|
||||
# lazy import here to avoid circular deps
|
||||
from tvm.script import tirx as _tirx_dsl # pylint: disable=import-outside-toplevel
|
||||
from tvm.script.parser import (
|
||||
ir, # pylint: disable=import-outside-toplevel
|
||||
relax, # pylint: disable=import-outside-toplevel
|
||||
)
|
||||
from tvm.script.parser import tirx as _tirx_parser # pylint: disable=import-outside-toplevel
|
||||
from tvm.script.tirx import tile as _tirx_tile # pylint: disable=import-outside-toplevel
|
||||
from tvm.tirx import layout as _tirx_layout # pylint: disable=import-outside-toplevel
|
||||
|
||||
# Expose the layout `Axis` class so printed layout sugar like
|
||||
# `4 @ Axis.laneid` round-trips without per-script imports. Injecting just
|
||||
# `Axis` (one short symbol) avoids name collisions with common user shape
|
||||
# vars like `m`, `P`, `F` that registered axes happen to share names with.
|
||||
return {
|
||||
"tvm": tvm,
|
||||
"I": ir,
|
||||
"ir": ir,
|
||||
"T": _tirx_parser,
|
||||
"tir": _tirx_parser,
|
||||
"R": relax,
|
||||
"relax": relax,
|
||||
"Tx": _tirx_tile,
|
||||
"tirx": _tirx_dsl,
|
||||
"Axis": _tirx_layout.Axis,
|
||||
}
|
||||
|
||||
|
||||
def scan_macro(program: Any | str, extra_vars: dict[str, Any] | None = None) -> Any:
|
||||
"""Generate the AST, and the source code for __repr__."""
|
||||
# The AST will be converted into TIR at the time of expansion.
|
||||
source = Source(program)
|
||||
closure_vars = extra_vars or _default_globals()
|
||||
return source, closure_vars
|
||||
|
||||
|
||||
def parse(
|
||||
program: doc.AST | Any | str,
|
||||
extra_vars: dict[str, Any] | None = None,
|
||||
check_well_formed: bool = True,
|
||||
s_tir: bool = False,
|
||||
) -> Any:
|
||||
"""Register a method for a operand type, AST operator node and operand index.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
program : Union[doc.AST, Any, str]
|
||||
The TVMScript code to parse.
|
||||
|
||||
extra_vars : Dict[str, Any]
|
||||
The extra variable table for parsing.
|
||||
|
||||
check_well_formed : bool
|
||||
Whether to check well-formedness after parsing.
|
||||
|
||||
Returns
|
||||
-------
|
||||
func : Any
|
||||
The parsed TVMScript program.
|
||||
"""
|
||||
if extra_vars is None:
|
||||
extra_vars = _default_globals()
|
||||
|
||||
ann = {}
|
||||
all_pyfuncs = {}
|
||||
if inspect.isfunction(program):
|
||||
ann = {program.__name__: program.__annotations__}
|
||||
elif inspect.isclass(program):
|
||||
for name, func in program.__dict__.items():
|
||||
if inspect.isfunction(func):
|
||||
ann[name] = func.__annotations__
|
||||
all_pyfuncs[name] = func
|
||||
|
||||
source = Source(program)
|
||||
parser = Parser(source, ann)
|
||||
with IRBuilder() as builder:
|
||||
try:
|
||||
parser.parse(extra_vars=extra_vars)
|
||||
except ParserError as err:
|
||||
parser.report_error(err.node, err.args[0])
|
||||
ret = builder.get()
|
||||
# Attach pyfuncs to the IRModule
|
||||
if inspect.isclass(program) and isinstance(ret, IRModule):
|
||||
_attach_pyfuncs_to_irmodule(ret, all_pyfuncs)
|
||||
|
||||
# check well-formedness in both Relax and TIR
|
||||
if check_well_formed:
|
||||
check_ret = ret
|
||||
if not isinstance(check_ret, IRModule):
|
||||
check_ret = IRModule.from_expr(ret)
|
||||
|
||||
source_ast = source.as_ast()
|
||||
|
||||
if isinstance(
|
||||
ret, IRModule | tvm.relax.Function
|
||||
) and not tvm.relax.analysis.check_well_formed(ret):
|
||||
parser.report_error(source_ast, err=WELL_FORMED_ERROR_MESSAGE)
|
||||
|
||||
try:
|
||||
if s_tir:
|
||||
tvm.tirx.analysis.verify_well_formed(check_ret)
|
||||
else:
|
||||
tvm.tirx.analysis.verify_tirx_well_formed(check_ret)
|
||||
except Exception as err: # pylint: disable=broad-exception-caught
|
||||
parser.report_error(
|
||||
source_ast,
|
||||
err=f"{WELL_FORMED_ERROR_MESSAGE}\n\nTraceback: {err!s}",
|
||||
)
|
||||
return ret
|
||||
|
||||
|
||||
def _create_python_packed_func(pyfunc):
|
||||
"""Create a PackedFunc wrapper for a Python function.
|
||||
|
||||
This function creates a PackedFunc that can be called from TVM runtime
|
||||
and will execute the original Python function.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
pyfunc : Callable
|
||||
The Python function to wrap.
|
||||
|
||||
Returns
|
||||
-------
|
||||
PackedFunc
|
||||
A PackedFunc that wraps the Python function.
|
||||
"""
|
||||
|
||||
def packed_func_wrapper(*args, **kwargs):
|
||||
"""Wrapper function that calls the original Python function."""
|
||||
try:
|
||||
result = pyfunc(*args, **kwargs)
|
||||
return result
|
||||
except Exception as error:
|
||||
print(f"Error calling Python function {pyfunc.__name__}: {error}")
|
||||
raise
|
||||
|
||||
return packed_func_wrapper
|
||||
|
||||
|
||||
def _attach_pyfuncs_to_irmodule(irmodule, all_pyfuncs):
|
||||
"""Attach Python functions to IRModule with reduced nesting."""
|
||||
if not all_pyfuncs:
|
||||
return
|
||||
|
||||
if not hasattr(irmodule, "pyfuncs"):
|
||||
irmodule.pyfuncs = {}
|
||||
|
||||
for global_var, func in irmodule.functions_items():
|
||||
if not isinstance(func, tvm.relax.ExternFunc):
|
||||
continue
|
||||
if not func.attrs.get("is_pyfunc", False):
|
||||
continue
|
||||
|
||||
pyfunc_name = global_var.name_hint
|
||||
if pyfunc_name not in all_pyfuncs:
|
||||
continue
|
||||
|
||||
pyfunc = all_pyfuncs[pyfunc_name]
|
||||
irmodule.pyfuncs[pyfunc_name] = pyfunc
|
||||
|
||||
try:
|
||||
source_code = inspect.getsource(pyfunc)
|
||||
func = func.with_attr("python_source", source_code)
|
||||
except (OSError, TypeError):
|
||||
func = func.with_attr("python_source", f"# Source unavailable for {pyfunc_name}")
|
||||
|
||||
packed_func = _create_python_packed_func(pyfunc)
|
||||
func = func.with_attr("python_packed_func", packed_func)
|
||||
|
||||
irmodule[global_var] = func
|
||||
Reference in New Issue
Block a user