chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,378 @@
|
||||
# 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: F401, RUF005
|
||||
|
||||
# pylint: disable=invalid-name,too-many-locals
|
||||
|
||||
"""Utility functions for Relax"""
|
||||
|
||||
import itertools
|
||||
import string
|
||||
from collections.abc import Callable
|
||||
from typing import Any, Optional
|
||||
|
||||
import tvm_ffi
|
||||
from tvm_ffi import Array, Map
|
||||
|
||||
import tvm
|
||||
|
||||
from .. import tirx
|
||||
from ..ir import Attrs, Type, VDevice
|
||||
from ..te import Tensor as te_Tensor
|
||||
from ..te import create_prim_func
|
||||
from . import _ffi_api
|
||||
from .expr import Expr, Function, ShapeExpr, StringImm, te_tensor
|
||||
from .expr import Tuple as rx_Tuple
|
||||
from .type import ShapeType, TensorType
|
||||
|
||||
|
||||
def metadata_partitioner(rx_txt: str) -> list[str]:
|
||||
"""Extract Relax program and metadata section.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
rx_txt : str
|
||||
The input relax text.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : List[str]
|
||||
The result list of partitioned text, the first element
|
||||
is the relax program, and the second is metadata section.
|
||||
"""
|
||||
partitions = []
|
||||
left_curly = 0
|
||||
meta_start = 0
|
||||
meta_end = 0
|
||||
for i, char in enumerate(rx_txt):
|
||||
if i < 0:
|
||||
raise ValueError("The program is invalid.")
|
||||
if char == "{":
|
||||
if meta_start == 0:
|
||||
meta_start = i
|
||||
left_curly += 1
|
||||
elif char == "}":
|
||||
left_curly -= 1
|
||||
if left_curly == 0:
|
||||
meta_end = i + 1
|
||||
break
|
||||
|
||||
if meta_end == 0:
|
||||
raise ValueError("The metadata section was not found.")
|
||||
metadata = rx_txt[meta_start:meta_end]
|
||||
rx_program = rx_txt[meta_end:-1]
|
||||
|
||||
partitions.append(rx_program)
|
||||
partitions.append(metadata)
|
||||
|
||||
return partitions
|
||||
|
||||
|
||||
def convert_to_expr(value: Any) -> Expr:
|
||||
"""Helper function to convert the input to Expr, which follows the rules:
|
||||
1. Return the input itself if it's already a `relax.Expr`;
|
||||
2. Return `Expr` if the input is a primitive scalar;
|
||||
3. Return `relax.StringImm` if the input is `tvm.String` or `str`;
|
||||
4. Return `relax.Tuple` if the input is a tuple/list of `Expr`.
|
||||
|
||||
Notes
|
||||
-----
|
||||
1. `tvm.tirx.StringImm` is not allowed because of ambiguity,
|
||||
which can be either `relax.StringImm` or `Expr`.
|
||||
"""
|
||||
if isinstance(value, int):
|
||||
return tirx.IntImm("int64", value)
|
||||
|
||||
if isinstance(value, float):
|
||||
return tirx.FloatImm("float64", value)
|
||||
|
||||
tvm_value = tvm_ffi.convert(value)
|
||||
# Case 1
|
||||
if tvm.ir.is_prim_expr(tvm_value):
|
||||
return tvm_value
|
||||
# Note`` 1
|
||||
if isinstance(tvm_value, tirx.StringImm):
|
||||
raise TypeError(
|
||||
"Cannot convert `tirx.StringImm` to `relax.Expr` because of ambiguity,"
|
||||
"which can be either `relax.StringImm` or `Expr` "
|
||||
)
|
||||
# Case 2
|
||||
if isinstance(tvm_value, Expr):
|
||||
return tvm_value
|
||||
# Case 3
|
||||
if isinstance(tvm_value, str):
|
||||
return StringImm(value)
|
||||
# Case 4
|
||||
if isinstance(value, tuple | list):
|
||||
# `convert_to_expr` ensures that all elements are `Expr` if no exception raises
|
||||
return rx_Tuple([convert_to_expr(v) for v in value])
|
||||
raise TypeError(f"Cannot convert {value} with type {type(value)} to `relax.Expr`")
|
||||
|
||||
|
||||
def copy_with_new_vars(func: Function) -> Function:
|
||||
"""Copy the given function. All variables that are bound inside the original function
|
||||
would be copied to satisfy the restriction in the well-formed check: Variables in
|
||||
Relax must be bound exactly once. This also ensures that both the function and its copy
|
||||
can be inserted into the same IRModule, and be asserted on the structural equality
|
||||
agaisnt IRModule created by TVMScript.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func : Function
|
||||
The relax function to copy.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : Function
|
||||
The copied function.
|
||||
"""
|
||||
return _ffi_api.CopyWithNewVars(func) # type: ignore
|
||||
|
||||
|
||||
def gen_call_tir_inputs(
|
||||
func: Callable, *args: Any, **kwargs: Any
|
||||
) -> tuple[tirx.PrimFunc, Expr, list[TensorType], ShapeExpr | None]:
|
||||
"""Generate the inputs for call_tir 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 keyword args 'primfunc_attrs' is reserved for passing func
|
||||
attributes to be added to the PrimFunc that gets created.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : Tuple[tirx.PrimFunc, Expr, List[TensorType], Optional[ShapeExpr]]
|
||||
ret contains the inputs for call_tir, including a tirx prim_func, args,
|
||||
out_ty, and tir_vars.
|
||||
"""
|
||||
|
||||
tir_var_map: dict[tirx.Var, tirx.Expr] = {}
|
||||
|
||||
call_tir_args = []
|
||||
create_primfunc_args = []
|
||||
# extra list of tirx expression arguments
|
||||
# that are not covered by Tensor
|
||||
extra_tir_args_list = []
|
||||
|
||||
def _copy_undefined_var(expr: tirx.Expr):
|
||||
def _visit_expr(e: tirx.Expr):
|
||||
if isinstance(e, tirx.Var) and e not in tir_var_map:
|
||||
new_var = tirx.Var(e.name, e.ty)
|
||||
tir_var_map[e] = new_var
|
||||
|
||||
tirx.stmt_functor.post_order_visit(expr, _visit_expr)
|
||||
|
||||
def _convert_te_arg(te_args: Any) -> Any:
|
||||
"""Helper function used to convert Relax expressions to TE tensor.
|
||||
|
||||
In the common case, the type of te_args is a Relax expression and is converted
|
||||
into a TE tensor.
|
||||
If te_args is a nested or recursive datatype (i.e list, dict, tvm_ffi.Map, tvm_ffi.Array),
|
||||
we recursive and convert any value of type Relax expression into a TE tensor.
|
||||
Common values of type int, float, and str are preserved.
|
||||
|
||||
In dynamic shape cases, the passed in arguments may contain TIR variable.
|
||||
For example, the argument can be a Relax Var with TensorType, which
|
||||
has symbolic shape, or the argument can be a ShapeExpr with symbolic variables.
|
||||
To make the PrimFunc generated has independent variables with
|
||||
the caller Relax function, we will substitute the TIR variables in the input
|
||||
arguments with fresh ones, which is done by maintaining a TIR variable mapping.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
te_args : Any
|
||||
Argument to convert to TE
|
||||
|
||||
tir_var_map : Dict[tirx.Var, tirx.Expr]
|
||||
The TIR variable mapping, which maps TIR variables on the Relax function
|
||||
side to the new set of variables used on the PrimFunc side.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : (Any, [tvm.te.Tensor])
|
||||
A tuple of the converted te_args, and a list of te tensors for each converted
|
||||
Relax expression
|
||||
"""
|
||||
|
||||
def _convert_te_arg_helper(arg):
|
||||
if isinstance(arg, tvm.relax.Var) and tvm.ir.is_prim_expr(arg):
|
||||
name = arg.name_hint or f"scalar_input_{len(create_primfunc_args)}"
|
||||
tir_param = tirx.Var(name, arg.ty.dtype)
|
||||
call_tir_args.append(arg)
|
||||
create_primfunc_args.append(tir_param)
|
||||
return tir_param
|
||||
|
||||
if tvm.ir.is_prim_expr(arg):
|
||||
_copy_undefined_var(arg)
|
||||
new_arg = tirx.stmt_functor.substitute(arg, tir_var_map)
|
||||
extra_tir_args_list.append(new_arg)
|
||||
return new_arg
|
||||
|
||||
if isinstance(arg, Expr): # type: ignore
|
||||
if isinstance(arg.ty, TensorType):
|
||||
assert isinstance(arg.ty.shape, ShapeExpr), (
|
||||
"emit_te now only supports Tensor that has ShapeExpr shape"
|
||||
)
|
||||
for shape_value in arg.ty.shape.values:
|
||||
_copy_undefined_var(shape_value)
|
||||
|
||||
n_args = len(create_primfunc_args)
|
||||
if isinstance(arg, tvm.relax.Var):
|
||||
name = arg.name_hint
|
||||
elif n_args < len(string.ascii_uppercase):
|
||||
name = string.ascii_uppercase[n_args]
|
||||
else:
|
||||
name = f"tensor_input_{n_args}"
|
||||
|
||||
te_arg = te_tensor(arg, tir_var_map, name)
|
||||
|
||||
call_tir_args.append(arg)
|
||||
create_primfunc_args.append(te_arg)
|
||||
|
||||
return te_arg
|
||||
|
||||
if isinstance(arg.ty, ShapeType):
|
||||
assert isinstance(arg, ShapeExpr), (
|
||||
"For Expr having ShapeType, emit_te now only supports ShapeExpr"
|
||||
)
|
||||
return [_convert_te_arg_helper(val) for val in arg.values]
|
||||
|
||||
elif isinstance(arg, list | Array):
|
||||
return [_convert_te_arg_helper(x) for x in arg]
|
||||
elif isinstance(arg, tuple):
|
||||
return tuple(_convert_te_arg_helper(x) for x in arg)
|
||||
elif isinstance(arg, dict | Map):
|
||||
for key in arg:
|
||||
assert isinstance(key, str), (
|
||||
"emit_te only supports dict with string as the key currently"
|
||||
)
|
||||
return {k: _convert_te_arg_helper(arg[k]) for k in arg}
|
||||
elif isinstance(arg, int | float | str | Type | Attrs) or arg is None:
|
||||
return arg
|
||||
raise TypeError(f"not supported type in emit_te: {type(arg)}")
|
||||
|
||||
new_arg = _convert_te_arg_helper(te_args)
|
||||
return new_arg
|
||||
|
||||
def _get_unbound_tir_vars(args: list[te_Tensor], extra_tir_args: list[Expr]) -> list[tirx.Var]:
|
||||
"""get unbound TIR vars (i.e TIR vars used in the shape but is not
|
||||
itself a dimension of a shape)"""
|
||||
|
||||
bound_vars = set()
|
||||
used_vars = set()
|
||||
|
||||
def _populate_bound_vars(expr):
|
||||
if isinstance(expr, te_Tensor):
|
||||
for dim in expr.shape:
|
||||
_populate_bound_vars(dim)
|
||||
elif isinstance(expr, tirx.Var):
|
||||
bound_vars.add(expr)
|
||||
|
||||
def _populate_used_vars(expr):
|
||||
if isinstance(expr, te_Tensor):
|
||||
for dim in expr.shape:
|
||||
_populate_used_vars(dim)
|
||||
elif tvm.ir.is_prim_expr(expr):
|
||||
used_vars.update(tirx.analysis.undefined_vars(expr))
|
||||
|
||||
for arg in itertools.chain(args, extra_tir_args):
|
||||
_populate_used_vars(arg)
|
||||
|
||||
for arg in args:
|
||||
_populate_bound_vars(arg)
|
||||
|
||||
diff = used_vars - bound_vars
|
||||
return list(diff)
|
||||
|
||||
def _get_vdevice(arg: Any) -> VDevice | None:
|
||||
"""get the virtual device from arguments."""
|
||||
vdevice = None
|
||||
if isinstance(arg, Expr): # type: ignore
|
||||
if isinstance(arg.ty, TensorType):
|
||||
vdevice = arg.ty.vdevice
|
||||
elif isinstance(arg, list | Array | tuple):
|
||||
for x in arg:
|
||||
vdevice = _get_vdevice(x)
|
||||
if vdevice is not None:
|
||||
return vdevice
|
||||
elif isinstance(arg, dict | Map):
|
||||
for k in arg:
|
||||
vdevice = _get_vdevice(arg[k])
|
||||
if vdevice is not None:
|
||||
return vdevice
|
||||
return vdevice
|
||||
|
||||
def _shape_with_old_tir_var(
|
||||
shape_values: list[tirx.Expr], tir_var_inverse_map: dict[tirx.Var, tirx.Expr]
|
||||
):
|
||||
return ShapeExpr(
|
||||
[tirx.stmt_functor.substitute(value, tir_var_inverse_map) for value in shape_values]
|
||||
)
|
||||
|
||||
primfunc_attrs = kwargs.pop("primfunc_attrs", None)
|
||||
custom_out_ty = kwargs.pop("ty_args", [])
|
||||
|
||||
te_args = _convert_te_arg(args)
|
||||
te_kwargs = _convert_te_arg(kwargs)
|
||||
|
||||
te_out = func(*te_args, **te_kwargs)
|
||||
assert isinstance(te_out, te_Tensor) or (
|
||||
isinstance(te_out, tuple | list | Array) and all(isinstance(t, te_Tensor) for t in te_out)
|
||||
), "only support te.tensor or tuple/list/Array of te.tensor as function output"
|
||||
|
||||
outs = [te_out] if isinstance(te_out, te_Tensor) else list(te_out)
|
||||
unbound_tir_vars = _get_unbound_tir_vars([*create_primfunc_args, *outs], extra_tir_args_list)
|
||||
|
||||
inputs = [*create_primfunc_args] + outs + unbound_tir_vars
|
||||
tir_func = create_prim_func(inputs, "int64")
|
||||
|
||||
if primfunc_attrs:
|
||||
tir_func = tir_func.with_attrs(primfunc_attrs)
|
||||
|
||||
tir_func = tir_func.without_attr("global_symbol")
|
||||
|
||||
# Invert the TIR variable mapping, to convert the output shape back
|
||||
# with old set of variables.
|
||||
tir_var_inverse_map = {v: k for k, v in tir_var_map.items()}
|
||||
|
||||
if len(custom_out_ty) == 1:
|
||||
output_ty = custom_out_ty[0]
|
||||
else:
|
||||
output_ty = [
|
||||
TensorType(
|
||||
_shape_with_old_tir_var(out.shape, tir_var_inverse_map),
|
||||
out.dtype,
|
||||
_get_vdevice(args),
|
||||
)
|
||||
for out in outs
|
||||
]
|
||||
|
||||
tir_vars = None
|
||||
if len(unbound_tir_vars) > 0:
|
||||
tir_vars = _shape_with_old_tir_var(unbound_tir_vars, tir_var_inverse_map)
|
||||
|
||||
return (tir_func, call_tir_args, output_ty, tir_vars)
|
||||
Reference in New Issue
Block a user