wehub-resource-sync
370 lines
12 KiB
Python
370 lines
12 KiB
Python
# 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=invalid-name
|
|
# ruff: noqa: RUF005
|
|
"""Default legalization function for manipulate operators."""
|
|
|
|
import tvm
|
|
from tvm import DataTypeCode, relax, s_tir, te, tirx, topi
|
|
from tvm.ir import Call
|
|
from tvm.relax.op.base import call_tir
|
|
from tvm.relax.type import TensorType
|
|
from tvm.relax.utils import gen_call_tir_inputs
|
|
from tvm.tirx.expr import IntImm
|
|
|
|
from ...block_builder import BlockBuilder
|
|
from ...expr import Expr, ShapeExpr, Tuple, TupleGetItem, Var
|
|
from .common import LegalizeFunc, TEFunc, register_legalize
|
|
|
|
|
|
def _reshape(
|
|
te_func: TEFunc, primfunc_name: str, is_collapse_sum_like: bool = False
|
|
) -> LegalizeFunc:
|
|
def reshape_call_te(bb: BlockBuilder, call: Call):
|
|
tgt_shape = call.args[1].ty.shape if is_collapse_sum_like else call.args[1]
|
|
# If target shape is Var, pass its bound expr only when it is ShapeExpr
|
|
if isinstance(tgt_shape, Var):
|
|
tgt_shape = bb.lookup_binding(tgt_shape)
|
|
assert isinstance(tgt_shape, ShapeExpr)
|
|
return bb.call_te(te_func, call.args[0], tgt_shape, primfunc_name_hint=primfunc_name)
|
|
|
|
return reshape_call_te
|
|
|
|
|
|
register_legalize("relax.broadcast_to", _reshape(topi.broadcast_to, "broadcast_to"))
|
|
register_legalize("relax.reshape", _reshape(topi.reshape, "reshape"))
|
|
register_legalize(
|
|
"relax.collapse_sum_like",
|
|
_reshape(topi.collapse_sum, "collapse_sum", is_collapse_sum_like=True),
|
|
)
|
|
|
|
register_legalize("relax.collapse_sum_to", _reshape(topi.collapse_sum, "collapse_sum"))
|
|
|
|
|
|
@register_legalize("relax.concat")
|
|
def _concat(bb: BlockBuilder, call: Call) -> Expr:
|
|
t = call.args[0]
|
|
n_field = len(t.ty.fields)
|
|
while isinstance(t, Var):
|
|
binding = bb.lookup_binding(t)
|
|
if not isinstance(binding, Tuple | Var):
|
|
break
|
|
t = binding
|
|
|
|
assert isinstance(t, Tuple | Var)
|
|
fields = (
|
|
t.fields if isinstance(t, Tuple) else [bb.emit(TupleGetItem(t, i)) for i in range(n_field)]
|
|
)
|
|
return bb.call_te(
|
|
topi.concatenate, fields, None if call.attrs.axis is None else call.attrs.axis
|
|
)
|
|
|
|
|
|
@register_legalize("relax.expand_dims")
|
|
def _expand_dims(bb: BlockBuilder, call: Call) -> Expr:
|
|
def te_expand_dims(data, axis):
|
|
data_relax = relax.Var("data", relax.TensorType(data.shape))
|
|
f_infer_ty = call.op.get_attr("FInferType")
|
|
output_shape = f_infer_ty(relax.op.expand_dims(data_relax, axis), bb).shape
|
|
output_ndim = len(output_shape)
|
|
|
|
data_dims = []
|
|
for i in range(output_ndim):
|
|
if i not in axis and (i - output_ndim) not in axis:
|
|
data_dims.append(i)
|
|
return te.compute(
|
|
output_shape,
|
|
lambda *idx: data(*[idx[dim] for dim in data_dims]),
|
|
name="expand_dims",
|
|
)
|
|
|
|
return bb.call_te(
|
|
te_expand_dims, call.args[0], call.attrs.axis, primfunc_name_hint="expand_dims"
|
|
)
|
|
|
|
|
|
@register_legalize("relax.flatten")
|
|
def _flatten(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(topi.reshape, call.args[0], call.ty.shape.values)
|
|
|
|
|
|
@register_legalize("relax.permute_dims")
|
|
def _permute_dims(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(topi.transpose, call.args[0], call.attrs.axes)
|
|
|
|
|
|
@register_legalize("relax.split")
|
|
def _split(bb: BlockBuilder, call: Call) -> Expr:
|
|
if isinstance(call.attrs.indices_or_sections, tirx.IntImm):
|
|
indices_or_sections = call.attrs.indices_or_sections.value
|
|
else:
|
|
indices_or_sections = call.attrs.indices_or_sections
|
|
return bb.call_te(topi.split, call.args[0], indices_or_sections, call.attrs.axis)
|
|
|
|
|
|
@register_legalize("relax.squeeze")
|
|
def _squeeze(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(topi.squeeze, call.args[0], call.attrs.axis)
|
|
|
|
|
|
@register_legalize("relax.stack")
|
|
def _stack(bb: BlockBuilder, call: Call) -> Expr:
|
|
t = call.args[0]
|
|
n_field = len(t.ty.fields)
|
|
|
|
# Follow bindings to find the actual tuple
|
|
while isinstance(t, Var):
|
|
binding = bb.lookup_binding(t)
|
|
if not isinstance(binding, Tuple | Var):
|
|
break
|
|
t = binding
|
|
|
|
assert isinstance(t, Tuple | Var)
|
|
|
|
# Extract fields from either Tuple or bound Var
|
|
fields = (
|
|
t.fields if isinstance(t, Tuple) else [bb.emit(TupleGetItem(t, i)) for i in range(n_field)]
|
|
)
|
|
|
|
return bb.call_te(topi.stack, fields, 0 if call.attrs.axis is None else call.attrs.axis)
|
|
|
|
|
|
@register_legalize("relax.repeat")
|
|
def _repeat(bb: BlockBuilder, call: Call) -> Expr:
|
|
def te_repeat(data: te.Tensor, repeats: IntImm, axis: IntImm | None):
|
|
if axis is None:
|
|
# flatten data
|
|
out_shape = data.shape[0]
|
|
for i in data.shape[1:]:
|
|
out_shape *= i
|
|
data = topi.reshape(data, (out_shape,))
|
|
axis = 0
|
|
# topi only receives int repeats and axis
|
|
return topi.repeat(data, int(repeats), int(axis))
|
|
|
|
return bb.call_te(
|
|
te_repeat, call.args[0], call.attrs.repeats, call.attrs.axis, primfunc_name_hint="repeat"
|
|
)
|
|
|
|
|
|
@register_legalize("relax.tile")
|
|
def _tile(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(topi.tile, call.args[0], call.attrs.repeats)
|
|
|
|
|
|
@register_legalize("relax.flip")
|
|
def _flip(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(topi.flip, call.args[0], int(call.attrs.axis))
|
|
|
|
|
|
@register_legalize("relax.reverse_sequence")
|
|
def _reverse_sequence(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(
|
|
topi.reverse_sequence,
|
|
call.args[0],
|
|
call.args[1],
|
|
int(call.attrs.seq_axis),
|
|
int(call.attrs.batch_axis),
|
|
primfunc_name_hint="reverse_sequence",
|
|
)
|
|
|
|
|
|
@register_legalize("relax.gather_elements")
|
|
def _gather_elements(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(topi.gather, call.args[0], int(call.attrs.axis), call.args[1])
|
|
|
|
|
|
@register_legalize("relax.gather_nd")
|
|
def _gather_nd(bb: BlockBuilder, call: Call) -> Expr:
|
|
def te_gather_nd(data, indices, batch_dims):
|
|
indices_ndim = len(indices.shape)
|
|
axes = [indices_ndim - 1] + list(range(indices_ndim - 1))
|
|
indices = topi.transpose(indices, axes)
|
|
return topi.gather_nd(data, indices, batch_dims)
|
|
|
|
return bb.call_te(te_gather_nd, call.args[0], call.args[1], int(call.attrs.batch_dims))
|
|
|
|
|
|
@register_legalize("relax.index_tensor")
|
|
def _index_tensor(bb: BlockBuilder, call: Call) -> Expr:
|
|
t = call.args[1]
|
|
n_field = len(t.ty.fields)
|
|
fields = [bb.emit(TupleGetItem(t, i)) for i in range(n_field)]
|
|
return bb.call_te(topi.index_tensor, call.args[0], fields)
|
|
|
|
|
|
@register_legalize("relax.index_put")
|
|
def _index_put(bb: BlockBuilder, call: Call) -> Expr:
|
|
data = call.args[0]
|
|
indices = call.args[1]
|
|
values = call.args[2]
|
|
accumulate = call.attrs.accumulate
|
|
|
|
# If indices is a Tuple, unpack it into individual tensors
|
|
if isinstance(indices, relax.Tuple):
|
|
indices_list = [indices.fields[i] for i in range(len(indices.fields))]
|
|
else:
|
|
indices_list = [indices]
|
|
|
|
return bb.call_te(
|
|
topi.index_put,
|
|
data,
|
|
indices_list,
|
|
values,
|
|
accumulate=accumulate,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.meshgrid")
|
|
def _meshgrid(bb: BlockBuilder, call: Call) -> Expr:
|
|
t = call.args[0]
|
|
n_field = len(t.ty.fields)
|
|
while isinstance(t, Var):
|
|
binding = bb.lookup_binding(t)
|
|
if not isinstance(binding, Tuple | Var):
|
|
break
|
|
t = binding
|
|
|
|
assert isinstance(t, Tuple | Var)
|
|
fields = (
|
|
t.fields if isinstance(t, Tuple) else [bb.emit(TupleGetItem(t, i)) for i in range(n_field)]
|
|
)
|
|
return bb.call_te(
|
|
topi.meshgrid, fields, "ij" if call.attrs.indexing is None else call.attrs.indexing
|
|
)
|
|
|
|
|
|
def _is_gpu_target():
|
|
target = tvm.target.Target.current(allow_none=True)
|
|
return target is not None and "gpu" in target.keys
|
|
|
|
|
|
@register_legalize("relax.scatter_elements")
|
|
def _scatter_elements(bb: BlockBuilder, call: Call) -> Expr:
|
|
te_func = topi.gpu.scatter_elements if _is_gpu_target() else topi.scatter_elements
|
|
return bb.call_te(
|
|
te_func,
|
|
call.args[0],
|
|
call.args[1],
|
|
call.args[2],
|
|
call.attrs.axis,
|
|
call.attrs.reduction,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.scatter_nd")
|
|
def _scatter_nd(bb: BlockBuilder, call: Call) -> Expr:
|
|
# TODO(relax-team): Support native scatter_nd without te extern
|
|
base_te = topi.gpu.scatter_nd if _is_gpu_target() else topi.scatter_nd
|
|
|
|
def scatter_nd(data, indices, updates, reduction):
|
|
axes = list(range(len(indices.shape)))
|
|
indices = topi.transpose(indices, axes[-1:] + axes[:-1])
|
|
return base_te(data, indices, updates, reduction)
|
|
|
|
return bb.call_te(
|
|
scatter_nd,
|
|
call.args[0],
|
|
call.args[1],
|
|
call.args[2],
|
|
call.attrs.reduction,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.slice_scatter")
|
|
def _slice_scatter(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(
|
|
topi.slice_scatter,
|
|
call.args[0],
|
|
call.args[1],
|
|
call.args[2],
|
|
call.args[3],
|
|
call.args[4],
|
|
call.attrs.axis,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.one_hot")
|
|
def _one_hot(bb: BlockBuilder, call: Call) -> Expr:
|
|
indices, on_value, off_value = call.args
|
|
if not (tvm.ir.is_prim_expr(on_value) and tvm.ir.is_prim_expr(off_value)):
|
|
raise ValueError("on_value and off_value must be Expr")
|
|
if on_value.ty != off_value.ty:
|
|
raise ValueError("on_value and off_value must have the same dtype")
|
|
return bb.call_te(
|
|
topi.one_hot,
|
|
indices,
|
|
on_value,
|
|
off_value,
|
|
call.attrs.depth,
|
|
call.attrs.axis,
|
|
on_value.ty,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.layout_transform")
|
|
def _layout_transform(bb: BlockBuilder, call: Call) -> Expr:
|
|
def te_layout_transform(data, name):
|
|
"""
|
|
Returns a passthrough TE compute with appropriate name. This is needed to generate
|
|
TIR function, output shape info, TIR vars from gen_call_tir_inputs function.
|
|
"""
|
|
return te.compute(
|
|
data.shape,
|
|
data,
|
|
name=name,
|
|
)
|
|
|
|
def set_axis_sep(axis_sep: list, sch: s_tir.schedule, buffer_type: str):
|
|
sch.set_axis_separator(primfunc_name, (buffer_type, 0), axis_separators=axis_sep)
|
|
|
|
index_map: tvm.tirx.IndexMap = call.attrs.index_map
|
|
pad_value = call.attrs.pad_value
|
|
if pad_value is not None:
|
|
pad_value = pad_value.value
|
|
else:
|
|
if call.args[0].ty.dtype.matches_code(DataTypeCode.INT, DataTypeCode.UINT):
|
|
pad_value = 0
|
|
else:
|
|
pad_value = 0.0
|
|
|
|
axis_separators: tvm.tirx.IndexMap.AXIS_SEPARATOR = call.attrs.axis_separators
|
|
input_axis_separators: tvm.tirx.IndexMap.AXIS_SEPARATOR = call.attrs.input_axis_separators
|
|
|
|
# Convert to list from array
|
|
axis_separators = [int(sep) for sep in axis_separators]
|
|
primfunc_name = "te_layout_transform"
|
|
_, padding_predicate = index_map.non_surjective_inverse(call.args[0].ty.shape)
|
|
if not isinstance(padding_predicate, tvm.tirx.expr.IntImm):
|
|
primfunc_name += "_with_pad"
|
|
if len(axis_separators) != 0:
|
|
primfunc_name += "_axis_separator"
|
|
tir_func, call_args, _, tir_vars = gen_call_tir_inputs(
|
|
te_layout_transform, call.args[0], primfunc_name
|
|
)
|
|
# Create TIR schedule to apply layout changes with axis separators
|
|
sch = tvm.s_tir.Schedule(tir_func)
|
|
sch.transform_layout(primfunc_name, ("write", 0), index_map, pad_value)
|
|
set_axis_sep(axis_separators, sch, "write")
|
|
if input_axis_separators is not None:
|
|
set_axis_sep(input_axis_separators, sch, "read")
|
|
gvar = bb.add_func(sch.mod["main"], primfunc_name)
|
|
output_shape = index_map.map_shape(list(call_args[0].ty.shape))
|
|
output_dtype = call_args[0].ty.dtype
|
|
output_ty = [TensorType(output_shape, output_dtype)]
|
|
return call_tir(gvar, call_args, output_ty, tir_vars)
|