chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,52 @@
|
||||
# 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.
|
||||
# ruff: noqa: RUF005
|
||||
"""Initial impl of relax parser for sugars"""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from tvm.relax.script.builder import * # pylint: disable=redefined-builtin
|
||||
from tvm.relax.script.builder import ir as _relax
|
||||
|
||||
from . import parser as _parser
|
||||
from .entry import Any, Callable, Object, Prim, Shape, Tensor, Tuple, match_cast
|
||||
|
||||
from . import dist
|
||||
from .dist import * # pylint: disable=wildcard-import,redefined-builtin
|
||||
|
||||
if TYPE_CHECKING:
|
||||
# pylint: disable=invalid-name
|
||||
# Define prim_func and make it type check as static method
|
||||
# so most tvmscript won't trigger pylint error here.
|
||||
function = staticmethod
|
||||
else:
|
||||
from .entry import function, macro
|
||||
|
||||
__all__ = _relax.__all__ + [
|
||||
"dist",
|
||||
"Any",
|
||||
"Callable",
|
||||
"Object",
|
||||
"Prim",
|
||||
"Shape",
|
||||
"Tensor",
|
||||
"Tuple",
|
||||
"function",
|
||||
"macro",
|
||||
"match_cast",
|
||||
]
|
||||
@@ -0,0 +1,106 @@
|
||||
# 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=redefined-builtin,missing-docstring, invalid-name, unused-import, redefined-outer-name
|
||||
# ruff: noqa: F401
|
||||
|
||||
from typing import Any, Optional, Union
|
||||
|
||||
from tvm.ir import Range
|
||||
from tvm.relax import TensorType
|
||||
from tvm.relax.distributed import DeviceMesh, DTensorType, Placement, device_mesh
|
||||
from tvm.relax.script.builder.distributed import (
|
||||
annotate_sharding,
|
||||
call_tir,
|
||||
call_tir_local_view,
|
||||
const,
|
||||
redistribute,
|
||||
redistribute_replica_to_shard,
|
||||
)
|
||||
from tvm.script.ir_builder import IRBuilder
|
||||
from tvm.script.ir_builder.ir import IRModuleFrame
|
||||
from tvm.tirx import Expr
|
||||
|
||||
from .entry import TensorProxy, TypeProxy
|
||||
|
||||
############################### R.DTensor ###############################
|
||||
|
||||
|
||||
class DTensorProxy(TypeProxy):
|
||||
tensor_ty_proxy: TensorProxy
|
||||
device_mesh: DeviceMesh
|
||||
placement: Placement
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
tensor_ty_proxy: TensorProxy,
|
||||
device_mesh: DeviceMesh,
|
||||
placement: Placement,
|
||||
) -> None:
|
||||
self.device_mesh = device_mesh
|
||||
self.placement = placement
|
||||
self.tensor_ty_proxy = tensor_ty_proxy
|
||||
super().__init__()
|
||||
|
||||
def get_symbolic_vars(self) -> set[str]:
|
||||
return self.tensor_ty_proxy.get_symbolic_vars()
|
||||
|
||||
def as_ty(self, dict_globals: dict[str, Any] | None = None) -> DTensorType:
|
||||
return DTensorType(
|
||||
self.tensor_ty_proxy.as_ty(dict_globals),
|
||||
self.device_mesh,
|
||||
self.placement,
|
||||
)
|
||||
|
||||
|
||||
def DTensor(
|
||||
shape: list[Expr | str] | None = None,
|
||||
dtype: str | None = None,
|
||||
device_mesh: DeviceMesh | str = DeviceMesh([], Range(0, 1)),
|
||||
placement: Placement | str = "",
|
||||
*,
|
||||
ndim: int = -1,
|
||||
) -> DTensorProxy:
|
||||
# scalar tensor case
|
||||
if shape is not None 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):
|
||||
raise ValueError(f"shape must be a list or tuple, but got: {shape}")
|
||||
if isinstance(device_mesh, str):
|
||||
if not IRBuilder.is_in_scope():
|
||||
return (
|
||||
DTensorProxy(
|
||||
TensorProxy(shape, dtype, None, ndim), DeviceMesh([], Range(0, 1)), ""
|
||||
),
|
||||
)
|
||||
name, index = device_mesh.split("[")
|
||||
index = int(index[:-1])
|
||||
frames = IRBuilder.current().frames
|
||||
for f in frames:
|
||||
if isinstance(f, IRModuleFrame):
|
||||
device_mesh = f.global_infos[name][index]
|
||||
break
|
||||
assert isinstance(device_mesh, DeviceMesh)
|
||||
if isinstance(placement, str):
|
||||
placement = Placement.from_text(placement)
|
||||
return DTensorProxy(TensorProxy(shape, dtype, None, ndim), device_mesh, placement)
|
||||
|
||||
|
||||
__all__ = ["DTensor", "device_mesh"]
|
||||
@@ -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)
|
||||
@@ -0,0 +1,456 @@
|
||||
# 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, unused-argument
|
||||
|
||||
import functools
|
||||
import numbers
|
||||
from typing import Any
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
import tvm
|
||||
from tvm import relax, tirx
|
||||
from tvm.ir import GlobalVar
|
||||
from tvm.relax import Expr, Type
|
||||
from tvm.relax.script import builder as R
|
||||
from tvm.relax.script.builder.frame import BindingBlockFrame
|
||||
from tvm.relax.utils import convert_to_expr
|
||||
from tvm.script.ir_builder import ir as I
|
||||
from tvm.script.ir_builder.base import IRBuilder
|
||||
from tvm.script.parser._core import Parser, dispatch, doc
|
||||
|
||||
from .entry import (
|
||||
MatchCastPair,
|
||||
TypeProxy,
|
||||
_normalize_ty,
|
||||
_normalize_ty_proxy,
|
||||
)
|
||||
|
||||
relax.Expr._dispatch_type = relax.Expr # pylint: disable=protected-access
|
||||
dispatch.register_op(relax.Expr, doc.GtE, 0)(lambda lhs, rhs: lhs >= rhs)
|
||||
dispatch.register_op(relax.Expr, doc.Gt, 0)(lambda lhs, rhs: lhs > rhs)
|
||||
dispatch.register_op(relax.Expr, doc.LtE, 0)(lambda lhs, rhs: lhs <= rhs)
|
||||
dispatch.register_op(relax.Expr, doc.Lt, 0)(lambda lhs, rhs: lhs < rhs)
|
||||
|
||||
|
||||
def bind_assign_value(
|
||||
self: Parser,
|
||||
node: doc.expr,
|
||||
var_name: str,
|
||||
value: Any,
|
||||
anno_ty: Type | None = None,
|
||||
emit_prim_expr: bool = False,
|
||||
) -> Any:
|
||||
var_table = self.var_table.get()
|
||||
|
||||
if isinstance(value, tirx.Var):
|
||||
if value.name and var_name != value.name:
|
||||
self.report_error(
|
||||
node,
|
||||
"Cannot define TIR variables with different names. The LHS of binding should "
|
||||
"has the same name provided in RHS.",
|
||||
)
|
||||
if var_name in var_table:
|
||||
prev_value = var_table[var_name]
|
||||
if not isinstance(prev_value, tirx.Var):
|
||||
self.report_error(
|
||||
node,
|
||||
"Cannot redefine a non-TIR-variable object to a TIR variable. Please "
|
||||
"define the TIR variable with another name.",
|
||||
)
|
||||
if prev_value.ty != value.ty:
|
||||
self.report_error(
|
||||
node,
|
||||
f"Expected the same dtype for TIR vars but got {value.ty} vs {prev_value.ty}",
|
||||
)
|
||||
if not isinstance(value, type(prev_value)):
|
||||
self.report_error(
|
||||
node,
|
||||
f"Expected the same IR type for TIR vars "
|
||||
f"but existing value {type(value)} is mismatched "
|
||||
f"to previous {type(prev_value)}",
|
||||
)
|
||||
value = prev_value
|
||||
IRBuilder.name(var_name, value)
|
||||
return value
|
||||
|
||||
if tvm.ir.is_prim_expr(value):
|
||||
if not emit_prim_expr:
|
||||
return value
|
||||
|
||||
if isinstance(value, tuple):
|
||||
value = convert_to_expr(value)
|
||||
if isinstance(value, numbers.Number):
|
||||
value = R.const(value)
|
||||
|
||||
if isinstance(value, relax.Expr):
|
||||
var = R.emit(value, anno_ty)
|
||||
elif isinstance(value, MatchCastPair):
|
||||
if anno_ty is not None and not tvm_ffi.structural_equal(anno_ty, value.ty):
|
||||
self.report_error(
|
||||
node, "Cannot specify inconsistent annotation for a match cast pair. "
|
||||
)
|
||||
var = R.emit_match_cast(value.value, value.ty)
|
||||
else:
|
||||
return value
|
||||
|
||||
IRBuilder.name(var_name, var)
|
||||
return var
|
||||
|
||||
|
||||
def is_prim_value_call(node: doc.expr) -> bool:
|
||||
return isinstance(node, doc.Call) and getattr(node.func, "attr", None) == "prim_value"
|
||||
|
||||
|
||||
def eval_ty_proxy(self: Parser, node: doc.expr) -> TypeProxy:
|
||||
try:
|
||||
annotation = self.eval_expr(node)
|
||||
return _normalize_ty_proxy(annotation)
|
||||
except Exception as err: # pylint: disable=broad-except
|
||||
self.report_error(node, err)
|
||||
raise
|
||||
|
||||
|
||||
def eval_ty(self: Parser, node: doc.expr, eval_str: bool = False) -> Type:
|
||||
var_table = self.var_table.get() if eval_str else None
|
||||
try:
|
||||
ty = self.eval_expr(node)
|
||||
return _normalize_ty(ty, var_table)
|
||||
except Exception as err: # pylint: disable=broad-except
|
||||
self.report_error(node, err)
|
||||
raise
|
||||
|
||||
|
||||
def is_called(node: Any, func_name: str) -> bool:
|
||||
# Check if it calls into a func
|
||||
if isinstance(node, doc.Call):
|
||||
# Recursive call was found
|
||||
if isinstance(node.func, doc.Name) and node.func.id == func_name:
|
||||
return True
|
||||
elif isinstance(node, list | tuple):
|
||||
for stmt in node:
|
||||
if is_called(stmt, func_name):
|
||||
return True
|
||||
elif isinstance(node, doc.AnnAssign | doc.Assign | doc.Return | doc.Expr):
|
||||
return is_called(node.value, func_name)
|
||||
elif isinstance(node, doc.With):
|
||||
return is_called(node.body, func_name)
|
||||
elif isinstance(node, doc.If):
|
||||
smts = []
|
||||
if node.body is not None:
|
||||
smts = smts + list(node.body)
|
||||
if node.orelse is not None:
|
||||
smts = smts + list(node.orelse)
|
||||
return is_called(smts, func_name)
|
||||
return False
|
||||
|
||||
|
||||
def is_recursive(node: doc.FunctionDef) -> bool:
|
||||
# Check if it is a recursive function
|
||||
for stmt in node.body:
|
||||
if is_called(stmt, node.name):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def collect_symbolic_var_from_prelude(
|
||||
self: Parser, node: doc.FunctionDef, symbolic_vars: dict[str, tirx.Var]
|
||||
) -> dict[str, tirx.Var]:
|
||||
prelude_vars = {}
|
||||
for stmt in node.body:
|
||||
if isinstance(stmt, doc.Assign) and all(
|
||||
isinstance(target, doc.Name) and target.id in symbolic_vars for target in stmt.targets
|
||||
):
|
||||
values = self.eval_expr(stmt.value)
|
||||
|
||||
try:
|
||||
iter(values)
|
||||
except TypeError:
|
||||
values = [values]
|
||||
|
||||
assert len(stmt.targets) == len(values)
|
||||
for target, value in zip(stmt.targets, values):
|
||||
name = target.id
|
||||
prelude_vars[name] = value
|
||||
|
||||
return {**symbolic_vars, **prelude_vars}
|
||||
|
||||
|
||||
def collect_symbolic_var_from_params(self: Parser, node: doc.FunctionDef) -> None:
|
||||
# Collect symbolic vars from parameters
|
||||
symbolic_vars = {}
|
||||
for arg in node.args.args:
|
||||
if arg.annotation is None:
|
||||
self.report_error(arg, "Type annotation is required for function parameters.")
|
||||
param_ty_proxy = eval_ty_proxy(self, arg.annotation)
|
||||
|
||||
for var_name in param_ty_proxy.get_symbolic_vars():
|
||||
if var_name not in symbolic_vars:
|
||||
symbolic_vars[var_name] = tirx.Var(var_name, "int64")
|
||||
|
||||
# Update symbolic vars based on
|
||||
symbolic_vars = collect_symbolic_var_from_prelude(self, node, symbolic_vars)
|
||||
|
||||
# Define symbolic vars to the current var_table frame
|
||||
for var_name, var in symbolic_vars.items():
|
||||
self.var_table.add(var_name, var, allow_shadowing=False)
|
||||
|
||||
|
||||
@dispatch.register(token="relax", type_name="FunctionDef")
|
||||
def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
|
||||
is_inner_function = self.inside_function
|
||||
self.inside_function = True
|
||||
|
||||
# reserve a var for local function
|
||||
func_val = self.var_table.get().get(node.name)
|
||||
if not func_val and is_recursive(node):
|
||||
collect_symbolic_var_from_params(self, node)
|
||||
if node.returns is None:
|
||||
ret_ty = relax.TupleType([])
|
||||
else:
|
||||
ret_ty = eval_ty(self, node.returns, eval_str=True)
|
||||
params_ty = []
|
||||
for arg in node.args.args:
|
||||
if arg.annotation is None:
|
||||
self.report_error(arg, "Type annotation is required for function parameters.")
|
||||
param_ty = eval_ty(self, arg.annotation, eval_str=True)
|
||||
params_ty.append(param_ty)
|
||||
# created a var for the local function, the same var could be used for recursive call
|
||||
local_func_var = relax.Var(node.name, relax.FuncType(params_ty, ret_ty))
|
||||
self.var_table.add(node.name, local_func_var)
|
||||
|
||||
purity = find_decorator_annotation(node, "pure")
|
||||
# treat the function as private if we are inside another function
|
||||
# or if it has a privacy annotation
|
||||
privacy = is_inner_function or find_decorator_annotation(node, "private", default=False)
|
||||
|
||||
with self.var_table.with_frame():
|
||||
with self.with_dispatch_token("relax"):
|
||||
with R.function(is_pure=purity, is_private=privacy):
|
||||
R.func_name(node.name)
|
||||
collect_symbolic_var_from_params(self, node)
|
||||
|
||||
if node.returns is not None:
|
||||
ann_ty = eval_ty(self, node.returns, eval_str=True)
|
||||
R.func_ret_ty(ann_ty)
|
||||
|
||||
self.visit(node.args)
|
||||
|
||||
for stmt in node.body:
|
||||
if isinstance(stmt, doc.FunctionDef):
|
||||
if not stmt.decorator_list:
|
||||
self.report_error(stmt, "Function must be decorated")
|
||||
dec = self.eval_expr(stmt.decorator_list[-1])
|
||||
# inline prim_func was found
|
||||
if dec.dispatch_token == "tirx":
|
||||
self.report_error(stmt, "inline prim_func is disallowed in Relax IR")
|
||||
|
||||
self.visit_body(node.body)
|
||||
self.inside_function = is_inner_function
|
||||
|
||||
|
||||
def find_decorator_annotation(node: doc.FunctionDef, annotation: str, default: bool = True) -> bool:
|
||||
"""
|
||||
Check the value of given annotation (argument name) in the function decorator.
|
||||
Returns the value of the annotation if present, otherwise giving the default value.
|
||||
"""
|
||||
# look for the named argument in the function decorator
|
||||
for dec in node.decorator_list:
|
||||
if not isinstance(dec, doc.Call) or dec.func.attr != "function":
|
||||
continue
|
||||
for keyword in dec.keywords:
|
||||
if keyword.arg == annotation:
|
||||
return keyword.value.value
|
||||
return default
|
||||
|
||||
|
||||
@dispatch.register(token="relax", type_name="tvm_declare_function")
|
||||
def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar:
|
||||
with self.var_table.with_frame():
|
||||
collect_symbolic_var_from_params(self, node)
|
||||
|
||||
if node.returns is None:
|
||||
# Use AnyType as unknown return type
|
||||
# NOTE: Cannot use VoidType here because the return type can be refined later.
|
||||
ret_ty = relax.AnyType()
|
||||
else:
|
||||
ret_ty = eval_ty(self, node.returns, eval_str=True)
|
||||
params = []
|
||||
for arg in node.args.args:
|
||||
if arg.annotation is None:
|
||||
self.report_error(arg, "Type annotation is required for function parameters.")
|
||||
param_ty = eval_ty(self, arg.annotation, eval_str=True)
|
||||
params.append(relax.Var(arg.arg, param_ty))
|
||||
|
||||
is_pure = find_decorator_annotation(node, "pure")
|
||||
|
||||
func_signature = relax.Function.create_empty(params, ret_ty, is_pure=is_pure)
|
||||
return I.decl_function(node.name, func_signature)
|
||||
|
||||
|
||||
@dispatch.register(token="relax", type_name="pre_visit_local_function")
|
||||
def pre_visit_local_function(self: Parser, node: doc.Expr) -> None:
|
||||
ir_builder = IRBuilder()
|
||||
ir_builder.__enter__()
|
||||
|
||||
|
||||
@dispatch.register(token="relax", type_name="post_visit_local_function")
|
||||
def post_visit_local_function(self: Parser, node: doc.Expr) -> None:
|
||||
ir_builder = IRBuilder.current()
|
||||
result = ir_builder.get()
|
||||
ir_builder.__exit__(None, None, None)
|
||||
# reuse var if it is reserved
|
||||
reserved_var = self.var_table.get().get(node.name)
|
||||
if reserved_var:
|
||||
var = R.emit_var_binding(relax.VarBinding(reserved_var, result))
|
||||
else:
|
||||
var = R.emit(result)
|
||||
IRBuilder.name(node.name, var)
|
||||
self.var_table.add(node.name, var, allow_shadowing=False)
|
||||
|
||||
|
||||
@dispatch.register(token="relax", type_name="Expr")
|
||||
def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
|
||||
value = self.eval_expr(node.value)
|
||||
if isinstance(value, relax.Expr):
|
||||
var = R.emit(value)
|
||||
IRBuilder.name("_", var)
|
||||
is_void_value = isinstance(var.ty, relax.TupleType) and len(var.ty.fields) == 0
|
||||
|
||||
if not is_void_value:
|
||||
self.report_error(
|
||||
node,
|
||||
f"Non-void relax expressions must be bound to a variable, "
|
||||
f"but expression of type {var.ty} was used as a statement.",
|
||||
)
|
||||
|
||||
elif value is not None:
|
||||
self.report_error(node, f"Unsupported Expr stmt type {value}.")
|
||||
|
||||
|
||||
@dispatch.register(token="relax", type_name="arguments")
|
||||
def visit_arguments(self: Parser, node: doc.arguments) -> None:
|
||||
arg: doc.arg
|
||||
for arg in node.args:
|
||||
if arg.annotation is None:
|
||||
self.report_error(arg, "Type annotation is required for function parameters.")
|
||||
param_ty = eval_ty(self, arg.annotation, eval_str=True)
|
||||
param = R.arg(arg.arg, param_ty)
|
||||
|
||||
self.var_table.add(arg.arg, param)
|
||||
|
||||
|
||||
@dispatch.register(token="relax", type_name="tvm_annotation")
|
||||
def visit_tvm_annotation(self: Parser, node: doc.expr) -> Type:
|
||||
return eval_ty(self, node, eval_str=False)
|
||||
|
||||
|
||||
@dispatch.register(token="relax", type_name="With")
|
||||
def visit_with(self: Parser, node: doc.With) -> None:
|
||||
# Currently only `with R.dataflow()` is supported
|
||||
if len(node.items) != 1:
|
||||
self.report_error(node, "Only one item is allowed.")
|
||||
item = node.items[0]
|
||||
if item.optional_vars is not None:
|
||||
self.report_error(
|
||||
item.context_expr,
|
||||
"Relax syntax doesn't allow binding expressions in `with` to variables",
|
||||
)
|
||||
frame = self.eval_expr(item.context_expr)
|
||||
with self.var_table.with_frame():
|
||||
with frame:
|
||||
self.visit(node.body)
|
||||
if isinstance(frame, BindingBlockFrame) and frame.is_dataflow:
|
||||
output_vars = frame.output_vars
|
||||
for var in output_vars:
|
||||
self.var_table.add(var.name_hint, var, allow_shadowing=True)
|
||||
|
||||
|
||||
@dispatch.register(token="relax", 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=functools.partial(
|
||||
bind_assign_value,
|
||||
emit_prim_expr=is_prim_value_call(node.value),
|
||||
),
|
||||
allow_shadowing=True,
|
||||
)
|
||||
|
||||
|
||||
@dispatch.register(token="relax", type_name="AnnAssign")
|
||||
def visit_ann_assign(self: Parser, node: doc.AnnAssign) -> None:
|
||||
lhs = node.target
|
||||
rhs = self.eval_expr(node.value)
|
||||
anno_ty = self.visit_tvm_annotation(node.annotation)
|
||||
self.eval_assign(
|
||||
target=lhs,
|
||||
source=rhs,
|
||||
bind_value=functools.partial(
|
||||
bind_assign_value,
|
||||
anno_ty=anno_ty,
|
||||
emit_prim_expr=is_prim_value_call(node.value),
|
||||
),
|
||||
allow_shadowing=True,
|
||||
)
|
||||
|
||||
|
||||
@dispatch.register(token="relax", type_name="Return")
|
||||
def visit_return(self: Parser, node: doc.Assign) -> None:
|
||||
value = self.eval_expr(node.value)
|
||||
value = convert_to_expr(value)
|
||||
R.func_ret_value(value)
|
||||
|
||||
|
||||
@dispatch.register(token="relax", type_name="If")
|
||||
def visit_if(self: Parser, node: doc.If) -> None:
|
||||
if node.orelse is None:
|
||||
raise ValueError("Else statements are required for relax dialect.")
|
||||
with R.If(self.eval_expr(node.test)) as if_frame:
|
||||
with self.var_table.with_frame():
|
||||
with R.Then():
|
||||
self.visit_body(node.body)
|
||||
with self.var_table.with_frame():
|
||||
with R.Else():
|
||||
self.visit_body(node.orelse)
|
||||
self.var_table.add(if_frame.var_name, if_frame.var, allow_shadowing=True)
|
||||
|
||||
|
||||
@dispatch.register(token="relax", type_name="enter_token")
|
||||
def enter_token(self: Parser) -> dict[str, Any]:
|
||||
def relax_call(self, *args) -> Expr:
|
||||
args = [convert_to_expr(arg) if isinstance(arg, tuple) else arg for arg in args]
|
||||
|
||||
if all(isinstance(x, Expr) for x in args):
|
||||
return relax.Call(self, args)
|
||||
arg_types = [type(x) for x in args]
|
||||
raise RuntimeError(f"Do not know how to handle GlobalVar.__call__ for types {arg_types}")
|
||||
|
||||
context = {"GlobalVar.__call__": GlobalVar.__call__}
|
||||
GlobalVar.__call__ = relax_call
|
||||
return context
|
||||
|
||||
|
||||
@dispatch.register(token="relax", type_name="exit_token")
|
||||
def exit_token(self: Parser, context: dict[str, Any]) -> None:
|
||||
assert "GlobalVar.__call__" in context
|
||||
GlobalVar.__call__ = context.get("GlobalVar.__call__")
|
||||
Reference in New Issue
Block a user