wehub-resource-sync
827 lines
28 KiB
Python
827 lines
28 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,unused-argument
|
|
"""Default legalization function for neural network operators."""
|
|
|
|
import logging
|
|
import math
|
|
|
|
from tvm import s_tir, te, tirx, topi
|
|
from tvm.ir import Call
|
|
|
|
from ...block_builder import BlockBuilder
|
|
from ...expr import Expr
|
|
from .common import _call_topi_without_attr, register_legalize
|
|
|
|
|
|
@register_legalize("relax.nn.conv1d")
|
|
def _nn_conv1d(bb: BlockBuilder, call: Call) -> Expr:
|
|
if call.attrs.out_layout != call.attrs.data_layout:
|
|
logging.info(
|
|
"TOPI conv1d does not support different input-output "
|
|
"layouts, and thus cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
if len(call.attrs.data_layout) != 3 or len(call.attrs.kernel_layout) != 3:
|
|
logging.info(
|
|
"Conv1D where data layout or kernel layout have channel chunk "
|
|
"cannot be legalized by TOPI at this moment."
|
|
)
|
|
return call
|
|
if call.attrs.groups != 1:
|
|
data_layout = s_tir.slayout(call.attrs.data_layout)
|
|
kernel_layout = s_tir.slayout(call.attrs.kernel_layout)
|
|
ic = call.args[0].ty.shape.values[data_layout.index_of("C")]
|
|
oc = call.args[1].ty.shape.values[kernel_layout.index_of("O")]
|
|
if not isinstance(ic, tirx.IntImm) or not isinstance(oc, tirx.IntImm):
|
|
logging.info(
|
|
"Conv1D where number of groups is more than one and input or output "
|
|
"channel size is symbolic cannot be legalized by TOPI at this moment."
|
|
)
|
|
return call
|
|
|
|
return bb.call_te(
|
|
topi.nn.conv1d,
|
|
data=call.args[0],
|
|
kernel=call.args[1],
|
|
strides=call.attrs.strides,
|
|
padding=call.attrs.padding,
|
|
dilation=call.attrs.dilation,
|
|
groups=call.attrs.groups,
|
|
data_layout=call.attrs.data_layout,
|
|
kernel_layout=call.attrs.kernel_layout,
|
|
out_dtype=call.attrs.out_dtype if call.attrs.out_dtype != "" else None,
|
|
primfunc_name_hint="conv1d",
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.conv2d")
|
|
def _nn_conv2d(bb: BlockBuilder, call: Call) -> Expr:
|
|
if call.attrs.out_layout != call.attrs.data_layout:
|
|
logging.info(
|
|
"TOPI conv2d does not support different input-output "
|
|
"layouts, and thus cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
if len(call.attrs.data_layout) != 4 or len(call.attrs.kernel_layout) != 4:
|
|
logging.info(
|
|
"Conv2D where data layout or kernel layout have channel chunk "
|
|
"cannot be legalized by TOPI at this moment."
|
|
)
|
|
return call
|
|
if call.attrs.groups != 1:
|
|
data_layout = s_tir.slayout(call.attrs.data_layout)
|
|
kernel_layout = s_tir.slayout(call.attrs.kernel_layout)
|
|
ic = call.args[0].ty.shape.values[data_layout.index_of("C")]
|
|
oc = call.args[1].ty.shape.values[kernel_layout.index_of("O")]
|
|
if not isinstance(ic, tirx.IntImm) or not isinstance(oc, tirx.IntImm):
|
|
logging.info(
|
|
"Conv2D where number of groups is more than one and input or output "
|
|
"channel size is symbolic cannot be legalized by TOPI at this moment."
|
|
)
|
|
return call
|
|
|
|
return bb.call_te(
|
|
topi.nn.conv,
|
|
inp=call.args[0],
|
|
filt=call.args[1],
|
|
stride=call.attrs.strides,
|
|
padding=call.attrs.padding,
|
|
dilation=call.attrs.dilation,
|
|
groups=call.attrs.groups,
|
|
data_layout=call.attrs.data_layout,
|
|
kernel_layout=call.attrs.kernel_layout,
|
|
out_dtype=call.attrs.out_dtype if call.attrs.out_dtype != "" else None,
|
|
primfunc_name_hint="conv2d",
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.conv3d")
|
|
def _nn_conv3d(bb: BlockBuilder, call: Call) -> Expr:
|
|
if call.attrs.out_layout != call.attrs.data_layout:
|
|
logging.info(
|
|
"TOPI conv3d does not support different input-output "
|
|
"layouts, and thus cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
if len(call.attrs.data_layout) != 5 or len(call.attrs.kernel_layout) != 5:
|
|
logging.info(
|
|
"Conv3D where data layout or kernel layout have channel chunk "
|
|
"cannot be legalized by TOPI at this moment."
|
|
)
|
|
return call
|
|
if call.attrs.groups != 1:
|
|
data_layout = s_tir.slayout(call.attrs.data_layout)
|
|
kernel_layout = s_tir.slayout(call.attrs.kernel_layout)
|
|
ic = call.args[0].ty.shape.values[data_layout.index_of("C")]
|
|
oc = call.args[1].ty.shape.values[kernel_layout.index_of("O")]
|
|
if not isinstance(ic, tirx.IntImm) or not isinstance(oc, tirx.IntImm):
|
|
logging.info(
|
|
"Conv3D where number of groups is more than one and input or output "
|
|
"channel size is symbolic cannot be legalized by TOPI at this moment."
|
|
)
|
|
return call
|
|
|
|
return bb.call_te(
|
|
topi.nn.conv,
|
|
inp=call.args[0],
|
|
filt=call.args[1],
|
|
stride=call.attrs.strides,
|
|
padding=call.attrs.padding,
|
|
dilation=call.attrs.dilation,
|
|
groups=call.attrs.groups,
|
|
data_layout=call.attrs.data_layout,
|
|
kernel_layout=call.attrs.kernel_layout,
|
|
out_dtype=call.attrs.out_dtype if call.attrs.out_dtype != "" else None,
|
|
primfunc_name_hint="conv3d",
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.conv1d_transpose")
|
|
def _nn_conv1d_transpose(bb: BlockBuilder, call: Call) -> Expr:
|
|
if call.attrs.out_layout != call.attrs.data_layout:
|
|
logging.info(
|
|
"TOPI conv1d_transpose does not support different input-output "
|
|
"layouts, and thus cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
if call.attrs.data_layout != "NCW" or call.attrs.kernel_layout != "IOW":
|
|
logging.info(
|
|
"TOPI conv1d_transpose does not support input layout other than NCW, "
|
|
"and kernel layout other than IOW, so cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
strides = [int(s) for s in call.attrs.strides]
|
|
padding = [int(p) for p in call.attrs.padding]
|
|
output_padding = [int(o) for o in call.attrs.output_padding]
|
|
groups = int(call.attrs.groups)
|
|
out_dtype = call.ty.dtype
|
|
dilation = [int(d) for d in call.attrs.dilation]
|
|
|
|
def te_conv1d_transpose(data, kernel):
|
|
# Dilated transposed conv == transposed conv with a spatially dilated (zero-filled) kernel.
|
|
if any(d != 1 for d in dilation):
|
|
kernel = topi.nn.dilate(kernel, [1, 1, dilation[0]], name="kernel_dilate")
|
|
return topi.nn.group_conv1d_transpose_ncw(
|
|
data, kernel, strides, padding, out_dtype, output_padding, groups
|
|
)
|
|
|
|
return bb.call_te(
|
|
te_conv1d_transpose, call.args[0], call.args[1], primfunc_name_hint="conv1d_transpose"
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.conv2d_transpose")
|
|
def _nn_conv2d_transpose(bb: BlockBuilder, call: Call) -> Expr:
|
|
if call.attrs.out_layout != call.attrs.data_layout:
|
|
logging.info(
|
|
"TOPI conv2d_transpose does not support different input-output "
|
|
"layouts, and thus cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
if call.attrs.data_layout != "NCHW" or call.attrs.kernel_layout != "IOHW":
|
|
logging.info(
|
|
"TOPI conv2d_transpose does not support input layout other than NCHW, "
|
|
"and kernel layout other than IOHW, so cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
strides = [int(s) for s in call.attrs.strides]
|
|
padding = [int(p) for p in call.attrs.padding]
|
|
output_padding = [int(o) for o in call.attrs.output_padding]
|
|
groups = int(call.attrs.groups)
|
|
out_dtype = call.ty.dtype
|
|
dilation = [int(d) for d in call.attrs.dilation]
|
|
|
|
def te_conv2d_transpose(data, kernel):
|
|
# Dilated transposed conv == transposed conv with a spatially dilated (zero-filled) kernel.
|
|
if any(d != 1 for d in dilation):
|
|
kernel = topi.nn.dilate(kernel, [1, 1, dilation[0], dilation[1]], name="kernel_dilate")
|
|
return topi.nn.group_conv2d_transpose_nchw(
|
|
data, kernel, strides, padding, out_dtype, output_padding, groups
|
|
)
|
|
|
|
return bb.call_te(
|
|
te_conv2d_transpose, call.args[0], call.args[1], primfunc_name_hint="conv2d_transpose"
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.conv3d_transpose")
|
|
def _nn_conv3d_transpose(bb: BlockBuilder, call: Call) -> Expr:
|
|
# Keep policy in sync with _nn_conv2d_transpose: only lower when TOPI supports
|
|
# the layout/dilation.
|
|
if call.attrs.out_layout != call.attrs.data_layout:
|
|
logging.info(
|
|
"TOPI conv3d_transpose does not support different input-output "
|
|
"layouts, and thus cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
if call.attrs.data_layout != "NCDHW" or call.attrs.kernel_layout != "IODHW":
|
|
logging.info(
|
|
"TOPI conv3d_transpose does not support input layout other than NCDHW, "
|
|
"and kernel layout other than IODHW, so cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
strides = [int(s) for s in call.attrs.strides]
|
|
padding = [int(p) for p in call.attrs.padding]
|
|
output_padding = [int(o) for o in call.attrs.output_padding]
|
|
groups = int(call.attrs.groups)
|
|
out_dtype = call.ty.dtype
|
|
dilation = [int(d) for d in call.attrs.dilation]
|
|
|
|
def te_conv3d_transpose(data, kernel):
|
|
# Dilated transposed conv == transposed conv with a spatially dilated (zero-filled) kernel.
|
|
if any(d != 1 for d in dilation):
|
|
kernel = topi.nn.dilate(
|
|
kernel, [1, 1, dilation[0], dilation[1], dilation[2]], name="kernel_dilate"
|
|
)
|
|
return topi.nn.group_conv3d_transpose_ncdhw(
|
|
data, kernel, strides, padding, out_dtype, output_padding, groups
|
|
)
|
|
|
|
return bb.call_te(
|
|
te_conv3d_transpose, call.args[0], call.args[1], primfunc_name_hint="conv3d_transpose"
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.pad")
|
|
def _nn_pad(bb: BlockBuilder, call: Call) -> Expr:
|
|
pad_mode = call.attrs.pad_mode
|
|
pad_widths = call.attrs.pad_width
|
|
pad_before = pad_widths[::2]
|
|
pad_after = pad_widths[1::2]
|
|
if pad_mode == "reflect":
|
|
return bb.call_te(
|
|
topi.nn.reflect_pad, call.args[0], pad_before=pad_before, pad_after=pad_after
|
|
)
|
|
elif pad_mode == "replicate":
|
|
return bb.call_te(
|
|
topi.nn.replicate_pad, call.args[0], pad_before=pad_before, pad_after=pad_after
|
|
)
|
|
elif pad_mode == "circular":
|
|
return bb.call_te(
|
|
topi.nn.circular_pad, call.args[0], pad_before=pad_before, pad_after=pad_after
|
|
)
|
|
else:
|
|
return bb.call_te(
|
|
topi.nn.pad,
|
|
call.args[0],
|
|
pad_before=pad_before,
|
|
pad_after=pad_after,
|
|
pad_value=call.attrs.pad_value,
|
|
primfunc_name_hint="pad",
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.pixel_shuffle")
|
|
def _nn_pixel_shuffle(bb: BlockBuilder, call: Call) -> Expr:
|
|
upscale_factor = call.attrs.upscale_factor
|
|
return bb.call_te(topi.nn.pixel_shuffle, call.args[0], upscale_factor=upscale_factor)
|
|
|
|
|
|
@register_legalize("relax.nn.max_pool1d")
|
|
def _nn_max_pool1d(bb: BlockBuilder, call: Call) -> Expr:
|
|
if call.attrs.out_layout != call.attrs.layout:
|
|
logging.info(
|
|
"TOPI max_pool1d does not support different input-output "
|
|
"layouts, and thus cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
|
|
return bb.call_te(
|
|
topi.nn.pool1d,
|
|
call.args[0],
|
|
kernel=call.attrs.pool_size,
|
|
stride=call.attrs.strides,
|
|
dilation=call.attrs.dilation,
|
|
padding=call.attrs.padding,
|
|
pool_type="max",
|
|
ceil_mode=call.attrs.ceil_mode,
|
|
layout=call.attrs.layout,
|
|
primfunc_name_hint="max_pool1d",
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.max_pool2d")
|
|
def _nn_max_pool2d(bb: BlockBuilder, call: Call) -> Expr:
|
|
if call.attrs.out_layout != call.attrs.layout:
|
|
logging.info(
|
|
"TOPI max_pool2d does not support different input-output "
|
|
"layouts, and thus cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
|
|
return bb.call_te(
|
|
topi.nn.pool2d,
|
|
call.args[0],
|
|
kernel=call.attrs.pool_size,
|
|
stride=call.attrs.strides,
|
|
dilation=call.attrs.dilation,
|
|
padding=call.attrs.padding,
|
|
pool_type="max",
|
|
ceil_mode=call.attrs.ceil_mode,
|
|
layout=call.attrs.layout,
|
|
primfunc_name_hint="max_pool2d",
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.max_pool3d")
|
|
def _nn_max_pool3d(bb: BlockBuilder, call: Call) -> Expr:
|
|
if call.attrs.out_layout != call.attrs.layout:
|
|
logging.info(
|
|
"TOPI max_pool3d does not support different input-output "
|
|
"layouts, and thus cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
|
|
return bb.call_te(
|
|
topi.nn.pool3d,
|
|
call.args[0],
|
|
kernel=call.attrs.pool_size,
|
|
stride=call.attrs.strides,
|
|
dilation=call.attrs.dilation,
|
|
padding=call.attrs.padding,
|
|
pool_type="max",
|
|
ceil_mode=call.attrs.ceil_mode,
|
|
layout=call.attrs.layout,
|
|
primfunc_name_hint="max_pool3d",
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.avg_pool1d")
|
|
def _nn_avg_pool1d(bb: BlockBuilder, call: Call) -> Expr:
|
|
if call.attrs.out_layout != call.attrs.layout:
|
|
logging.info(
|
|
"TOPI avg_pool1d does not support different input-output "
|
|
"layouts, and thus cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
|
|
return bb.call_te(
|
|
topi.nn.pool1d,
|
|
call.args[0],
|
|
kernel=call.attrs.pool_size,
|
|
stride=call.attrs.strides,
|
|
dilation=call.attrs.dilation,
|
|
padding=call.attrs.padding,
|
|
pool_type="avg",
|
|
ceil_mode=call.attrs.ceil_mode,
|
|
layout=call.attrs.layout,
|
|
count_include_pad=call.attrs.count_include_pad,
|
|
primfunc_name_hint="avg_pool1d",
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.avg_pool2d")
|
|
def _nn_avg_pool2d(bb: BlockBuilder, call: Call) -> Expr:
|
|
if call.attrs.out_layout != call.attrs.layout:
|
|
logging.info(
|
|
"TOPI avg_pool2d does not support different input-output "
|
|
"layouts, and thus cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
|
|
return bb.call_te(
|
|
topi.nn.pool2d,
|
|
call.args[0],
|
|
kernel=call.attrs.pool_size,
|
|
stride=call.attrs.strides,
|
|
dilation=call.attrs.dilation,
|
|
padding=call.attrs.padding,
|
|
pool_type="avg",
|
|
ceil_mode=call.attrs.ceil_mode,
|
|
layout=call.attrs.layout,
|
|
count_include_pad=call.attrs.count_include_pad,
|
|
primfunc_name_hint="avg_pool2d",
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.avg_pool3d")
|
|
def _nn_avg_pool3d(bb: BlockBuilder, call: Call) -> Expr:
|
|
if call.attrs.out_layout != call.attrs.layout:
|
|
logging.info(
|
|
"TOPI avg_pool3d does not support different input-output "
|
|
"layouts, and thus cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
|
|
return bb.call_te(
|
|
topi.nn.pool3d,
|
|
call.args[0],
|
|
kernel=call.attrs.pool_size,
|
|
stride=call.attrs.strides,
|
|
dilation=call.attrs.dilation,
|
|
padding=call.attrs.padding,
|
|
pool_type="avg",
|
|
ceil_mode=call.attrs.ceil_mode,
|
|
layout=call.attrs.layout,
|
|
count_include_pad=call.attrs.count_include_pad,
|
|
primfunc_name_hint="avg_pool3d",
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.adaptive_avg_pool1d")
|
|
def _nn_adaptive_avg_pool1d(bb: BlockBuilder, call: Call) -> Expr:
|
|
if call.attrs.out_layout != call.attrs.layout:
|
|
logging.info(
|
|
"TOPI adaptive_avg_pool1d does not support different input-output "
|
|
"layouts, and thus cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
|
|
def te_adaptive_avg_pool1d(data, output_size, layout_str):
|
|
if output_size is None:
|
|
layout = s_tir.slayout(layout_str)
|
|
idx_W = layout.index_of("W")
|
|
assert idx_W != -1
|
|
output_size = data.shape[idx_W]
|
|
|
|
return topi.nn.adaptive_pool1d(data, output_size, "avg", layout_str)
|
|
|
|
return bb.call_te(
|
|
te_adaptive_avg_pool1d,
|
|
call.args[0],
|
|
call.attrs.output_size,
|
|
call.attrs.layout,
|
|
primfunc_name_hint="adaptive_avg_pool1d",
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.adaptive_avg_pool2d")
|
|
def _nn_adaptive_avg_pool2d(bb: BlockBuilder, call: Call) -> Expr:
|
|
if call.attrs.out_layout != call.attrs.layout:
|
|
logging.info(
|
|
"TOPI adaptive_avg_pool2d does not support different input-output "
|
|
"layouts, and thus cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
|
|
def te_adaptive_avg_pool2d(data, output_size, layout_str):
|
|
if output_size is None:
|
|
layout = s_tir.slayout(layout_str)
|
|
idx_H = layout.index_of("H")
|
|
idx_W = layout.index_of("W")
|
|
assert idx_H != -1 and idx_W != -1
|
|
output_size = (data.shape[idx_H], data.shape[idx_W])
|
|
|
|
return topi.nn.adaptive_pool(data, output_size, "avg", layout_str)
|
|
|
|
return bb.call_te(
|
|
te_adaptive_avg_pool2d,
|
|
call.args[0],
|
|
call.attrs.output_size,
|
|
call.attrs.layout,
|
|
primfunc_name_hint="adaptive_avg_pool2d",
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.adaptive_avg_pool3d")
|
|
def _nn_adaptive_avg_pool3d(bb: BlockBuilder, call: Call) -> Expr:
|
|
if call.attrs.out_layout != call.attrs.layout:
|
|
logging.info(
|
|
"TOPI adaptive_avg_pool3d does not support different input-output "
|
|
"layouts, and thus cannot be legalized by TOPI"
|
|
)
|
|
return call
|
|
|
|
def te_adaptive_avg_pool3d(data, output_size, layout_str):
|
|
if output_size is None:
|
|
layout = s_tir.slayout(layout_str)
|
|
idx_D = layout.index_of("D")
|
|
idx_H = layout.index_of("H")
|
|
idx_W = layout.index_of("W")
|
|
assert idx_D != -1 and idx_H != -1 and idx_W != -1
|
|
output_size = (data.shape[idx_D], data.shape[idx_H], data.shape[idx_W])
|
|
|
|
return topi.nn.adaptive_pool3d(data, output_size, "avg", layout_str)
|
|
|
|
return bb.call_te(
|
|
te_adaptive_avg_pool3d,
|
|
call.args[0],
|
|
call.attrs.output_size,
|
|
call.attrs.layout,
|
|
primfunc_name_hint="adaptive_avg_pool3d",
|
|
)
|
|
|
|
|
|
register_legalize("relax.nn.relu", _call_topi_without_attr(topi.nn.relu))
|
|
|
|
|
|
@register_legalize("relax.nn.leakyrelu")
|
|
def _nn_leakyrelu(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(topi.nn.leaky_relu, call.args[0], call.attrs.alpha)
|
|
|
|
|
|
@register_legalize("relax.nn.prelu")
|
|
def _nn_prelu(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(topi.nn.prelu, call.args[0], call.args[1], call.attrs.axis)
|
|
|
|
|
|
@register_legalize("relax.nn.gelu")
|
|
def _nn_gelu(bb: BlockBuilder, call: Call) -> Expr:
|
|
def te_gelu(x: te.Tensor):
|
|
dtype = x.dtype
|
|
erf_inp = x * tirx.const(0.5**0.5, dtype)
|
|
|
|
if dtype == "float16":
|
|
erf = topi.math.cast(topi.erf(topi.math.cast(erf_inp, "float32")), "float16")
|
|
else:
|
|
erf = topi.erf(erf_inp)
|
|
|
|
return x * (tirx.const(0.5, dtype) + erf * tirx.const(0.5, dtype))
|
|
|
|
return bb.call_te(te_gelu, call.args[0], primfunc_name_hint="gelu")
|
|
|
|
|
|
@register_legalize("relax.nn.gelu_tanh")
|
|
def _nn_gelu_tanh(bb: BlockBuilder, call: Call) -> Expr:
|
|
def te_gelu_tanh(x: te.Tensor):
|
|
dtype = x.dtype
|
|
return (
|
|
tirx.const(0.5, dtype)
|
|
* x
|
|
* (
|
|
tirx.const(1.0, dtype)
|
|
+ topi.tanh(
|
|
tirx.const(math.sqrt(2.0 / math.pi), dtype)
|
|
* x
|
|
* (1 + tirx.const(0.044715, dtype) * x * x)
|
|
)
|
|
)
|
|
)
|
|
|
|
return bb.call_te(te_gelu_tanh, call.args[0], primfunc_name_hint="gelu_tanh")
|
|
|
|
|
|
@register_legalize("relax.nn.selu")
|
|
def _nn_selu(bb: BlockBuilder, call: Call) -> Expr:
|
|
def te_selu(x: te.Tensor):
|
|
dtype = x.dtype
|
|
alpha = tirx.const(1.6732632423543772848170429916717, dtype)
|
|
scale = tirx.const(1.0507009873554804934193349852946, dtype)
|
|
|
|
# Compute SELU
|
|
# SELU(x) = scale*(max(0,x)+min(0,a*(exp(x)-1)))
|
|
positive_part = topi.maximum(x, tirx.const(0, dtype))
|
|
negative_part = topi.minimum(
|
|
tirx.const(0, dtype), alpha * (topi.exp(x) - tirx.const(1, dtype))
|
|
)
|
|
return scale * (positive_part + negative_part)
|
|
|
|
return bb.call_te(te_selu, call.args[0], primfunc_name_hint="selu")
|
|
|
|
|
|
@register_legalize("relax.nn.silu")
|
|
def _nn_silu(bb: BlockBuilder, call: Call) -> Expr:
|
|
def te_silu(x: te.Tensor):
|
|
return topi.multiply(x, topi.sigmoid(x))
|
|
|
|
return bb.call_te(te_silu, call.args[0], primfunc_name_hint="silu")
|
|
|
|
|
|
@register_legalize("relax.nn.softplus")
|
|
def _nn_softplus(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(
|
|
topi.nn.softplus,
|
|
call.args[0],
|
|
call.attrs.beta,
|
|
call.attrs.threshold,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.softmax")
|
|
def _nn_softmax(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(topi.nn.softmax, call.args[0], call.attrs.axis)
|
|
|
|
|
|
@register_legalize("relax.nn.log_softmax")
|
|
def _nn_log_softmax(bb: BlockBuilder, call: Call):
|
|
return bb.call_te(topi.nn.log_softmax, call.args[0], call.attrs.axis)
|
|
|
|
|
|
@register_legalize("relax.nn.cross_entropy_with_logits")
|
|
def _nn_cross_entropy_with_logits(bb: BlockBuilder, call: Call):
|
|
def te_cross_entropy_with_logits(x, y):
|
|
if len(x.shape) > 1:
|
|
return -topi.sum(x * y) / x.shape[0]
|
|
return -topi.sum(x * y)
|
|
|
|
return bb.call_te(
|
|
te_cross_entropy_with_logits,
|
|
call.args[0],
|
|
call.args[1],
|
|
primfunc_name_hint="cross_entropy_with_logits",
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.batch_norm")
|
|
def _nn_batch_norm(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(
|
|
topi.nn.batch_norm,
|
|
data=call.args[0],
|
|
gamma=call.args[1],
|
|
beta=call.args[2],
|
|
moving_mean=call.args[3],
|
|
moving_var=call.args[4],
|
|
axis=call.attrs.axis,
|
|
epsilon=call.attrs.epsilon,
|
|
center=call.attrs.center,
|
|
scale=call.attrs.scale,
|
|
training=call.attrs.training,
|
|
momentum=call.attrs.momentum,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.layer_norm")
|
|
def _nn_layer_norm(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(
|
|
topi.nn.layer_norm,
|
|
call.args[0],
|
|
call.args[1],
|
|
call.args[2],
|
|
axis=call.attrs.axes,
|
|
epsilon=call.attrs.epsilon,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.group_norm")
|
|
def _nn_group_norm(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(
|
|
topi.nn.group_norm,
|
|
call.args[0],
|
|
call.args[1],
|
|
call.args[2],
|
|
call.attrs.num_groups,
|
|
call.attrs.channel_axis,
|
|
call.attrs.axes,
|
|
call.attrs.epsilon,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.instance_norm")
|
|
def _nn_instance_norm(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(
|
|
topi.nn.instance_norm,
|
|
data=call.args[0],
|
|
gamma=call.args[1],
|
|
beta=call.args[2],
|
|
channel_axis=call.attrs.channel_axis,
|
|
axis=call.attrs.axes,
|
|
epsilon=call.attrs.epsilon,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.rms_norm")
|
|
def _nn_rms_norm(bb: BlockBuilder, call: Call) -> Expr:
|
|
return bb.call_te(
|
|
topi.nn.rms_norm,
|
|
call.args[0],
|
|
call.args[1],
|
|
axis=call.attrs.axes,
|
|
epsilon=call.attrs.epsilon,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.dropout")
|
|
def _nn_dropout(bb: BlockBuilder, call: Call) -> Expr:
|
|
# Dropout is a no-op at inference: pass the input through and return an all-ones mask.
|
|
return bb.call_te(
|
|
lambda x: [topi.identity(x), topi.full_like(x, 1.0)],
|
|
call.args[0],
|
|
primfunc_name_hint="dropout",
|
|
)
|
|
|
|
|
|
def _te_attention(
|
|
q: te.Tensor,
|
|
k: te.Tensor,
|
|
v: te.Tensor,
|
|
bias: te.Tensor,
|
|
scale: tirx.FloatImm,
|
|
causal_mask: str | None,
|
|
) -> te.Tensor:
|
|
batch_size, seq_len, num_head, head_dim = q.shape
|
|
_, seq_len_kv, _, head_dim_v = v.shape
|
|
q = topi.transpose(q, [0, 2, 1, 3])
|
|
k = topi.transpose(k, [0, 2, 1, 3])
|
|
v = topi.transpose(v, [0, 2, 1, 3])
|
|
bs = batch_size * num_head
|
|
q = topi.reshape(q, [bs, seq_len, head_dim])
|
|
k = topi.reshape(k, [bs, seq_len_kv, head_dim])
|
|
v = topi.reshape(v, [bs, seq_len_kv, head_dim_v])
|
|
p = topi.nn.batch_matmul(q, k, oshape=[bs, seq_len, seq_len_kv])
|
|
if scale is not None:
|
|
p = topi.multiply(p, scale)
|
|
else:
|
|
p = topi.divide(p, tirx.sqrt(tirx.Cast(p.dtype, head_dim)))
|
|
if bias is not None:
|
|
p = topi.reshape(p, [batch_size, num_head, seq_len, seq_len_kv])
|
|
p = topi.add(p, bias)
|
|
p = topi.reshape(p, [bs, seq_len, seq_len_kv])
|
|
if causal_mask is None:
|
|
s = topi.nn.softmax(p)
|
|
else:
|
|
if causal_mask == "TopLeft":
|
|
offset = tirx.IntImm("int32", 0)
|
|
elif causal_mask == "BottomRight":
|
|
offset = tirx.abs(seq_len - seq_len_kv).astype("int32")
|
|
else:
|
|
raise NotImplementedError()
|
|
p_masked = topi.trilu(p, k=offset, upper=False)
|
|
p_masked_exp = topi.trilu(
|
|
topi.exp(p_masked - topi.max(p_masked, axis=-1, keepdims=True)), k=offset, upper=False
|
|
)
|
|
p_masked_sum = topi.sum(p_masked_exp, axis=-1, keepdims=True)
|
|
s = topi.divide(p_masked_exp, p_masked_sum)
|
|
o = topi.nn.batch_matmul(s, v, transpose_b=False, oshape=[bs, seq_len, head_dim_v])
|
|
o = topi.reshape(o, [batch_size, num_head, seq_len, head_dim_v])
|
|
return topi.transpose(o, [0, 2, 1, 3])
|
|
|
|
|
|
@register_legalize("relax.nn.attention")
|
|
def _nn_attention(bb: BlockBuilder, call: Call) -> Expr:
|
|
assert call.attrs.window_size is None, (
|
|
"Legalization for sliding-window attention is not supported yet."
|
|
)
|
|
return bb.call_te(
|
|
_te_attention,
|
|
call.args[0],
|
|
call.args[1],
|
|
call.args[2],
|
|
None,
|
|
call.attrs.scale,
|
|
call.attrs.causal_mask,
|
|
primfunc_name_hint="attention",
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.attention_bias")
|
|
def _nn_attention_bias(bb: BlockBuilder, call: Call) -> Expr:
|
|
assert call.attrs.window_size is None, (
|
|
"Legalization for sliding-window attention is not supported yet."
|
|
)
|
|
return bb.call_te(
|
|
_te_attention,
|
|
call.args[0],
|
|
call.args[1],
|
|
call.args[2],
|
|
call.args[3],
|
|
call.attrs.scale,
|
|
call.attrs.causal_mask,
|
|
primfunc_name_hint="attention_bias",
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.attention_var_len")
|
|
def _nn_attention_var_len(bb: BlockBuilder, call: Call) -> Expr:
|
|
raise RuntimeError("Legalization of attention_var_len op is not supported yet.")
|
|
|
|
|
|
@register_legalize("relax.nn.nll_loss")
|
|
def _nn_nll_loss(bb: BlockBuilder, call: Call) -> Expr:
|
|
def nll_loss_without_weight(predictions, targets, reduction, ignore_index):
|
|
weight = topi.full(
|
|
(predictions.shape[1] if len(predictions.shape) > 1 else predictions.shape[0],),
|
|
predictions.dtype,
|
|
1.0,
|
|
)
|
|
return topi.nn.nll_loss(predictions, targets, weight, reduction, ignore_index)
|
|
|
|
if len(call.args) == 2:
|
|
return bb.call_te(
|
|
nll_loss_without_weight,
|
|
call.args[0],
|
|
call.args[1],
|
|
reduction=call.attrs.reduction,
|
|
ignore_index=call.attrs.ignore_index,
|
|
)
|
|
|
|
return bb.call_te(
|
|
topi.nn.nll_loss,
|
|
call.args[0],
|
|
call.args[1],
|
|
call.args[2],
|
|
reduction=call.attrs.reduction,
|
|
ignore_index=call.attrs.ignore_index,
|
|
)
|
|
|
|
|
|
@register_legalize("relax.nn.batch_flatten")
|
|
def _nn_batch_flatten(bb: BlockBuilder, call: Call) -> Expr:
|
|
if call.ty.shape is None:
|
|
return call
|
|
return bb.call_te(topi.reshape, call.args[0], call.ty.shape.values)
|