chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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 *
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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): ...
|
||||
@@ -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",
|
||||
]
|
||||
Reference in New Issue
Block a user