chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,521 @@
|
||||
# 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=missing-docstring, invalid-name
|
||||
import inspect
|
||||
from collections.abc import Callable as _Callable
|
||||
from typing import Any, TypeVar
|
||||
|
||||
import tvm
|
||||
from tvm.ir import PrimType
|
||||
from tvm.relax import (
|
||||
AnyType,
|
||||
Expr,
|
||||
Function,
|
||||
FuncType,
|
||||
SeqExpr,
|
||||
ShapeExpr,
|
||||
ShapeType,
|
||||
TensorType,
|
||||
TupleType,
|
||||
Type,
|
||||
)
|
||||
from tvm.relax.expr import Var
|
||||
from tvm.relax.script import builder as R
|
||||
from tvm.runtime import ObjectConvertible
|
||||
from tvm.script.ir_builder.ir import lookup_vdevice
|
||||
from tvm.script.parser._core import doc, parse, utils
|
||||
from tvm.script.parser.core.entry import scan_macro
|
||||
from tvm.script.parser.core.parser import Parser, ScriptMacro
|
||||
|
||||
FType = TypeVar("FType", bound=_Callable)
|
||||
|
||||
############################## R.function ##############################
|
||||
|
||||
|
||||
# this formulation allows us to support having @R.function
|
||||
# appear as a decorator by itself or to have optional arguments
|
||||
# like @R.function(pure=False)
|
||||
def function(
|
||||
f: FType | None = None, pure: bool = True, private: bool = False, check_well_formed=True
|
||||
) -> Function | FType:
|
||||
# pylint: disable=unused-argument
|
||||
# (pure and private aren't used here, but are used later in parsing)
|
||||
|
||||
# need to inspect the stack first because is_defined_in_class expects the outer class
|
||||
# to be in a particular position in the stack
|
||||
orig_stack = inspect.stack()
|
||||
|
||||
def decorator_wrapper(f):
|
||||
if not inspect.isfunction(f):
|
||||
raise TypeError(f"Expect a function, but got: {f}")
|
||||
if utils.is_defined_in_class(orig_stack, f):
|
||||
return f
|
||||
return parse(f, utils.inspect_function_capture(f), check_well_formed=check_well_formed)
|
||||
|
||||
if f is not None:
|
||||
# if there are no optional args given, this will directly invoke the wrapper
|
||||
return decorator_wrapper(f)
|
||||
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", "relax")
|
||||
return decorator_wrapper
|
||||
|
||||
|
||||
setattr(function, "dispatch_token", "relax")
|
||||
|
||||
|
||||
############################## R.macro ##############################
|
||||
|
||||
|
||||
class RelaxMacro(ScriptMacro):
|
||||
"""Specialization of the ScriptMacro class for Relax."""
|
||||
|
||||
def parse_macro(self, parser: Parser) -> Expr:
|
||||
macro_def = self.get_macro_def()
|
||||
ret_value = None
|
||||
|
||||
with R.SeqExpr() as seq:
|
||||
for idx, stmt in enumerate(macro_def.body):
|
||||
# Normally, a "return" statement is only allowed in a R.function. We don't
|
||||
# want to parse the macro's body as if it was a body of a function, because
|
||||
# the latter imposes some constraints that we want to avoid.
|
||||
# At the same time, we want to use "return" to indicate the value of the
|
||||
# macro (since in Relax everything is an expression), so add special handling
|
||||
# of "return".
|
||||
if isinstance(stmt, doc.Return):
|
||||
ret_value = parser.eval_expr(stmt.value)
|
||||
if idx + 1 != len(macro_def.body):
|
||||
parser.report_error(macro_def, "'return' should be the last statement")
|
||||
break
|
||||
parser.visit(stmt)
|
||||
|
||||
if ret_value is None:
|
||||
parser.report_error(macro_def, "Macros must end with a return statement")
|
||||
|
||||
return SeqExpr(seq.binding_blocks, ret_value)
|
||||
|
||||
|
||||
def macro(*args, hygienic: bool = True) -> _Callable:
|
||||
"""Decorator for macro definitions.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
hygienic: bool
|
||||
Specifies whether the macro is hygienic or not.
|
||||
A macro is hygienic if all symbols used in the macro's body are resolved
|
||||
to values from the location of the macro definition. A non-hygienic macro
|
||||
will have its symbols resolved to values at the time of the macro's use.
|
||||
"""
|
||||
|
||||
def _decorator(func: _Callable) -> ScriptMacro:
|
||||
source, closure_vars = scan_macro(func, utils.inspect_function_capture(func))
|
||||
obj = RelaxMacro(source, closure_vars, func, hygienic)
|
||||
|
||||
def wrapper(*args, **kwargs):
|
||||
return obj(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
if len(args) == 0:
|
||||
return _decorator
|
||||
if len(args) == 1 and inspect.isfunction(args[0]):
|
||||
return _decorator(args[0])
|
||||
|
||||
raise ValueError(
|
||||
"Invalid use of R.macro. Usage: @R.macro, @R.macro(), @R.macro(hygienic=[True|False])"
|
||||
)
|
||||
|
||||
|
||||
############################# Type ##############################
|
||||
|
||||
|
||||
class TypeProxy(ObjectConvertible):
|
||||
def as_ty(self, dict_globals: dict[str, Any] | None = None) -> Type:
|
||||
raise NotImplementedError()
|
||||
|
||||
def get_symbolic_vars(self) -> set[str]:
|
||||
return {}
|
||||
|
||||
def asobject(self):
|
||||
return self.as_ty(None)
|
||||
|
||||
|
||||
############################### R.Any ################################
|
||||
|
||||
|
||||
class AnyProxy(TypeProxy):
|
||||
"""The proxy for AnyType.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
values : Optional[List[Expr]]
|
||||
The symbolic shape values if known.
|
||||
|
||||
ndim : Optional[int]
|
||||
The size of the shape.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def get_symbolic_vars(self) -> set[str]:
|
||||
return set()
|
||||
|
||||
def as_ty(self, dict_globals: dict[str, Any] | None = None) -> AnyType:
|
||||
return AnyType()
|
||||
|
||||
|
||||
def Any() -> AnyProxy:
|
||||
return AnyProxy()
|
||||
|
||||
|
||||
ObjectProxy = AnyProxy
|
||||
|
||||
|
||||
def Object() -> AnyProxy:
|
||||
return AnyProxy()
|
||||
|
||||
|
||||
############################### R.Tensor ###############################
|
||||
|
||||
|
||||
def _eval_shape(expr: str | Expr, dict_globals: dict[str, Any] | None) -> Expr:
|
||||
if isinstance(expr, str):
|
||||
code = compile(expr, "<string>", "eval")
|
||||
return eval(code, dict_globals or {}) # pylint: disable=eval-used
|
||||
else:
|
||||
return expr
|
||||
|
||||
|
||||
class TensorProxy(TypeProxy):
|
||||
shape: list[str | Expr] | None
|
||||
dtype: str
|
||||
vdevice: str | None
|
||||
ndim: int
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
shape: list[Expr | str] | Expr | None = None,
|
||||
dtype: str | None = None,
|
||||
vdevice: str | None = None,
|
||||
ndim: int = -1,
|
||||
) -> None:
|
||||
if isinstance(shape, Expr):
|
||||
if not isinstance(shape, ShapeExpr | Var):
|
||||
raise ValueError(
|
||||
"When the shape is an Expr, it must be a ShapeExpr or a Var with ShapeExpr "
|
||||
f"value. But got: {shape} with type: {type(shape)}"
|
||||
)
|
||||
if isinstance(shape, Var) and not isinstance(shape.ty, ShapeType):
|
||||
raise ValueError(
|
||||
"When the shape is a Var, it must have shape ty. But got "
|
||||
f"{shape} with ty: {shape.ty}"
|
||||
)
|
||||
self.shape = shape
|
||||
self.dtype = dtype
|
||||
self.vdevice = vdevice
|
||||
self.ndim = ndim
|
||||
|
||||
def get_symbolic_vars(self) -> set[str]:
|
||||
if self.shape is None or isinstance(self.shape, Expr):
|
||||
return {}
|
||||
else:
|
||||
return {s for s in self.shape if isinstance(s, str) and s.isidentifier()}
|
||||
|
||||
def as_ty(self, dict_globals: dict[str, Any] | None = None) -> TensorType:
|
||||
vdev = self.vdevice
|
||||
if isinstance(self.vdevice, str):
|
||||
if ":" in self.vdevice:
|
||||
split_vdev = self.vdevice.split(":")
|
||||
vdev = lookup_vdevice(split_vdev[0], int(split_vdev[1]))
|
||||
else:
|
||||
vdev = lookup_vdevice(self.vdevice, 0)
|
||||
|
||||
if self.shape is None:
|
||||
return TensorType(None, self.dtype, vdev, self.ndim)
|
||||
elif isinstance(self.shape, ShapeExpr | Var):
|
||||
return TensorType(self.shape, self.dtype, vdev, self.ndim)
|
||||
else:
|
||||
if dict_globals is None and any([isinstance(s, str) for s in self.shape]):
|
||||
raise ValueError(
|
||||
"String-defined shape expr is only allowed when parsing function parameters "
|
||||
"and return annotations for TVMScript."
|
||||
)
|
||||
shape = [_eval_shape(s, dict_globals) for s in self.shape]
|
||||
return TensorType(shape, self.dtype, vdev, self.ndim)
|
||||
|
||||
|
||||
def Tensor(
|
||||
shape: list[Expr | str] | Expr | None = None,
|
||||
dtype: str | None = None,
|
||||
vdevice: str | None = None,
|
||||
ndim: int = -1,
|
||||
) -> TensorProxy:
|
||||
# scalar tensor case
|
||||
if shape is not None and not isinstance(shape, Var) and len(shape) == 0:
|
||||
shape = []
|
||||
if isinstance(shape, str) and dtype is None:
|
||||
dtype = shape
|
||||
shape = None
|
||||
|
||||
if shape is not None and not isinstance(shape, tuple | list) and not isinstance(shape, Expr):
|
||||
raise ValueError(f"shape must be a list/tuple or an Expr, but got: {shape}")
|
||||
return TensorProxy(shape, dtype, vdevice, ndim)
|
||||
|
||||
|
||||
############################## R.Callable ##############################
|
||||
|
||||
|
||||
class CallableProxy(TypeProxy):
|
||||
params: list[TypeProxy]
|
||||
ret: TypeProxy
|
||||
purity: bool
|
||||
derive_func: str | tvm.ir.EnvFunc | None
|
||||
|
||||
"""Function type.
|
||||
|
||||
A function type consists of a list of type parameters to enable
|
||||
the definition of generic functions,
|
||||
a set of type constraints which we omit for the time being,
|
||||
a sequence of argument types, the purity of the function, and a return type.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
params : List[TypeProxy]
|
||||
The argument TypeProxy
|
||||
|
||||
ret : TypeProxy
|
||||
The return TypeProxy.
|
||||
|
||||
purity : bool
|
||||
Whether the callable is pure.
|
||||
|
||||
derive_func: Optional[Union[str, tvm.ir.EnvFunc]]
|
||||
The derivation function to determine the output Type,
|
||||
based on the arguments provided to the function. The
|
||||
specified function should be accessible using
|
||||
`tvm.get_global_func`, and should have a signature
|
||||
`Callable[[relax.Call, relax.BlockBuilder], relax.Type]`.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
params: TypeProxy | list[TypeProxy] | None = None,
|
||||
ret: TypeProxy | None = None,
|
||||
purity: bool | None = None,
|
||||
derive_func: str | tvm.ir.EnvFunc | None = None,
|
||||
) -> None:
|
||||
if params is None:
|
||||
self.params = params
|
||||
else:
|
||||
if not isinstance(params, list | tuple):
|
||||
params = [params]
|
||||
# convert `R.Callable` to `R.Callable()`
|
||||
self.params = [param() if callable(param) else param for param in params]
|
||||
|
||||
# Mimic the C++ defaults, where an opaque function is assumed
|
||||
# to be impure, and a non-opaque function is assumed to be
|
||||
# pure.
|
||||
if purity is None:
|
||||
purity = params is not None
|
||||
|
||||
self.ret = ret() if callable(ret) else ret
|
||||
self.purity = purity
|
||||
self.derive_func = derive_func
|
||||
|
||||
def get_symbolic_vars(self) -> set[str]:
|
||||
if self.params is None:
|
||||
return set()
|
||||
else:
|
||||
return set().union(*[p.get_symbolic_vars() for p in self.params])
|
||||
|
||||
def as_ty(self, dict_globals: dict[str, Any] | None = None) -> FuncType:
|
||||
if self.ret is None:
|
||||
ret = None
|
||||
else:
|
||||
ret = self.ret.as_ty(dict_globals)
|
||||
|
||||
if self.params is None:
|
||||
params = None
|
||||
else:
|
||||
params = [param.as_ty(dict_globals) for param in self.params]
|
||||
|
||||
if params is None:
|
||||
return FuncType.opaque_func(ret=ret, derive_func=self.derive_func, purity=self.purity)
|
||||
else:
|
||||
return FuncType(params, ret, purity=self.purity)
|
||||
|
||||
|
||||
def Callable(
|
||||
params: TypeProxy | list[TypeProxy] | None = None,
|
||||
ret: TypeProxy | None = None,
|
||||
purity: bool | None = None,
|
||||
derive_func: str | tvm.ir.EnvFunc | None = None,
|
||||
) -> CallableProxy:
|
||||
return CallableProxy(params, ret, purity=purity, derive_func=derive_func)
|
||||
|
||||
|
||||
############################### R.Tuple ################################
|
||||
|
||||
|
||||
class TupleProxy(TypeProxy):
|
||||
fields: list[TypeProxy]
|
||||
"""The type of tuple values.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
fields : List[TypeProxy]
|
||||
The fields in the tuple
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*fields: list[TypeProxy],
|
||||
) -> None:
|
||||
if len(fields) == 1 and isinstance(fields[0], tuple | list):
|
||||
fields = fields[0]
|
||||
# convert `R.Tensor` to `R.Tensor()`
|
||||
self.fields = [field() if callable(field) else field for field in fields]
|
||||
|
||||
def get_symbolic_vars(self) -> set[str]:
|
||||
return set().union(*[f.get_symbolic_vars() for f in self.fields])
|
||||
|
||||
def as_ty(self, dict_globals: dict[str, Any] | None = None) -> TupleType:
|
||||
fields = [field.as_ty(dict_globals) for field in self.fields]
|
||||
return TupleType(fields)
|
||||
|
||||
|
||||
def Tuple(*fields: list[TypeProxy]) -> TupleProxy:
|
||||
return TupleProxy(*fields)
|
||||
|
||||
|
||||
############################### R.Shape ################################
|
||||
|
||||
|
||||
class ShapeProxy(TypeProxy):
|
||||
values: list[Expr] | None
|
||||
ndim: int
|
||||
"""The type of shape values.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
values : Optional[List[Expr]]
|
||||
The symbolic shape values if known.
|
||||
|
||||
ndim : Optional[int]
|
||||
The size of the shape.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
values: list[Expr] | None = None,
|
||||
ndim: int = -1,
|
||||
) -> None:
|
||||
self.values = values
|
||||
self.ndim = ndim
|
||||
|
||||
def get_symbolic_vars(self) -> set[str]:
|
||||
if self.values is None:
|
||||
return set()
|
||||
else:
|
||||
return {v for v in self.values if isinstance(v, str) and v.isidentifier()}
|
||||
|
||||
def as_ty(self, dict_globals: dict[str, Any] | None = None) -> ShapeType:
|
||||
values = [_eval_shape(v, dict_globals) for v in self.values] if self.values else None
|
||||
return ShapeType(values, self.ndim)
|
||||
|
||||
|
||||
def Shape(values: list[Expr] | None = None, ndim: int = -1) -> ShapeProxy:
|
||||
return ShapeProxy(values, ndim)
|
||||
|
||||
|
||||
################################ R.Prim ################################
|
||||
|
||||
|
||||
class PrimProxy(TypeProxy):
|
||||
dtype: str
|
||||
|
||||
"""The type of TIR-representable values.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
dtype : str
|
||||
The data type.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dtype: str,
|
||||
) -> None:
|
||||
self.dtype = dtype
|
||||
|
||||
def get_symbolic_vars(self) -> set[str]:
|
||||
return set()
|
||||
|
||||
def as_ty(self, dict_globals: dict[str, Any] | None = None) -> PrimType:
|
||||
return PrimType(self.dtype)
|
||||
|
||||
|
||||
def Prim(
|
||||
dtype: str,
|
||||
) -> PrimProxy:
|
||||
return PrimProxy(dtype)
|
||||
|
||||
|
||||
############################ R.match_cast #############################
|
||||
class MatchCastPair:
|
||||
value: Expr
|
||||
ty: Type
|
||||
|
||||
def __init__(self, value: Expr, ty: Type) -> None:
|
||||
self.value = value
|
||||
self.ty = ty
|
||||
|
||||
|
||||
def match_cast(value: Expr, ty: Type):
|
||||
ty = _normalize_ty(ty)
|
||||
|
||||
if value is None:
|
||||
raise ValueError("value of match_cast cannot be None")
|
||||
if ty is None:
|
||||
raise ValueError("ty of match_cast cannot be None")
|
||||
return MatchCastPair(value, ty)
|
||||
|
||||
|
||||
def _normalize_ty_proxy(annotation) -> TypeProxy:
|
||||
if annotation is None:
|
||||
return TupleProxy([])
|
||||
elif callable(annotation):
|
||||
annotation = annotation()
|
||||
if tvm.ir.is_prim_expr(annotation):
|
||||
return PrimProxy(annotation.ty.dtype)
|
||||
return annotation
|
||||
elif isinstance(annotation, TypeProxy):
|
||||
return annotation
|
||||
else:
|
||||
raise TypeError(f"Expected TypeProxy but got {type(annotation)}.")
|
||||
|
||||
|
||||
def _normalize_ty(ty, dict_globals: dict[str, Any] | None = None) -> Type:
|
||||
if isinstance(ty, Type):
|
||||
return ty
|
||||
else:
|
||||
proxy = _normalize_ty_proxy(ty)
|
||||
return proxy.as_ty(dict_globals)
|
||||
Reference in New Issue
Block a user