chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,35 @@
|
||||
# isort: skip_file
|
||||
# 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 ir module parser"""
|
||||
|
||||
from tvm.ir import Range
|
||||
from ...ir_builder.ir import * # pylint: disable=redefined-builtin
|
||||
from . import parser as _parser
|
||||
from .entry import ir_module, pyfunc
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Range",
|
||||
"dummy_global_info",
|
||||
"ir_module",
|
||||
"lookup_vdevice",
|
||||
"module_attrs",
|
||||
"module_global_infos",
|
||||
"pyfunc",
|
||||
"vdevice",
|
||||
]
|
||||
@@ -0,0 +1,163 @@
|
||||
# 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.
|
||||
# pylint: disable=import-outside-toplevel
|
||||
"""The entry point of TVM parser for ir module."""
|
||||
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
|
||||
from tvm import cpu, ir
|
||||
from tvm.ir import GlobalVar, IRModule
|
||||
|
||||
from .._core import parse, utils
|
||||
|
||||
|
||||
# this formulation allows us to support having @I.ir_module
|
||||
# appear as a decorator by itself or to have optional arguments
|
||||
# like @I.ir_module(check_well_formed=False)
|
||||
def ir_module(
|
||||
mod: type | None = None, check_well_formed: bool = True, s_tir: bool = False
|
||||
) -> IRModule:
|
||||
"""The parsing method for ir module, by using `@ir_module` as decorator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mod : Type
|
||||
The class to be parsed as ir module.
|
||||
|
||||
check_well_formed : bool
|
||||
Whether to check well-formedness during parsing.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ir_module : IRModule
|
||||
The parsed ir module.
|
||||
"""
|
||||
|
||||
# Capture stack outside wrapper (wrapper adds to the stack)
|
||||
outer_stack = inspect.stack()
|
||||
|
||||
def decorator_wrapper(mod):
|
||||
if not inspect.isclass(mod):
|
||||
raise TypeError(f"Expect a class, but got: {mod}")
|
||||
|
||||
# Check BasePyModule inheritance
|
||||
base_py_module_inherited = any(base.__name__ == "BasePyModule" for base in mod.__bases__)
|
||||
|
||||
extra_vars = utils.inspect_class_capture(mod)
|
||||
# Resolve closure variables hidden by PEP 563 (annotation-only names)
|
||||
utils.resolve_closure_vars(mod, extra_vars, outer_stack)
|
||||
m = parse(mod, extra_vars, check_well_formed=check_well_formed, s_tir=s_tir)
|
||||
|
||||
if base_py_module_inherited:
|
||||
# Lazy import: tvm.relax cannot be imported at module level in tvm.script.parser
|
||||
# because tvm.script is loaded before tvm.relax during tvm initialization.
|
||||
from tvm.relax.base_py_module import BasePyModule
|
||||
from tvm.relax.expr import ExternFunc # pylint: disable=import-outside-toplevel
|
||||
|
||||
# Collect pyfunc methods
|
||||
pyfunc_methods = [
|
||||
name
|
||||
for name, attr in mod.__dict__.items()
|
||||
if hasattr(attr, "dispatch_token") and attr.dispatch_token == "pyfunc"
|
||||
]
|
||||
|
||||
mod._pyfunc_methods = pyfunc_methods
|
||||
|
||||
# Create ExternFunc nodes
|
||||
|
||||
for method_name in pyfunc_methods:
|
||||
try:
|
||||
existing_gvars = [
|
||||
global_var
|
||||
for global_var in m.get_global_vars()
|
||||
if global_var.name_hint == method_name
|
||||
]
|
||||
|
||||
extern_func = ExternFunc(method_name)
|
||||
extern_func = extern_func.with_attr("is_pyfunc", True)
|
||||
extern_func = extern_func.with_attr("function_type", "python")
|
||||
extern_func = extern_func.with_attr("python_function_name", method_name)
|
||||
extern_func = extern_func.with_attr(
|
||||
"python_source", f"# Source for {method_name}"
|
||||
)
|
||||
extern_func = extern_func.with_attr("python_packed_func", None)
|
||||
|
||||
if existing_gvars:
|
||||
m[existing_gvars[0]] = extern_func
|
||||
else:
|
||||
m[GlobalVar(method_name)] = extern_func
|
||||
|
||||
except Exception: # pylint: disable=broad-exception-caught
|
||||
continue
|
||||
|
||||
class ModuleFactory:
|
||||
"""Factory class for creating BasePyModule instances with Python functions."""
|
||||
|
||||
def __init__(self, module, pyfunc_methods, original_class):
|
||||
self.ir_module = module
|
||||
self.pyfunc_methods = pyfunc_methods
|
||||
self.original_class = original_class
|
||||
|
||||
def __call__(self, device=None, target=None):
|
||||
if device is None:
|
||||
device = cpu(0)
|
||||
|
||||
instance_ir_mod = ir.IRModule()
|
||||
for global_var, func in self.ir_module.functions_items():
|
||||
instance_ir_mod[global_var] = func
|
||||
|
||||
instance = BasePyModule(instance_ir_mod, device, target)
|
||||
|
||||
for method_name in self.pyfunc_methods:
|
||||
if hasattr(self.original_class, method_name):
|
||||
method = getattr(self.original_class, method_name)
|
||||
instance.add_python_function(method_name, method)
|
||||
|
||||
return instance
|
||||
|
||||
def __getattr__(self, name):
|
||||
if hasattr(self.ir_module, name):
|
||||
return getattr(self.ir_module, name)
|
||||
raise AttributeError(
|
||||
f"'{self.__class__.__name__}' object has no attribute '{name}'"
|
||||
)
|
||||
|
||||
factory = ModuleFactory(m, pyfunc_methods, mod)
|
||||
setattr(factory, "__name__", mod.__name__)
|
||||
return factory
|
||||
|
||||
setattr(m, "__name__", mod.__name__)
|
||||
return m
|
||||
|
||||
if mod is not None:
|
||||
# if there are no optional args given, this will directly invoke the wrapper
|
||||
return decorator_wrapper(mod)
|
||||
else:
|
||||
# if there is a optional arg given, it returns the wrapper function
|
||||
# as a new decorator and applies it
|
||||
setattr(decorator_wrapper, "dispatch_token", "ir")
|
||||
return decorator_wrapper
|
||||
|
||||
|
||||
def pyfunc(func: Callable):
|
||||
# Set the dispatch_token on the decorated function
|
||||
setattr(func, "dispatch_token", "pyfunc")
|
||||
return func
|
||||
|
||||
|
||||
setattr(pyfunc, "dispatch_token", "pyfunc")
|
||||
@@ -0,0 +1,212 @@
|
||||
# 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.
|
||||
# pylint: disable=unused-argument
|
||||
"""The base parser for ir module"""
|
||||
|
||||
from tvm.ir import GlobalVar
|
||||
|
||||
from ...ir_builder import ir as I
|
||||
from .._core import Parser, dispatch, doc
|
||||
|
||||
|
||||
class ModuleWithGlobalVars:
|
||||
"""A Module that can add global vars during parsing, to support `Module.function` syntax."""
|
||||
|
||||
def __getattr__(self, attr):
|
||||
# Customize the error message.
|
||||
# NOTE: `__getattr__` is only called when the attribute access fails with an AttributeError
|
||||
raise AttributeError(f"Cannot find the function `{attr}` in the current IRModule")
|
||||
|
||||
|
||||
@dispatch.register(token="ir", type_name="ClassDef")
|
||||
def _visit_class_def(self: Parser, node: doc.ClassDef) -> None:
|
||||
"""The class definition visiting method for ir module.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.ClassDef
|
||||
The doc AST class definition node.
|
||||
"""
|
||||
|
||||
with self.var_table.with_frame():
|
||||
with I.ir_module():
|
||||
# Step 0. Add the class name to the var table
|
||||
fake_module = ModuleWithGlobalVars()
|
||||
self.var_table.add(node.name, fake_module)
|
||||
|
||||
# Step 1: Check if this class inherits from BasePyModule
|
||||
is_base_py_module = _check_base_py_module_inheritance(node)
|
||||
if is_base_py_module:
|
||||
# Store this information in the IRModule for later use
|
||||
I.module_attrs({"base_py_module": True})
|
||||
# Set the parser context to allow Python functions
|
||||
self.set_class_context(node.name, True)
|
||||
else:
|
||||
# Set the parser context to disallow Python functions
|
||||
self.set_class_context(node.name, False)
|
||||
|
||||
# Step 2. Visit non-function stmts, including but not limited to
|
||||
# 1. `I.module_attrs`
|
||||
# 2. `I.module_global_infos`
|
||||
with self.with_dispatch_token("ir"):
|
||||
for stmt in node.body:
|
||||
if not isinstance(stmt, doc.FunctionDef):
|
||||
self.visit(stmt)
|
||||
|
||||
# Step 3. Visit function stmts to declare the global vars
|
||||
for stmt in node.body:
|
||||
if isinstance(stmt, doc.FunctionDef):
|
||||
global_var = self.visit_tvm_declare_function(stmt)
|
||||
fake_module.__setattr__(stmt.name, global_var)
|
||||
|
||||
# Step 4. Visit and parse the functions
|
||||
with self.with_dispatch_token("ir"):
|
||||
for stmt in node.body:
|
||||
if isinstance(stmt, doc.FunctionDef):
|
||||
self.visit(stmt)
|
||||
|
||||
|
||||
@dispatch.register(token="ir", type_name="Assign")
|
||||
def _visit_assign(self: Parser, node: doc.Assign) -> None:
|
||||
"""The assign visiting method for ir module.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.ClassDef
|
||||
The doc AST assign node.
|
||||
"""
|
||||
if len(node.targets) != 1:
|
||||
self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
|
||||
lhs = node.targets[0].id
|
||||
rhs = self.eval_expr(node.value)
|
||||
|
||||
I.decl_function(lhs, rhs)
|
||||
I.def_function(lhs, rhs)
|
||||
|
||||
|
||||
@dispatch.register(token="ir", type_name="Expr")
|
||||
def _visit_expr(self: Parser, node: doc.Expr) -> None:
|
||||
"""The expression visiting method for ir module.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.ClassDef
|
||||
The doc AST expression node.
|
||||
"""
|
||||
self.eval_expr(node.value)
|
||||
|
||||
|
||||
@dispatch.register(token="default", type_name="Assign")
|
||||
def visit_assign(self: Parser, node: doc.Assign) -> None:
|
||||
if len(node.targets) != 1:
|
||||
self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
|
||||
lhs = node.targets[0]
|
||||
rhs = self.eval_expr(node.value)
|
||||
self.eval_assign(
|
||||
target=lhs, source=rhs, bind_value=lambda _a, _b, _c, value: value, allow_shadowing=True
|
||||
)
|
||||
|
||||
|
||||
@dispatch.register(token="default", type_name="pre_visit_local_function")
|
||||
def pre_visit_local_function(self: Parser, node: doc.Expr) -> None:
|
||||
pass
|
||||
|
||||
|
||||
@dispatch.register(token="default", type_name="post_visit_local_function")
|
||||
def post_visit_local_function(self: Parser, node: doc.Expr) -> None:
|
||||
pass
|
||||
|
||||
|
||||
@dispatch.register(token="pyfunc", type_name="tvm_declare_function")
|
||||
def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar:
|
||||
"""Declare a Python function as an ExternFunc in the IRModule."""
|
||||
# Check if Python functions are allowed in this context
|
||||
# We need to check if we're in a class that inherits from BasePyModule
|
||||
current_class = self._get_current_class_context()
|
||||
if current_class and not self._is_base_py_module_context():
|
||||
self.report_error(
|
||||
node,
|
||||
"@I.pyfunc are only allowed in classes that inherit from BasePyModule. "
|
||||
f"Class '{current_class}' does not inherit from BasePyModule.",
|
||||
)
|
||||
|
||||
# Lazy import: tvm.relax cannot be imported at module level in tvm.script.parser
|
||||
# because tvm.script is loaded before tvm.relax during tvm initialization.
|
||||
from tvm.relax import ExternFunc # pylint: disable=import-outside-toplevel
|
||||
|
||||
# Create ExternFunc with proper attributes for Python functions
|
||||
func = ExternFunc(node.name)
|
||||
func = func.with_attr("is_pyfunc", True)
|
||||
func = func.with_attr("function_type", "python")
|
||||
func = func.with_attr("python_function_name", node.name)
|
||||
|
||||
# Add placeholder attributes that will be filled in later
|
||||
func = func.with_attr("python_source", f"# Source will be filled for {node.name}")
|
||||
func = func.with_attr("python_packed_func", None) # Will be filled in entry.py
|
||||
|
||||
# Store the function name for later retrieval
|
||||
return I.decl_function(node.name, func)
|
||||
|
||||
|
||||
@dispatch.register(token="pyfunc", type_name="FunctionDef")
|
||||
def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
|
||||
"""Visit Python function definition - no need to parse the body."""
|
||||
# Python function body is not parsed in TVMScript
|
||||
|
||||
|
||||
def _check_base_py_module_inheritance(node: doc.ClassDef) -> bool:
|
||||
"""Check if a class inherits from BasePyModule.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.ClassDef
|
||||
The class definition node to check.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bool
|
||||
True if the class inherits from BasePyModule, False otherwise.
|
||||
"""
|
||||
if not node.bases:
|
||||
return False
|
||||
|
||||
# Check each base class
|
||||
for base in node.bases:
|
||||
if hasattr(base, "id"):
|
||||
if base.id == "BasePyModule":
|
||||
return True
|
||||
elif hasattr(base, "attr"):
|
||||
if base.attr == "BasePyModule":
|
||||
return True
|
||||
elif hasattr(base, "value") and hasattr(base.value, "id"):
|
||||
if (
|
||||
base.value.id in ["BasePyModule", "tvm", "relax"]
|
||||
and hasattr(base, "attr")
|
||||
and base.attr == "BasePyModule"
|
||||
):
|
||||
return True
|
||||
|
||||
return False
|
||||
Reference in New Issue
Block a user