wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tvm/relax/script/__init__.py
+27
View File
@@ -0,0 +1,27 @@
# 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.
"""Relax-layer TVMScript pieces (parser, builder).
After the per-dialect TVMScript restructure, the Relax layer owns its own
``script/{parser,builder}`` subpackages. ``tvm.script.relax`` resolves to
this module via the dialect registry, so the public parser surface
(``function``, ``Tensor``, ``match_cast``, etc.) is re-exported here.
"""
# pylint: disable=redefined-builtin,wildcard-import,unused-wildcard-import
from .parser import *
from .parser import dist
python/tvm/relax/script/builder/__init__.py
+25
View File
@@ -0,0 +1,25 @@
# 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.
"""Package tvm.relax.script.builder.
Holds the per-dialect ir_builder API for Relax. The legacy path
``tvm.script.ir_builder.relax`` resolves here via the dialect registry.
"""
# pylint: disable=wildcard-import,redefined-builtin
from . import distributed, frame, ir
from .ir import *
python/tvm/relax/script/builder/_ffi_api.py
+21
View File
@@ -0,0 +1,21 @@
# 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.
"""FFI APIs for tvm.script.ir_builder.relax"""
import tvm_ffi
tvm_ffi.init_ffi_api("script.ir_builder.relax", __name__) # pylint: disable=protected-access
python/tvm/relax/script/builder/distributed/__init__.py
+21
View File
@@ -0,0 +1,21 @@
# 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.
# pylint: disable=unused-import
"""Package tvm.script.ir_builder.relax.distributed"""
from .ir import * # pylint: disable=wildcard-import,redefined-builtin
python/tvm/relax/script/builder/distributed/_ffi_api.py
+21
View File
@@ -0,0 +1,21 @@
# 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.
"""FFI APIs for tvm.script.ir_builder.relax.distributed"""
import tvm_ffi
tvm_ffi.init_ffi_api("script.ir_builder.relax.distributed", __name__) # pylint: disable=protected-access
python/tvm/relax/script/builder/distributed/ir.py
+170
View File
@@ -0,0 +1,170 @@
# 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, wrong-import-order, no-member, invalid-name, unused-import
# ruff: noqa: F401
"""IRBuilder for distributed Relax dialect"""
from numbers import Number
from typing import Optional, Union
import numpy as _np # type: ignore
import tvm
from tvm import base as _base
from tvm.ir import Call
from tvm.relax.distributed import DeviceMesh, DTensorType, Placement
from tvm.relax.expr import Constant, Expr, ExternFunc, ShapeExpr
from tvm.relax.expr import Tuple as RxTuple
from tvm.relax.op.distributed import (
annotate_sharding as _annotate_sharding,
)
from tvm.relax.op.distributed import (
call_tir_local_view,
redistribute_replica_to_shard,
)
from tvm.relax.op.distributed import (
redistribute as _redistribute,
)
from tvm.relax.script.builder.ir import py_str
from tvm.relax.utils import convert_to_expr
from tvm.runtime import _tensor
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder.ir import IRModuleFrame
from . import _ffi_api
def call_tir(
func: str | Expr,
args: Expr,
out_ty: DTensorType | list[DTensorType],
tir_vars: ShapeExpr | tuple[Expr] | list[Expr] | None = None,
) -> Call:
"""Distributed version of call_tir
Parameters:
----------
func : Union[str, Expr]
The destination-passing-style function, can be ExternFunc or PrimFunc.
args : Expr
The input arguments.
out_ty : Union[DTensorType, List[DTensorType]]
The type information of the call_tir output.
It should be a single or a list of DTensorType. Each one denotes the
type information of a returned distributed tensor.
tir_vars : Optional[Union[ShapeExpr, Tuple[Expr], List[Expr]]]
ShapeExpr representing a tuple of integers to unpack when calling func. Is null if not used
Returns
-------
ret: Call
A call node for the call_tir operator.
"""
if isinstance(func, str):
func = ExternFunc(func)
if isinstance(args, tuple | list):
args = RxTuple([convert_to_expr(a) for a in args])
elif isinstance(args, Expr) and not isinstance(args, RxTuple): # type: ignore
args = RxTuple((args,))
if not isinstance(out_ty, list):
out_ty = [out_ty]
if isinstance(tir_vars, list | tuple):
tir_vars = ShapeExpr(tir_vars)
return _ffi_api.call_tir_dist(func, args, out_ty, tir_vars) # type: ignore
def const(
value: bool | int | float | _np.ndarray | tvm.runtime.Tensor,
ty: DTensorType,
) -> Constant:
"""Create a constant value.
Parameters
----------
value: Union[bool, int, float, numpy.ndarray, tvm.runtime.Tensor]
The constant value.
dtype: Optional[str]
The data type of the resulting constant.
Note
----
When dtype is None, we use the following rule:
- int maps to "int32"
- float maps to "float32"
- bool maps to "bool"
- other using the same default rule as numpy.
"""
ty = tvm.runtime.convert(ty)
if not isinstance(ty, DTensorType):
raise TypeError("ty needs to be an instance of DTensorType. ")
dtype = str(ty.tensor_ty.dtype)
if isinstance(value, Number | (bool | list)):
value = _np.array(value, dtype=dtype)
if isinstance(value, _np.ndarray | _np.generic):
if dtype is not None:
value = value.astype(dtype)
value = _tensor.tensor(value)
if not isinstance(value, _tensor.Tensor):
raise ValueError("value has to be scalar or Tensor")
return Constant(value, ty)
def _lookup_device_mesh(device_mesh_str: py_str) -> DeviceMesh:
if not IRBuilder.is_in_scope():
raise ValueError("device_mesh cannot be found in global info")
name, index_str = device_mesh_str.split("[")
index = int(index_str[:-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)
return device_mesh
def annotate_sharding(
value: Expr, device_mesh: py_str | DeviceMesh, placement: py_str | Placement
) -> Expr:
if isinstance(device_mesh, py_str):
device_mesh = _lookup_device_mesh(device_mesh)
if isinstance(placement, py_str):
placement = Placement.from_text(placement)
return _annotate_sharding(value, device_mesh, placement)
def redistribute(
value: Expr, device_mesh: py_str | DeviceMesh, placement: py_str | Placement
) -> Expr:
if isinstance(device_mesh, py_str):
device_mesh = _lookup_device_mesh(device_mesh)
if isinstance(placement, py_str):
placement = Placement.from_text(placement)
return _redistribute(value, device_mesh, placement)
python/tvm/relax/script/builder/frame.py
+52
View File
@@ -0,0 +1,52 @@
# 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.
"""IR Builder Frame for Relax dialect"""
from tvm_ffi import register_object as _register_object
from tvm.script.ir_builder.base import IRBuilderFrame
@_register_object("script.ir_builder.relax.RelaxFrame")
class RelaxFrame(IRBuilderFrame):
"""The base ir_builder frame for the relax dialect."""
@_register_object("script.ir_builder.relax.SeqExprFrame")
class SeqExprFrame(RelaxFrame): ...
@_register_object("script.ir_builder.relax.FunctionFrame")
class FunctionFrame(SeqExprFrame):
"""The ir_builder frame for the relax function."""
@_register_object("script.ir_builder.relax.BindingBlockFrame")
class BindingBlockFrame(RelaxFrame):
"""The ir_builder frame for relax binding blocks."""
@_register_object("script.ir_builder.relax.IfFrame")
class IfFrame(RelaxFrame): ...
@_register_object("script.ir_builder.relax.ThenFrame")
class ThenFrame(SeqExprFrame): ...
@_register_object("script.ir_builder.relax.ElseFrame")
class ElseFrame(SeqExprFrame): ...
python/tvm/relax/script/builder/ir.py
+999
View File
@@ -0,0 +1,999 @@
# 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, wrong-import-order, no-member, invalid-name
"""IRBuilder for Relax dialect"""
import builtins
import functools
import inspect
from collections.abc import Callable
from typing import Any
import tvm
from tvm import DataType, relax
from tvm.ir import IRModule, VDevice
from tvm.relax import (
Call,
Expr,
ExternFunc,
ShapeExpr,
StringImm,
TupleGetItem,
Var,
VarBinding,
const,
)
from tvm.relax.dpl import PatternMatchingRewriter
############################### Operators ###############################
from tvm.relax.op import (
abs,
acos,
acosh,
add,
arange,
argmax,
argmin,
argsort,
asin,
asinh,
assert_op,
astype,
atan,
atan2,
atanh,
bitwise_and,
bitwise_not,
bitwise_or,
bitwise_xor,
broadcast_to,
bucketize,
builtin,
call_builtin_with_ctx,
call_dps_packed,
call_inplace_packed,
call_pure_packed,
call_tir,
call_tir_inplace,
call_tir_with_grad,
ccl,
ceil,
clip,
collapse_sum_like,
collapse_sum_to,
concat,
cos,
cosh,
cumprod,
cumsum,
dequantize,
divide,
dynamic_strided_slice,
einsum,
equal,
erf,
ewise_fma,
exp,
expand_dims,
eye,
eye_like,
flatten,
flip,
floor,
floor_divide,
floor_mod,
full,
full_like,
gather_elements,
gather_nd,
grad,
greater,
greater_equal,
hamming_window,
hint_on_device,
image,
index_put,
index_tensor,
invoke_closure,
invoke_pure_closure,
isfinite,
isinf,
isnan,
layout_transform,
left_shift,
less,
less_equal,
linear,
log,
log_add_exp,
logical_and,
logical_not,
logical_or,
logical_xor,
make_closure,
matmul,
max,
maximum,
mean,
median,
memory,
meshgrid,
min,
minimum,
mod,
multinomial_from_uniform,
multiply,
negative,
nn,
nonzero,
not_equal,
null_value,
one_hot,
ones,
ones_like,
outer,
permute_dims,
power,
print,
prod,
quantize,
repeat,
reshape,
reverse_sequence,
right_shift,
round,
rsqrt,
scatter_elements,
scatter_nd,
shape_of,
shape_to_tensor,
sigmoid,
sign,
sin,
sinh,
size,
slice_scatter,
sort,
split,
sqrt,
square,
squeeze,
stack,
std,
strided_slice,
subtract,
sum,
take,
tan,
tanh,
tensor_to_shape,
tile,
topk,
tril,
triu,
trunc,
unique,
variance,
vision,
vm,
where,
wrap_param,
zeros,
zeros_like,
)
from tvm.relax.op import (
call_py_func as _call_py_func,
)
from tvm.relax.op.builtin import stop_lift_params
from tvm.relax.type import Type
from tvm.relax.utils import convert_to_expr, gen_call_tir_inputs
from tvm.runtime import Object as tvm_Object
from tvm.runtime import ObjectConvertible
from tvm.runtime._tensor import (
cpu,
cuda,
device,
ext_dev,
hexagon,
metal,
opencl,
rocm,
vpi,
vulkan,
webgpu,
)
from tvm.script.ir_builder.ir import decl_function, lookup_vdevice
from . import _ffi_api, frame
##################### Python Native Function Alias ######################
py_print = builtins.print
py_tuple = tuple # pylint: disable=used-before-assignment
py_str = str # pylint: disable=used-before-assignment
################################ Device ################################
def to_vdevice(data: Expr, dst_vdevice: py_str | VDevice) -> Expr:
"""Copy data to the destination device.
Parameters
----------
data : Expr
The tensor to be copied.
dst_device : Union[py_str, VDevice]
The destination device where the data is copied to.
Returns
-------
result : Expr
The copied result.
"""
if isinstance(dst_vdevice, py_str):
if ":" in dst_vdevice:
split_vdev = dst_vdevice.split(":")
dst_vdevice = lookup_vdevice(split_vdev[0], int(split_vdev[1]))
else:
dst_vdevice = lookup_vdevice(dst_vdevice, 0)
return tvm.relax.op.to_vdevice(data, dst_vdevice)
############################### Function ################################
def function(is_pure: bool = True, is_private: bool = False) -> frame.FunctionFrame:
"""Start a function frame.
Parameters
----------
is_pure: bool
Whether the function is annotated as pure.
is_private : bool
Whether the function is annotated as private.
Returns
-------
frame: FunctionFrame
The constructed function frame.
"""
return _ffi_api.Function( # type: ignore[attr-defined] # pylint: disable=no-member
is_pure, is_private
)
def arg(name: py_str, ty: Type) -> Var:
"""Add a parameter to the last function frame.
Parameters
----------
name: str
The name of the parameter.
ty: Type
The type of the parameter
Returns
-------
var: Var
The created function parameter var.
"""
return _ffi_api.Arg(name, ty) # type: ignore[attr-defined] # pylint: disable=no-member
def func_name(name: py_str) -> None:
"""Specify the name of the last function frame.
Parameters
----------
name: str
The function name.
"""
return _ffi_api.FuncName(name) # type: ignore[attr-defined] # pylint: disable=no-member
def func_attr(attrs: dict[py_str, tvm_Object]) -> None:
"""Specify the attrs of the last function frame.
Parameters
----------
attrs: Dict[str, Object]
The function attrs.
"""
return _ffi_api.FuncAttrs(attrs) # type: ignore[attr-defined] # pylint: disable=no-member
def func_ret_type(ret_ty: Type) -> None:
"""Specify the return type of the last function frame.
Parameters
----------
ret_ty: Type
The function return type.
"""
return _ffi_api.FuncRetType(ret_ty) # type: ignore[attr-defined] # pylint: disable=no-member
def func_ret_ty(ret_ty: Type) -> None:
"""Backward-compatible alias for `func_ret_type`."""
return func_ret_type(ret_ty)
def func_ret_value(value: Expr) -> None:
"""Specify the return value of the last function frame.
Parameters
----------
value: Expr
The function return value.
"""
return _ffi_api.FuncRetValue(value) # type: ignore[attr-defined] # pylint: disable=no-member
def rewriter(rewriter_mod: IRModule | type) -> PatternMatchingRewriter:
"""Define a pattern-rewrite rule
The IRModule must have two publicly-exposed functions, `pattern`
and `replacement`, where `pattern` and `replacement` have the same
function signature.
.. code-block:: python
@R.rewriter
class RewriteAddIntoMultiply:
@R.function
def pattern(A: R.Tensor):
B = A + A
return B
@R.function
def replacement(A: R.Tensor):
B = A * 2
return B
Parameters
----------
rewriter_mod: Union[IRModule, Type]
Either an IRModule that defines a rewrite pattern, or a
TVMScript class that can be parsed into an IRModule.
Returns
-------
rewriter: PatternMatchingRewriter
A rewriter object, which can be applied either to a Relax
function or to an entire IRModule.
"""
if not isinstance(rewriter_mod, IRModule):
rewriter_mod = tvm.script.ir_module(rewriter_mod)
return PatternMatchingRewriter.from_module(rewriter_mod)
############################# BindingBlock ##############################
def dataflow() -> frame.BindingBlockFrame:
"""Start a dataflow binding block frame.
Returns
-------
frame: frame.BindingBlockFrame
The created ir_builder Block frame.
"""
return _ffi_api.Dataflow() # type: ignore[attr-defined] # pylint: disable=no-member
def output(*vars: tuple[Var]) -> None:
"""Expose the dataflow block output variables as global ones.
Parameters
----------
vars: Tuple[Var]
The output variables of a dataflow block.
"""
return _ffi_api.DataflowBlockOutput(vars) # type: ignore[attr-defined] # pylint: disable=no-member
################################## Ops #################################
def call_packed(
func: py_str,
*args: Expr,
ty_args: Type | list[Type] | None = None,
**kwargs: Any,
) -> Call:
"""Create a relax Call, which calls a packed function.
Parameters
----------
func: str
The name of extern function.
*args : Expr
The arguments.
ty_args: Optional[Union[Type, List[Type]]]
The list of type information arguments.
kwargs: Expr
The keyword arguments.
Returns
-------
call: Call
The created Relax Call
"""
op = ExternFunc(func)
args = py_tuple(convert_to_expr(a) for a in args)
if ty_args is None:
ty_args = []
if isinstance(ty_args, py_tuple): # type: ignore
ty_args = list(ty_args)
elif not isinstance(ty_args, list):
ty_args = [ty_args]
ty_args = [
(ty() if callable(ty) else ty.asobject() if isinstance(ty, ObjectConvertible) else ty)
for ty in ty_args
]
is_default = False
if "attrs_type_key" in kwargs:
attrs_type_key = kwargs["attrs_type_key"]
kwargs.pop("attrs_type_key")
else:
attrs_type_key = "ir.DictAttrs"
is_default = True
attrs = None
if kwargs or not is_default:
attrs = tvm.ir.attrs.make_node(attrs_type_key, **kwargs)
return Call(op, args, attrs=attrs, ty_args=ty_args)
def call_py_func(
py_func_name: py_str,
*args: Expr,
out_ty: Type | list[Type],
) -> Call:
"""Create a relax Call, which calls a Python function.
Parameters
----------
py_func_name: str
The name of the Python function to call. This should correspond to a function
in the IRModule's pyfuncs attribute.
*args : Expr
The arguments.
out_ty: Union[Type, List[Type]]
The type information of the call_py_func output.
It should be a single or a list of TensorType. Each one denotes the
type information of a returned tensor.
Returns
-------
call: Call
The created Relax Call for call_py_func operator.
"""
args = py_tuple(convert_to_expr(a) for a in args)
if isinstance(out_ty, py_tuple): # type: ignore
out_ty = list(out_ty)
elif not isinstance(out_ty, list):
out_ty = [out_ty]
out_ty = [
(ty() if callable(ty) else ty.asobject() if isinstance(ty, ObjectConvertible) else ty)
for ty in out_ty
]
# Convert string to StringImm
try:
func_name_imm = (
StringImm(py_func_name) if isinstance(py_func_name, py_str) else py_func_name
)
except (TypeError, ValueError, AttributeError):
func_name_imm = StringImm(py_func_name)
return _call_py_func(func_name_imm, args, out_ty)
def _ty_arg_wrapper(func):
"""A wrapper to convert TypeProxies to Type for builtin operators with ty_args"""
def _convert_tensor_type(args):
if isinstance(args, list | py_tuple): # type: ignore
new_args = [_convert_tensor_type(x) for x in args]
return type(args)(new_args)
if isinstance(args, dict):
return {_convert_tensor_type(k): _convert_tensor_type(v) for k, v in args.items()}
if inspect.isfunction(args):
args = args()
if isinstance(args, ObjectConvertible):
args = args.asobject()
return args
@functools.wraps(func)
def wrapped(*args, **kwargs):
return func(*_convert_tensor_type(args), **_convert_tensor_type(kwargs))
return wrapped # type: ignore
invoke_closure = _ty_arg_wrapper(invoke_closure) # pylint: disable=invalid-name
call_builtin_with_ctx = _ty_arg_wrapper(call_builtin_with_ctx) # pylint: disable=invalid-name
############################### Emits ###############################
def emit(value: Expr, annotate_ty: Type | None = None) -> Var:
"""Emit a binding to the last binding block frame.
Parameters
----------
value: Expr
The right side value of the bindings to be emitted.
annotate_ty: Optional[Type]
The optional type annotation for the emitted value.
Returns
-------
var: Var
The left side var of the emitted binding.
"""
return _ffi_api.Emit(value, annotate_ty) # type: ignore[attr-defined] # pylint: disable=no-member
def emit_te(func: Callable, *args: Any, **kwargs: Any) -> Call:
"""Emit a call node according to the te function.
This function converts arguments from relax expression to te tensor,
The callback func should return a te tensor or a list of te tensors.
Parameters
----------
func : Callable
A function that returns a te tensor or a list of te tensors.
args : Any, optional
arguments passed to the function.
kwargs : Any, optional
The keyword arguments passed to the function.
Note that the following keyword args are reserved:
- 'primfunc_name_hint' for passing name hint to the PrimFunc
that gets generated.
- 'primfunc_attrs' is reserved for passing func attributes to
be added to the PrimFunc that gets created.
Returns
-------
call : Call
A newly created call that calls into a tirx function.
"""
primfunc_name_hint = kwargs.pop("primfunc_name_hint", None)
tir_func, call_args, out_ty, tir_vars = gen_call_tir_inputs(func, *args, **kwargs)
if not primfunc_name_hint:
primfunc_name_hint = func.__name__
gvar = decl_function(primfunc_name_hint, tir_func) # type: ignore
return call_tir(gvar, call_args, out_ty, tir_vars)
def emit_match_cast(value: Expr, ty: Type) -> Var:
"""Emit a match_cast binding to the last binding block frame.
Parameters
----------
value: Expr
The value of the MatchCast to be emitted.
ty: Type
The ty of the MatchCast to be emitted.
Returns
-------
var: Var
The left side var of the emitted binding.
"""
return _ffi_api.EmitMatchCast(value, ty) # type: ignore
def emit_var_binding(value: VarBinding) -> Var:
"""Emit a binding to the last binding block frame.
Parameters
----------
value: VarBinding
The binding to be emitted.
Returns
-------
var: Var
The left side var of the emitted binding.
"""
return _ffi_api.EmitVarBinding(value) # type: ignore
def emit_with_type(
op: str,
args: Expr,
ty_args: Type | list[Type] | None = None,
) -> Call:
"""Create a Relax Call with type arguments.
Parameters
----------
op: Expr
The relax op for which type args are to be appended
args : Expr
The arguments.
ty_args: Optional[Union[Type, List[Type]]]
The list of type arguments.
Returns
-------
call: Call
The created Relax Call
"""
builtin_call = tvm.ir.Op.get(op)
return Call(builtin_call, args, attrs=None, ty_args=ty_args)
def emit_with_ty(
op: str,
args: Expr,
ty_args: Type | list[Type] | None = None,
) -> Call:
"""Backward-compatible alias for `emit_with_type`."""
return emit_with_type(op, args, ty_args)
############################### SeqExpr ###############################
def SeqExpr() -> frame.SeqExprFrame: # pylint: disable=invalid-name
"""Create a SeqExpr frame.
Returns
-------
res : frame.SeqExprFrame
The result SeqExprFrame
"""
return _ffi_api.SeqExpr() # type: ignore[attr-defined] # pylint: disable=no-member
############################# If Then Else #############################
def If(condition: Expr) -> frame.IfFrame: # pylint: disable=invalid-name
"""Create an if frame.
Parameters
----------
condition : Expr
The condition of if statement, executes the true branch if the
condition is true, otherwise jump into the false branch.
Returns
-------
res : frame.IfFrame
The result IfFrame.
"""
if not isinstance(condition, Expr):
condition = relax.prim_value(condition)
return _ffi_api.If(condition) # type: ignore[attr-defined] # pylint: disable=no-member
def Then() -> frame.ThenFrame: # pylint: disable=invalid-name
"""Create a then frame.
Returns
-------
res : frame.ThenFrame
The result ThenFrame.
"""
return _ffi_api.Then() # type: ignore[attr-defined] # pylint: disable=no-member
def Else() -> frame.ElseFrame: # pylint: disable=invalid-name
"""Create an else frame.
Returns
-------
res : frame.ElseFrame
The result ElseFrame.
"""
return _ffi_api.Else() # type: ignore[attr-defined] # pylint: disable=no-member
############################### R.tuple ################################
def tuple(*fields: Expr) -> Expr:
"""Create a tuple expression.
Parameters
----------
*fields : Expr
The fields of the tuple.
Returns
-------
res : Expr
The result tuple.
"""
if len(fields) == 0:
fields = py_tuple()
return relax.Tuple(fields) # type: ignore[attr-defined] # pylint: disable=no-member
############################### R.shape ################################
def shape(value: list[Expr]) -> Expr:
"""Create a ShapeExpr.
Parameters
----------
value : List[Expr]
The fields of the tuple.
Returns
-------
res : Expr
The result tuple.
"""
return relax.ShapeExpr(value) # pylint: disable=no-member # type: ignore
############################### Expr ###############################
def prim_value(value: Expr | int | float) -> Expr:
"""Convert a value to a primitive expression.
Parameters
----------
value : Expr | int | float
The value to convert.
Returns
-------
res : Expr
The primitive expression.
"""
return relax.prim_value(value) # type: ignore[attr-defined] # pylint: disable=no-member
def str(value: py_str) -> Expr:
"""Create a string imm expression.
Parameters
----------
value : str
The value of the str.
Returns
-------
res : Expr
The result str.
"""
return relax.StringImm(value) # type: ignore[attr-defined] # pylint: disable=no-member
def dtype(value: py_str | DataType) -> Expr:
"""Create a dtype imm expression.
Parameters
----------
value : dtype
The value of the dtype.
Returns
-------
res : Expr
The result dtype.
"""
return relax.DataTypeImm(value) # type: ignore[attr-defined] # pylint: disable=no-member
############################### Importer ###############################
__all__ = [
"Else",
"ExternFunc",
"If",
"SeqExpr",
"ShapeExpr",
"Then",
"TupleGetItem",
"abs",
"acos",
"acosh",
"add",
"arange",
"arg",
"argmax",
"argmin",
"argsort",
"asin",
"asinh",
"assert_op",
"astype",
"atan",
"atan2",
"atanh",
"bitwise_and",
"bitwise_not",
"bitwise_or",
"bitwise_xor",
"broadcast_to",
"bucketize",
"builtin",
"call_builtin_with_ctx",
"call_dps_packed",
"call_inplace_packed",
"call_packed",
"call_pure_packed",
"call_py_func",
"call_tir",
"call_tir_inplace",
"call_tir_with_grad",
"ccl",
"ceil",
"clip",
"collapse_sum_like",
"collapse_sum_to",
"concat",
"const",
"cos",
"cosh",
"cpu",
"cuda",
"cumprod",
"cumsum",
"dataflow",
"dequantize",
"device",
"divide",
"dtype",
"dynamic_strided_slice",
"einsum",
"emit",
"emit_match_cast",
"emit_te",
"emit_var_binding",
"emit_with_ty",
"emit_with_type",
"equal",
"erf",
"ewise_fma",
"exp",
"expand_dims",
"ext_dev",
"eye",
"eye_like",
"flatten",
"flip",
"floor",
"floor_divide",
"floor_mod",
"full",
"full_like",
"func_attr",
"func_name",
"func_ret_ty",
"func_ret_type",
"func_ret_value",
"function",
"gather_elements",
"gather_nd",
"grad",
"greater",
"greater_equal",
"hamming_window",
"hexagon",
"hint_on_device",
"image",
"index_put",
"index_tensor",
"invoke_closure",
"invoke_pure_closure",
"isfinite",
"isinf",
"isnan",
"layout_transform",
"left_shift",
"less",
"less_equal",
"linear",
"log",
"log_add_exp",
"logical_and",
"logical_not",
"logical_or",
"logical_xor",
"make_closure",
"matmul",
"max",
"maximum",
"mean",
"median",
"memory",
"meshgrid",
"metal",
"min",
"minimum",
"mod",
"multinomial_from_uniform",
"multiply",
"negative",
"nn",
"nonzero",
"not_equal",
"null_value",
"one_hot",
"ones",
"ones_like",
"opencl",
"outer",
"output",
"permute_dims",
"power",
"prim_value",
"print",
"prod",
"quantize",
"repeat",
"reshape",
"reverse_sequence",
"rewriter",
"right_shift",
"rocm",
"round",
"rsqrt",
"scatter_elements",
"scatter_nd",
"shape",
"shape_of",
"shape_to_tensor",
"sigmoid",
"sign",
"sin",
"sinh",
"size",
"slice_scatter",
"sort",
"split",
"sqrt",
"square",
"squeeze",
"stack",
"std",
"stop_lift_params",
"str",
"str",
"strided_slice",
"subtract",
"sum",
"take",
"tan",
"tanh",
"tensor_to_shape",
"tile",
"to_vdevice",
"topk",
"tril",
"triu",
"trunc",
"tuple",
"unique",
"variance",
"vision",
"vm",
"vpi",
"vulkan",
"webgpu",
"where",
"wrap_param",
"zeros",
"zeros_like",
]
python/tvm/relax/script/parser/__init__.py
+52
View File
@@ -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",
]
python/tvm/relax/script/parser/dist.py
+106
View File
@@ -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"]
python/tvm/relax/script/parser/entry.py
+521
View File
@@ -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)
python/tvm/relax/script/parser/parser.py
+456
View File
@@ -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__")