chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,57 @@
|
||||
# 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=wildcard-import
|
||||
"""Neural network operators"""
|
||||
|
||||
from .conv1d import *
|
||||
from .conv2d import *
|
||||
from .conv3d import *
|
||||
from .correlation import *
|
||||
from .deformable_conv2d import *
|
||||
from .depthwise_conv2d import *
|
||||
from .elemwise import *
|
||||
from .dilate import *
|
||||
from .flatten import *
|
||||
from .dense import *
|
||||
from .mapping import *
|
||||
from .pooling import *
|
||||
from .softmax import *
|
||||
from .conv3d_transpose import *
|
||||
from .conv2d_transpose import *
|
||||
from .conv1d_transpose import *
|
||||
from .bnn import *
|
||||
from .qnn import *
|
||||
from .upsampling import *
|
||||
from .instance_norm import instance_norm
|
||||
from .layer_norm import layer_norm
|
||||
from .group_norm import group_norm
|
||||
from .rms_norm import rms_norm
|
||||
from .local_response_norm import *
|
||||
from .bitserial_conv2d import *
|
||||
from .bitserial_dense import *
|
||||
from .batch_matmul import *
|
||||
from .batch_norm import *
|
||||
from .pad import *
|
||||
from .fifo_buffer import *
|
||||
from .depth_to_space import *
|
||||
from .space_to_depth import *
|
||||
from .space_to_batch_nd import *
|
||||
from .batch_to_space_nd import *
|
||||
from .loss import *
|
||||
from .lstm import *
|
||||
@@ -0,0 +1,152 @@
|
||||
# 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: E731
|
||||
"""Batch matrix multiplication"""
|
||||
|
||||
# pylint: disable=invalid-name
|
||||
import logging
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
|
||||
from ..utils import get_const_tuple
|
||||
|
||||
logger = logging.getLogger("topi")
|
||||
|
||||
|
||||
def batch_matmul(
|
||||
tensor_a,
|
||||
tensor_b,
|
||||
oshape=None,
|
||||
out_dtype=None,
|
||||
transpose_a=False,
|
||||
transpose_b=True,
|
||||
auto_scheduler_rewritten_layout="",
|
||||
meta_schedule_original_shape=None,
|
||||
):
|
||||
"""Compute batch matrix multiplication of `tensor_a` and `tensor_b`.
|
||||
|
||||
Both `tensor_a` and `tensor_b` can be transposed. For legacy reason, we use NT format
|
||||
(transpose_a=False, transpose_b=True) by default.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tensor_a : tvm.te.Tensor
|
||||
3-D with shape [batch, M, K] or [batch, K, M].
|
||||
|
||||
tensor_b : tvm.te.Tensor
|
||||
3-D with shape [batch, K, N] or [batch, N, K].
|
||||
|
||||
oshape : List[Optional]
|
||||
Explicit intended output shape of the computation. Can be useful in cases
|
||||
with dynamic input shapes.
|
||||
|
||||
out_dtype : Optional[str]
|
||||
Specifies the output data type for mixed precision batch matmul.
|
||||
|
||||
transpose_a : Optional[bool] = False
|
||||
Whether the first tensor is in transposed format.
|
||||
|
||||
transpose_b : Optional[bool] = True
|
||||
Whether the second tensor is in transposed format.
|
||||
|
||||
auto_scheduler_rewritten_layout: Optional[str] = ""
|
||||
The layout after auto-scheduler's layout rewrite pass.
|
||||
|
||||
meta_schedule_original_shape: Optional[List[Expr]] = None
|
||||
The original shape of the tensor
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
3-D with shape [batch, M, N]
|
||||
"""
|
||||
assert len(tensor_a.shape) == 3, "tensor_a only support 3-dim"
|
||||
if transpose_a:
|
||||
XB, XK, XI = get_const_tuple(tensor_a.shape)
|
||||
else:
|
||||
XB, XI, XK = get_const_tuple(tensor_a.shape)
|
||||
if auto_scheduler_rewritten_layout:
|
||||
raise RuntimeError("LEGACY-FLOW triggered, to be removed")
|
||||
if meta_schedule_original_shape:
|
||||
raise RuntimeError("LEGACY-FLOW triggered, to be removed")
|
||||
|
||||
assert len(tensor_b.shape) == 3, "tensor_b only support 3-dim"
|
||||
if transpose_b:
|
||||
YB, YJ, YK = get_const_tuple(tensor_b.shape)
|
||||
else:
|
||||
YB, YK, YJ = get_const_tuple(tensor_b.shape)
|
||||
|
||||
assert XK == YK or isinstance(YK, tvm.tirx.expr.Var), "shapes of x and y are inconsistent"
|
||||
k = te.reduce_axis((0, XK), name="k")
|
||||
if oshape is None:
|
||||
assert XB == YB or XB == 1 or YB == 1, "batch dimension doesn't match"
|
||||
batch = (
|
||||
tvm.tirx.expr.Var("batch", "int32")
|
||||
if isinstance(XB, tvm.tirx.expr.Var) or isinstance(YB, tvm.tirx.expr.Var)
|
||||
else te.max(XB, YB)
|
||||
)
|
||||
oshape = (batch, XI, YJ)
|
||||
if out_dtype is None:
|
||||
out_dtype = tensor_a.dtype
|
||||
if tensor_a.dtype != tensor_b.dtype:
|
||||
logger.warning(
|
||||
"tensor_a has different data type with tensor_b: %s, %s",
|
||||
tensor_a.dtype,
|
||||
tensor_b.dtype,
|
||||
)
|
||||
|
||||
if (transpose_a, transpose_b) == (True, True):
|
||||
compute_lambda = lambda b, i, j: te.sum(
|
||||
tensor_a[b if XB != 1 else 0, k, i].astype(out_dtype)
|
||||
* tensor_b[b if YB != 1 else 0, j, k].astype(out_dtype),
|
||||
axis=k,
|
||||
)
|
||||
compute_name = "T_batch_matmul_TT"
|
||||
elif (transpose_a, transpose_b) == (True, False):
|
||||
compute_lambda = lambda b, i, j: te.sum(
|
||||
tensor_a[b if XB != 1 else 0, k, i].astype(out_dtype)
|
||||
* tensor_b[b if YB != 1 else 0, k, j].astype(out_dtype),
|
||||
axis=k,
|
||||
)
|
||||
compute_name = "T_batch_matmul_TN"
|
||||
elif (transpose_a, transpose_b) == (False, True):
|
||||
compute_lambda = lambda b, i, j: te.sum(
|
||||
tensor_a[b if XB != 1 else 0, i, k].astype(out_dtype)
|
||||
* tensor_b[b if YB != 1 else 0, j, k].astype(out_dtype),
|
||||
axis=k,
|
||||
)
|
||||
compute_name = "T_batch_matmul_NT"
|
||||
else: # (transpose_a, transpose_b) == (False, False):
|
||||
compute_lambda = lambda b, i, j: te.sum(
|
||||
tensor_a[b if XB != 1 else 0, i, k].astype(out_dtype)
|
||||
* tensor_b[b if YB != 1 else 0, k, j].astype(out_dtype),
|
||||
axis=k,
|
||||
)
|
||||
compute_name = "T_batch_matmul_NN"
|
||||
|
||||
output = te.compute(
|
||||
oshape,
|
||||
compute_lambda,
|
||||
name=compute_name,
|
||||
tag="batch_matmul",
|
||||
attrs={"layout_free_placeholders": [tensor_b]},
|
||||
)
|
||||
if auto_scheduler_rewritten_layout:
|
||||
raise RuntimeError("LEGACY-FLOW triggered, to be removed")
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,146 @@
|
||||
# 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.
|
||||
"""Batch normalization."""
|
||||
|
||||
from functools import reduce
|
||||
|
||||
from tvm import te, topi
|
||||
|
||||
|
||||
def batch_norm(
|
||||
data: te.Tensor,
|
||||
gamma: te.Tensor,
|
||||
beta: te.Tensor,
|
||||
moving_mean: te.Tensor,
|
||||
moving_var: te.Tensor,
|
||||
axis: int | None = None,
|
||||
epsilon: float | None = None,
|
||||
center: bool | None = None,
|
||||
scale: bool | None = None,
|
||||
training: bool | None = None,
|
||||
momentum: float | None = None,
|
||||
) -> list[te.Tensor]:
|
||||
"""Batch normalization layer (Ioffe and Szegedy, 2014).
|
||||
|
||||
Normalizes the input at each batch, i.e. applies a transformation
|
||||
that maintains the mean activation close to 0 and the activation
|
||||
standard deviation close to 1.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
Input to be batch-normalized.
|
||||
|
||||
gamma : tvm.te.Tensor
|
||||
Scale factor to be applied to the normalized tensor.
|
||||
|
||||
beta : tvm.te.Tensor
|
||||
Offset to be applied to the normalized tensor.
|
||||
|
||||
moving_mean : tvm.te.Tensor
|
||||
Running mean of input.
|
||||
|
||||
moving_var : tvm.te.Tensor
|
||||
Running variance of input.
|
||||
|
||||
axis : int, optional, default=1
|
||||
Specify along which shape axis the normalization should occur.
|
||||
|
||||
epsilon : float, optional, default=1e-5
|
||||
Small float added to variance to avoid dividing by zero.
|
||||
|
||||
center : bool, optional, default=True
|
||||
If True, add offset of beta to normalized tensor, If False,
|
||||
beta is ignored.
|
||||
|
||||
scale : bool, optional, defualt=True
|
||||
If True, scale normalized tensor by gamma. If False, gamma
|
||||
is ignored.
|
||||
|
||||
training : bool, optional, defualt=False
|
||||
Indicating whether it is in training mode. If True, update
|
||||
moving_mean and moving_var.
|
||||
|
||||
momentum : float, optional, default=0.1
|
||||
The value used for the moving_mean and moving_var update.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : list of tvm.te.Tensor
|
||||
Normalized data with same shape as input
|
||||
|
||||
moving_mean : tvm.te.Tensor
|
||||
Running mean of input.
|
||||
|
||||
moving_var : tvm.te.Tensor
|
||||
Running variance of input.
|
||||
"""
|
||||
if axis is None:
|
||||
axis = 1
|
||||
|
||||
if epsilon is None:
|
||||
epsilon = 1e-5
|
||||
|
||||
if center is None:
|
||||
center = True
|
||||
|
||||
if scale is None:
|
||||
scale = True
|
||||
|
||||
if training is None:
|
||||
training = False
|
||||
|
||||
if momentum is None:
|
||||
momentum = 0.1
|
||||
|
||||
shape = [1] * len(data.shape)
|
||||
shape[axis] = data.shape[axis]
|
||||
data_mean = None
|
||||
data_var = None
|
||||
|
||||
if training:
|
||||
reduce_axes = list(range(len(data.shape)))
|
||||
reduce_axes.remove(axis)
|
||||
shape_prod = reduce(lambda x, y: x * y, [data.shape[ax] for ax in reduce_axes], 1)
|
||||
data_mean = topi.sum(data, axis=reduce_axes) / shape_prod
|
||||
data_mean_rs = topi.reshape(data_mean, shape)
|
||||
data_var = (
|
||||
topi.sum((data - data_mean_rs) * (data - data_mean_rs), axis=reduce_axes) / shape_prod
|
||||
)
|
||||
data_var_rs = topi.reshape(data_var, shape)
|
||||
out = (data - data_mean_rs) / topi.math.sqrt(data_var_rs + epsilon)
|
||||
else:
|
||||
moving_mean_rs = topi.reshape(moving_mean, shape)
|
||||
moving_var_rs = topi.reshape(moving_var, shape)
|
||||
out = (data - moving_mean_rs) / topi.math.sqrt(moving_var_rs + epsilon)
|
||||
|
||||
if scale:
|
||||
out = out * topi.reshape(gamma, shape)
|
||||
if center:
|
||||
out = out + topi.reshape(beta, shape)
|
||||
|
||||
if training:
|
||||
assert 0 <= momentum <= 1, "the valid momentum range is [0, 1]."
|
||||
return [
|
||||
out,
|
||||
(1 - momentum) * moving_mean + momentum * data_mean,
|
||||
(1 - momentum) * moving_var + momentum * data_var,
|
||||
]
|
||||
|
||||
# Moving mean and var aren't updated during test. To avoid
|
||||
# placeholder reuse, we multiply by 1 and return them.
|
||||
return [out, moving_mean * 1, moving_var * 1]
|
||||
@@ -0,0 +1,49 @@
|
||||
# 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
|
||||
"""TVM operator batch_to_space_nd compute."""
|
||||
|
||||
from . import cpp
|
||||
|
||||
|
||||
def batch_to_space_nd(data, block_shape, crop_begin_list, crop_end_list):
|
||||
"""Perform space to batch transformation on the data
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
N-D Tensor with shape [batch, spatial_shape, remaining_shapes],
|
||||
where spatial_shape has M dimensions.
|
||||
|
||||
block_shape : list of ints
|
||||
list of size [M] where M is number of spatial dims, specifies block
|
||||
size for each spatial dimension.
|
||||
|
||||
crop_begin_list : list of ints
|
||||
list of shape [M] where M is number of spatial dims, specifies
|
||||
begin crop size for each spatial dimension.
|
||||
|
||||
crop_end_list : list of ints
|
||||
list of shape [M] where M is number of spatial dims, specifies
|
||||
end crop size for each spatial dimension.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
"""
|
||||
|
||||
return cpp.nn.batch_to_space_nd(data, block_shape, crop_begin_list, crop_end_list)
|
||||
@@ -0,0 +1,276 @@
|
||||
# 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, too-many-locals, too-many-arguments
|
||||
# pylint: disable=unused-argument, redefined-builtin
|
||||
"""Bitserial Conv2D operators"""
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
|
||||
from ..utils import get_const_tuple
|
||||
from .bitserial_util import bitpack
|
||||
from .pad import pad
|
||||
from .utils import get_pad_tuple
|
||||
|
||||
|
||||
def bitserial_conv2d_nchw(
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
padding,
|
||||
activation_bits,
|
||||
weight_bits,
|
||||
pack_dtype="uint32",
|
||||
out_dtype="int16",
|
||||
unipolar=True,
|
||||
):
|
||||
"""Bitserial Conv2D operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
kernel : tvm.te.Tensor
|
||||
4-D with shape [num_filter, in_channel, filter_height, filter_width]
|
||||
|
||||
stride : int or a list/tuple of two ints
|
||||
stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : int or a list/tuple of two or four ints
|
||||
padding size, [pad_height, pad_width], [pad_top, pad_left, pad_down, pad_right]
|
||||
|
||||
activation_bits: int
|
||||
number of bits used for activations/input elements
|
||||
|
||||
weight_bits: int
|
||||
number of bits used for weight elements
|
||||
|
||||
out_dtype: str
|
||||
return type of convolution
|
||||
|
||||
pack_dtype: str
|
||||
bit packing type
|
||||
|
||||
unipolar: bool
|
||||
if binarization style is in unipolar 1/0 format, instead of bipolar -1/+1 format
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
assert isinstance(stride, int) or len(stride) == 2
|
||||
Input_q = bitpack(data, activation_bits, pack_axis=1, bit_axis=2, pack_type=pack_dtype)
|
||||
if len(kernel.shape) == 4:
|
||||
Filter_q = bitpack(kernel, weight_bits, pack_axis=1, bit_axis=4, pack_type=pack_dtype)
|
||||
else:
|
||||
Filter_q = kernel
|
||||
batch, in_channel, activation_bits, in_height, in_width = Input_q.shape
|
||||
num_filter, _, kernel_h, kernel_w, weight_bits = Filter_q.shape
|
||||
|
||||
if isinstance(padding, int) or (isinstance(padding, tuple | list) and len(padding) == 2):
|
||||
TPAD, LPAD, DPAD, RPAD = get_pad_tuple(padding, kernel)
|
||||
else:
|
||||
TPAD, LPAD, DPAD, RPAD = padding
|
||||
pad_before = [0, 0, 0, TPAD, LPAD]
|
||||
pad_after = [0, 0, 0, DPAD, RPAD]
|
||||
|
||||
PadInput_q = pad(Input_q, pad_before, pad_after, name="pad_temp")
|
||||
# compute the output shape
|
||||
if isinstance(stride, int):
|
||||
stride_h = stride_w = stride
|
||||
else:
|
||||
stride_h, stride_w = stride
|
||||
out_channel = num_filter
|
||||
out_height = (in_height - kernel_h + TPAD + DPAD) // stride_h + 1
|
||||
out_width = (in_width - kernel_w + LPAD + RPAD) // stride_w + 1
|
||||
|
||||
rc = te.reduce_axis((0, in_channel), name="rc")
|
||||
ry = te.reduce_axis((0, kernel_h), name="ry")
|
||||
rx = te.reduce_axis((0, kernel_w), name="rx")
|
||||
b1 = te.reduce_axis((0, activation_bits), name="b1")
|
||||
b2 = te.reduce_axis((0, weight_bits), name="b2")
|
||||
|
||||
if unipolar:
|
||||
|
||||
def _conv(nn, ff, yy, xx):
|
||||
b1b2 = (b1 + b2).astype(out_dtype)
|
||||
return te.sum(
|
||||
(
|
||||
(
|
||||
tvm.tirx.popcount(
|
||||
PadInput_q[nn, rc, b1, yy * stride_h + ry, xx * stride_w + rx]
|
||||
& Filter_q[ff, rc, ry, rx, b2]
|
||||
)
|
||||
- tvm.tirx.popcount(
|
||||
PadInput_q[nn, rc, b1, yy * stride_h + ry, xx * stride_w + rx]
|
||||
& ~Filter_q[ff, rc, ry, rx, b2]
|
||||
)
|
||||
)
|
||||
<< (b1b2)
|
||||
).astype(out_dtype),
|
||||
axis=[rc, ry, rx, b2, b1],
|
||||
).astype(out_dtype)
|
||||
|
||||
else:
|
||||
|
||||
def _conv(nn, ff, yy, xx):
|
||||
b1b2 = (b1 + b2).astype(out_dtype)
|
||||
return te.sum(
|
||||
(
|
||||
tvm.tirx.popcount(
|
||||
PadInput_q[nn, rc, b1, yy * stride_h + ry, xx * stride_w + rx]
|
||||
& Filter_q[ff, rc, ry, rx, b2]
|
||||
)
|
||||
<< (b1b2)
|
||||
).astype(out_dtype),
|
||||
axis=[rc, ry, rx, b2, b1],
|
||||
).astype(out_dtype)
|
||||
|
||||
return te.compute(
|
||||
(batch, out_channel, out_height, out_width),
|
||||
_conv,
|
||||
name="Conv2dOutput",
|
||||
tag="bitserial_conv2d_nchw",
|
||||
)
|
||||
|
||||
|
||||
def bitserial_conv2d_nhwc(
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
padding,
|
||||
activation_bits,
|
||||
weight_bits,
|
||||
pack_dtype="uint32",
|
||||
out_dtype="int16",
|
||||
unipolar=True,
|
||||
):
|
||||
"""Bitserial Conv2D operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
4-D with shape [batch, in_height, in_width, in_channel]
|
||||
|
||||
kernel : tvm.te.Tensor
|
||||
4-D with shape [filter_height, filter_width, in_channel, num_filter]
|
||||
|
||||
stride : int or a list/tuple of two ints
|
||||
stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : int or a list/tuple of two or four ints
|
||||
padding size, [pad_height, pad_width], [pad_top, pad_left, pad_down, pad_right]
|
||||
|
||||
activation_bits: int
|
||||
number of bits used for activations/input elements
|
||||
|
||||
weight_bits: int
|
||||
number of bits used for weight elements
|
||||
|
||||
out_dtype: str
|
||||
return type of convolution
|
||||
|
||||
pack_dtype: str
|
||||
bit packing type
|
||||
|
||||
unipolar: bool
|
||||
if binarization style is in unipolar 1/0 format, instead of bipolar -1/+1 format
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
4-D with shape [batch, out_height, out_width, out_channel]
|
||||
"""
|
||||
assert isinstance(stride, int) or len(stride) == 2
|
||||
Input_q = bitpack(data, activation_bits, pack_axis=3, bit_axis=4, pack_type=pack_dtype)
|
||||
if len(kernel.shape) == 4:
|
||||
Filter_q = bitpack(kernel, weight_bits, pack_axis=2, bit_axis=4, pack_type=pack_dtype)
|
||||
kernel_h, kernel_w, _, num_filter, _ = get_const_tuple(Filter_q.shape)
|
||||
else:
|
||||
Filter_q = kernel
|
||||
kernel_h, kernel_w, _, _, num_filter = get_const_tuple(Filter_q.shape)
|
||||
batch, in_height, in_width, in_channel_q, _ = get_const_tuple(Input_q.shape)
|
||||
|
||||
if isinstance(padding, int) or (isinstance(padding, tuple | list) and len(padding) == 2):
|
||||
TPAD, LPAD, DPAD, RPAD = get_pad_tuple(padding, kernel)
|
||||
else:
|
||||
TPAD, LPAD, DPAD, RPAD = padding
|
||||
pad_before = [0, TPAD, LPAD, 0, 0]
|
||||
pad_after = [0, DPAD, RPAD, 0, 0]
|
||||
|
||||
# compute the output shape
|
||||
if isinstance(stride, int):
|
||||
stride_h = stride_w = stride
|
||||
else:
|
||||
stride_h, stride_w = stride
|
||||
out_channel = num_filter
|
||||
out_height = (in_height - kernel_h + TPAD + DPAD) // stride_h + 1
|
||||
out_width = (in_width - kernel_w + LPAD + RPAD) // stride_w + 1
|
||||
PadInput_q = pad(Input_q, pad_before, pad_after, name="PaddedInput")
|
||||
|
||||
rc = te.reduce_axis((0, in_channel_q), name="rc")
|
||||
ry = te.reduce_axis((0, kernel_h), name="ry")
|
||||
rx = te.reduce_axis((0, kernel_w), name="rx")
|
||||
b1 = te.reduce_axis((0, activation_bits), name="b1")
|
||||
b2 = te.reduce_axis((0, weight_bits), name="b2")
|
||||
|
||||
if unipolar:
|
||||
|
||||
def _conv(nn, yy, xx, ff):
|
||||
b1b2 = (b1 + b2).astype(out_dtype)
|
||||
return te.sum(
|
||||
(
|
||||
(
|
||||
tvm.tirx.popcount(
|
||||
PadInput_q[nn, yy * stride_h + ry, xx * stride_w + rx, rc, b1]
|
||||
& Filter_q[ry, rx, rc, ff, b2]
|
||||
)
|
||||
- tvm.tirx.popcount(
|
||||
PadInput_q[nn, yy * stride_h + ry, xx * stride_w + rx, rc, b1]
|
||||
& ~Filter_q[ry, rx, rc, ff, b2]
|
||||
)
|
||||
)
|
||||
<< b1b2
|
||||
).astype(out_dtype),
|
||||
axis=[rc, ry, rx, b2, b1],
|
||||
)
|
||||
|
||||
else:
|
||||
|
||||
def _conv(nn, yy, xx, ff):
|
||||
b1b2 = (b1 + b2).astype(out_dtype)
|
||||
return te.sum(
|
||||
(
|
||||
tvm.tirx.popcount(
|
||||
PadInput_q[nn, yy * stride_h + ry, xx * stride_w + rx, rc, b1]
|
||||
& Filter_q[ry, rx, rc, ff, b2]
|
||||
)
|
||||
<< b1b2
|
||||
).astype(out_dtype),
|
||||
axis=[rc, ry, rx, b2, b1],
|
||||
)
|
||||
|
||||
conv = te.compute(
|
||||
(batch, out_height, out_width, out_channel),
|
||||
_conv,
|
||||
name="Conv2dOutput",
|
||||
tag="bitserial_conv2d_nhwc",
|
||||
)
|
||||
|
||||
return conv
|
||||
@@ -0,0 +1,82 @@
|
||||
# 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, too-many-locals, too-many-arguments
|
||||
"""Bitserial Dense operator."""
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
from tvm.topi.utils import get_const_tuple
|
||||
|
||||
from .bitserial_util import bitpack
|
||||
|
||||
|
||||
def bitserial_dense(
|
||||
data, weight, data_bits, weight_bits, pack_dtype="uint32", out_dtype="int16", unipolar=True
|
||||
):
|
||||
"""The default implementation of bitserial dense in topi.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
2-D with shape [batch, in_dim]
|
||||
weight : tvm.te.Tensor
|
||||
2-D with shape [out_dim, in_dim] or
|
||||
3-D with shape [out_dim, weight_bits, in_dim]
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
2-D with shape [batch, out_dim]
|
||||
"""
|
||||
data_packed = bitpack(data, data_bits, pack_axis=1, bit_axis=1, pack_type=pack_dtype)
|
||||
if len(weight.shape) == 2:
|
||||
weight_packed = bitpack(weight, weight_bits, pack_axis=1, bit_axis=1, pack_type=pack_dtype)
|
||||
else:
|
||||
weight_packed = weight
|
||||
Y, DB, K = get_const_tuple(data_packed.shape)
|
||||
X, WB, _ = get_const_tuple(weight_packed.shape)
|
||||
|
||||
oshape = (Y, X)
|
||||
k = te.reduce_axis((0, K), name="k")
|
||||
db = te.reduce_axis((0, DB), name="db")
|
||||
wb = te.reduce_axis((0, WB), name="wb")
|
||||
|
||||
matmul_unipolar = te.compute(
|
||||
oshape,
|
||||
lambda i, j: te.sum(
|
||||
(
|
||||
tvm.tirx.popcount(weight_packed[j, wb, k] & data_packed[i, db, k])
|
||||
- tvm.tirx.popcount(~weight_packed[j, wb, k] & data_packed[i, db, k])
|
||||
).astype(out_dtype)
|
||||
<< (db + wb).astype(out_dtype),
|
||||
axis=[wb, db, k],
|
||||
),
|
||||
tag="bitserial_dense_unipolar",
|
||||
)
|
||||
|
||||
matmul = te.compute(
|
||||
oshape,
|
||||
lambda i, j: te.sum(
|
||||
tvm.tirx.popcount(weight_packed[j, wb, k] & data_packed[i, db, k]).astype(out_dtype)
|
||||
<< (db + wb).astype(out_dtype),
|
||||
axis=[wb, db, k],
|
||||
),
|
||||
tag="bitserial_dense",
|
||||
)
|
||||
|
||||
if unipolar:
|
||||
return matmul_unipolar
|
||||
return matmul
|
||||
@@ -0,0 +1,110 @@
|
||||
# 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, too-many-locals, too-many-arguments
|
||||
"""Utility functions for bitserial operators"""
|
||||
|
||||
import numpy as np
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
from tvm.topi.transform import concatenate
|
||||
|
||||
from ..utils import get_const_int
|
||||
|
||||
|
||||
def bitpack(data, bits, pack_axis, bit_axis, pack_type, name="QuantizeInput"):
|
||||
"""Packs data into format necessary for bitserial computation
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
The input tvm tensor
|
||||
bits : int
|
||||
Number of bits to use for packing
|
||||
pack_axis : int
|
||||
index of the axis to pack in data
|
||||
bit_axis : int
|
||||
index of axis to place bit axis in resulting packed data
|
||||
pack_type : str
|
||||
Data type for packing, must be one of: ['uint8', 'uint16', 'uint32', 'uint64']
|
||||
name : Optional[str] = "QuantizeInput"
|
||||
Name for the operation
|
||||
"""
|
||||
ishape = data.shape
|
||||
n = len(ishape)
|
||||
if pack_type == "uint8":
|
||||
data_width = 8
|
||||
elif pack_type == "uint16":
|
||||
data_width = 16
|
||||
elif pack_type == "uint32":
|
||||
data_width = 32
|
||||
elif pack_type == "uint64":
|
||||
data_width = 64
|
||||
|
||||
# Data must be in multiples of the data_width
|
||||
assert get_const_int(ishape[pack_axis]) % data_width == 0, "Not a multiple of word size"
|
||||
|
||||
shape_vec = list(ishape)
|
||||
shape_vec[pack_axis] = shape_vec[pack_axis] // data_width
|
||||
shape_vec.insert(bit_axis, 1)
|
||||
bitserial_oshape = tuple(shape_vec)
|
||||
masks = np.array([0x1, 0x2, 0x4, 0x8, 0x10, 0x20, 0x40, 0x80])
|
||||
|
||||
# pack axis shifts if bit axis comes before
|
||||
if bit_axis <= pack_axis:
|
||||
pack_axis += 1
|
||||
|
||||
def _bitpack(*indices):
|
||||
packed_data = [tvm.tirx.const(0, pack_type)] * bits
|
||||
for k in range(data_width):
|
||||
# Translate indices for packed data back to original
|
||||
idx = [0] * n
|
||||
j = 0
|
||||
for i in range(n + 1):
|
||||
if i == bit_axis:
|
||||
continue
|
||||
if i == pack_axis:
|
||||
idx[j] = indices[i] * data_width + k
|
||||
else:
|
||||
idx[j] = indices[i]
|
||||
j += 1
|
||||
|
||||
element = data(*idx)
|
||||
for b in range(bits):
|
||||
extracted_bit = ((element & tvm.tirx.const(masks[b], "int32")) >> b).astype(
|
||||
pack_type
|
||||
)
|
||||
packed_data[b] = packed_data[b] | extracted_bit
|
||||
if k < data_width - 1:
|
||||
packed_data[b] = packed_data[b] << 1
|
||||
|
||||
if k == data_width - 1:
|
||||
return tuple(packed_data)
|
||||
return tuple(packed_data)
|
||||
|
||||
output_tuple = te.compute(bitserial_oshape, _bitpack, name=name, tag="bitpack")
|
||||
|
||||
if bits > 1:
|
||||
return concatenate(output_tuple, axis=bit_axis)
|
||||
return output_tuple
|
||||
|
||||
|
||||
def binary_op_multiplier(pack_dtype):
|
||||
""" "Returns number of bits packed into
|
||||
pack_dtype: string
|
||||
pack type for the operator (must be a uint)"""
|
||||
return int(pack_dtype[4:])
|
||||
@@ -0,0 +1,99 @@
|
||||
# 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.
|
||||
"""Binary Neural Network (BNN) Operators"""
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
|
||||
from .. import tag
|
||||
from ..utils import get_const_int, simplify
|
||||
|
||||
|
||||
def binarize_pack(data, axis=None, name="PackedInput"):
|
||||
"""Binarization and bit-packing along a certain axis.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
n-D input, can be any layout.
|
||||
|
||||
axis : None or int
|
||||
The axis along which to do binarization and bit-packing,
|
||||
default is the last axis.
|
||||
|
||||
name : str, optional
|
||||
The name prefix operators generate.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
n-D, the same layout as input, dtype is uint32.
|
||||
"""
|
||||
ishape = data.shape
|
||||
if axis is None:
|
||||
axis = len(ishape) - 1
|
||||
assert get_const_int(ishape[axis]) % 32 == 0
|
||||
n = len(ishape)
|
||||
oshape = tuple(simplify(ishape[i] // 32) if i == axis else ishape[i] for i in range(n))
|
||||
|
||||
def _binarize_pack(*indices):
|
||||
start_idx = [indices[i] * 32 if i == axis else indices[i] for i in range(n)]
|
||||
packed = tvm.tirx.const(0, "uint32")
|
||||
for j in range(32):
|
||||
idx = [start_idx[i] + j if i == axis else start_idx[i] for i in range(n)]
|
||||
sign = (data(*idx) >= 0).astype("uint32")
|
||||
packed = packed | sign
|
||||
if j == 31:
|
||||
return packed
|
||||
packed = packed << 1
|
||||
raise RuntimeError("not resach")
|
||||
|
||||
return te.compute(oshape, _binarize_pack, name=name, tag="binarize_pack")
|
||||
|
||||
|
||||
def binary_dense(data, weight):
|
||||
"""Binary matrix multiplication using xor and bit-count.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
2-D with shape [batch, in_dim], dtype is uint32.
|
||||
|
||||
weight : tvm.te.Tensor
|
||||
2-D with shape [out_dim, in_dim], dtype is uint32.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
2-D with shape [batch, out_dim], dtype is float32.
|
||||
"""
|
||||
assert data.dtype == "uint32" and weight.dtype == "uint32", (
|
||||
"dtype of data and weight should be uint32"
|
||||
)
|
||||
assert len(data.shape) == 2 and len(weight.shape) == 2, "only support 2-dim binary dense"
|
||||
batch, in_dim = data.shape
|
||||
out_dim, _ = weight.shape
|
||||
k = te.reduce_axis((0, in_dim), name="k")
|
||||
matmul = te.compute(
|
||||
(batch, out_dim),
|
||||
lambda i, j: te.sum(tvm.tirx.popcount(data[i, k] ^ weight[j, k]), axis=k),
|
||||
tag="binary_dense",
|
||||
)
|
||||
|
||||
return te.compute(
|
||||
(batch, out_dim), lambda i, j: 32 * in_dim - 2.0 * matmul(i, j), tag=tag.ELEMWISE
|
||||
)
|
||||
@@ -0,0 +1,141 @@
|
||||
# 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-variable, unused-argument
|
||||
"""1D convolution operators."""
|
||||
|
||||
from .conv2d import conv
|
||||
|
||||
|
||||
def conv1d(
|
||||
data,
|
||||
kernel,
|
||||
strides=1,
|
||||
padding="VALID",
|
||||
dilation=1,
|
||||
groups=1,
|
||||
data_layout="NCW",
|
||||
kernel_layout="",
|
||||
out_dtype=None,
|
||||
):
|
||||
"""1D convolution forward operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
3-D input shape [batch, in_channel, in_width] for data_layout == 'NCW'
|
||||
and [batch, in_width, in_channel] for data_layout == 'NWC'
|
||||
|
||||
kernel : tvm.te.Tensor
|
||||
3-D kernel with shape [num_filter, in_channel, filter_size] for kernel_layout == 'OIW'
|
||||
and [filter_size, in_channel, num_filter] for kernel_layout == 'WIO'
|
||||
|
||||
strides : int or tuple
|
||||
The spatial stride along width
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
dilation : int or tuple
|
||||
Dilation rate if convolution should be dilated.
|
||||
|
||||
data_layout : str
|
||||
How input data is laid out, must be one of ['NCW', 'NWC']
|
||||
|
||||
kernel_layout: Optiona[str]
|
||||
The layout of the kernel. If unspecified, use default layout. "OIW" if data_layout == "NCW",
|
||||
"WIO" if data_layout == "NWC".
|
||||
|
||||
out_dtype : str
|
||||
The output data type. If None then output is same type as input.
|
||||
"""
|
||||
return conv(
|
||||
data, kernel, strides, padding, dilation, groups, data_layout, kernel_layout, out_dtype
|
||||
)
|
||||
|
||||
|
||||
def conv1d_nwc(data, kernel, strides=1, padding="VALID", dilation=1, out_dtype=None):
|
||||
"""1D convolution in NWC layout. See :py:func:`conv` for details on parameters"""
|
||||
return conv(data, kernel, strides, padding, dilation, 1, "NWC", "WIO", out_dtype=out_dtype)
|
||||
|
||||
|
||||
def conv1d_ncw(data, kernel, strides=1, padding="VALID", dilation=1, out_dtype=None):
|
||||
"""1D convolution in NCW layout. See :py:func:`conv` for details on parameters"""
|
||||
return conv(data, kernel, strides, padding, dilation, 1, "NCW", "OIW", out_dtype=out_dtype)
|
||||
|
||||
|
||||
def group_conv1d_nwc(
|
||||
data, kernel, strides=1, padding="VALID", dilation=1, groups=1, out_dtype=None
|
||||
):
|
||||
"""1D convolution forward operator for NWC layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
3-D with shape [batch, in_width, in_channel]
|
||||
|
||||
kernel : tvm.te.Tensor
|
||||
3-D with shape [filter_size, in_channel, num_filter]
|
||||
|
||||
strides : int or tuple
|
||||
The spatial stride along width
|
||||
|
||||
padding : int, tuple, or str
|
||||
Padding size can be an integer for equal padding,
|
||||
a tuple of (left, right) or a string in ['VALID', 'SAME'].
|
||||
|
||||
dilation : int or tuple
|
||||
Dilation rate if convolution should be dilated.
|
||||
|
||||
groups : int
|
||||
Number of groups
|
||||
|
||||
out_dtype : str
|
||||
The output data type. If None then output is same type as input.
|
||||
"""
|
||||
return conv(data, kernel, strides, padding, dilation, groups, "NWC", "WIO", out_dtype=out_dtype)
|
||||
|
||||
|
||||
def group_conv1d_ncw(
|
||||
data, kernel, strides=1, padding="VALID", dilation=1, groups=1, out_dtype=None
|
||||
):
|
||||
"""1D convolution forward operator for NCW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
3-D with shape [batch, in_channel, in_width]
|
||||
|
||||
kernel : tvm.te.Tensor
|
||||
3-D with shape [num_filter, in_channel, filter_size]
|
||||
|
||||
strides : int or tuple
|
||||
The spatial stride along width
|
||||
|
||||
padding : int, tuple, or str
|
||||
Padding size can be an integer for equal padding,
|
||||
a tuple of (left, right) or a string in ['VALID', 'SAME'].
|
||||
|
||||
dilation : int or tuple
|
||||
Dilation rate if convolution should be dilated.
|
||||
|
||||
groups : int
|
||||
Number of groups
|
||||
|
||||
out_dtype : str
|
||||
The output data type. If None then output is same type as input.
|
||||
"""
|
||||
return conv(data, kernel, strides, padding, dilation, groups, "NCW", "OIW", out_dtype=out_dtype)
|
||||
@@ -0,0 +1,217 @@
|
||||
# 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-variable, unused-argument
|
||||
"""Transposed 1D convolution operators (sometimes called Deconvolution)."""
|
||||
|
||||
from tvm import te
|
||||
|
||||
from ..utils import simplify
|
||||
from .dilate import dilate
|
||||
from .pad import pad
|
||||
from .utils import get_pad_tuple1d
|
||||
|
||||
|
||||
def _conv1d_transpose_ncw_preprocess(data, kernel, stride, padding, out_dtype, output_padding):
|
||||
"""Preprocess data and kernel to make the compute pattern
|
||||
of conv1d_transpose the same as conv1d.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
3-D with shape [batch, in_channel, in_width]
|
||||
|
||||
kernel : tvm.te.Tensor
|
||||
3-D with shape [in_channel, num_filter, filter_width]
|
||||
|
||||
stride : ints
|
||||
The spatial stride along width
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
out_dtype : str
|
||||
The output data type. This is used for mixed precision.
|
||||
|
||||
output_padding : ints
|
||||
Used to recover the actual output shape in case there are more
|
||||
than one possible shape. Must be smaller than stride.
|
||||
|
||||
Returns
|
||||
-------
|
||||
data_pad : tvm.te.Tensor
|
||||
Padded input data. 3-D with shape [batch, in_channel, in_width]
|
||||
|
||||
kernel: tvm.te.Tensor
|
||||
Transformed kernel. 3-D with shape [num_filter, in_channel, filter_width]
|
||||
"""
|
||||
# some pre-processing and prelimnary checks
|
||||
if out_dtype is None:
|
||||
out_dtype = data.dtype
|
||||
|
||||
# dilate and pad
|
||||
if isinstance(stride, tuple | list):
|
||||
stride = stride[0]
|
||||
if isinstance(output_padding, tuple | list):
|
||||
output_padding = output_padding[0]
|
||||
|
||||
_, channels_in, _ = data.shape
|
||||
_, channels_out, kernel_width = kernel.shape
|
||||
assert output_padding < stride
|
||||
channels_out = simplify(channels_out)
|
||||
data_dilate = dilate(data, [1, 1, stride], name="data_dilate")
|
||||
pad_left, pad_right = get_pad_tuple1d(padding, (kernel_width,))
|
||||
pad_left = kernel_width - 1 - pad_left
|
||||
pad_right = kernel_width - 1 - pad_right + output_padding
|
||||
data_pad = pad(data_dilate, [0, 0, pad_left], [0, 0, pad_right], name="data_pad")
|
||||
|
||||
# transform kernel layout from IOW to OIW, and rotate kernel by 180 degrees
|
||||
kernel = te.compute(
|
||||
(channels_out, channels_in, kernel_width),
|
||||
lambda o, i, w: kernel[i][o][kernel_width - 1 - w],
|
||||
name="kernel",
|
||||
)
|
||||
return data_pad, kernel
|
||||
|
||||
|
||||
def conv1d_transpose_ncw(data, kernel, stride, padding, out_dtype, output_padding):
|
||||
"""Transposed 1D convolution ncw forward operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
3-D with shape [batch, in_channel, in_width]
|
||||
|
||||
kernel : tvm.te.Tensor
|
||||
3-D with shape [in_channel, num_filter, filter_width]
|
||||
|
||||
stride : ints
|
||||
The spatial stride along width
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
out_dtype : str
|
||||
The output data type. This is used for mixed precision.
|
||||
|
||||
output_padding : ints
|
||||
Used to recover the actual output shape in case there are more
|
||||
than one possible shape. Must be smaller than stride.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
3-D with shape [batch, out_channel, out_width]
|
||||
|
||||
"""
|
||||
|
||||
batch, channels_in, _ = data.shape
|
||||
_, channels_out, kernel_width = kernel.shape
|
||||
|
||||
data_pad, transformed_kernel = _conv1d_transpose_ncw_preprocess(
|
||||
data, kernel, stride, padding, out_dtype, output_padding
|
||||
)
|
||||
|
||||
# convolution
|
||||
_, _, data_width = data_pad.shape
|
||||
out_w = simplify(data_width - kernel_width + 1)
|
||||
dc = te.reduce_axis((0, channels_in), name="dc")
|
||||
dw = te.reduce_axis((0, kernel_width), name="dw")
|
||||
output = te.compute(
|
||||
(batch, channels_out, out_w),
|
||||
lambda b, c, w: te.sum(
|
||||
data_pad[b, dc, w + dw].astype(out_dtype)
|
||||
* transformed_kernel[c, dc, dw].astype(out_dtype),
|
||||
axis=[dc, dw],
|
||||
),
|
||||
tag="conv1d_transpose_ncw",
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def group_conv1d_transpose_ncw(data, kernel, stride, padding, out_dtype, output_padding, groups):
|
||||
"""Transposed 1D group convolution ncw forward operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
3-D with shape [batch, in_channel, in_width]
|
||||
|
||||
kernel : tvm.te.Tensor
|
||||
3-D with shape [in_channel, num_filter, filter_width]
|
||||
|
||||
stride : ints
|
||||
The spatial stride along width
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
out_dtype : str
|
||||
The output data type. This is used for mixed precision.
|
||||
|
||||
output_padding : ints
|
||||
Used to recover the actual output shape in case there are more
|
||||
than one possible shape. Must be smaller than stride.
|
||||
|
||||
groups : int
|
||||
number of groups
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
3-D with shape [batch, out_channel, out_width]
|
||||
|
||||
"""
|
||||
if groups == 1:
|
||||
return conv1d_transpose_ncw(data, kernel, stride, padding, out_dtype, output_padding)
|
||||
|
||||
_, in_channels, _ = data.shape
|
||||
|
||||
assert in_channels % groups == 0, (
|
||||
f"input channels {in_channels} must divide group size {groups}"
|
||||
)
|
||||
|
||||
data_pad, transformed_kernel = _conv1d_transpose_ncw_preprocess(
|
||||
data, kernel, stride, padding, out_dtype, output_padding
|
||||
)
|
||||
|
||||
batch, in_channels, in_w = data_pad.shape
|
||||
out_c, _, filter_w = transformed_kernel.shape
|
||||
|
||||
# convolution stage
|
||||
out_channels = simplify(out_c * groups)
|
||||
out_w = simplify(in_w - filter_w + 1)
|
||||
dc = te.reduce_axis((0, in_channels // groups), name="dc")
|
||||
dw = te.reduce_axis((0, filter_w), name="dw")
|
||||
|
||||
# data: batch, in_channels, out_w
|
||||
# weight: out_channels // G, in_channels, out_w
|
||||
return te.compute(
|
||||
(batch, out_channels, out_w),
|
||||
lambda b, c, w: te.sum(
|
||||
data_pad[
|
||||
b, c // (out_channels // groups) * (in_channels // groups) + dc, w + dw
|
||||
].astype(out_dtype)
|
||||
* transformed_kernel[
|
||||
c % (out_channels // groups),
|
||||
c // (out_channels // groups) * (in_channels // groups) + dc,
|
||||
dw,
|
||||
].astype(out_dtype),
|
||||
axis=[dc, dw],
|
||||
),
|
||||
tag="group_conv1d_transpose_ncw",
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,246 @@
|
||||
# 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-variable, unused-argument
|
||||
# ruff: noqa: F821
|
||||
"""Transposed 2D convolution operators (sometimes called Deconvolution)."""
|
||||
|
||||
import collections
|
||||
|
||||
from tvm import te
|
||||
|
||||
from ..utils import simplify
|
||||
from .dilate import dilate
|
||||
from .pad import pad
|
||||
from .utils import get_pad_tuple
|
||||
|
||||
|
||||
def _ntuple(n):
|
||||
def parse(x):
|
||||
if isinstance(x, collections.abc.Iterable):
|
||||
assert len(x) == n, f"Input can only have {n} elements, but got {len(x)} instead: {x}."
|
||||
return x
|
||||
return tuple(repeat(x, n))
|
||||
|
||||
return parse
|
||||
|
||||
|
||||
_single = _ntuple(1)
|
||||
_pair = _ntuple(2)
|
||||
_triple = _ntuple(3)
|
||||
_quadruple = _ntuple(4)
|
||||
|
||||
|
||||
def conv2d_transpose_nchw(Input, Filter, strides, padding, out_dtype, output_padding):
|
||||
"""Transposed 2D convolution nchw forward operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
Input : tvm.te.Tensor
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
Filter : tvm.te.Tensor
|
||||
4-D with shape [in_channel, num_filter, filter_height, filter_width]
|
||||
|
||||
strides : tuple of two ints
|
||||
The spatial stride along height and width
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
out_dtype : str
|
||||
The output data type. This is used for mixed precision.
|
||||
|
||||
output_padding : tuple of ints
|
||||
Used to get the right output shape for gradients
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
return declaration_conv2d_transpose_impl(
|
||||
Input, Filter, strides, padding, out_dtype, output_padding=output_padding
|
||||
)
|
||||
|
||||
|
||||
def conv2d_transpose_nchw_preprocess(data, kernel, strides, padding, out_dtype, output_padding):
|
||||
"""Preprocess data and kernel to make the compute pattern
|
||||
of conv2d_transpose the same as conv2d"""
|
||||
batch, in_c, in_h, in_w = data.shape
|
||||
_, out_c, filter_h, filter_w = kernel.shape
|
||||
stride_h, stride_w = strides
|
||||
opad_h, opad_w = output_padding
|
||||
assert opad_h < stride_h and opad_w < stride_w
|
||||
# dilate data
|
||||
data_dilate = dilate(data, [1, 1, stride_h, stride_w], name="data_dilate")
|
||||
# pad data
|
||||
fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple(padding, (filter_h, filter_w))
|
||||
bpad_top = filter_h - 1 - fpad_top
|
||||
bpad_bottom = filter_h - 1 - fpad_bottom + opad_h
|
||||
bpad_left = filter_w - 1 - fpad_left
|
||||
bpad_right = filter_w - 1 - fpad_right + opad_w
|
||||
data_pad = pad(
|
||||
data_dilate, [0, 0, bpad_top, bpad_left], [0, 0, bpad_bottom, bpad_right], name="data_pad"
|
||||
)
|
||||
# transform kernel layout from IOHW to OIHW, and rotate kernel by 180 degrees
|
||||
kernel_transform = te.compute(
|
||||
(out_c, in_c, filter_h, filter_w),
|
||||
lambda o, i, h, w: kernel[i][o][filter_h - 1 - h][filter_w - 1 - w],
|
||||
name="kernel_transform",
|
||||
)
|
||||
return data_pad, kernel_transform
|
||||
|
||||
|
||||
def declaration_conv2d_transpose_impl(data, kernel, strides, padding, out_dtype, output_padding):
|
||||
"""Implementation of conv2d transpose"""
|
||||
data_pad, kernel_transform = conv2d_transpose_nchw_preprocess(
|
||||
data, kernel, strides, padding, out_dtype, output_padding
|
||||
)
|
||||
batch, in_c, in_h, in_w = data_pad.shape
|
||||
out_c, _, filter_h, filter_w = kernel_transform.shape
|
||||
|
||||
# convolution stage
|
||||
out_c = simplify(out_c)
|
||||
|
||||
out_h = simplify(in_h - filter_h + 1)
|
||||
out_w = simplify(in_w - filter_w + 1)
|
||||
dc = te.reduce_axis((0, in_c), name="dc")
|
||||
dh = te.reduce_axis((0, filter_h), name="dh")
|
||||
dw = te.reduce_axis((0, filter_w), name="dw")
|
||||
|
||||
Output = te.compute(
|
||||
(batch, out_c, out_h, out_w),
|
||||
lambda b, c, h, w: te.sum(
|
||||
data_pad[b, dc, h + dh, w + dw].astype(out_dtype)
|
||||
* kernel_transform[c, dc, dh, dw].astype(out_dtype),
|
||||
axis=[dc, dh, dw],
|
||||
),
|
||||
tag="conv2d_transpose_nchw",
|
||||
)
|
||||
|
||||
return Output
|
||||
|
||||
|
||||
def group_conv2d_transpose_nchw(data, kernel, stride, padding, out_dtype, output_padding, groups):
|
||||
"""Group convolution operator in NCHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
kernel : tvm.te.Tensor
|
||||
4-D with shape [in_channel, out_channel // groups, filter_height, filter_width]
|
||||
|
||||
stride : int or a list/tuple of two ints
|
||||
Stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : int or a list/tuple of 2 or 4 ints
|
||||
padding size, or
|
||||
[pad_height, pad_width] for 2 ints, or
|
||||
[pad_top, pad_left, pad_bottom, pad_right] for 4 ints
|
||||
|
||||
out_dtype : str
|
||||
The output data type. This is used for mixed precision.
|
||||
|
||||
output_padding : tuple of ints
|
||||
Used to get the right output shape for gradients
|
||||
|
||||
groups : int
|
||||
number of groups
|
||||
|
||||
out_dtype : str
|
||||
The output type. This is used for mixed precision.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
if groups == 1:
|
||||
return conv2d_transpose_nchw(data, kernel, stride, padding, out_dtype, output_padding)
|
||||
|
||||
# some pre-processing and prelimnary checks
|
||||
if out_dtype is None:
|
||||
out_dtype = data.dtype
|
||||
|
||||
batch, in_channels, in_h, in_w = data.shape
|
||||
_, out_c, filter_h, filter_w = kernel.shape
|
||||
assert in_channels % groups == 0, (
|
||||
f"input channels {in_channels} must divide group size {groups}"
|
||||
)
|
||||
# assert out_c % groups == 0, f"output channels {in_c} must divide group size {groups}"
|
||||
|
||||
strides = _pair(stride)
|
||||
# padding = _pair(padding)
|
||||
# output_padding = _pair(output_padding)
|
||||
# dilation = _pair(dilation)
|
||||
|
||||
stride_h, stride_w = strides
|
||||
opad_h, opad_w = output_padding
|
||||
assert opad_h < stride_h and opad_w < stride_w, (
|
||||
f"[{output_padding}] opad_h:{opad_h} < stride_h:{stride_h} \
|
||||
and opad_w:{opad_w} < stride_w:{stride_w} does not satisfy."
|
||||
)
|
||||
# dilate data
|
||||
data_dilate = dilate(data, [1, 1, stride_h, stride_w], name="data_dilate")
|
||||
# pad data
|
||||
fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple(padding, (filter_h, filter_w))
|
||||
bpad_top = filter_h - 1 - fpad_top
|
||||
bpad_bottom = filter_h - 1 - fpad_bottom + opad_h
|
||||
bpad_left = filter_w - 1 - fpad_left
|
||||
bpad_right = filter_w - 1 - fpad_right + opad_w
|
||||
data_pad = pad(
|
||||
data_dilate, [0, 0, bpad_top, bpad_left], [0, 0, bpad_bottom, bpad_right], name="data_pad"
|
||||
)
|
||||
# transform kernel layout from IOHW to OIHW, and rotate kernel by 180 degrees
|
||||
kernel_transform = te.compute(
|
||||
(out_c, in_channels, filter_h, filter_w),
|
||||
lambda i, o, h, w: kernel[o][i][filter_h - 1 - h][filter_w - 1 - w],
|
||||
name="kernel_transform",
|
||||
)
|
||||
|
||||
batch, in_channels, in_h, in_w = data_pad.shape
|
||||
out_c, _, filter_h, filter_w = kernel_transform.shape
|
||||
|
||||
# convolution stage
|
||||
out_channels = simplify(out_c * groups)
|
||||
|
||||
out_h = simplify(in_h - filter_h + 1)
|
||||
out_w = simplify(in_w - filter_w + 1)
|
||||
dc = te.reduce_axis((0, in_channels // groups), name="dc")
|
||||
dh = te.reduce_axis((0, filter_h), name="dh")
|
||||
dw = te.reduce_axis((0, filter_w), name="dw")
|
||||
|
||||
# data: batch, in_channels, out_h, out_w
|
||||
# weight: out_channels // G, in_channels, out_h, out_w
|
||||
return te.compute(
|
||||
(batch, out_channels, out_h, out_w),
|
||||
lambda b, c, h, w: te.sum(
|
||||
data_pad[
|
||||
b, c // (out_channels // groups) * (in_channels // groups) + dc, h + dh, w + dw
|
||||
].astype(out_dtype)
|
||||
* kernel_transform[
|
||||
c % (out_channels // groups),
|
||||
c // (out_channels // groups) * (in_channels // groups) + dc,
|
||||
dh,
|
||||
dw,
|
||||
].astype(out_dtype),
|
||||
axis=[dc, dh, dw],
|
||||
),
|
||||
tag="group_conv2d_transpose_nchw",
|
||||
)
|
||||
@@ -0,0 +1,169 @@
|
||||
# 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-variable, too-many-locals
|
||||
# pylint: disable=unused-argument, redefined-builtin, no-else-return
|
||||
"""Conv3D operators"""
|
||||
|
||||
from tvm import te
|
||||
|
||||
from ..utils import get_const_tuple
|
||||
from .conv2d import conv
|
||||
from .winograd_util import winograd_transform_matrices
|
||||
|
||||
|
||||
def conv3d_ncdhw(Input, Filter, stride, padding, dilation, groups, out_dtype=None):
|
||||
"""Conv3D operator in NCDHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
Input : tvm.te.Tensor
|
||||
5-D with shape [batch, in_channel, in_depth, in_height, in_width]
|
||||
|
||||
Filter : tvm.te.Tensor
|
||||
5-D with shape [num_filter, in_channel, filter_depth, filter_height, filter_width]
|
||||
|
||||
stride : int or a list/tuple of three ints
|
||||
Stride size, or [strid_depth, stride_height, stride_width]
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
dilation: int or a list/tuple of three ints
|
||||
dilation size, or [dilation_depth, dilation_height, dilation_width]
|
||||
|
||||
groups: int
|
||||
Number of groups.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
5-D with shape [batch, out_channel, out_depth, out_height, out_width]
|
||||
"""
|
||||
return conv(Input, Filter, stride, padding, dilation, groups, "NCDHW", "OIDHW", out_dtype)
|
||||
|
||||
|
||||
def conv3d_ndhwc(
|
||||
Input,
|
||||
Filter,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
groups,
|
||||
out_dtype="float32",
|
||||
auto_scheduler_rewritten_layout="",
|
||||
meta_schedule_origin_shape=None,
|
||||
):
|
||||
"""Convolution operator in NDHWC layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
Input : tvm.te.Tensor
|
||||
5-D with shape [batch, in_depth, in_height, in_width, in_channel]
|
||||
|
||||
Filter : tvm.te.Tensor
|
||||
5-D with shape [filter_depth, filter_height, filter_width, in_channel, num_filter]
|
||||
|
||||
stride : int or a list/tuple of three ints
|
||||
Stride size, or [stride_depth, stride_height, stride_width]
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
dilation: int or a list/tuple of three ints
|
||||
dilation size, or [dilation_depth, dilation_height, dilation_width]
|
||||
|
||||
groups: int
|
||||
Number of groups.
|
||||
|
||||
out_dtype: str = "float32",
|
||||
The type of output tensor
|
||||
|
||||
auto_scheduler_rewritten_layout: str = ""
|
||||
The layout after auto-scheduler's layout rewrite pass.
|
||||
|
||||
meta_schedule_origin_shape: Optional[List[Expr]] = None
|
||||
The original shape of the input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
5-D with shape [batch, out_depth, out_height, out_width, out_channel]
|
||||
"""
|
||||
return conv(
|
||||
Input,
|
||||
Filter,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
groups,
|
||||
"NDHWC",
|
||||
"DHWIO",
|
||||
out_dtype,
|
||||
auto_scheduler_rewritten_layout,
|
||||
meta_schedule_origin_shape,
|
||||
)
|
||||
|
||||
|
||||
def conv3d_winograd_weight_transform(kernel, tile_size):
|
||||
"""Weight transformation for 3D winograd
|
||||
|
||||
Parameters
|
||||
----------
|
||||
kernel: Tensor
|
||||
The raw kernel tensor with layout "NCDHW".
|
||||
tile_size: int
|
||||
Tile size of winograd transform. e.g. 2 for F(2x2, 3x3) and 4 for F(4x4, 3x3)
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
5-D with shape [alpha, alpha, alpha, CO, CI]
|
||||
"""
|
||||
CO, CI, KD, KH, KW = get_const_tuple(kernel.shape)
|
||||
|
||||
depth_transform = 2 < KD < 8 and KD == KH
|
||||
|
||||
if depth_transform:
|
||||
assert KD == KH == KW, "Only support NxNxN kernel"
|
||||
else:
|
||||
assert KH == KW, "Only supports DxNxN kernel"
|
||||
|
||||
r = tile_size + KH - 1
|
||||
|
||||
r_kh = te.reduce_axis((0, KH), name="r_kh")
|
||||
r_kw = te.reduce_axis((0, KW), name="r_kw")
|
||||
_, _, G = winograd_transform_matrices(tile_size, KH, kernel.dtype)
|
||||
if depth_transform:
|
||||
shape = (r, r, r, CO, CI)
|
||||
r_kd = te.reduce_axis((0, KD), name="r_kd")
|
||||
return te.compute(
|
||||
shape,
|
||||
lambda omg, eps, nu, co, ci: te.sum(
|
||||
kernel[co][ci][r_kd][r_kh][r_kw] * G[omg][r_kd] * G[eps][r_kh] * G[nu][r_kw],
|
||||
axis=[r_kd, r_kh, r_kw],
|
||||
),
|
||||
name="transform_weight",
|
||||
)
|
||||
else:
|
||||
shape = (r, r, KD, CO, CI)
|
||||
return te.compute(
|
||||
shape,
|
||||
lambda eps, nu, d, co, ci: te.sum(
|
||||
kernel[co][ci][d][r_kh][r_kw] * G[eps][r_kh] * G[nu][r_kw], axis=[r_kh, r_kw]
|
||||
),
|
||||
name="transform_weight",
|
||||
)
|
||||
@@ -0,0 +1,200 @@
|
||||
# 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-variable, unused-argument
|
||||
"""Transposed 3D convolution operators (sometimes called Deconvolution)."""
|
||||
|
||||
from tvm import te
|
||||
|
||||
from ..utils import simplify
|
||||
from .dilate import dilate
|
||||
from .pad import pad
|
||||
from .utils import get_pad_tuple3d
|
||||
|
||||
|
||||
def conv3d_transpose_ncdhw(Input, Filter, strides, padding, out_dtype, output_padding):
|
||||
"""Transposed 3D convolution ncdhw forward operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
Input : tvm.te.Tensor
|
||||
5-D with shape [batch, in_channel, in_depth, in_height, in_width]
|
||||
|
||||
Filter : tvm.te.Tensor
|
||||
5-D with shape [in_channel, num_filter, filter_depth, filter_height, filter_width]
|
||||
|
||||
strides : int or a list/tuple of three ints
|
||||
The spatial stride along depth,height and width
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
out_dtype : str
|
||||
The output data type. This is used for mixed precision.
|
||||
|
||||
output_padding : tuple of ints
|
||||
Used to get the right output shape for gradients
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
5-D with shape [batch, out_channel, out_depth, out_height, out_width]
|
||||
"""
|
||||
return declaration_conv3d_transpose_impl(
|
||||
Input, Filter, strides, padding, out_dtype, output_padding
|
||||
)
|
||||
|
||||
|
||||
def conv3d_transpose_ncdhw_preprocess(data, kernel, strides, padding, out_dtype, output_padding):
|
||||
"""Preprocess data and kernel to make the compute pattern
|
||||
of conv3d_transpose the same as conv3d"""
|
||||
batch, in_c, in_d, in_h, in_w = data.shape
|
||||
_, out_c, filter_d, filter_h, filter_w = kernel.shape
|
||||
stride_d, stride_h, stride_w = strides
|
||||
opad_d, opad_h, opad_w = output_padding
|
||||
assert opad_d < stride_d and opad_h < stride_h and opad_w < stride_w
|
||||
# dilate data
|
||||
data_dilate = dilate(data, [1, 1, stride_d, stride_h, stride_w], name="data_dilate")
|
||||
# pad data
|
||||
fpad_front, fpad_top, fpad_left, fpad_back, fpad_bottom, fpad_right = get_pad_tuple3d(
|
||||
padding, (filter_d, filter_h, filter_w)
|
||||
)
|
||||
bpad_front = filter_d - 1 - fpad_front
|
||||
bpad_back = filter_d - 1 - fpad_back + opad_d
|
||||
bpad_top = filter_h - 1 - fpad_top
|
||||
bpad_bottom = filter_h - 1 - fpad_bottom + opad_h
|
||||
bpad_left = filter_w - 1 - fpad_left
|
||||
bpad_right = filter_w - 1 - fpad_right + opad_w
|
||||
data_pad = pad(
|
||||
data_dilate,
|
||||
[0, 0, bpad_front, bpad_top, bpad_left],
|
||||
[0, 0, bpad_back, bpad_bottom, bpad_right],
|
||||
name="data_pad",
|
||||
)
|
||||
# transform kernel layout from IODHW to OIDHW, and rotate kernel by 180 degrees
|
||||
kernel_transform = te.compute(
|
||||
(out_c, in_c, filter_d, filter_h, filter_w),
|
||||
lambda o, i, d, h, w: kernel[i][o][filter_d - 1 - d][filter_h - 1 - h][filter_w - 1 - w],
|
||||
name="kernel_transform",
|
||||
)
|
||||
return data_pad, kernel_transform
|
||||
|
||||
|
||||
def declaration_conv3d_transpose_impl(data, kernel, strides, padding, out_dtype, output_padding):
|
||||
"""Implementation of conv3d transpose"""
|
||||
data_pad, kernel_transform = conv3d_transpose_ncdhw_preprocess(
|
||||
data, kernel, strides, padding, out_dtype, output_padding
|
||||
)
|
||||
batch, in_c, in_d, in_h, in_w = data_pad.shape
|
||||
out_c, _, filter_d, filter_h, filter_w = kernel_transform.shape
|
||||
stride_d, stride_h, stride_w = strides
|
||||
|
||||
# convolution stage
|
||||
out_c = simplify(out_c)
|
||||
out_d = simplify(in_d - filter_d + 1)
|
||||
out_h = simplify(in_h - filter_h + 1)
|
||||
out_w = simplify(in_w - filter_w + 1)
|
||||
dc = te.reduce_axis((0, in_c), name="dc")
|
||||
dd = te.reduce_axis((0, filter_d), name="dd")
|
||||
dh = te.reduce_axis((0, filter_h), name="dh")
|
||||
dw = te.reduce_axis((0, filter_w), name="dw")
|
||||
|
||||
Output = te.compute(
|
||||
(batch, out_c, out_d, out_h, out_w),
|
||||
lambda b, c, d, h, w: te.sum(
|
||||
data_pad[b, dc, d + dd, h + dh, w + dw].astype(out_dtype)
|
||||
* kernel_transform[c, dc, dd, dh, dw].astype(out_dtype),
|
||||
axis=[dc, dd, dh, dw],
|
||||
),
|
||||
tag="conv3d_transpose_ncdhw",
|
||||
)
|
||||
|
||||
return Output
|
||||
|
||||
|
||||
def group_conv3d_transpose_ncdhw(data, kernel, strides, padding, out_dtype, output_padding, groups):
|
||||
"""Transposed group 3D convolution ncdhw forward operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
5-D with shape [batch, in_channel, in_depth, in_height, in_width]
|
||||
|
||||
kernel : tvm.te.Tensor
|
||||
5-D with shape [in_channel, num_filter, filter_depth, filter_height, filter_width]
|
||||
|
||||
strides : int or a list/tuple of three ints
|
||||
The spatial stride along depth,height and width
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
out_dtype : str
|
||||
The output data type. This is used for mixed precision.
|
||||
|
||||
output_padding : tuple of ints
|
||||
Used to get the right output shape for gradients
|
||||
|
||||
groups : int
|
||||
number of groups
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
5-D with shape [batch, out_channel, out_depth, out_height, out_width]
|
||||
"""
|
||||
if not isinstance(strides, tuple | list):
|
||||
strides = (strides, strides, strides)
|
||||
|
||||
if groups == 1:
|
||||
return conv3d_transpose_ncdhw(data, kernel, strides, padding, out_dtype, output_padding)
|
||||
|
||||
data_pad, kernel_transform = conv3d_transpose_ncdhw_preprocess(
|
||||
data, kernel, strides, padding, out_dtype, output_padding
|
||||
)
|
||||
batch, in_c, in_d, in_h, in_w = data_pad.shape
|
||||
out_c, _, filter_d, filter_h, filter_w = kernel_transform.shape
|
||||
assert in_c % groups == 0, f"input channels {in_c} must divide group size {groups}"
|
||||
|
||||
# convolution stage
|
||||
out_c = simplify(out_c * groups)
|
||||
out_d = simplify(in_d - filter_d + 1)
|
||||
out_h = simplify(in_h - filter_h + 1)
|
||||
out_w = simplify(in_w - filter_w + 1)
|
||||
dc = te.reduce_axis((0, in_c // groups), name="dc")
|
||||
dd = te.reduce_axis((0, filter_d), name="dd")
|
||||
dh = te.reduce_axis((0, filter_h), name="dh")
|
||||
dw = te.reduce_axis((0, filter_w), name="dw")
|
||||
|
||||
# data: batch, in_channels, out_d, out_h, out_w
|
||||
# weight: out_channels // G, in_channels, out_d, out_h, out_w
|
||||
return te.compute(
|
||||
(batch, out_c, out_d, out_h, out_w),
|
||||
lambda b, c, d, h, w: te.sum(
|
||||
data_pad[
|
||||
b, c // (out_c // groups) * (in_c // groups) + dc, d + dd, h + dh, w + dw
|
||||
].astype(out_dtype)
|
||||
* kernel_transform[
|
||||
c % (out_c // groups),
|
||||
c // (out_c // groups) * (in_c // groups) + dc,
|
||||
dd,
|
||||
dh,
|
||||
dw,
|
||||
].astype(out_dtype),
|
||||
axis=[dc, dd, dh, dw],
|
||||
),
|
||||
tag="group_conv3d_transpose_ncdhw",
|
||||
)
|
||||
@@ -0,0 +1,124 @@
|
||||
# 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: E731
|
||||
"""Correlation operators"""
|
||||
|
||||
from tvm import te
|
||||
|
||||
from ..utils import get_const_tuple
|
||||
from .pad import pad
|
||||
|
||||
|
||||
def correlation_nchw(
|
||||
data1, data2, kernel_size, max_displacement, stride1, stride2, padding, is_multiply
|
||||
):
|
||||
"""Correlation operator in NCHW layout.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data1 : tvm.te.Tensor
|
||||
4-D with shape [batch, channel, height, width]
|
||||
|
||||
data2 : tvm.te.Tensor
|
||||
4-D with shape [batch, channel, height, width]
|
||||
|
||||
kernel_size: int
|
||||
Kernel size for correlation, must be an odd number
|
||||
|
||||
max_displacement: int
|
||||
Max displacement of Correlation
|
||||
|
||||
stride1: int
|
||||
Stride for data1
|
||||
|
||||
stride2: int
|
||||
Stride for data2 within the neightborhood centered around data1
|
||||
|
||||
padding : int or a list/tuple of 2 or 4 ints
|
||||
Padding size, or
|
||||
[pad_height, pad_width] for 2 ints, or
|
||||
[pad_top, pad_left, pad_bottom, pad_right] for 4 ints
|
||||
|
||||
is_multiply: bool
|
||||
operation type is either multiplication or substraction
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
# pylint: disable=unnecessary-lambda, invalid-name
|
||||
data_shape = get_const_tuple(data1.shape)
|
||||
assert get_const_tuple(data2.shape) == data_shape, "data1 and data2 should have the same shape"
|
||||
assert kernel_size > 0 and kernel_size % 2, "kernel_size should be non-negative odd number"
|
||||
if isinstance(padding, tuple | list):
|
||||
if len(padding) == 2:
|
||||
pad_before_h = pad_after_h = padding[0]
|
||||
pad_before_w = pad_after_w = padding[1]
|
||||
elif len(padding) == 4:
|
||||
pad_before_h, pad_before_w, pad_after_h, pad_after_w = padding
|
||||
else:
|
||||
raise ValueError("invalid padding")
|
||||
elif isinstance(padding, int):
|
||||
pad_before_h = pad_after_h = pad_before_w = pad_after_w = padding
|
||||
else:
|
||||
raise ValueError("invalid padding")
|
||||
pad_before = [0, 0, pad_before_h, pad_before_w]
|
||||
pad_after = [0, 0, pad_after_h, pad_after_w]
|
||||
padded_data1 = pad(data1, pad_before, pad_after)
|
||||
padded_data2 = pad(data2, pad_before, pad_after)
|
||||
|
||||
batch, channel, height, width = data_shape
|
||||
|
||||
kernel_radius = (kernel_size - 1) // 2
|
||||
border_size = max_displacement + kernel_radius
|
||||
displacement_radius = max_displacement // stride2
|
||||
displacement_size = 2 * displacement_radius + 1
|
||||
|
||||
padded_width = width + pad_before_w + pad_after_w
|
||||
padded_height = height + pad_before_h + pad_after_h
|
||||
out_channel = displacement_size * displacement_size
|
||||
out_height = (padded_height - 2 * border_size + stride1 - 1) // stride1
|
||||
out_width = (padded_width - 2 * border_size + stride1 - 1) // stride1
|
||||
|
||||
rc = te.reduce_axis((0, channel), name="rc")
|
||||
ry = te.reduce_axis((0, kernel_size), name="ry")
|
||||
rx = te.reduce_axis((0, kernel_size), name="rx")
|
||||
|
||||
if is_multiply:
|
||||
corr_func = lambda x, y: x * y
|
||||
else:
|
||||
corr_func = lambda x, y: te.abs(x - y)
|
||||
|
||||
def _compute_correlation(n, q, i, j):
|
||||
# location in data1
|
||||
y1 = i * stride1 + max_displacement
|
||||
x1 = j * stride1 + max_displacement
|
||||
# location in data2
|
||||
y2 = y1 + (te.indexdiv(q, displacement_size) - displacement_radius) * stride2
|
||||
x2 = x1 + (te.indexmod(q, displacement_size) - displacement_radius) * stride2
|
||||
return te.sum(
|
||||
corr_func(padded_data1[n, rc, y1 + ry, x1 + rx], padded_data2[n, rc, y2 + ry, x2 + rx]),
|
||||
axis=[rc, ry, rx],
|
||||
)
|
||||
|
||||
correlation = te.compute(
|
||||
(batch, out_channel, out_height, out_width),
|
||||
lambda n, q, i, j: _compute_correlation(n, q, i, j),
|
||||
tag="correlation_nchw",
|
||||
)
|
||||
return correlation / (kernel_size * kernel_size * channel)
|
||||
@@ -0,0 +1,241 @@
|
||||
# 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, too-many-locals, too-many-arguments
|
||||
"""Deformable Conv2D operators"""
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
|
||||
from ..cpp.utils import bilinear_sample_nchw, bilinear_sample_nhwc
|
||||
from ..utils import get_const_tuple
|
||||
from .utils import get_pad_tuple
|
||||
|
||||
|
||||
def deformable_conv2d_nchw(
|
||||
data, offset, kernel, strides, padding, dilation, deformable_groups, groups, out_dtype
|
||||
):
|
||||
"""Deformable conv2D operator in NCHW layout.
|
||||
|
||||
The deformable convolution operation is described in https://arxiv.org/abs/1703.06211
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
offset : tvm.te.Tensor
|
||||
4-D with shape [batch, deformable_groups * filter_height * filter_width * 2,
|
||||
out_height, out_width].
|
||||
|
||||
kernel : tvm.te.Tensor
|
||||
4-D with shape [num_filter, in_channel, filter_height, filter_width]
|
||||
|
||||
strides : int or a list/tuple of two ints
|
||||
stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : int or a list/tuple of two ints
|
||||
padding size, or [pad_height, pad_width]
|
||||
|
||||
dilation : int or a list/tuple of two ints
|
||||
dilation size, or [dilation_height, dilation_width]
|
||||
|
||||
deformable_groups : int
|
||||
number of deformable groups
|
||||
|
||||
groups : int
|
||||
number of groups
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
if out_dtype is None:
|
||||
out_dtype = data.dtype
|
||||
|
||||
if isinstance(strides, int):
|
||||
stride_h = stride_w = strides
|
||||
else:
|
||||
stride_h, stride_w = strides
|
||||
|
||||
if isinstance(dilation, int):
|
||||
dilation_h = dilation_w = dilation
|
||||
else:
|
||||
dilation_h, dilation_w = dilation
|
||||
|
||||
batch, in_channel, in_height, in_width = get_const_tuple(data.shape)
|
||||
out_channel, channel, kernel_h, kernel_w = get_const_tuple(kernel.shape)
|
||||
_, _, out_height, out_width = get_const_tuple(offset.shape)
|
||||
assert in_channel % deformable_groups == 0, "Input cahnnels must divide deformable group size"
|
||||
assert groups == 1, "deformable_conv2d_nchw does not support groups > 1"
|
||||
|
||||
ic_per_dgroup = channel // deformable_groups
|
||||
|
||||
dilated_kernel_h = (kernel_h - 1) * dilation_h + 1
|
||||
dilated_kernel_w = (kernel_w - 1) * dilation_w + 1
|
||||
pad_top, pad_left, _, _ = get_pad_tuple(padding, (dilated_kernel_h, dilated_kernel_w))
|
||||
rc = te.reduce_axis((0, in_channel), name="rc")
|
||||
ry = te.reduce_axis((0, kernel_h), name="ry")
|
||||
rx = te.reduce_axis((0, kernel_w), name="rx")
|
||||
|
||||
zero = tvm.tirx.const(0.0, data.dtype)
|
||||
|
||||
def _bilinear(n, c, h, w):
|
||||
outside = tvm.tirx.any(h < 0, w < 0, h >= in_height, w >= in_width)
|
||||
val = bilinear_sample_nchw(data, (n, c, h, w), in_height - 1, in_width - 1)
|
||||
return tvm.tirx.if_then_else(outside, zero, val)
|
||||
|
||||
data_deform = te.compute(
|
||||
(batch, in_channel, kernel_h, kernel_w, out_height, out_width),
|
||||
lambda n, c, kh, kw, y, x: _bilinear(
|
||||
n,
|
||||
c,
|
||||
y * stride_h
|
||||
- pad_top
|
||||
+ kh * dilation_h
|
||||
+ offset[
|
||||
n, c // ic_per_dgroup * (kernel_w * kernel_h * 2) + (kh * kernel_w + kw) * 2, y, x
|
||||
],
|
||||
x * stride_w
|
||||
- pad_left
|
||||
+ kw * dilation_w
|
||||
+ offset[
|
||||
n,
|
||||
c // ic_per_dgroup * (kernel_w * kernel_h * 2) + (kh * kernel_w + kw) * 2 + 1,
|
||||
y,
|
||||
x,
|
||||
],
|
||||
),
|
||||
tag="data_deform",
|
||||
)
|
||||
return te.compute(
|
||||
(batch, out_channel, out_height, out_width),
|
||||
lambda n, f, y, x: te.sum(
|
||||
data_deform[n, rc, ry, rx, y, x].astype(out_dtype)
|
||||
* kernel[f, rc, ry, rx].astype(out_dtype),
|
||||
axis=[rc, ry, rx],
|
||||
),
|
||||
tag="deformable_conv2d_nchw",
|
||||
)
|
||||
|
||||
|
||||
def deformable_conv2d_nhwc(
|
||||
data, offset, kernel, strides, padding, dilation, deformable_groups, groups, out_dtype
|
||||
):
|
||||
"""Deformable conv2D operator in NHWC layout.
|
||||
|
||||
The deformable convolution operation is described in https://arxiv.org/abs/1703.06211
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
4-D with shape [batch, in_height, in_width, in_channel]
|
||||
|
||||
offset : tvm.te.Tensor
|
||||
4-D with shape [batch, out_height, out_width,
|
||||
deformable_groups * filter_height * filter_width * 2].
|
||||
|
||||
kernel : tvm.te.Tensor
|
||||
4-D with shape [filter_height, filter_width, in_channel, num_filter]
|
||||
|
||||
strides : int or a list/tuple of two ints
|
||||
stride size, or [stride_height, stride_width]
|
||||
|
||||
padding : int or a list/tuple of two ints
|
||||
padding size, or [pad_height, pad_width]
|
||||
|
||||
dilation : int or a list/tuple of two ints
|
||||
dilation size, or [dilation_height, dilation_width]
|
||||
|
||||
deformable_groups : int
|
||||
number of deformable groups
|
||||
|
||||
groups : int
|
||||
number of groups
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
4-D with shape [batch, out_height, out_width, out_channel]
|
||||
"""
|
||||
if out_dtype is None:
|
||||
out_dtype = data.dtype
|
||||
|
||||
if isinstance(strides, int):
|
||||
stride_h = stride_w = strides
|
||||
else:
|
||||
stride_h, stride_w = strides
|
||||
|
||||
if isinstance(dilation, int):
|
||||
dilation_h = dilation_w = dilation
|
||||
else:
|
||||
dilation_h, dilation_w = dilation
|
||||
|
||||
batch, in_height, in_width, in_channel = get_const_tuple(data.shape)
|
||||
kernel_h, kernel_w, channel, out_channel = get_const_tuple(kernel.shape)
|
||||
_, out_height, out_width, _ = get_const_tuple(offset.shape)
|
||||
assert in_channel % deformable_groups == 0, "Input cahnnels must divide deformable group size"
|
||||
assert groups == 1, "deformable_conv2d_nchw does not support groups > 1"
|
||||
|
||||
ic_per_dgroup = channel // deformable_groups
|
||||
|
||||
dilated_kernel_h = (kernel_h - 1) * dilation_h + 1
|
||||
dilated_kernel_w = (kernel_w - 1) * dilation_w + 1
|
||||
pad_top, pad_left, _, _ = get_pad_tuple(padding, (dilated_kernel_h, dilated_kernel_w))
|
||||
rc = te.reduce_axis((0, in_channel), name="rc")
|
||||
ry = te.reduce_axis((0, kernel_h), name="ry")
|
||||
rx = te.reduce_axis((0, kernel_w), name="rx")
|
||||
|
||||
zero = tvm.tirx.const(0.0, data.dtype)
|
||||
|
||||
def _bilinear(n, h, w, c):
|
||||
outside = tvm.tirx.any(h < 0, w < 0, h >= in_height, w >= in_width)
|
||||
val = bilinear_sample_nhwc(data, (n, h, w, c), in_height - 1, in_width - 1)
|
||||
return tvm.tirx.if_then_else(outside, zero, val)
|
||||
|
||||
data_deform = te.compute(
|
||||
(batch, kernel_h, kernel_w, in_channel, out_height, out_width),
|
||||
lambda n, kh, kw, c, y, x: _bilinear(
|
||||
n,
|
||||
y * stride_h
|
||||
- pad_top
|
||||
+ kh * dilation_h
|
||||
+ offset[
|
||||
n, y, x, c // ic_per_dgroup * (kernel_w * kernel_h * 2) + (kh * kernel_w + kw) * 2
|
||||
],
|
||||
x * stride_w
|
||||
- pad_left
|
||||
+ kw * dilation_w
|
||||
+ offset[
|
||||
n,
|
||||
y,
|
||||
x,
|
||||
c // ic_per_dgroup * (kernel_w * kernel_h * 2) + (kh * kernel_w + kw) * 2 + 1,
|
||||
],
|
||||
c,
|
||||
),
|
||||
tag="data_deform",
|
||||
)
|
||||
return te.compute(
|
||||
(batch, out_height, out_width, out_channel),
|
||||
lambda n, y, x, f: te.sum(
|
||||
data_deform[n, ry, rx, rc, y, x].astype(out_dtype)
|
||||
* kernel[ry, rx, rc, f].astype(out_dtype),
|
||||
axis=[ry, rx, rc],
|
||||
),
|
||||
tag="deformable_conv2d_nhwc",
|
||||
)
|
||||
@@ -0,0 +1,261 @@
|
||||
# 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
|
||||
# ruff: noqa: E741, F821
|
||||
"""TVM operator fully connected compute."""
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
|
||||
from .. import add, tag
|
||||
|
||||
|
||||
def matmul(
|
||||
tensor_a,
|
||||
tensor_b,
|
||||
bias=None,
|
||||
out_dtype=None,
|
||||
transpose_a=False,
|
||||
transpose_b=False,
|
||||
auto_scheduler_rewritten_layout="",
|
||||
meta_schedule_original_shape=None,
|
||||
):
|
||||
"""The default implementation of matmul in topi.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tensor_a : tvm.te.Tensor
|
||||
2-D with shape [batch, in_dim]
|
||||
|
||||
tensor_b : tvm.te.Tensor
|
||||
2-D with shape [out_dim, in_dim]
|
||||
|
||||
bias : Optional[tvm.te.Tensor]
|
||||
1-D with shape [out_dim]
|
||||
|
||||
out_dtype : Optional[str]
|
||||
The output type. This is used for mixed precision.
|
||||
|
||||
transpose_a : Optional[bool] = False
|
||||
Whether the tensor_a is in transposed format.
|
||||
|
||||
transpose_b : Optional[bool] = False
|
||||
Whether the tensor_b is in transposed format.
|
||||
|
||||
auto_scheduler_rewritten_layout: Optional[str] = ""
|
||||
The layout after auto-scheduler's layout rewrite pass.
|
||||
|
||||
meta_schedule_original_shape: Optional[List[Expr]] = None
|
||||
The original shape of the input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
2-D with shape [batch, out_dim]
|
||||
"""
|
||||
# TODO(yixin): support cases for 1-dim input
|
||||
# TODO(yixin): adding support and further check for >2-dim input in autotvm template
|
||||
assert len(tensor_a.shape) >= 2 and len(tensor_b.shape) >= 2, (
|
||||
"1-dim matmul is not supported yet."
|
||||
)
|
||||
|
||||
if bias is not None:
|
||||
assert len(bias.shape) == 1
|
||||
if out_dtype is None:
|
||||
out_dtype = tensor_a.dtype
|
||||
if transpose_a:
|
||||
reduce_dim_a, in_dim = tensor_a.shape[-2:]
|
||||
else:
|
||||
in_dim, reduce_dim_a = tensor_a.shape[-2:]
|
||||
batch_dims_a = tensor_a.shape[:-2]
|
||||
|
||||
if auto_scheduler_rewritten_layout:
|
||||
# Infer shape for the rewritten layout
|
||||
raise RuntimeError("LEGACY-FLOW triggered, to be removed")
|
||||
if meta_schedule_original_shape:
|
||||
raise RuntimeError("LEGACY-FLOW triggered, to be removed")
|
||||
|
||||
if transpose_b:
|
||||
out_dim, reduce_dim_b = tensor_b.shape[-2:]
|
||||
else:
|
||||
reduce_dim_b, out_dim = tensor_b.shape[-2:]
|
||||
batch_dims_b = tensor_b.shape[:-2]
|
||||
|
||||
if not isinstance(reduce_dim_a, tvm.tirx.Var) and not isinstance(reduce_dim_b, tvm.tirx.Var):
|
||||
assert int(reduce_dim_a) == int(reduce_dim_b), (
|
||||
f"Reduction dimensions of dense do not match. {reduce_dim_a} vs {reduce_dim_b}."
|
||||
)
|
||||
|
||||
result_ndim = max(len(batch_dims_a), len(batch_dims_b))
|
||||
batch_dims_a = [1] * (result_ndim - len(batch_dims_a)) + batch_dims_a
|
||||
batch_dims_b = [1] * (result_ndim - len(batch_dims_b)) + batch_dims_b
|
||||
|
||||
for idx, (l, r) in enumerate(zip(batch_dims_a, batch_dims_b)):
|
||||
if (
|
||||
not isinstance(l, tvm.tirx.Var)
|
||||
and not isinstance(r, tvm.tirx.Var)
|
||||
and int(l) != 1
|
||||
and int(r) != 1
|
||||
):
|
||||
assert int(l) == int(r), (
|
||||
"Batch dimensions of dense do not match: "
|
||||
f"{tensor_a.shape[:-2]} vs {tensor_b.shape[:-2]}."
|
||||
)
|
||||
if not isinstance(l, tvm.tirx.Var) and int(l) == 1:
|
||||
batch_dims_a[idx] = batch_dims_b[idx]
|
||||
|
||||
k = te.reduce_axis((0, reduce_dim_a), name="k")
|
||||
|
||||
def compute(*indices):
|
||||
batch_indices_a = indices[-len(tensor_a.shape) : -2]
|
||||
batch_indices_a = [
|
||||
i if isinstance(dim, tvm.tirx.Var) or int(dim) != 1 else 0
|
||||
for i, dim in zip(batch_indices_a, tensor_a.shape[:-2])
|
||||
]
|
||||
batch_indices_b = indices[-len(tensor_b.shape) : -2]
|
||||
batch_indices_b = [
|
||||
i if isinstance(dim, tvm.tirx.Var) or int(dim) != 1 else 0
|
||||
for i, dim in zip(batch_indices_b, tensor_b.shape[:-2])
|
||||
]
|
||||
i, j = indices[-2:]
|
||||
a_indices = (*batch_indices_a, k, i) if transpose_a else (*batch_indices_a, i, k)
|
||||
b_indices = (*batch_indices_b, j, k) if transpose_b else (*batch_indices_b, k, j)
|
||||
return te.sum(
|
||||
tensor_a[a_indices].astype(out_dtype) * tensor_b[b_indices].astype(out_dtype), axis=k
|
||||
)
|
||||
|
||||
compute_name = {
|
||||
(True, True): "T_matmul_TT",
|
||||
(True, False): "T_matmul_TN",
|
||||
(False, True): "T_matmul_NT",
|
||||
(False, False): "T_matmul_NN",
|
||||
}[(transpose_a, transpose_b)]
|
||||
|
||||
# TODO(jcf94): Remove `dense` when `matmul` is finally ready
|
||||
compute_tag = "dense" if (transpose_a, transpose_b) == (False, True) else "matmul"
|
||||
|
||||
mat = te.compute(
|
||||
(*batch_dims_a, in_dim, out_dim),
|
||||
compute,
|
||||
name=compute_name,
|
||||
tag=compute_tag,
|
||||
attrs={"layout_free_placeholders": [tensor_b]},
|
||||
)
|
||||
|
||||
if bias is not None:
|
||||
mat = add(mat, bias.astype(out_dtype))
|
||||
|
||||
if auto_scheduler_rewritten_layout:
|
||||
raise RuntimeError("LEGACY-FLOW triggered, to be removed")
|
||||
|
||||
return mat
|
||||
|
||||
|
||||
def dense(
|
||||
data,
|
||||
weight,
|
||||
bias=None,
|
||||
out_dtype=None,
|
||||
auto_scheduler_rewritten_layout="",
|
||||
meta_schedule_original_shape=None,
|
||||
):
|
||||
"""The default implementation of dense in topi.
|
||||
This is an alias of matmul_nt operator for data tensor in non-transposed format and weight
|
||||
tensor in transposed format.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
2-D with shape [batch, in_dim]
|
||||
|
||||
weight : tvm.te.Tensor
|
||||
2-D with shape [out_dim, in_dim]
|
||||
|
||||
bias : Optional[tvm.te.Tensor]
|
||||
1-D with shape [out_dim]
|
||||
|
||||
out_dtype : Optional[str]
|
||||
The output type. This is used for mixed precision.
|
||||
|
||||
auto_scheduler_rewritten_layout: str = ""
|
||||
The layout after auto-scheduler's layout rewrite pass.
|
||||
|
||||
meta_schedule_original_shape: Optional[List[Expr]] = None
|
||||
The original shape of the input tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
2-D with shape [batch, out_dim]
|
||||
"""
|
||||
|
||||
return matmul(
|
||||
data,
|
||||
weight,
|
||||
bias,
|
||||
out_dtype,
|
||||
False,
|
||||
True,
|
||||
auto_scheduler_rewritten_layout,
|
||||
meta_schedule_original_shape,
|
||||
)
|
||||
|
||||
|
||||
def dense_pack(data, weight, bias=None, out_dtype=None):
|
||||
"""The default implementation of dense_pack in topi.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
2-D with shape [batch, in_dim]
|
||||
|
||||
weight : tvm.te.Tensor
|
||||
2-D with shape [out_dim, in_dim]
|
||||
|
||||
bias : Optional[tvm.te.Tensor]
|
||||
1-D with shape [out_dim]
|
||||
|
||||
out_dtype : Optional[str]
|
||||
The output type. This is used for mixed precision.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
2-D with shape [batch, out_dim]
|
||||
"""
|
||||
if out_dtype is None:
|
||||
out_dtype = data.dtype
|
||||
M, K = get_const_tuple(data.shape) # batch, in_dim
|
||||
N, _, packw_bn = get_const_tuple(weight.shape) # out_dim
|
||||
N = N * packw_bn
|
||||
|
||||
idxdiv = tvm.tirx.indexdiv
|
||||
idxmod = tvm.tirx.indexmod
|
||||
k = te.reduce_axis((0, K), name="k")
|
||||
C = te.compute(
|
||||
(M, N),
|
||||
lambda y, x: te.sum(
|
||||
data[y, k].astype(out_dtype)
|
||||
* weight[idxdiv(x, packw_bn), k, idxmod(x, packw_bn)].astype(out_dtype),
|
||||
axis=k,
|
||||
),
|
||||
name="T_dense_pack",
|
||||
tag="dense_pack",
|
||||
)
|
||||
if bias is not None:
|
||||
C = te.compute((M, N), lambda i, j: C[i, j] + bias[j].astype(out_dtype), tag=tag.BROADCAST)
|
||||
return C
|
||||
@@ -0,0 +1,88 @@
|
||||
# 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
|
||||
"""TVM operator depth_to_space compute."""
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
|
||||
from .. import tag
|
||||
|
||||
|
||||
def depth_to_space(data, block_size, layout="NCHW", mode="DCR"):
|
||||
"""Perform depth to space transformation on the data
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
4-D tensor in either NCHW or NHWC layout.
|
||||
|
||||
block_size : int
|
||||
Size of blocks to compose from channel dimension.
|
||||
|
||||
layout : string
|
||||
Either NCHW or NHWC, indicating data layout.
|
||||
|
||||
mode : string
|
||||
Either DCR or CDR, indicates how channels should be accessed.
|
||||
In DCR, channels are interwoven in the Tensorflow style while
|
||||
in CDR channels are accessed sequentially as in Pytorch.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
Output of shape [N, C / block_size**2, H * block_size, W * block_size]
|
||||
"""
|
||||
if layout == "NCHW":
|
||||
in_n, in_c, in_h, in_w = data.shape
|
||||
channel_factor = tvm.tirx.truncdiv(in_c, (block_size * block_size))
|
||||
output_shape = [in_n, channel_factor, in_h * block_size, in_w * block_size]
|
||||
elif layout == "NHWC":
|
||||
in_n, in_h, in_w, in_c = data.shape
|
||||
channel_factor = tvm.tirx.truncdiv(in_c, (block_size * block_size))
|
||||
output_shape = [in_n, in_h * block_size, in_w * block_size, channel_factor]
|
||||
else:
|
||||
raise ValueError("Only NCHW and NHWC layouts are currently supported.")
|
||||
|
||||
def _get_indices(*indices):
|
||||
if layout == "NCHW":
|
||||
n, c, y, x = indices
|
||||
elif layout == "NHWC":
|
||||
n, y, x, c = indices
|
||||
return n, c, y, x
|
||||
|
||||
def _get_pixel(n, c, y, x):
|
||||
block_x = tvm.tirx.truncdiv(x, block_size)
|
||||
block_y = tvm.tirx.truncdiv(y, block_size)
|
||||
idx_x = tvm.tirx.truncmod(x, block_size)
|
||||
idx_y = tvm.tirx.truncmod(y, block_size)
|
||||
if mode == "DCR":
|
||||
channel_idx = channel_factor * ((block_size * idx_y) + idx_x) + c
|
||||
else:
|
||||
channel_idx = (c * block_size * block_size) + ((block_size * idx_y) + idx_x)
|
||||
|
||||
if layout == "NCHW":
|
||||
output = data(n, channel_idx, block_y, block_x)
|
||||
else:
|
||||
output = data(n, block_y, block_x, channel_idx)
|
||||
return output
|
||||
|
||||
def _compute(*indices):
|
||||
n, c, y, x = _get_indices(*indices)
|
||||
return _get_pixel(n, c, y, x)
|
||||
|
||||
return te.compute(output_shape, _compute, name="depth_to_space", tag=tag.INJECTIVE)
|
||||
@@ -0,0 +1,462 @@
|
||||
# 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-variable, too-many-locals, unused-argument
|
||||
# ruff: noqa: F841
|
||||
"""Depthwise convolution operators"""
|
||||
|
||||
from collections import namedtuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
|
||||
from ..utils import get_const_tuple, simplify
|
||||
from .dilate import dilate
|
||||
from .pad import pad
|
||||
from .utils import get_pad_tuple
|
||||
|
||||
# workload description of depthwise-conv2d
|
||||
Workload = namedtuple(
|
||||
"Workload",
|
||||
[
|
||||
"in_dtype",
|
||||
"out_dtype",
|
||||
"height",
|
||||
"width",
|
||||
"in_filter",
|
||||
"out_filter",
|
||||
"kernel_h",
|
||||
"kernel_w",
|
||||
"padt",
|
||||
"padl",
|
||||
"padb",
|
||||
"padr",
|
||||
"dilation_h",
|
||||
"dilation_w",
|
||||
"stride_h",
|
||||
"stride_w",
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _get_workload(data, kernel, stride, padding, dilation, out_dtype, data_layout="NCHW"):
|
||||
"""Get the workload structure for a depthwise conv2d.
|
||||
|
||||
Input data and filter should use NCHW layout.
|
||||
"""
|
||||
if data_layout == "NCHW":
|
||||
_, in_channel, height, width = get_const_tuple(data.shape)
|
||||
filter_channel, channel_multiplier, kh, kw = get_const_tuple(kernel.shape)
|
||||
elif data_layout == "NHWC":
|
||||
_, height, width, in_channel = get_const_tuple(data.shape)
|
||||
kh, kw, filter_channel, channel_multiplier = get_const_tuple(kernel.shape)
|
||||
elif data_layout == "NCHWc":
|
||||
_, in_channel_chunk, height, width, in_channel_block = get_const_tuple(data.shape)
|
||||
in_channel = in_channel_chunk * in_channel_block
|
||||
(filter_channel_chunk, cm_chunk, kh, kw, cm_block, filter_channel_block) = get_const_tuple(
|
||||
kernel.shape
|
||||
)
|
||||
filter_channel = filter_channel_chunk * filter_channel_block
|
||||
channel_multiplier = cm_chunk * cm_block
|
||||
|
||||
assert in_channel_block == filter_channel_block, (
|
||||
f"Incorrect dimensions, data has block size {in_channel_block}, but filter has "
|
||||
f"block size {filter_channel_block}"
|
||||
)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Data layout {data_layout} not supported")
|
||||
|
||||
assert in_channel == filter_channel, (
|
||||
f"Incorrect dimensions, data has {in_channel} channels but filter expects "
|
||||
f"{filter_channel} channels"
|
||||
)
|
||||
|
||||
out_channel = filter_channel * channel_multiplier
|
||||
dilation_h, dilation_w = (
|
||||
dilation if isinstance(dilation, tuple | list) else (dilation, dilation)
|
||||
)
|
||||
if isinstance(stride, tuple | list):
|
||||
HSTR, WSTR = stride
|
||||
else:
|
||||
HSTR, WSTR = stride, stride
|
||||
assert (data.dtype == kernel.dtype) or (data.dtype == "uint8" and kernel.dtype == "int8"), (
|
||||
f"Do not support inputs with different data types now. {data.dtype} vs. {kernel.dtype}"
|
||||
)
|
||||
dilated_kernel_h = (kh - 1) * dilation_h + 1
|
||||
dilated_kernel_w = (kw - 1) * dilation_w + 1
|
||||
pt, pl, pb, pr = get_pad_tuple(padding, (dilated_kernel_h, dilated_kernel_w))
|
||||
return Workload(
|
||||
data.dtype,
|
||||
out_dtype,
|
||||
height,
|
||||
width,
|
||||
in_channel,
|
||||
out_channel,
|
||||
kh,
|
||||
kw,
|
||||
pt,
|
||||
pl,
|
||||
pb,
|
||||
pr,
|
||||
dilation_h,
|
||||
dilation_w,
|
||||
HSTR,
|
||||
WSTR,
|
||||
)
|
||||
|
||||
|
||||
def depthwise_conv2d_nchw(Input, Filter, stride, padding, dilation, out_dtype=None):
|
||||
"""Depthwise convolution nchw forward operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
Input : tvm.te.Tensor
|
||||
4-D with shape [batch, in_channel, in_height, in_width]
|
||||
|
||||
Filter : tvm.te.Tensor
|
||||
4-D with shape [in_channel, channel_multiplier, filter_height, filter_width]
|
||||
|
||||
stride : int or a list/tuple of two ints
|
||||
The spatial stride, or (stride_height, stride_width).
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
dilation: int or a list/tuple of two ints
|
||||
dilation size, or [dilation_height, dilation_width]
|
||||
|
||||
out_dtype: str, optional
|
||||
Output data type
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
4-D with shape [batch, out_channel, out_height, out_width]
|
||||
"""
|
||||
out_dtype = Input.dtype if out_dtype is None else out_dtype
|
||||
|
||||
if isinstance(stride, int):
|
||||
stride_h = stride_w = stride
|
||||
else:
|
||||
stride_h, stride_w = stride
|
||||
|
||||
if isinstance(dilation, int):
|
||||
dilation_h = dilation_w = dilation
|
||||
else:
|
||||
dilation_h, dilation_w = dilation
|
||||
|
||||
batch, in_channel, in_height, in_width = Input.shape
|
||||
# shape of dilated kernel
|
||||
filter_channel, channel_multiplier, filter_height, filter_width = Filter.shape
|
||||
|
||||
dilated_kernel_h = (filter_height - 1) * dilation_h + 1
|
||||
dilated_kernel_w = (filter_width - 1) * dilation_w + 1
|
||||
pad_top, pad_left, pad_down, pad_right = get_pad_tuple(
|
||||
padding, (dilated_kernel_h, dilated_kernel_w)
|
||||
)
|
||||
out_channel = simplify(in_channel * channel_multiplier)
|
||||
out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1)
|
||||
out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1)
|
||||
|
||||
# padding stage
|
||||
pad_before = [0, 0, pad_top, pad_left]
|
||||
pad_after = [0, 0, pad_down, pad_right]
|
||||
PaddedInput = pad(Input, pad_before, pad_after, name="PaddedInput")
|
||||
# depthconv stage
|
||||
idxdiv = tvm.tirx.indexdiv
|
||||
idxmod = tvm.tirx.indexmod
|
||||
di = te.reduce_axis((0, filter_height), name="di")
|
||||
dj = te.reduce_axis((0, filter_width), name="dj")
|
||||
Output = te.compute(
|
||||
(batch, out_channel, out_height, out_width),
|
||||
lambda b, c, i, j: te.sum(
|
||||
(
|
||||
PaddedInput[
|
||||
b,
|
||||
idxdiv(c, channel_multiplier),
|
||||
i * stride_h + di * dilation_h,
|
||||
j * stride_w + dj * dilation_w,
|
||||
].astype(out_dtype)
|
||||
* Filter[
|
||||
idxdiv(c, channel_multiplier), idxmod(c, channel_multiplier), di, dj
|
||||
].astype(out_dtype)
|
||||
),
|
||||
axis=[di, dj],
|
||||
),
|
||||
name="DepthwiseConv2d",
|
||||
tag="depthwise_conv2d_nchw",
|
||||
)
|
||||
return Output
|
||||
|
||||
|
||||
def depthwise_conv2d_nhwc(
|
||||
Input, Filter, stride, padding, dilation, kernel_layout="HWOI", out_dtype=None
|
||||
):
|
||||
"""Depthwise convolution nhwc forward operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
Input : tvm.te.Tensor
|
||||
4-D with shape [batch, in_height, in_width, in_channel]
|
||||
|
||||
Filter : tvm.te.Tensor
|
||||
4-D with shape [filter_height, filter_width, in_channel, channel_multiplier]
|
||||
|
||||
stride : tuple of two ints
|
||||
The spatial stride along height and width
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
dilation: int or a list/tuple of two ints
|
||||
dilation size, or [dilation_height, dilation_width]
|
||||
|
||||
out_dtype: str, optional
|
||||
Output data type
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
4-D with shape [batch, out_height, out_width, out_channel]
|
||||
"""
|
||||
out_dtype = Input.dtype if out_dtype is None else out_dtype
|
||||
|
||||
if isinstance(stride, int):
|
||||
stride_h = stride_w = stride
|
||||
else:
|
||||
stride_h, stride_w = stride
|
||||
|
||||
if isinstance(dilation, int):
|
||||
dilation_h = dilation_w = dilation
|
||||
else:
|
||||
dilation_h, dilation_w = dilation
|
||||
|
||||
batch, in_height, in_width, in_channel = Input.shape
|
||||
|
||||
# shape of dilated kernel
|
||||
if kernel_layout == "HWIO":
|
||||
filter_height, filter_width, channel_multiplier, filter_channel = Filter.shape
|
||||
kernel_permutation = [0, 1, 3, 2]
|
||||
else:
|
||||
filter_height, filter_width, filter_channel, channel_multiplier = Filter.shape
|
||||
kernel_permutation = [0, 1, 2, 3]
|
||||
|
||||
dilated_kernel_h = (filter_height - 1) * dilation_h + 1
|
||||
dilated_kernel_w = (filter_width - 1) * dilation_w + 1
|
||||
pad_top, pad_left, pad_down, pad_right = get_pad_tuple(
|
||||
padding, (dilated_kernel_h, dilated_kernel_w)
|
||||
)
|
||||
out_channel = simplify(in_channel * channel_multiplier)
|
||||
out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1)
|
||||
out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1)
|
||||
|
||||
# padding stage
|
||||
pad_before = [0, pad_top, pad_left, 0]
|
||||
pad_after = [0, pad_down, pad_right, 0]
|
||||
PaddedInput = pad(Input, pad_before, pad_after, name="PaddedInput")
|
||||
# depthconv stage
|
||||
idxdiv = tvm.tirx.indexdiv
|
||||
idxmod = tvm.tirx.indexmod
|
||||
|
||||
di = te.reduce_axis((0, filter_height), name="di")
|
||||
dj = te.reduce_axis((0, filter_width), name="dj")
|
||||
Output = te.compute(
|
||||
(batch, out_height, out_width, out_channel),
|
||||
lambda b, i, j, c: te.sum(
|
||||
(
|
||||
PaddedInput[
|
||||
b,
|
||||
i * stride_h + di * dilation_h,
|
||||
j * stride_w + dj * dilation_w,
|
||||
idxdiv(c, channel_multiplier),
|
||||
].astype(out_dtype)
|
||||
* Filter[
|
||||
tuple(
|
||||
np.array(
|
||||
[di, dj, idxdiv(c, channel_multiplier), idxmod(c, channel_multiplier)]
|
||||
)[kernel_permutation]
|
||||
)
|
||||
].astype(out_dtype)
|
||||
),
|
||||
axis=[di, dj],
|
||||
),
|
||||
name="DepthwiseConv2d",
|
||||
tag="depthwise_conv2d_nhwc",
|
||||
)
|
||||
return Output
|
||||
|
||||
|
||||
def depthwise_conv2d_backward_input_nhwc(Filter, Out_grad, oshape, ishape, stride, padding):
|
||||
"""Depthwise convolution nhwc backward wrt input operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
Filter : tvm.te.Tensor
|
||||
4-D with shape [filter_height, filter_width, in_channel, channel_multiplier]
|
||||
|
||||
Out_grad : tvm.te.Tensor
|
||||
4-D with shape [batch, out_height, out_width, out_channel]
|
||||
|
||||
stride : tuple of two ints
|
||||
The spatial stride along height and width
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
4-D with shape [batch, in_height, in_width, in_channel]
|
||||
"""
|
||||
batch, in_h, in_w, in_c = ishape
|
||||
_, out_h, out_w, out_c = oshape
|
||||
filter_h, filter_w, _, channel_multiplier = Filter.shape
|
||||
if isinstance(stride, int):
|
||||
stride_h = stride_w = stride
|
||||
else:
|
||||
stride_h, stride_w = stride
|
||||
|
||||
dilated_out_grad = dilate(Out_grad, [1, stride_h, stride_w, 1], name="dilated_out_grad")
|
||||
|
||||
# padding params in forward propagation
|
||||
fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple(padding, (filter_h, filter_w))
|
||||
# padding params in backward propagation
|
||||
bpad_top = filter_h - 1 - fpad_top
|
||||
bpad_bottom = (filter_h - 1 - fpad_bottom) + (stride_h - 1)
|
||||
bpad_left = filter_w - 1 - fpad_left
|
||||
bpad_right = (filter_w - 1 - fpad_right) + (stride_w - 1)
|
||||
|
||||
padded_out_grad = pad(
|
||||
dilated_out_grad,
|
||||
[0, bpad_top, bpad_left, 0],
|
||||
[0, bpad_bottom, bpad_right, 0],
|
||||
name="padded_out_grad",
|
||||
)
|
||||
|
||||
dh = te.reduce_axis((0, filter_h), name="dh")
|
||||
dw = te.reduce_axis((0, filter_w), name="dw")
|
||||
dc = te.reduce_axis((0, channel_multiplier), name="dc")
|
||||
|
||||
In_grad = te.compute(
|
||||
(batch, in_h, in_w, in_c),
|
||||
lambda b, h, w, c: te.sum(
|
||||
padded_out_grad[b, h + dh, w + dw, c * channel_multiplier + dc]
|
||||
* Filter[filter_h - 1 - dh, filter_w - 1 - dw, c, dc],
|
||||
axis=[dh, dw, dc],
|
||||
),
|
||||
tag="depthwise_conv2d_backward_input_nhwc",
|
||||
)
|
||||
|
||||
return In_grad
|
||||
|
||||
|
||||
def depthwise_conv2d_backward_weight_nhwc(Input, Out_grad, oshape, fshape, stride, padding):
|
||||
"""Depthwise convolution nhwc backward wrt weight operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
Input : tvm.te.Tensor
|
||||
4-D with shape [batch, in_height, in_width, in_channel]
|
||||
|
||||
Out_grad : tvm.te.Tensor
|
||||
4-D with shape [batch, out_height, out_width, out_channel]
|
||||
|
||||
stride : tuple of two ints
|
||||
The spatial stride along height and width
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
4-D with shape [filter_height, filter_width, in_channel, channel_multiplier]
|
||||
"""
|
||||
batch, out_h, out_w, out_c = oshape
|
||||
filter_h, filter_w, _, channel_multiplier = fshape
|
||||
in_c = Input.shape[3].value
|
||||
if isinstance(stride, int):
|
||||
stride_h = stride_w = stride
|
||||
else:
|
||||
stride_h, stride_w = stride
|
||||
|
||||
pad_top, pad_left, pad_bottom, pad_right = get_pad_tuple(padding, (filter_h, filter_w))
|
||||
|
||||
padded_in = pad(
|
||||
Input, [0, pad_top, pad_left, 0], [0, pad_bottom, pad_right, 0], name="padded_in"
|
||||
)
|
||||
|
||||
dh = te.reduce_axis((0, Out_grad.shape[1].value), name="dh")
|
||||
dw = te.reduce_axis((0, Out_grad.shape[2].value), name="dw")
|
||||
db = te.reduce_axis((0, batch), name="db")
|
||||
idxdiv = tvm.tirx.indexdiv
|
||||
idxmod = tvm.tirx.indexmod
|
||||
|
||||
Weight_grad = te.compute(
|
||||
(filter_h, filter_w, in_c, channel_multiplier),
|
||||
lambda fh, fw, c, m: te.sum(
|
||||
Out_grad[db, dh, dw, c * channel_multiplier + idxmod(m, channel_multiplier)]
|
||||
* padded_in[db, fh + dh * stride_h, fw + dw * stride_w, c],
|
||||
axis=[db, dh, dw],
|
||||
),
|
||||
tag="depthwise_conv2d_backward_weight_nhwc",
|
||||
)
|
||||
|
||||
return Weight_grad
|
||||
|
||||
|
||||
def depthwise_conv2d_NCHWc(
|
||||
Input, Filter, stride, padding, dilation, layout, out_layout, out_dtype=None
|
||||
):
|
||||
"""Depthwise convolution NCHW[x]c forward operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
Input : tvm.te.Tensor
|
||||
5-D with shape [batch, in_channel_chunk, in_height, in_width, in_channel_block]
|
||||
|
||||
Filter : tvm.te.Tensor
|
||||
6-D with shape [out_channel_chunk, 1, filter_height, filter_width, 1, out_channel_block]
|
||||
In NCHWc depthwise convolution,
|
||||
we group kernel's in_channel and channel_multiplier together then do the tiling.
|
||||
|
||||
stride : tuple of two ints
|
||||
The spatial stride along height and width
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
dilation: int or a list/tuple of two ints
|
||||
dilation size, or [dilation_height, dilation_width]
|
||||
|
||||
layout : str
|
||||
Input data layout
|
||||
|
||||
out_layout : str
|
||||
Output data layout
|
||||
|
||||
out_dtype: str, optional
|
||||
Output data type
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
5-D with shape [batch, out_channel_chunk, out_height, out_width, out_channel_block]
|
||||
"""
|
||||
raise ValueError("missing register for topi.nn.depthwise_conv2d_NCHWc")
|
||||
@@ -0,0 +1,73 @@
|
||||
# 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
|
||||
"""Dilation operators"""
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
|
||||
from .. import tag, utils
|
||||
|
||||
|
||||
@te.tag_scope(tag=tag.INJECTIVE + ",dilate")
|
||||
def dilate(data, strides, dilation_value=0.0, name="DilatedInput"):
|
||||
"""Dilate data with given dilation value (0 by default).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
n-D, can be any layout.
|
||||
|
||||
strides : list / tuple of n ints
|
||||
Dilation stride on each dimension, 1 means no dilation.
|
||||
|
||||
dilation_value : int/float, optional
|
||||
Value used to dilate the input.
|
||||
|
||||
name : str, optional
|
||||
The name prefix operators generated
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
n-D, the same layout as data.
|
||||
"""
|
||||
n = len(data.shape)
|
||||
if len(strides) != n:
|
||||
raise ValueError(f"data dimension and strides size dismatch : {n} vs {len(strides)}")
|
||||
ana = tvm.arith.Analyzer()
|
||||
out_shape = tuple(ana.simplify((data.shape[i] - 1) * strides[i] + 1) for i in range(n))
|
||||
|
||||
def _dilate(*indices):
|
||||
not_zero = []
|
||||
index_tuple = []
|
||||
idxdiv = tvm.tirx.indexdiv
|
||||
idxmod = tvm.tirx.indexmod
|
||||
for i in range(n):
|
||||
if not utils.equal_const_int(strides[i], 1):
|
||||
index_tuple.append(idxdiv(indices[i], strides[i]))
|
||||
not_zero.append(idxmod(indices[i], strides[i]).equal(0))
|
||||
else:
|
||||
index_tuple.append(indices[i])
|
||||
if not_zero:
|
||||
not_zero = tvm.tirx.all(*not_zero)
|
||||
return tvm.tirx.if_then_else(
|
||||
not_zero, data(*index_tuple), tvm.tirx.const(dilation_value, data.dtype)
|
||||
)
|
||||
return data(*index_tuple)
|
||||
|
||||
return te.compute(out_shape, _dilate, name=name)
|
||||
@@ -0,0 +1,143 @@
|
||||
# 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.
|
||||
"""Elementwise operators"""
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
|
||||
from .. import tag
|
||||
from ..utils import get_const_int
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
||||
def relu(x):
|
||||
"""Take relu of input x.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : tvm.te.Tensor
|
||||
Input argument.
|
||||
|
||||
Returns
|
||||
-------
|
||||
y : tvm.te.Tensor
|
||||
The result.
|
||||
"""
|
||||
return te.compute(x.shape, lambda *i: tvm.te.max(x(*i), tvm.tirx.const(0, x.dtype)))
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
||||
def leaky_relu(x, alpha):
|
||||
"""Take leaky relu of input x.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : tvm.te.Tensor
|
||||
Input argument.
|
||||
|
||||
alpha : float
|
||||
The slope for the small gradient when x < 0
|
||||
|
||||
Returns
|
||||
-------
|
||||
y : tvm.te.Tensor
|
||||
The result.
|
||||
"""
|
||||
|
||||
def _compute(*indices):
|
||||
value = x(*indices)
|
||||
calpha = tvm.tirx.const(alpha, value.ty)
|
||||
return tvm.tirx.Select(value > 0, value, value * calpha)
|
||||
|
||||
return te.compute(x.shape, _compute)
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag=tag.ELEMWISE)
|
||||
def softplus(x, beta=1.0, threshold=20.0):
|
||||
"""Compute Softplus activation for input x with numerical stability.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : tvm.te.Tensor
|
||||
Input tensor.
|
||||
|
||||
beta : float, optional
|
||||
The scaling factor β in the Softplus formula (default is 1.0).
|
||||
|
||||
threshold : float, optional
|
||||
The threshold value for numerical stability (default is 20.0).
|
||||
|
||||
Returns
|
||||
-------
|
||||
y : tvm.te.Tensor
|
||||
The result.
|
||||
"""
|
||||
|
||||
def _compute(*indices):
|
||||
value = x(*indices)
|
||||
b = tvm.tirx.const(beta, value.ty)
|
||||
t = tvm.tirx.const(threshold, value.ty)
|
||||
|
||||
return tvm.tirx.Select(
|
||||
b * value > t, value, (1 / b) * tvm.tirx.log(1 + tvm.tirx.exp(b * value))
|
||||
)
|
||||
|
||||
return te.compute(x.shape, _compute)
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag=tag.BROADCAST)
|
||||
def prelu(x, slope, axis=1):
|
||||
"""PReLU.
|
||||
It accepts two arguments: an input ``x`` and a weight array ``W``
|
||||
and computes the output as :math:`PReLU(x) y = x > 0 ? x : W * x`,
|
||||
where :math:`*` is an elementwise multiplication for each sample in the
|
||||
batch.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : tvm.te.Tensor
|
||||
Input argument.
|
||||
|
||||
slope : tvm.te.Tensor
|
||||
Channelised slope tensor for prelu
|
||||
|
||||
axis : int
|
||||
The axis where the channel data needs to be applied
|
||||
|
||||
Returns
|
||||
-------
|
||||
y : tvm.te.Tensor
|
||||
The result.
|
||||
|
||||
Links
|
||||
-----
|
||||
[http://arxiv.org/pdf/1502.01852v1.pdf]
|
||||
"""
|
||||
|
||||
assert len(slope.shape) == 1
|
||||
assert axis < len(x.shape)
|
||||
if slope.shape[0] == 1:
|
||||
slope = te.compute(
|
||||
(get_const_int(x.shape[axis]),), lambda c: slope[0], name="slope_broadcasted"
|
||||
)
|
||||
assert get_const_int(slope.shape[0]) == get_const_int(x.shape[axis])
|
||||
|
||||
def _compute_channelwise(*indices):
|
||||
xval = x(*indices)
|
||||
return tvm.tirx.Select(xval > 0, xval, xval * slope(indices[axis]))
|
||||
|
||||
return te.compute(x.shape, _compute_channelwise)
|
||||
@@ -0,0 +1,188 @@
|
||||
# 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: E741
|
||||
|
||||
"""FIFO buffer op"""
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
|
||||
from .. import tag
|
||||
from ..transform import concatenate, strided_slice
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag=tag.INJECTIVE + ",fifo_buffer")
|
||||
def fifo_buffer(data, buffer, axis):
|
||||
"""
|
||||
FIFO buffer to enable computation reuse in CNNs with sliding indow input
|
||||
|
||||
Compute equivalent of
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
concat(buffer, data, axis=axis)
|
||||
.slice_axis(axis=axis,
|
||||
begin=data.shape[axis],
|
||||
end=data.shape[axis]+buffer.shape[axis])
|
||||
|
||||
Useful for
|
||||
|
||||
* Encoding explicit re-use of computation in convolution ops operated on a sliding window input
|
||||
* Implementing a FIFO queue to cache intermediate results, e.g. as in Fast WaveNet.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
The input data
|
||||
buffer : tvm.te.Tensor
|
||||
Previous value of the FIFO buffer
|
||||
axis : int
|
||||
Specify which axis should be used for buffering
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : tvm.te.Tensor
|
||||
Updated value for the buffer
|
||||
"""
|
||||
assert len(data.shape) == len(buffer.shape), (
|
||||
f"buffer and data must have same number of dimensions, "
|
||||
f"buffer.shape = {buffer.shape}, data.shape = {data.shape}"
|
||||
)
|
||||
assert len(buffer.shape) >= 1, "Zero-dimension tensor not supported"
|
||||
assert 0 <= axis < len(buffer.shape), "buffer axis out of range"
|
||||
for i in range(len(data.shape)):
|
||||
if i == axis:
|
||||
assert int(str(data.shape[i])) <= int(str(buffer.shape[i]))
|
||||
else:
|
||||
assert int(str(data.shape[i])) == int(str(buffer.shape[i]))
|
||||
|
||||
buflen = buffer.shape[axis]
|
||||
data_size = data.shape[axis]
|
||||
|
||||
# Explicitly write out formula up to 4D, and then use concat+slice combo for 5D and higher
|
||||
if len(buffer.shape) == 1:
|
||||
return te.compute(
|
||||
buffer.shape,
|
||||
lambda i: tvm.tirx.if_then_else(
|
||||
i < buflen - data_size, buffer[i + data_size], data[i - buflen + data_size]
|
||||
),
|
||||
name="new_buffer",
|
||||
)
|
||||
if len(buffer.shape) == 2:
|
||||
if axis == 0:
|
||||
return te.compute(
|
||||
buffer.shape,
|
||||
lambda i, j: tvm.tirx.if_then_else(
|
||||
i < buflen - data_size,
|
||||
buffer[i + data_size, j],
|
||||
data[i - buflen + data_size, j],
|
||||
),
|
||||
name="new_buffer",
|
||||
)
|
||||
if axis == 1:
|
||||
return te.compute(
|
||||
buffer.shape,
|
||||
lambda i, j: tvm.tirx.if_then_else(
|
||||
j < buflen - data_size,
|
||||
buffer[i, j + data_size],
|
||||
data[i, j - buflen + data_size],
|
||||
),
|
||||
name="new_buffer",
|
||||
)
|
||||
assert False, f"Invalid value for axis; it should be at most {len(buffer.shape)}"
|
||||
elif len(buffer.shape) == 3:
|
||||
if axis == 0:
|
||||
return te.compute(
|
||||
buffer.shape,
|
||||
lambda i, j, k: tvm.tirx.if_then_else(
|
||||
i < buflen - data_size,
|
||||
buffer[i + data_size, j, k],
|
||||
data[i - buflen + data_size, j, k],
|
||||
),
|
||||
name="new_buffer",
|
||||
)
|
||||
if axis == 1:
|
||||
return te.compute(
|
||||
buffer.shape,
|
||||
lambda i, j, k: tvm.tirx.if_then_else(
|
||||
j < buflen - data_size,
|
||||
buffer[i, j + data_size, k],
|
||||
data[i, j - buflen + data_size, k],
|
||||
),
|
||||
name="new_buffer",
|
||||
)
|
||||
if axis == 2:
|
||||
return te.compute(
|
||||
buffer.shape,
|
||||
lambda i, j, k: tvm.tirx.if_then_else(
|
||||
k < buflen - data_size,
|
||||
buffer[i, j, k + data_size],
|
||||
data[i, j, k - buflen + data_size],
|
||||
),
|
||||
name="new_buffer",
|
||||
)
|
||||
assert False, f"Invalid value for axis; it should be at most {len(buffer.shape)}"
|
||||
elif len(buffer.shape) == 4:
|
||||
if axis == 0:
|
||||
return te.compute(
|
||||
buffer.shape,
|
||||
lambda i, j, k, l: tvm.tirx.if_then_else(
|
||||
i < buflen - data_size,
|
||||
buffer[i + data_size, j, k, l],
|
||||
data[i - buflen + data_size, j, k, l],
|
||||
),
|
||||
name="new_buffer",
|
||||
)
|
||||
if axis == 1:
|
||||
return te.compute(
|
||||
buffer.shape,
|
||||
lambda i, j, k, l: tvm.tirx.if_then_else(
|
||||
j < buflen - data_size,
|
||||
buffer[i, j + data_size, k, l],
|
||||
data[i, j - buflen + data_size, k, l],
|
||||
),
|
||||
name="new_buffer",
|
||||
)
|
||||
if axis == 2:
|
||||
return te.compute(
|
||||
buffer.shape,
|
||||
lambda i, j, k, l: tvm.tirx.if_then_else(
|
||||
k < buflen - data_size,
|
||||
buffer[i, j, k + data_size, l],
|
||||
data[i, j, k - buflen + data_size, l],
|
||||
),
|
||||
name="new_buffer",
|
||||
)
|
||||
if axis == 3:
|
||||
return te.compute(
|
||||
buffer.shape,
|
||||
lambda i, j, k, l: tvm.tirx.if_then_else(
|
||||
l < buflen - data_size,
|
||||
buffer[i, j, k, l + data_size],
|
||||
data[i, j, k, l - buflen + data_size],
|
||||
),
|
||||
name="new_buffer",
|
||||
)
|
||||
assert False, f"Invalid value for axis; it should be at most {len(buffer.shape)}"
|
||||
else:
|
||||
# Implement FIFO buffer as combination of concat and slice
|
||||
begin = [0] * len(buffer.shape)
|
||||
begin[axis] = data.shape[axis]
|
||||
end = list(buffer.shape[:])
|
||||
end[axis] += data.shape[axis]
|
||||
return strided_slice(concatenate((buffer, data), axis=axis), begin=begin, end=end)
|
||||
return None
|
||||
@@ -0,0 +1,54 @@
|
||||
# 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.
|
||||
"""TVM operator flatten compute."""
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
|
||||
from .. import tag
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag=tag.INJECTIVE)
|
||||
def flatten(data):
|
||||
"""Flattens the input array into a 2-D array by collapsing the higher dimensions.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
Input array.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
2-D array with collapsed higher dimensions.
|
||||
"""
|
||||
ishape = data.shape
|
||||
dim = 1
|
||||
for i in range(1, len(ishape)):
|
||||
dim = dim * ishape[i]
|
||||
oshape = [ishape[0], dim]
|
||||
idxdiv = tvm.tirx.indexdiv
|
||||
idxmod = tvm.tirx.indexmod
|
||||
|
||||
def unwrap(idx, shape):
|
||||
index = []
|
||||
for s in reversed(shape):
|
||||
index.append(idxmod(idx, s))
|
||||
idx = idxdiv(idx, s)
|
||||
return list(reversed(index))
|
||||
|
||||
return te.compute(oshape, lambda i, j: data(i, *unwrap(j, ishape[1:])))
|
||||
@@ -0,0 +1,55 @@
|
||||
# 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.
|
||||
"""Layer normalization operator."""
|
||||
|
||||
from .. import cpp
|
||||
|
||||
|
||||
def group_norm(data, gamma, beta, num_groups, channel_axis, axes, epsilon=1e-5):
|
||||
"""Group normalization operator.
|
||||
It accepts fp16 and fp32 as input data type. It will cast the input to fp32
|
||||
to perform the computation. The output will have the same data type as input.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
N-D with shape (d_0, d_1, ..., d_{N-1})
|
||||
|
||||
gamma: tvm.te.Tensor
|
||||
1-D with shape (r_0) where r_0 == d_{channel_axis}
|
||||
|
||||
beta: tvm.te.Tensor
|
||||
Optional, 1-D with shape (r_0) where r_0 == d_{channel_axis}
|
||||
|
||||
num_groups : int
|
||||
The number of groups
|
||||
|
||||
channel_axis : int
|
||||
The channel axis
|
||||
|
||||
axes : list of int
|
||||
Axis over the normalization applied, excluding the channel axis
|
||||
|
||||
epsilon : float
|
||||
The epsilon value to avoid division by zero.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : tvm.te.Tensor
|
||||
N-D with shape (d_0, d_1, ..., d_{N-1})
|
||||
"""
|
||||
return cpp.nn.group_norm(data, gamma, beta, num_groups, channel_axis, axes, epsilon)
|
||||
@@ -0,0 +1,48 @@
|
||||
# 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.
|
||||
"""Instance normalization operator."""
|
||||
|
||||
from .. import cpp
|
||||
|
||||
|
||||
def instance_norm(data, gamma, beta, channel_axis, axis, epsilon=1e-5):
|
||||
"""Instance normalization operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
N-D with shape (d_0, d_1, ..., d_{N-1})
|
||||
|
||||
gamma: tvm.te.Tensor
|
||||
K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
|
||||
|
||||
beta: tvm.te.Tensor
|
||||
Optional, K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
|
||||
|
||||
axis : list of int
|
||||
Axis over the normalization applied (the axis along which the mean and variance are
|
||||
computed)
|
||||
|
||||
epsilon : float
|
||||
The epsilon value to avoid division by zero.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : tvm.te.Tensor
|
||||
N-D with shape (d_0, d_1, ..., d_{N-1})
|
||||
"""
|
||||
return cpp.nn.instance_norm(data, gamma, beta, channel_axis, axis, epsilon)
|
||||
@@ -0,0 +1,49 @@
|
||||
# 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.
|
||||
"""Layer normalization operator."""
|
||||
|
||||
from .. import cpp
|
||||
|
||||
|
||||
def layer_norm(data, gamma, beta, axis, epsilon=1e-5):
|
||||
"""Layer normalization operator.
|
||||
It accepts fp16 and fp32 as input data type. It will cast the input to fp32
|
||||
to perform the computation. The output will have the same data type as input.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
N-D with shape (d_0, d_1, ..., d_{N-1})
|
||||
|
||||
gamma: tvm.te.Tensor
|
||||
K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
|
||||
|
||||
beta: tvm.te.Tensor
|
||||
Optional, K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
|
||||
|
||||
axis : list of int
|
||||
Axis over the normalization applied
|
||||
|
||||
epsilon : float
|
||||
The epsilon value to avoid division by zero.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : tvm.te.Tensor
|
||||
N-D with shape (d_0, d_1, ..., d_{N-1})
|
||||
"""
|
||||
return cpp.nn.layer_norm(data, gamma, beta, axis, epsilon)
|
||||
@@ -0,0 +1,59 @@
|
||||
# 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
|
||||
"""TVM operator for local response norm compute."""
|
||||
|
||||
from .. import cpp
|
||||
|
||||
|
||||
def lrn(data, size, axis=1, alpha=0.0001, beta=0.75, bias=2):
|
||||
"""Perform the across channels local response normalisation
|
||||
on the input data.
|
||||
|
||||
sum_sqr_up^i{x, y} = (bias+((alpha/size)* \
|
||||
{sum_{j=max(0, i-size/2)}^{min(N-1,i+size/2)} \
|
||||
(data^j{x,y})^2}))^beta
|
||||
output^i{x, y} = data^i{x, y}/sum_sqr_up^i{x, y}
|
||||
N is the number for input channels
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
4-D with shape [batch, channel, height, width]
|
||||
|
||||
size : int
|
||||
normalisation window size
|
||||
|
||||
axis : int
|
||||
input data layout channel axis
|
||||
default value is 1 for NCHW format
|
||||
|
||||
bias : float
|
||||
offset to avoid dividing by 0
|
||||
|
||||
alpha : float
|
||||
to be divided
|
||||
|
||||
beta : float
|
||||
exponent
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
4-D output with same shape
|
||||
"""
|
||||
return cpp.nn.lrn(data, size, axis, alpha, beta, bias)
|
||||
@@ -0,0 +1,60 @@
|
||||
# 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
|
||||
"""Loss functions definitions."""
|
||||
|
||||
from . import cpp
|
||||
|
||||
|
||||
def nll_loss(predictions, targets, weights, reduction, ignore_index):
|
||||
"""Negative log likelihood loss on the input data.
|
||||
|
||||
output{n, i_1, i_2, ..., i_k} = -p * w
|
||||
where t = target{n, i_1, i_2, ..., i_k}
|
||||
p = predictions{n, t, i_1, i_2, i_k}
|
||||
w = weights{n, i_1, i_2, ..., i_k} if t != ignore_index else 0
|
||||
|
||||
result = reduction(output)
|
||||
|
||||
Parameters
|
||||
----------
|
||||
predictions : tvm.te.Tensor
|
||||
(k+2)-D with shape (N, C, d_1, d_2, ..., d_k),
|
||||
where C is the number of target classes
|
||||
|
||||
targets : tvm.te.Tensor
|
||||
(k+1)-D with shape (N, d_1, d_2, ..., d_k)
|
||||
The target value of the input.
|
||||
|
||||
weights : tvm.te.Tensor
|
||||
1-D with shape (C,)
|
||||
The weight of each target value.
|
||||
|
||||
reduction : string
|
||||
The reduction method to apply to output.
|
||||
Can be "mean", "sum" or "none".
|
||||
|
||||
ignore_index : int
|
||||
The target value to ignore.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
a scalar if the reduction type is "mean" or "sum",
|
||||
otherwise the same shape as `target`.
|
||||
"""
|
||||
return cpp.nn.nll_loss(predictions, targets, weights, reduction, ignore_index)
|
||||
@@ -0,0 +1,238 @@
|
||||
# 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: E731
|
||||
"""General LSTM implementation using TE scan."""
|
||||
|
||||
from tvm import te, tirx
|
||||
from tvm.topi import tag
|
||||
|
||||
|
||||
def lstm(
|
||||
Xs,
|
||||
Wi,
|
||||
Wh,
|
||||
Bi=None,
|
||||
Bh=None,
|
||||
h_init=None,
|
||||
c_init=None,
|
||||
proj=None,
|
||||
p_i=None,
|
||||
p_f=None,
|
||||
p_o=None,
|
||||
f_act=tirx.sigmoid,
|
||||
g_act=tirx.tanh,
|
||||
h_act=tirx.tanh,
|
||||
reverse=False,
|
||||
weight_layout: str = "IFGO",
|
||||
):
|
||||
"""General LSTM implemented using TE scan.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
Xs : te.Tensor
|
||||
Input sequence with shape `(seq_len, batch_size, in_dim)`
|
||||
Wi : te.Tensor
|
||||
Input weight matrix with shape `(4 * hidden_dim, in_dim)`. The weights are packed according
|
||||
to `weight_layout`.
|
||||
Wh : te.Tensor
|
||||
Hidden weight matrix with shape `(4 * hidden_dim, hidden_dim or proj_dim)`. Packed as `Wh`.
|
||||
Bi : te.Tensor, optional
|
||||
Input bias with shape `(4 * hidden_dim,)`, by default None. Packed as `Wh`.
|
||||
Bh : te.Tensor, optional
|
||||
Hidden bias with shape as `Bi`, by default None. Packed as `Wh`.
|
||||
h_init : te.Tensor, optional
|
||||
Initial hidden state with shape `(batch_size, hidden_dim or proj_dim)`, zero if None
|
||||
c_init : te.Tensor, optional
|
||||
Initial cell state with same shape as `h_init`, zero if None
|
||||
proj : te.Tensor, optional
|
||||
Projection matrix with shape `(proj_dim, hidden_dim)`, by default None
|
||||
p_i, p_f, p_o : te.Tensor, optional
|
||||
Peephole LSTM matrices with shape `(batch_size, hidden_dim)`, by default None
|
||||
f_act, g_act, h_act : F, optional
|
||||
Gate activation functions
|
||||
reverse : bool, optional
|
||||
Whether to process `Xs` in reverse, by default False
|
||||
weight_layout : str, optional
|
||||
The packed weight layout for gates, by default "IFGO". Note: I = input, F = forget,
|
||||
G = cell, O = output.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : te.Tensor, te.Tensor
|
||||
Tuple of hidden states (with shape `(seq_len, batch_size, hidden_dim or proj_dim)`), and
|
||||
cell states (with shape `(seq_len, batch_size, hidden_dim)`).
|
||||
"""
|
||||
assert len(weight_layout) == 4 and sorted(weight_layout) == sorted("IFGO"), (
|
||||
f'given weight layout "{weight_layout}" is not a permutation of "IFGO"'
|
||||
)
|
||||
|
||||
i_gate_idx = weight_layout.find("I")
|
||||
f_gate_idx = weight_layout.find("F")
|
||||
g_gate_idx = weight_layout.find("G")
|
||||
o_gate_idx = weight_layout.find("O")
|
||||
|
||||
seq_len, batch_size, in_dim = Xs.shape
|
||||
assert Wi.shape[0] % 4 == 0, (
|
||||
f"dim 0 of input weight should be 4 * hidden_dim, but {Wi.shape[0]} is not divisible by 4"
|
||||
)
|
||||
hidden_dim = Wi.shape[0] // 4
|
||||
proj_dim = hidden_dim
|
||||
if proj is not None:
|
||||
proj_dim = proj.shape[0]
|
||||
|
||||
# te.scan uses up 1 element for the initial value
|
||||
scan_len = seq_len + 1
|
||||
|
||||
# precompute input-hidden matmul outside the scan
|
||||
ki = te.reduce_axis((0, in_dim), name="ki2h")
|
||||
Xi2h = te.compute(
|
||||
(seq_len * batch_size, 4 * hidden_dim),
|
||||
lambda tb, ij: te.sum(Xs[(tb // batch_size), tb % batch_size, ki] * Wi[ij, ki], axis=ki),
|
||||
name="Xi2h",
|
||||
)
|
||||
if Bi is not None:
|
||||
Xi2h = te.compute(
|
||||
Xi2h.shape, lambda tb, ij: Xi2h[tb, ij] + Bi[ij], name="Xi2h_bias", tag=tag.INJECTIVE
|
||||
)
|
||||
|
||||
h_state = te.placeholder((scan_len, batch_size, proj_dim), name="h_state")
|
||||
c_state = te.placeholder((scan_len, batch_size, hidden_dim), name="c_state")
|
||||
h_init = te.compute(
|
||||
(1, batch_size, proj_dim),
|
||||
lambda _, b, i: h_init[b, i] if h_init is not None else 0.0,
|
||||
name="h_init",
|
||||
)
|
||||
c_init = te.compute(
|
||||
(1, batch_size, hidden_dim),
|
||||
lambda _, b, i: c_init[b, i] if c_init is not None else 0.0,
|
||||
name="c_init",
|
||||
)
|
||||
|
||||
# begin scan computations, first the (batched) hidden-hidden dense
|
||||
kh = te.reduce_axis((0, proj_dim), name="kh2h")
|
||||
s_h2h = te.compute(
|
||||
(scan_len, batch_size, 4, hidden_dim),
|
||||
lambda t, b, i, j: te.sum(h_state[t - 1, b, kh] * Wh[i * hidden_dim + j, kh], axis=kh),
|
||||
name="s_h2h",
|
||||
)
|
||||
if Bh is not None:
|
||||
s_h2h = te.compute(
|
||||
s_h2h.shape,
|
||||
lambda t, b, i, j: s_h2h[t, b, i, j] + Bh[i * hidden_dim + j],
|
||||
name="s_h2h_bias",
|
||||
tag=tag.INJECTIVE,
|
||||
)
|
||||
|
||||
# helper to reverse time if scanning backwards
|
||||
get_x_t = lambda t: seq_len - t if reverse else t - 1
|
||||
|
||||
gates = te.compute(
|
||||
(scan_len, batch_size, 4, hidden_dim),
|
||||
lambda t, b, i, j: (
|
||||
Xi2h[get_x_t(t) * batch_size + b, i * hidden_dim + j] + s_h2h[t, b, i, j]
|
||||
),
|
||||
name="gates",
|
||||
tag=tag.INJECTIVE,
|
||||
)
|
||||
|
||||
# helper to correctly read each gate dense from the batched output
|
||||
read_gate = lambda t, b, j, idx: gates[t, b, idx, j]
|
||||
|
||||
gate_shape = (scan_len, batch_size, hidden_dim)
|
||||
|
||||
# compute the activated gates (and do some extra stuff if peephole weights are present)
|
||||
if p_i is not None and p_f is not None:
|
||||
i_gate = te.compute(
|
||||
gate_shape,
|
||||
lambda t, b, j: f_act(
|
||||
read_gate(t, b, j, i_gate_idx) + p_i[b, j] * c_state[t - 1, b, j]
|
||||
),
|
||||
name="i_gate_p",
|
||||
tag=tag.INJECTIVE,
|
||||
)
|
||||
f_gate = te.compute(
|
||||
gate_shape,
|
||||
lambda t, b, j: f_act(
|
||||
read_gate(t, b, j, f_gate_idx) + p_f[b, j] * c_state[t - 1, b, j]
|
||||
),
|
||||
name="f_gate_p",
|
||||
tag=tag.INJECTIVE,
|
||||
)
|
||||
else:
|
||||
i_gate = te.compute(
|
||||
gate_shape,
|
||||
lambda *i: f_act(read_gate(*i, i_gate_idx)),
|
||||
name="i_gate",
|
||||
tag=tag.INJECTIVE,
|
||||
)
|
||||
f_gate = te.compute(
|
||||
gate_shape,
|
||||
lambda *i: f_act(read_gate(*i, f_gate_idx)),
|
||||
name="f_gate",
|
||||
tag=tag.INJECTIVE,
|
||||
)
|
||||
|
||||
g_gate = te.compute(
|
||||
gate_shape, lambda *i: g_act(read_gate(*i, g_gate_idx)), name="g_gate", tag=tag.INJECTIVE
|
||||
)
|
||||
|
||||
next_c = te.compute(
|
||||
gate_shape,
|
||||
lambda t, b, j: f_gate[t, b, j] * c_state[t - 1, b, j] + i_gate[t, b, j] * g_gate[t, b, j],
|
||||
name="next_c",
|
||||
)
|
||||
|
||||
if p_o is not None:
|
||||
o_gate = te.compute(
|
||||
gate_shape,
|
||||
lambda t, b, j: f_act(read_gate(t, b, j, o_gate_idx) + p_o[b, j] * next_c[t, b, j]),
|
||||
name="o_gate_p",
|
||||
tag=tag.INJECTIVE,
|
||||
)
|
||||
else:
|
||||
o_gate = te.compute(
|
||||
gate_shape,
|
||||
lambda *i: f_act(read_gate(*i, o_gate_idx)),
|
||||
name="o_gate",
|
||||
tag=tag.INJECTIVE,
|
||||
)
|
||||
|
||||
next_h = te.compute(gate_shape, lambda *i: o_gate(*i) * h_act(next_c(*i)), name="next_h")
|
||||
|
||||
# project hidden state back to proj_dim if projection matrix is present
|
||||
if proj is not None:
|
||||
kr = te.reduce_axis((0, hidden_dim), name="kh2p")
|
||||
next_h = te.compute(
|
||||
(scan_len, batch_size, proj_dim),
|
||||
lambda t, b, j: te.sum(next_h[t, b, kr] * proj[j, kr], axis=kr),
|
||||
name="next_h_proj",
|
||||
)
|
||||
|
||||
scan_h, scan_c = te.scan(
|
||||
[h_init, c_init], [next_h, next_c], [h_state, c_state], name="lstm_scan"
|
||||
)
|
||||
|
||||
# drop the initial values, TODO(@altanh): is there a better way?
|
||||
scan_h = te.compute(
|
||||
(seq_len, batch_size, proj_dim), lambda t, b, j: scan_h[t + 1, b, j], name="hidden_states"
|
||||
)
|
||||
scan_c = te.compute(
|
||||
(seq_len, batch_size, hidden_dim), lambda t, b, j: scan_c[t + 1, b, j], name="cell_states"
|
||||
)
|
||||
|
||||
return scan_h, scan_c
|
||||
@@ -0,0 +1,100 @@
|
||||
# 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, line-too-long
|
||||
"""Operators of one-to-one-mapping on the first input"""
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
|
||||
from .. import tag
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag=tag.BROADCAST)
|
||||
def scale_shift_nchw(Input, Scale, Shift):
|
||||
"""Batch normalization operator in inference.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
Input : tvm.te.Tensor
|
||||
4-D input tensor, NCHW layout [batch, channel, height, width]
|
||||
|
||||
Scale : tvm.te.Tensor
|
||||
Scale tensor, 1-D of size channel number
|
||||
|
||||
Shift : tvm.te.Tensor
|
||||
Shift tensor, 1-D of size channel number
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
Output tensor, layout is NCHW
|
||||
"""
|
||||
return te.compute(
|
||||
Input.shape, lambda b, c, i, j: Input[b, c, i, j] * Scale[c] + Shift[c], name="ScaleShift"
|
||||
)
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag=tag.BROADCAST)
|
||||
def scale_shift_nhwc(Input, Scale, Shift):
|
||||
"""Batch normalization operator in inference.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
Input : tvm.te.Tensor
|
||||
4-D input tensor, NHWC layout [batch, height, width, channel]
|
||||
|
||||
Scale : tvm.te.Tensor
|
||||
Scale tensor, 1-D of size channel number
|
||||
|
||||
Shift : tvm.te.Tensor
|
||||
Shift tensor, 1-D of size channel number
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
Output tensor, layout is NHWC
|
||||
"""
|
||||
return te.compute(
|
||||
Input.shape, lambda b, i, j, c: Input[b, i, j, c] * Scale[c] + Shift[c], name="ScaleShift"
|
||||
)
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag=tag.BROADCAST)
|
||||
def scale_shift_nchwc(Input, Scale, Shift):
|
||||
"""Batch normalization operator in inference.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
Input : tvm.te.Tensor
|
||||
5-D input tensor, NCHWc layout [batch, channel_chunk, height, width, channel_block]
|
||||
|
||||
Scale : tvm.te.Tensor
|
||||
Scale tensor, 2-D of size [channel_chunk, channel_block]
|
||||
|
||||
Shift : tvm.te.Tensor
|
||||
Shift tensor, 2-D of size [channel_chunk, channel_block]
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
Output tensor, layout is NHWC
|
||||
"""
|
||||
return te.compute(
|
||||
Input.shape,
|
||||
lambda b, cc, i, j, cb: Input[b, cc, i, j, cb] * Scale[cc, cb] + Shift[cc, cb],
|
||||
name="ScaleShift",
|
||||
)
|
||||
@@ -0,0 +1,316 @@
|
||||
# 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.
|
||||
"""Pad the data by constant value"""
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
from tvm.tirx import if_then_else
|
||||
|
||||
from .. import tag
|
||||
from ..utils import equal_const_int
|
||||
|
||||
|
||||
def get_padded_shape(data, pad_before, pad_after=None):
|
||||
"""
|
||||
Calculates the output shape of a tensor after applying padding.
|
||||
|
||||
Args:
|
||||
data (tvm.te.Tensor): The input tensor to which padding is applied.
|
||||
pad_before : list / tuple of n ints
|
||||
Pad width on each dimension to pad the before the axis begin.
|
||||
pad_after : list / tuple of n ints, optional
|
||||
Pad width each dimension to pad the after the axis end.
|
||||
|
||||
Raises:
|
||||
ValueError: If `pad_before` or `pad_after` lengths mismatch with `data` dimensions.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple representing the padded shape of the tensor.
|
||||
"""
|
||||
n = data.ndim
|
||||
pad_after = pad_after if pad_after else pad_before
|
||||
|
||||
if len(pad_before) != n:
|
||||
raise ValueError(f"pad_before length {len(pad_before)} != input dims {n}")
|
||||
if len(pad_after) != n:
|
||||
raise ValueError(f"pad_after length {len(pad_after)} != input dims {n}")
|
||||
|
||||
ana = tvm.arith.Analyzer()
|
||||
out_shape = tuple(ana.simplify(data.shape[i] + pad_before[i] + pad_after[i]) for i in range(n))
|
||||
|
||||
return out_shape
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag=tag.INJECTIVE + ",pad")
|
||||
def pad(data, pad_before, pad_after=None, pad_value=0.0, name="PadInput", attrs=None):
|
||||
"""Pad Input with using pad values.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
n-D input, can be any layout.
|
||||
|
||||
pad_before : list / tuple of n ints
|
||||
Pad width on each dimension to pad the before the axis begin.
|
||||
|
||||
pad_after : list / tuple of n ints, optional
|
||||
Pad width each dimension to pad the after the axis end.
|
||||
|
||||
pad_value : float, optional
|
||||
The value to be padded.
|
||||
|
||||
name : str, optional
|
||||
The name prefix operators generated
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
n-D, the same layout as Input.
|
||||
"""
|
||||
n = len(data.shape)
|
||||
pad_after = pad_after if pad_after else pad_before
|
||||
if len(pad_before) != n:
|
||||
raise ValueError(f"Input dimension and pad_before dismatch : {n} vs {len(pad_before)}")
|
||||
if len(pad_after) != n:
|
||||
raise ValueError(f"Input dimension and pad_after dismatch : {n} vs {len(pad_after)}")
|
||||
ana = tvm.arith.Analyzer()
|
||||
dshape = []
|
||||
for dim in data.shape:
|
||||
dshape.append(dim)
|
||||
out_shape = tuple(ana.simplify(dshape[i] + pad_before[i] + pad_after[i]) for i in range(n))
|
||||
pad_value = (
|
||||
pad_value if tvm.ir.is_prim_expr(pad_value) else tvm.tirx.const(pad_value, data.dtype)
|
||||
)
|
||||
|
||||
def _pad(*indices):
|
||||
not_zero = []
|
||||
index_tuple = []
|
||||
for i in range(n):
|
||||
if equal_const_int(pad_before[i], 0) and equal_const_int(pad_after[i], 0):
|
||||
index_tuple.append(indices[i])
|
||||
else:
|
||||
index_tuple.append(indices[i] - pad_before[i])
|
||||
not_zero.append(indices[i] >= pad_before[i])
|
||||
not_zero.append(indices[i] < data.shape[i] + pad_before[i])
|
||||
if not_zero:
|
||||
not_zero = tvm.tirx.all(*not_zero)
|
||||
return tvm.tirx.if_then_else(not_zero, data(*index_tuple), pad_value)
|
||||
return data(*index_tuple)
|
||||
|
||||
return te.compute(out_shape, _pad, name=name, attrs=attrs)
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag=tag.INJECTIVE + ",pad")
|
||||
def mirror_pad(data, pad_before, pad_after=None, mode="SYMMETRIC", name="MirrorPadInput"):
|
||||
"""Pad Input with mirroring either symmetric or reflected.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
n-D input, can be any layout.
|
||||
|
||||
pad_before : list / tuple of n ints
|
||||
Pad width on each dimension to pad the before the axis begin.
|
||||
|
||||
pad_after : list / tuple of n ints, optional
|
||||
Pad width each dimension to pad the after the axis end.
|
||||
|
||||
mode: str, optional
|
||||
Type of mirror padding to apply. Must be SYMMETRIC or REFLECT
|
||||
|
||||
name : str, optional
|
||||
The name prefix operators generated
|
||||
|
||||
Returns
|
||||
-------
|
||||
Output : tvm.te.Tensor
|
||||
n-D, the same layout as Input.
|
||||
"""
|
||||
n = len(data.shape)
|
||||
pad_after = pad_after if pad_after else pad_before
|
||||
if len(pad_before) != n:
|
||||
raise ValueError(f"Input dimension and pad_before dismatch : {n} vs {len(pad_before)}")
|
||||
if len(pad_after) != n:
|
||||
raise ValueError(f"Input dimension and pad_after dismatch : {n} vs {len(pad_after)}")
|
||||
ana = tvm.arith.Analyzer()
|
||||
out_shape = tuple(ana.simplify(data.shape[i] + pad_before[i] + pad_after[i]) for i in range(n))
|
||||
assert mode in ("SYMMETRIC", "REFLECT")
|
||||
mode = int(mode == "SYMMETRIC")
|
||||
|
||||
def _pad(*indices):
|
||||
index_tuple = []
|
||||
above = []
|
||||
below = []
|
||||
for i in range(n):
|
||||
if equal_const_int(pad_before[i], 0) and equal_const_int(pad_after[i], 0):
|
||||
index_tuple.append(indices[i])
|
||||
above.append(False)
|
||||
below.append(False)
|
||||
else:
|
||||
index_tuple.append(indices[i] - pad_before[i])
|
||||
above.append(indices[i] >= data.shape[i] + pad_before[i])
|
||||
below.append(indices[i] < pad_before[i])
|
||||
mapped_tuple = []
|
||||
for i, axis in enumerate(index_tuple):
|
||||
mapped_axis = tvm.tirx.if_then_else(below[i], -axis - mode, axis)
|
||||
mapped_axis = tvm.tirx.if_then_else(
|
||||
above[i], (2 * (data.shape[i] - 1)) - axis + mode, mapped_axis
|
||||
)
|
||||
mapped_tuple.append(mapped_axis)
|
||||
return data(*mapped_tuple)
|
||||
|
||||
return te.compute(out_shape, _pad, name=name)
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag=tag.INJECTIVE + ",pad")
|
||||
def reflect_pad(data, pad_before, pad_after=None, name="ReflectPadInput"):
|
||||
"""
|
||||
Apply reflect padding to the input tensor.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
Input tensor.
|
||||
|
||||
pad_before : List[int]
|
||||
Amount to pad before each dimension.
|
||||
|
||||
pad_after : List[int], optional
|
||||
Amount to pad after each dimension. If None, defaults to pad_before.
|
||||
|
||||
name : str
|
||||
Name of the resulting tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : tvm.te.Tensor
|
||||
Reflect-padded tensor.
|
||||
"""
|
||||
out_shape = get_padded_shape(data, pad_before, pad_after)
|
||||
|
||||
def _pad(*indices):
|
||||
index_tuple = []
|
||||
for i in range(data.ndim):
|
||||
idx = indices[i]
|
||||
size = data.shape[i]
|
||||
before = pad_before[i]
|
||||
|
||||
orig_idx = idx - before
|
||||
|
||||
reflected_idx = if_then_else(
|
||||
orig_idx < 0,
|
||||
-orig_idx, # reflect from start (no repeat)
|
||||
if_then_else(
|
||||
orig_idx >= size,
|
||||
(2 * size - 2) - orig_idx, # reflect from end
|
||||
orig_idx,
|
||||
),
|
||||
)
|
||||
index_tuple.append(reflected_idx)
|
||||
return data(*index_tuple)
|
||||
|
||||
return te.compute(out_shape, _pad, name=name)
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag=tag.INJECTIVE + ",pad")
|
||||
def replicate_pad(data, pad_before, pad_after=None, name="ReplicatePadInput"):
|
||||
"""
|
||||
Apply replicate padding (edge padding) to the input tensor.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
Input tensor.
|
||||
|
||||
pad_before : List[int]
|
||||
Amount to pad before each dimension.
|
||||
|
||||
pad_after : List[int], optional
|
||||
Amount to pad after each dimension. If None, defaults to pad_before.
|
||||
|
||||
name : str
|
||||
Name of the resulting tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : tvm.te.Tensor
|
||||
Replicate-padded tensor.
|
||||
"""
|
||||
out_shape = get_padded_shape(data, pad_before, pad_after)
|
||||
|
||||
def _pad(*indices):
|
||||
index_tuple = []
|
||||
for i in range(data.ndim):
|
||||
idx = indices[i]
|
||||
size = data.shape[i]
|
||||
before = pad_before[i]
|
||||
|
||||
orig_idx = idx - before
|
||||
clamped_idx = if_then_else(
|
||||
orig_idx < 0,
|
||||
tvm.tirx.const(0, "int32"), # replicate first element
|
||||
if_then_else(
|
||||
orig_idx >= size,
|
||||
size - 1, # replicate last element
|
||||
orig_idx,
|
||||
),
|
||||
)
|
||||
index_tuple.append(clamped_idx)
|
||||
return data(*index_tuple)
|
||||
|
||||
return te.compute(out_shape, _pad, name=name)
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag=tag.INJECTIVE + ",pad")
|
||||
def circular_pad(data, pad_before, pad_after=None, name="CircularPadInput"):
|
||||
"""
|
||||
Apply circular padding (wrap around) to the input tensor.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
Input tensor.
|
||||
|
||||
pad_before : List[int]
|
||||
Amount to pad before each dimension.
|
||||
|
||||
pad_after : List[int], optional
|
||||
Amount to pad after each dimension. If None, defaults to pad_before.
|
||||
|
||||
name : str
|
||||
Name of the resulting tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
out : tvm.te.Tensor
|
||||
Circular-padded tensor.
|
||||
"""
|
||||
out_shape = get_padded_shape(data, pad_before, pad_after)
|
||||
|
||||
def _pad(*indices):
|
||||
index_tuple = []
|
||||
for i in range(data.ndim):
|
||||
idx = indices[i]
|
||||
size = data.shape[i]
|
||||
before = pad_before[i]
|
||||
|
||||
orig_idx = idx - before
|
||||
wrapped_idx = tvm.tirx.indexmod(orig_idx + size, size)
|
||||
index_tuple.append(wrapped_idx)
|
||||
return data(*index_tuple)
|
||||
|
||||
return te.compute(out_shape, _pad, name=name)
|
||||
@@ -0,0 +1,75 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
# ruff: noqa: RUF005
|
||||
"""TVM operator pixel shuffle compute."""
|
||||
|
||||
import tvm
|
||||
|
||||
|
||||
def pixel_shuffle(data, upscale_factor, name="PixelShuffle"):
|
||||
"""PixelShuffle operator that rearranges elements in a tensor of shape
|
||||
[..., C * r * r, H, W] to [..., C, H * r, W * r].
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
N-D input tensor with at least 3 dimensions. Channel must be at index -3.
|
||||
|
||||
upscale_factor : int
|
||||
The upscale factor (r).
|
||||
|
||||
name : str
|
||||
Name of the output tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
Pixel shuffled tensor with shape [..., C, H*r, W*r]
|
||||
"""
|
||||
assert isinstance(upscale_factor, int) and upscale_factor > 0
|
||||
ndim = len(data.shape)
|
||||
assert ndim >= 3, "Input must be at least 3D"
|
||||
|
||||
upscale_factor_const = tvm.tirx.const(upscale_factor, "int32")
|
||||
c_in, h_in, w_in = data.shape[-3], data.shape[-2], data.shape[-1]
|
||||
|
||||
c_out = tvm.tirx.floordiv(c_in, upscale_factor_const * upscale_factor_const)
|
||||
h_out = h_in * upscale_factor_const
|
||||
w_out = w_in * upscale_factor_const
|
||||
|
||||
out_shape = list(data.shape[:-3]) + [c_out, h_out, w_out]
|
||||
|
||||
def _compute(*indices):
|
||||
batch_indices = indices[:-3]
|
||||
c_out_idx, h_out_idx, w_out_idx = indices[-3], indices[-2], indices[-1]
|
||||
|
||||
h_idx = tvm.tirx.floordiv(h_out_idx, upscale_factor_const)
|
||||
h_offset = h_out_idx % upscale_factor_const
|
||||
|
||||
w_idx = tvm.tirx.floordiv(w_out_idx, upscale_factor_const)
|
||||
w_offset = w_out_idx % upscale_factor_const
|
||||
|
||||
c_in_idx = (
|
||||
(c_out_idx * upscale_factor_const * upscale_factor_const)
|
||||
+ (h_offset * upscale_factor_const)
|
||||
+ w_offset
|
||||
)
|
||||
|
||||
index_tuple = batch_indices + (c_in_idx, h_idx, w_idx)
|
||||
return data[index_tuple]
|
||||
|
||||
return tvm.te.compute(out_shape, _compute, name=name)
|
||||
@@ -0,0 +1,406 @@
|
||||
# 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.
|
||||
"""TVM operator pooling compute."""
|
||||
|
||||
from .. import cpp
|
||||
|
||||
POOL_TYPE_CODE = {"avg": 0, "max": 1}
|
||||
|
||||
|
||||
def global_pool(data, pool_type, layout="NCHW"):
|
||||
"""Perform global pooling on height and width dimension of data.
|
||||
It decides the height and width dimension according to the layout string,
|
||||
in which 'W' and 'H' means width and height respectively.
|
||||
Width and height dimension cannot be split.
|
||||
For example, NCHW, NCHW16c, etc. are valid for pool,
|
||||
while NCHW16w, NCHW16h are not.
|
||||
See parameter `layout` for more information of the layout string convention.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
n-D with shape of layout
|
||||
|
||||
pool_type : str
|
||||
Pool type, 'max' or 'avg'
|
||||
|
||||
layout : str
|
||||
Layout of the input data.
|
||||
The layout is supposed to be composed of upper cases, lower cases and numbers,
|
||||
where upper case indicates a dimension and
|
||||
the corresponding lower case with factor size indicates the split dimension.
|
||||
For example, NCHW16c can describe a 5-D tensor of
|
||||
[batch_size, channel, height, width, channel_block],
|
||||
in which channel_block=16 is a split of dimension channel.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
n-D in same layout with height and width dimension size of 1.
|
||||
e.g., for NCHW, the output shape will be [batch, channel, 1, 1]
|
||||
"""
|
||||
return cpp.nn.global_pool(data, POOL_TYPE_CODE[pool_type], layout)
|
||||
|
||||
|
||||
def pool_grad(
|
||||
grads,
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
padding,
|
||||
pool_type,
|
||||
ceil_mode=False,
|
||||
count_include_pad=True,
|
||||
layout="NCHW",
|
||||
):
|
||||
"""Gradient of pooling on height and width dimension of data.
|
||||
It decides the height and width dimension according to the layout string,
|
||||
in which 'W' and 'H' means width and height respectively.
|
||||
Width and height dimension cannot be split.
|
||||
For example, NCHW, NCHW16c, etc. are valid for pool,
|
||||
while NCHW16w, NCHW16h are not.
|
||||
See parameter `layout` for more information of the layout string convention.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
grads : tvm.te.Tensor
|
||||
n-D with shape of layout
|
||||
|
||||
data : tvm.te.Tensor
|
||||
n-D with shape of layout
|
||||
|
||||
kernel : list/tuple of two ints
|
||||
Kernel size, [kernel_height, kernel_width]
|
||||
|
||||
stride : list/tuple of two ints
|
||||
Stride size, [stride_height, stride_width]
|
||||
|
||||
padding : list/tuple of four ints
|
||||
Pad size, [pad_top, pad_left, pad_bottom, pad_right]]
|
||||
|
||||
pool_type : str
|
||||
Pool type, 'max' or 'avg'
|
||||
|
||||
ceil_mode : bool
|
||||
Whether to use ceil when calculating output size.
|
||||
|
||||
count_include_pad: bool
|
||||
Whether include padding in the calculation when pool_type is 'avg'
|
||||
|
||||
layout: string
|
||||
Layout of the input data.
|
||||
The layout is supposed to be composed of upper cases, lower cases and numbers,
|
||||
where upper case indicates a dimension and
|
||||
the corresponding lower case with factor size indicates the split dimension.
|
||||
For example, NCHW16c can describe a 5-D tensor of
|
||||
[batch_size, channel, height, width, channel_block],
|
||||
in which channel_block=16 is a split of dimension channel.
|
||||
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
n-D in the same layout
|
||||
"""
|
||||
return cpp.nn.pool_grad(
|
||||
grads,
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
padding,
|
||||
POOL_TYPE_CODE[pool_type],
|
||||
ceil_mode,
|
||||
layout,
|
||||
count_include_pad,
|
||||
)
|
||||
|
||||
|
||||
def adaptive_pool(data, output_size, pool_type, layout="NCHW"):
|
||||
"""Perform pooling on height and width dimension of data.
|
||||
The pooling kernel and stride sizes are automatically chosen for desired
|
||||
output sizes.
|
||||
It decides the height and width dimension according to the layout string,
|
||||
in which 'W' and 'H' means width and height respectively.
|
||||
Width and height dimension cannot be split.
|
||||
For example, NCHW, NCHW16c, etc. are valid for pool,
|
||||
while NCHW16w, NCHW16h are not.
|
||||
See parameter `layout` for more information of the layout string convention.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
n-D with shape of layout
|
||||
|
||||
output_size : tuple of int
|
||||
output height and width.
|
||||
|
||||
pool_type : str
|
||||
Pool type, 'max' or 'avg'
|
||||
|
||||
layout: string
|
||||
Layout of the input data.
|
||||
The layout is supposed to be composed of upper cases, lower cases and numbers,
|
||||
where upper case indicates a dimension and
|
||||
the corresponding lower case with factor size indicates the split dimension.
|
||||
For example, NCHW16c can describe a 5-D tensor of
|
||||
[batch_size, channel, height, width, channel_block],
|
||||
in which channel_block=16 is a split of dimension channel.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
n-D in the same layout
|
||||
"""
|
||||
return cpp.nn.adaptive_pool(data, output_size, POOL_TYPE_CODE[pool_type], layout)
|
||||
|
||||
|
||||
def adaptive_pool1d(data, output_size, pool_type, layout="NCW"):
|
||||
"""Perform pooling on three dimensional data.
|
||||
See the two dimensional version above for details.
|
||||
"""
|
||||
return cpp.nn.adaptive_pool1d(data, output_size, POOL_TYPE_CODE[pool_type], layout)
|
||||
|
||||
|
||||
def adaptive_pool3d(data, output_size, pool_type, layout="NCDHW"):
|
||||
"""Perform pooling on three dimensional data.
|
||||
See the two dimensional version above for details.
|
||||
"""
|
||||
return cpp.nn.adaptive_pool3d(data, output_size, POOL_TYPE_CODE[pool_type], layout)
|
||||
|
||||
|
||||
def pool1d(
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
dilation,
|
||||
padding,
|
||||
pool_type,
|
||||
ceil_mode=False,
|
||||
layout="NCW",
|
||||
count_include_pad=True,
|
||||
):
|
||||
"""Perform pooling on width dimension of data.
|
||||
Width axis is determined according to the layout string.
|
||||
in which 'w' means width.
|
||||
Width dimension cannot be split.
|
||||
For example, NCW, NCW16c, etc. are valid for pool,
|
||||
while NCW16w is not.
|
||||
See parameter `layout` for more information of the layout string convention.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
n-D with shape of layout
|
||||
|
||||
kernel : list/tuple of one int or int
|
||||
Kernel size, [kernel_width]
|
||||
|
||||
stride : list/tuple of one int or int
|
||||
Stride size, [stride_width]
|
||||
|
||||
dilation: list/tuple of two ints
|
||||
Dilation size, [dilation_height, dilation_width]
|
||||
|
||||
padding : list/tuple of two ints
|
||||
Pad size, [pad_left, pad_right]
|
||||
|
||||
pool_type : str
|
||||
Pool type, 'max' or 'avg'
|
||||
|
||||
ceil_mode : bool
|
||||
Whether to use ceil when calculating output size.
|
||||
|
||||
layout: string
|
||||
Layout of the input data.
|
||||
The layout is supposed to be composed of upper cases, lower cases and numbers,
|
||||
where upper case indicates a dimension and
|
||||
the corresponding lower case with factor size indicates the split dimension.
|
||||
For example, NCW16c can describe a 4-D tensor of
|
||||
[batch_size, channel, width, channel_block],
|
||||
in which channel_block=16 is a split of dimension channel.
|
||||
|
||||
count_include_pad: bool
|
||||
Whether include padding in the calculation when pool_type is 'avg'
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
n-D in the same layout
|
||||
"""
|
||||
if isinstance(kernel, int):
|
||||
kernel = [
|
||||
kernel,
|
||||
]
|
||||
if isinstance(stride, int):
|
||||
stride = [
|
||||
stride,
|
||||
]
|
||||
return cpp.nn.pool1d(
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
dilation,
|
||||
padding,
|
||||
POOL_TYPE_CODE[pool_type],
|
||||
ceil_mode,
|
||||
layout,
|
||||
count_include_pad,
|
||||
)
|
||||
|
||||
|
||||
def pool2d(
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
dilation,
|
||||
padding,
|
||||
pool_type,
|
||||
ceil_mode=False,
|
||||
layout="NCHW",
|
||||
count_include_pad=True,
|
||||
):
|
||||
"""Perform pooling on height and width dimension of data.
|
||||
It decides the height and width dimension according to the layout string,
|
||||
in which 'W' and 'H' means width and height respectively.
|
||||
Width and height dimension cannot be split.
|
||||
For example, NCHW, NCHW16c, etc. are valid for pool,
|
||||
while NCHW16w, NCHW16h are not.
|
||||
See parameter `layout` for more information of the layout string convention.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
n-D with shape of layout
|
||||
|
||||
kernel : list/tuple of two ints
|
||||
Kernel size, [kernel_height, kernel_width]
|
||||
|
||||
stride : list/tuple of two ints
|
||||
Stride size, [stride_height, stride_width]
|
||||
|
||||
dilation: list/tuple of two ints
|
||||
Dilation size, [dilation_height, dilation_width]
|
||||
|
||||
padding : list/tuple of four ints
|
||||
Pad size, [pad_top, pad_left, pad_bottom, pad_right]]
|
||||
|
||||
pool_type : str
|
||||
Pool type, 'max' or 'avg'
|
||||
|
||||
ceil_mode : bool
|
||||
Whether to use ceil when calculating output size.
|
||||
|
||||
layout: string
|
||||
Layout of the input data.
|
||||
The layout is supposed to be composed of upper cases, lower cases and numbers,
|
||||
where upper case indicates a dimension and
|
||||
the corresponding lower case with factor size indicates the split dimension.
|
||||
For example, NCHW16c can describe a 5-D tensor of
|
||||
[batch_size, channel, height, width, channel_block],
|
||||
in which channel_block=16 is a split of dimension channel.
|
||||
|
||||
count_include_pad: bool
|
||||
Whether include padding in the calculation when pool_type is 'avg'
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
n-D in the same layout
|
||||
"""
|
||||
return cpp.nn.pool2d(
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
dilation,
|
||||
padding,
|
||||
POOL_TYPE_CODE[pool_type],
|
||||
ceil_mode,
|
||||
layout,
|
||||
count_include_pad,
|
||||
)
|
||||
|
||||
|
||||
def pool3d(
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
dilation,
|
||||
padding,
|
||||
pool_type,
|
||||
ceil_mode=False,
|
||||
layout="NCDHW",
|
||||
count_include_pad=True,
|
||||
):
|
||||
"""Perform pooling on depth, height and width dimension of data.
|
||||
It decides the depth, height and width dimension according to the layout string,
|
||||
in which 'D', 'W' and 'H' means depth, width and height respectively.
|
||||
Depth, width and height dimension cannot be split.
|
||||
For example, NCDHW, NCDHW16c, etc. are valid for pool,
|
||||
while NCDHW16d, NCDHW16w, NCDHW16h are not.
|
||||
See parameter `layout` for more information of the layout string convention.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
n-D with shape of layout
|
||||
|
||||
kernel : list/tuple of three ints
|
||||
Kernel size, [kernel_depth, kernel_height, kernel_width]
|
||||
|
||||
stride : list/tuple of three ints
|
||||
Stride size, [stride_depth, stride_height, stride_width]
|
||||
|
||||
dilation: list/tuple of two ints
|
||||
Dilation size, [dilation_height, dilation_width]
|
||||
|
||||
padding : list/tuple of six ints
|
||||
Pad size, [pad_front, pad_top, pad_left, pad_back, pad_bottom, pad_right]
|
||||
|
||||
pool_type : str
|
||||
Pool type, 'max' or 'avg'
|
||||
|
||||
ceil_mode : bool
|
||||
Whether to use ceil when calculating output size.
|
||||
|
||||
layout: string
|
||||
Layout of the input data.
|
||||
The layout is supposed to be composed of upper cases, lower cases and numbers,
|
||||
where upper case indicates a dimension and
|
||||
the corresponding lower case with factor size indicates the split dimension.
|
||||
For example, NCDHW16c can describe a 6-D tensor of
|
||||
[batch_size, channel, depth, height, width, channel_block],
|
||||
in which channel_block=16 is a split of dimension channel.
|
||||
|
||||
count_include_pad: bool
|
||||
Whether include padding in the calculation when pool_type is 'avg'
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
n-D in the same layout
|
||||
"""
|
||||
return cpp.nn.pool3d(
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
dilation,
|
||||
padding,
|
||||
POOL_TYPE_CODE[pool_type],
|
||||
ceil_mode,
|
||||
layout,
|
||||
count_include_pad,
|
||||
)
|
||||
@@ -0,0 +1,193 @@
|
||||
# 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.
|
||||
"""Quantized Neural Network (QNN) Operators"""
|
||||
|
||||
import tvm
|
||||
from tvm import te, tirx, topi
|
||||
|
||||
SQNN_DISABLE = 0
|
||||
SQNN_INT8 = 1
|
||||
SQNN_UINT8 = 2
|
||||
SQNN_INT32 = 3
|
||||
|
||||
SQNN_DTYPE_TO_CODE = {
|
||||
"disable": SQNN_DISABLE,
|
||||
"int8": SQNN_INT8,
|
||||
"uint8": SQNN_UINT8,
|
||||
"int32": SQNN_INT32,
|
||||
}
|
||||
|
||||
SQNN_CODE_TO_DTYPE = {v: k for k, v in SQNN_DTYPE_TO_CODE.items()}
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag=topi.tag.ELEMWISE)
|
||||
def simulated_quantize(data, out_dtype, output_scale=None, output_zero_point=None, axis=-1):
|
||||
"""Simulated QNN quantize operator that mimics QNN outputs without changing datatype.
|
||||
The benefit of this operator over true QNN quantize is that this operator allows dynamic
|
||||
datatype selection and can operate on both per-channel and scalar scales and zero points while
|
||||
QNN quantize requires both of these to be fixed at compile time.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data: tvm.te.Tensor
|
||||
An N-D input tensor to the operator.
|
||||
|
||||
out_dtype: tvm.te.Tensor
|
||||
A scalar variable that indicates which datatype to simulate quantization with. Use
|
||||
SQNN_DTYPE_TO_CODE to convert a dtype string into the corresponding variable
|
||||
value.
|
||||
|
||||
output_scale: tvm.te.Tensor, optional
|
||||
A scalar tensor representing the scale to use when quantizing to integer datatypes.
|
||||
When it contains more than a single value, N must match the number of channels in data.
|
||||
|
||||
output_zero_point: tvm.te.Tensor, optional
|
||||
A 1-D tensor representing the zero point to use when quantizing to integer datatypes.
|
||||
When it contains more than a single value, N must match the number of channels in data.
|
||||
|
||||
axis: int, optional
|
||||
The channel axis for quantization. Default value is -1 which corresponds to the last axis.
|
||||
|
||||
"""
|
||||
|
||||
# When disabled, just pass through the input values.
|
||||
def _compute_pass_through(value, *indices):
|
||||
return value[indices]
|
||||
|
||||
# Simulate quantization for arbitrary integer datatypes. The computation for all datatypes is:
|
||||
# Q_output = clip((round(input_tensor/output_scale) + output_zero_point),
|
||||
# out_dtype::min,
|
||||
# out_dtype::max)
|
||||
def _compute_intn(dtype, value, *indices):
|
||||
assert output_scale is not None and output_zero_point is not None
|
||||
const_min = tvm.tirx.min_value(dtype)
|
||||
const_max = tvm.tirx.max_value(dtype)
|
||||
# Use indexmod to handle both scalar and per-channel QNN parameters.
|
||||
scale_idx = tirx.indexmod(indices[axis], topi.shape(output_scale)[0])
|
||||
zp_idx = tirx.indexmod(indices[axis], topi.shape(output_zero_point)[0])
|
||||
return te.max(
|
||||
te.min(
|
||||
te.round(value[indices] / output_scale[scale_idx]) + output_zero_point[zp_idx],
|
||||
const_max,
|
||||
),
|
||||
const_min,
|
||||
)
|
||||
|
||||
# Use an if chain to dynamically return the proper quantization based on the input datatype.
|
||||
# This allows the op to compile once but apply different quantization approaches
|
||||
# using a variable datatype input.
|
||||
def _dispatch_sim_quantize(value):
|
||||
pass_through_value = te.compute(
|
||||
data.shape, lambda *indices: _compute_pass_through(value, *indices)
|
||||
)
|
||||
int8_value = te.compute(
|
||||
data.shape,
|
||||
lambda *indices: tirx.if_then_else(
|
||||
out_dtype.equal(SQNN_DTYPE_TO_CODE["int8"]),
|
||||
_compute_intn("int8", value, *indices),
|
||||
pass_through_value[indices],
|
||||
),
|
||||
)
|
||||
uint8_value = te.compute(
|
||||
data.shape,
|
||||
lambda *indices: tirx.if_then_else(
|
||||
out_dtype.equal(SQNN_DTYPE_TO_CODE["uint8"]),
|
||||
_compute_intn("uint8", value, *indices),
|
||||
int8_value[indices],
|
||||
),
|
||||
)
|
||||
int32_value = te.compute(
|
||||
data.shape,
|
||||
lambda *indices: tirx.if_then_else(
|
||||
out_dtype.equal(SQNN_DTYPE_TO_CODE["int32"]),
|
||||
_compute_intn("int32", value, *indices),
|
||||
uint8_value[indices],
|
||||
),
|
||||
)
|
||||
|
||||
return int32_value
|
||||
|
||||
return te.compute(data.shape, lambda *indices: _dispatch_sim_quantize(data)[indices])
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag=topi.tag.ELEMWISE)
|
||||
def simulated_dequantize(data, in_dtype, input_scale=None, input_zero_point=None, axis=-1):
|
||||
"""Simulated QNN dequantize operator that mimics QNN outputs without changing datatype.
|
||||
The benefit of this operator over true QNN dequantize is that this operator allows dynamic
|
||||
datatype selection and can operate on both per-channel and scalar scales and zero points while
|
||||
QNN dequantize requires both of these to be fixed at compile time.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data: tvm.te.Tensor
|
||||
An N-D input tensor to the operator.
|
||||
|
||||
in_dtype: tvm.te.Tensor
|
||||
A scalar variable that indicates which datatype to simulate dequantization with. Use
|
||||
SQNN_DTYPE_TO_CODE to convert a dtype string into the corresponding variable
|
||||
value.
|
||||
|
||||
input_scale: tvm.te.Tensor, optional
|
||||
A scalar tensor representing the scale to use when dequantizing from integer datatypes.
|
||||
When it contains more than a single value, N must match the number of channels in data.
|
||||
|
||||
input_zero_point: tvm.te.Tensor, optional
|
||||
A 1-D tensor representing the zero point to use when dequantizing from integer datatypes.
|
||||
When it contains more than a single value, N must match the number of channels in data.
|
||||
|
||||
axis: int, optional
|
||||
The channel axis for quantization. Default value is -1 which corresponds to the last axis.
|
||||
|
||||
"""
|
||||
|
||||
# When disabled simply return the input tensor.
|
||||
def _compute_pass_through(value, *indices):
|
||||
return value[indices]
|
||||
|
||||
# Simulate dequantization for arbitrary integer datatypes. The computation for all datatypes is:
|
||||
# DQ_output = (input - zero_point) * scale
|
||||
def _compute_intn(value, *indices):
|
||||
assert input_scale is not None and input_zero_point is not None
|
||||
# Use indexmod to handle both scalar and per-channel QNN parameters.
|
||||
scale_idx = tirx.indexmod(indices[axis], topi.shape(input_scale)[0])
|
||||
zp_idx = tirx.indexmod(indices[axis], topi.shape(input_zero_point)[0])
|
||||
return (value[indices] - input_zero_point[zp_idx]) * input_scale[scale_idx]
|
||||
|
||||
# Use an if chain to dynamically return the proper dequantization based on the input datatype.
|
||||
# This allows the op to compile once but apply different quantization approaches
|
||||
# using a variable datatype input.
|
||||
def _dispatch_sim_dequantize(value):
|
||||
pass_through_value = te.compute(
|
||||
data.shape, lambda *indices: _compute_pass_through(value, *indices)
|
||||
)
|
||||
intn_condition = tvm.te.any(
|
||||
in_dtype.equal(SQNN_DTYPE_TO_CODE["int8"]),
|
||||
in_dtype.equal(SQNN_DTYPE_TO_CODE["uint8"]),
|
||||
in_dtype.equal(SQNN_DTYPE_TO_CODE["int32"]),
|
||||
)
|
||||
intn_value = te.compute(
|
||||
data.shape,
|
||||
lambda *indices: tirx.if_then_else(
|
||||
intn_condition,
|
||||
_compute_intn(value, *indices),
|
||||
pass_through_value[indices],
|
||||
),
|
||||
)
|
||||
|
||||
return intn_value
|
||||
|
||||
return te.compute(data.shape, lambda *indices: _dispatch_sim_dequantize(data)[indices])
|
||||
@@ -0,0 +1,44 @@
|
||||
# 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.
|
||||
"""Root mean square normalization operator."""
|
||||
|
||||
from .. import cpp
|
||||
|
||||
|
||||
def rms_norm(data, weight, axis, epsilon=1e-5):
|
||||
"""Root mean square normalization operator. The output will have the same data type as input.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
N-D with shape (d_0, d_1, ..., d_{N-1})
|
||||
|
||||
weight: tvm.te.Tensor
|
||||
K-D with shape (r_0, r_1, ..., r_{K-1}) where K == len(axis) and d_{axis_k} == r_k
|
||||
|
||||
axis : list of int
|
||||
Axis over the normalization applied
|
||||
|
||||
epsilon : float
|
||||
The epsilon value to avoid division by zero.
|
||||
|
||||
Returns
|
||||
-------
|
||||
result : tvm.te.Tensor
|
||||
N-D with shape (d_0, d_1, ..., d_{N-1})
|
||||
"""
|
||||
return cpp.nn.rms_norm(data, weight, axis, epsilon)
|
||||
@@ -0,0 +1,177 @@
|
||||
# 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, pointless-exception-statement
|
||||
# ruff: noqa: RUF005
|
||||
"""TVM operator for softmax and log_softmax compute."""
|
||||
|
||||
import tvm
|
||||
from tvm import te, topi
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag="softmax_output")
|
||||
def softmax(x, axis=-1):
|
||||
"""Perform softmax activation on the data.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : tvm.te.Tensor
|
||||
can be any dimension
|
||||
|
||||
axis : int
|
||||
channel axis
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
output shape is the same as input
|
||||
"""
|
||||
return softmax_common(x, axis, False)
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag="fast_softmax_output")
|
||||
def fast_softmax(x, axis=-1):
|
||||
"""Perform softmax activation on the data.
|
||||
Use approximation to compute exponent for faster speed.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : tvm.te.Tensor
|
||||
can be any dimension
|
||||
|
||||
axis : int
|
||||
channel axis
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
output shape is the same as input
|
||||
"""
|
||||
return softmax_common(x, axis, True)
|
||||
|
||||
|
||||
def softmax_common(x, axis, use_fast_exp):
|
||||
"""The common part of softmax and fast_softmax"""
|
||||
shape = x.shape
|
||||
if axis < 0:
|
||||
axis = len(shape) + axis
|
||||
if axis >= len(shape):
|
||||
ValueError("axis parameter should be less than input dim")
|
||||
|
||||
k1 = te.reduce_axis((0, shape[axis]), name="k")
|
||||
k2 = te.reduce_axis((0, shape[axis]), name="k")
|
||||
|
||||
def insert_reduce_index(indices, reduce_index):
|
||||
return indices[:axis] + (reduce_index,) + indices[axis:]
|
||||
|
||||
def get_non_reduce_indices(indices):
|
||||
return tuple([var for (i, var) in enumerate(indices) if i != axis])
|
||||
|
||||
def _compute_max(*indices):
|
||||
eval_range = insert_reduce_index(indices, k1)
|
||||
return tvm.te.max(x[eval_range], axis=k1)
|
||||
|
||||
def _compute_delta(max_elem, *indices):
|
||||
non_reduce_indices = get_non_reduce_indices(indices)
|
||||
return x[indices] - max_elem[non_reduce_indices]
|
||||
|
||||
def _compute_exp(max_elem, *indices):
|
||||
non_reduce_indices = get_non_reduce_indices(indices)
|
||||
return te.exp(x[indices] - max_elem[non_reduce_indices])
|
||||
|
||||
def _compute_expsum(exp, *indices):
|
||||
eval_range = insert_reduce_index(indices, k2)
|
||||
return te.sum(exp[eval_range], axis=k2)
|
||||
|
||||
def _normalize(exp, expsum, *indices):
|
||||
non_reduce_indices = get_non_reduce_indices(indices)
|
||||
return exp[indices] / expsum[non_reduce_indices]
|
||||
|
||||
reduced_shape = tuple([dim for (i, dim) in enumerate(shape) if i != axis])
|
||||
max_elem = te.compute(reduced_shape, _compute_max, name="T_softmax_maxelem")
|
||||
|
||||
if use_fast_exp:
|
||||
delta = te.compute(
|
||||
shape, lambda *indices: _compute_delta(max_elem, *indices), name="T_softmax_delta"
|
||||
)
|
||||
exp = topi.math.fast_exp(delta)
|
||||
else:
|
||||
exp = te.compute(
|
||||
shape, lambda *indices: _compute_exp(max_elem, *indices), name="T_softmax_exp"
|
||||
)
|
||||
expsum = te.compute(
|
||||
reduced_shape, lambda *indices: _compute_expsum(exp, *indices), name="T_softmax_expsum"
|
||||
)
|
||||
return te.compute(
|
||||
shape,
|
||||
lambda *indices: _normalize(exp, expsum, *indices),
|
||||
name="T_softmax_norm",
|
||||
attrs={"axis": axis},
|
||||
)
|
||||
|
||||
|
||||
@tvm.te.tag_scope(tag="log_softmax_output")
|
||||
def log_softmax(x, axis=-1):
|
||||
"""Perform log softmax activation on the data
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : tvm.te.Tensor
|
||||
N-D input data
|
||||
|
||||
axis : int
|
||||
channel axis
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
N-D output with same shape
|
||||
"""
|
||||
shape = x.shape
|
||||
if axis < 0:
|
||||
axis = len(shape) + axis
|
||||
if axis >= len(shape):
|
||||
ValueError("axis parameter should be less than input dim")
|
||||
|
||||
k1 = te.reduce_axis((0, shape[axis]), name="k")
|
||||
k2 = te.reduce_axis((0, shape[axis]), name="k")
|
||||
|
||||
def insert_reduce_index(indices, reduce_index):
|
||||
return indices[:axis] + (reduce_index,) + indices[axis:]
|
||||
|
||||
def get_non_reduce_indices(indices):
|
||||
return tuple([var for (i, var) in enumerate(indices) if i != axis])
|
||||
|
||||
def _compute_max(*indices):
|
||||
eval_range = insert_reduce_index(indices, k1)
|
||||
return tvm.te.max(x[eval_range], axis=k1)
|
||||
|
||||
def _compute_expsum(max_elem, *indices):
|
||||
eval_range = insert_reduce_index(indices, k2)
|
||||
return te.sum(te.exp(x[eval_range] - max_elem[indices]), axis=k2)
|
||||
|
||||
def _normalize(max_elem, expsum, *indices):
|
||||
non_reduce_indices = get_non_reduce_indices(indices)
|
||||
return x[indices] - max_elem[non_reduce_indices] - te.log(expsum[non_reduce_indices])
|
||||
|
||||
reduced_shape = tuple([dim for (i, dim) in enumerate(shape) if i != axis])
|
||||
max_elem = te.compute(reduced_shape, _compute_max, name="T_softmax_maxelem")
|
||||
expsum = te.compute(reduced_shape, lambda *indices: _compute_expsum(max_elem, *indices))
|
||||
return te.compute(
|
||||
shape,
|
||||
lambda *indices: _normalize(max_elem, expsum, *indices),
|
||||
attrs={"axis": axis},
|
||||
)
|
||||
@@ -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.
|
||||
# pylint: disable=invalid-name
|
||||
"""TVM operator space_to_batch_nd compute."""
|
||||
|
||||
from . import cpp
|
||||
|
||||
|
||||
def space_to_batch_nd(data, block_shape, pad_before, pad_after, pad_value=0.0):
|
||||
"""Perform batch to space transformation on the data
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
N-D Tensor with shape [batch, spatial_shape, remaining_shapes],
|
||||
where spatial_shape has M dimensions.
|
||||
|
||||
block_shape : list of ints
|
||||
list of size [M] where M is number of spatial dims, specifies block
|
||||
size for each spatial dimension.
|
||||
|
||||
pad_before : list of ints
|
||||
list of shape [M] where M is number of spatial dims, specifies
|
||||
zero-padding size before each spatial dimension.
|
||||
|
||||
pad_after : list of ints
|
||||
list of shape [M] where M is number of spatial dims, specifies
|
||||
zero-padding size after each spatial dimension.
|
||||
|
||||
pad_value : float, optional
|
||||
The value used for padding.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
"""
|
||||
|
||||
return cpp.nn.space_to_batch_nd(data, block_shape, pad_before, pad_after, pad_value)
|
||||
@@ -0,0 +1,88 @@
|
||||
# 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
|
||||
"""TVM operator space_to_depth compute."""
|
||||
|
||||
import tvm
|
||||
from tvm import te
|
||||
|
||||
from .. import tag
|
||||
|
||||
|
||||
def space_to_depth(data, block_size, layout="NCHW"):
|
||||
"""Perform space to depth transformation on the data
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
4-D tensor in either NCHW or NHWC layout.
|
||||
|
||||
block_size : int
|
||||
Size of blocks to decompose into channel dimension.
|
||||
|
||||
layout : string
|
||||
Either NCHW or NHWC, indicating data layout.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
Output of shape [N, C * block_size**2, H / block_size, W / block_size]
|
||||
"""
|
||||
|
||||
if layout == "NCHW":
|
||||
in_n, in_c, in_h, in_w = data.shape
|
||||
output_shape = [
|
||||
in_n,
|
||||
in_c * block_size * block_size,
|
||||
tvm.tirx.truncdiv(in_h, block_size),
|
||||
tvm.tirx.truncdiv(in_w, block_size),
|
||||
]
|
||||
elif layout == "NHWC":
|
||||
in_n, in_h, in_w, in_c = data.shape
|
||||
output_shape = [
|
||||
in_n,
|
||||
tvm.tirx.truncdiv(in_h, block_size),
|
||||
tvm.tirx.truncdiv(in_w, block_size),
|
||||
in_c * block_size * block_size,
|
||||
]
|
||||
else:
|
||||
raise ValueError("Only NCHW and NHWC layouts are currently supported.")
|
||||
|
||||
def _get_indices(*indices):
|
||||
if layout == "NCHW":
|
||||
n, c, y, x = indices
|
||||
elif layout == "NHWC":
|
||||
n, y, x, c = indices
|
||||
return n, c, y, x
|
||||
|
||||
def _get_pixel(n, c, y, x):
|
||||
block_offset = tvm.tirx.truncdiv(c, in_c)
|
||||
channel_idx = tvm.tirx.truncmod(c, in_c)
|
||||
x_idx = tvm.tirx.truncmod(block_offset, block_size)
|
||||
y_idx = tvm.tirx.truncdiv(block_offset, block_size)
|
||||
|
||||
if layout == "NCHW":
|
||||
output = data(n, channel_idx, y_idx + (y * block_size), x_idx + (x * block_size))
|
||||
else:
|
||||
output = data(n, y_idx + (y * block_size), x_idx + (x * block_size), channel_idx)
|
||||
return output
|
||||
|
||||
def _compute(*indices):
|
||||
n, c, y, x = _get_indices(*indices)
|
||||
return _get_pixel(n, c, y, x)
|
||||
|
||||
return te.compute(output_shape, _compute, name="space_to_depth", tag=tag.INJECTIVE)
|
||||
@@ -0,0 +1,204 @@
|
||||
# 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.
|
||||
"""TVM operator upsampling compute."""
|
||||
|
||||
from tvm import te, topi
|
||||
|
||||
from ..utils import simplify
|
||||
|
||||
|
||||
def upsampling(
|
||||
data,
|
||||
scale_h,
|
||||
scale_w,
|
||||
layout="NCHW",
|
||||
method="nearest_neighbor",
|
||||
align_corners=False,
|
||||
output_shape=None,
|
||||
):
|
||||
"""Perform upsampling on the data.
|
||||
Nearest neighbor and bilinear upsampling are supported.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
inputs is a 4-D tensor with shape
|
||||
[batch, channel, in_height, in_width]
|
||||
or [batch, in_height, in_width, channel]
|
||||
|
||||
scale_h : float
|
||||
Scaling factor for height
|
||||
|
||||
scale_w : float
|
||||
Scaling factor for width
|
||||
|
||||
layout : string, optional
|
||||
either "NCHW" or "NHWC"
|
||||
|
||||
method : {"bilinear", "nearest_neighbor", "bicubic"}
|
||||
Method to be used for upsampling.
|
||||
|
||||
output_shape: tvm_ffi.Array, optional
|
||||
Shape to return. If left None will be inferred
|
||||
(If shape is determined dynamically, pass out_dtype.shape as output_shape)
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
4-D with shape [batch, channel, in_height*scale_h, in_width*scale_w]
|
||||
or [batch, in_height*scale, in_width*scale, channel]
|
||||
"""
|
||||
base_layout = layout[0:4]
|
||||
if base_layout == "NCHW":
|
||||
if not output_shape: # static case
|
||||
scaled_h = data.shape[2] * scale_h
|
||||
scaled_w = data.shape[3] * scale_w
|
||||
reshape_size = (
|
||||
simplify(topi.cast(te.round(scaled_h), data.shape[2].ty)),
|
||||
simplify(topi.cast(te.round(scaled_w), data.shape[3].ty)),
|
||||
)
|
||||
else: # dynamic case -- we don't need to scale; already done in shape func
|
||||
reshape_size = (
|
||||
simplify(topi.cast(te.round(output_shape[2]), output_shape[2].ty)),
|
||||
simplify(topi.cast(te.round(output_shape[3]), output_shape[3].ty)),
|
||||
)
|
||||
elif layout == "NHWC":
|
||||
if not output_shape: # static case
|
||||
scaled_h = data.shape[1] * scale_h
|
||||
scaled_w = data.shape[2] * scale_w
|
||||
reshape_size = (
|
||||
simplify(topi.cast(te.round(scaled_h), data.shape[1].ty)),
|
||||
simplify(topi.cast(te.round(scaled_w), data.shape[2].ty)),
|
||||
)
|
||||
else: # dynamic case
|
||||
reshape_size = (
|
||||
simplify(topi.cast(te.round(output_shape[1]), output_shape[1].ty)),
|
||||
simplify(topi.cast(te.round(output_shape[2]), output_shape[2].ty)),
|
||||
)
|
||||
|
||||
else:
|
||||
raise ValueError(f"not support this layout {layout} yet")
|
||||
coord_trans = "align_corners" if align_corners else "asymmetric"
|
||||
if method[0:2] == "bi":
|
||||
method = method[2:]
|
||||
return topi.image.resize2d(
|
||||
data,
|
||||
[0.0] * 4,
|
||||
reshape_size,
|
||||
layout=layout,
|
||||
method=method,
|
||||
coordinate_transformation_mode=coord_trans,
|
||||
output_shape=output_shape,
|
||||
)
|
||||
|
||||
|
||||
def upsampling3d(
|
||||
data,
|
||||
scale_d,
|
||||
scale_h,
|
||||
scale_w,
|
||||
layout="NCDHW",
|
||||
method="nearest_neighbor",
|
||||
coordinate_transformation_mode="half_pixel",
|
||||
output_shape=None,
|
||||
):
|
||||
"""Perform upsampling on the data.
|
||||
Nearest neighbor and bilinear upsampling are supported.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
inputs is a 5-D tensor with shape
|
||||
[batch, channel, in_depth, in_height, in_width]
|
||||
or [batch, in_depth, in_height, in_width, channel]
|
||||
|
||||
scale_d : float
|
||||
Scaling factor for depth
|
||||
|
||||
scale_h : float
|
||||
Scaling factor for height
|
||||
|
||||
scale_w : float
|
||||
Scaling factor for width
|
||||
|
||||
layout : string, optional
|
||||
either "NCDHW" or "NDHWC"
|
||||
|
||||
method : {"trilinear", "nearest_neighbor"}
|
||||
Method to be used for upsampling.
|
||||
|
||||
coordinate_transformation_mode: string, optional
|
||||
Describes how to transform the coordinate in the resized tensor
|
||||
to the coordinate in the original tensor.
|
||||
Refer to the ONNX Resize operator specification for details.
|
||||
Available options are "half_pixel", "align_corners" and "asymmetric".
|
||||
|
||||
output_shape: tvm_ffi.Array, optional
|
||||
Shape to return. If left None will be inferred
|
||||
(If shape is determined dynamically, pass out_dtype.shape as output_shape)
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
5-D with shape [batch, channel, in_depth*scale, in_height*scale, in_width*scale]
|
||||
or [batch, in_depth*scale, in_height*scale, in_width*scale, channel]
|
||||
"""
|
||||
base_layout = layout[0:5]
|
||||
if base_layout == "NCDHW":
|
||||
if not output_shape: # static case
|
||||
scaled_d = data.shape[2] * scale_d
|
||||
scaled_h = data.shape[3] * scale_h
|
||||
scaled_w = data.shape[4] * scale_w
|
||||
resize_shape = (
|
||||
simplify(topi.cast(te.round(scaled_d), data.shape[2].ty)),
|
||||
simplify(topi.cast(te.round(scaled_h), data.shape[3].ty)),
|
||||
simplify(topi.cast(te.round(scaled_w), data.shape[4].ty)),
|
||||
)
|
||||
else: # dynamic case -- don't need to scale; already done in shape func
|
||||
resize_shape = (
|
||||
simplify(topi.cast(te.round(output_shape[2]), data.shape[2].ty)),
|
||||
simplify(topi.cast(te.round(output_shape[3]), data.shape[3].ty)),
|
||||
simplify(topi.cast(te.round(output_shape[4]), data.shape[4].ty)),
|
||||
)
|
||||
elif layout == "NDHWC":
|
||||
if not output_shape: # static case
|
||||
scaled_d = data.shape[1] * scale_d
|
||||
scaled_h = data.shape[2] * scale_h
|
||||
scaled_w = data.shape[3] * scale_w
|
||||
resize_shape = (
|
||||
simplify(topi.cast(te.round(scaled_d), data.shape[1].ty)),
|
||||
simplify(topi.cast(te.round(scaled_h), data.shape[2].ty)),
|
||||
simplify(topi.cast(te.round(scaled_w), data.shape[3].ty)),
|
||||
)
|
||||
else: # dynamic case
|
||||
resize_shape = (
|
||||
simplify(topi.cast(te.round(output_shape[1]), data.shape[1].ty)),
|
||||
simplify(topi.cast(te.round(output_shape[2]), data.shape[2].ty)),
|
||||
simplify(topi.cast(te.round(output_shape[3]), data.shape[3].ty)),
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"not support this layout {layout} yet")
|
||||
if method[0:3] == "tri":
|
||||
method = method[3:]
|
||||
return topi.image.resize3d(
|
||||
data,
|
||||
[0.0] * 6,
|
||||
resize_shape,
|
||||
layout=layout,
|
||||
method=method,
|
||||
coordinate_transformation_mode=coordinate_transformation_mode,
|
||||
)
|
||||
@@ -0,0 +1,310 @@
|
||||
# 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-variable
|
||||
"""NN operator common utilities"""
|
||||
|
||||
import tvm
|
||||
|
||||
from ..utils import get_const_int
|
||||
|
||||
|
||||
def infer_pad(data, data_pad):
|
||||
"""Infer the padding from stages in reverse.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : Tensor
|
||||
data stage.
|
||||
|
||||
data_pad : Tensor
|
||||
pad stage.
|
||||
|
||||
Returns
|
||||
-------
|
||||
hpad : int
|
||||
padding size on height
|
||||
wpad : int
|
||||
padding size on width
|
||||
"""
|
||||
if data_pad is None:
|
||||
return 0, 0
|
||||
_, _, IH, IW = data.shape
|
||||
_, _, TH, TW = data_pad.shape
|
||||
hpad = (TH - IH) // 2
|
||||
wpad = (TW - IW) // 2
|
||||
return get_const_int(hpad), get_const_int(wpad)
|
||||
|
||||
|
||||
def infer_pad3d(data, data_pad, layout):
|
||||
"""Infer the padding from stages in reverse.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : Tensor
|
||||
data stage.
|
||||
|
||||
data_pad : Tensor
|
||||
pad stage.
|
||||
|
||||
Returns
|
||||
-------
|
||||
dpad : int
|
||||
padding depth
|
||||
hpad : int
|
||||
padding height
|
||||
wpad : int
|
||||
padding width
|
||||
"""
|
||||
if data_pad is None:
|
||||
return 0, 0, 0
|
||||
|
||||
if layout == "NDHWC":
|
||||
_, ID, IH, IW, _ = data.shape
|
||||
_, TD, TH, TW, _ = data_pad.shape
|
||||
elif layout == "NCDHW":
|
||||
_, _, ID, IH, IW = data.shape
|
||||
_, _, TD, TH, TW = data_pad.shape
|
||||
else:
|
||||
raise ValueError(f"Layout {layout} is not supported")
|
||||
dpad = TD - ID
|
||||
hpad = TH - IH
|
||||
wpad = TW - IW
|
||||
return get_const_int(dpad), get_const_int(hpad), get_const_int(wpad)
|
||||
|
||||
|
||||
def infer_stride(data, kernel, out):
|
||||
"""Infer the stride from stages in reverse.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : Tensor
|
||||
data stage.
|
||||
|
||||
kernel : Tensor
|
||||
kernel stage.
|
||||
|
||||
out : Tensor
|
||||
output stage.
|
||||
|
||||
Returns
|
||||
-------
|
||||
hstride : int
|
||||
stride size on height
|
||||
wstride : int
|
||||
stride size on width
|
||||
"""
|
||||
_, _, IH, IW = data.shape
|
||||
_, _, KH, KW = kernel.shape
|
||||
_, _, OH, OW = out.shape
|
||||
hstride = (IH - KH) // tvm.te.max(OH - 1, 1) + tvm.tirx.Select(OH == 1, 1, 0)
|
||||
wstride = (IW - KW) // tvm.te.max(OW - 1, 1) + tvm.tirx.Select(OW == 1, 1, 0)
|
||||
return get_const_int(hstride), get_const_int(wstride)
|
||||
|
||||
|
||||
def get_pad_tuple(padding, kernel):
|
||||
"""Common code to get the pad option
|
||||
|
||||
Parameters
|
||||
----------
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
kernel : tuple of int
|
||||
Conv kernel size
|
||||
|
||||
Returns
|
||||
-------
|
||||
pad_top : int
|
||||
Padding size on top
|
||||
|
||||
pad_left : int
|
||||
Padding size on left
|
||||
|
||||
pad_down : int
|
||||
Padding size on down.
|
||||
|
||||
pad_right : int
|
||||
Padding size on right.
|
||||
"""
|
||||
# compute the padding size
|
||||
if isinstance(padding, tuple | list):
|
||||
if len(padding) == 2:
|
||||
pad_h = padding[0] * 2
|
||||
pad_w = padding[1] * 2
|
||||
elif len(padding) == 4:
|
||||
return padding[0], padding[1], padding[2], padding[3]
|
||||
else:
|
||||
raise ValueError("Size of padding can only be 2 or 4")
|
||||
elif isinstance(padding, int):
|
||||
pad_h = pad_w = padding * 2
|
||||
elif padding == "VALID":
|
||||
pad_h = 0
|
||||
pad_w = 0
|
||||
elif padding == "SAME":
|
||||
pad_h = kernel[0] - 1
|
||||
pad_w = kernel[1] - 1
|
||||
else:
|
||||
raise ValueError(f"Unknown padding option {padding}")
|
||||
pad_top = (pad_h + 1) // 2
|
||||
pad_left = (pad_w + 1) // 2
|
||||
return pad_top, pad_left, pad_h - pad_top, pad_w - pad_left
|
||||
|
||||
|
||||
def get_pad_tuple_generic(padding, kernel):
|
||||
"""Common code to get the pad option
|
||||
|
||||
Parameters
|
||||
----------
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
kernel : tuple of int
|
||||
Conv kernel size
|
||||
|
||||
Returns
|
||||
-------
|
||||
pad_top : int
|
||||
Padding size on top
|
||||
|
||||
pad_down : int
|
||||
Padding size on down.
|
||||
|
||||
pad_left : int
|
||||
Padding size on left
|
||||
|
||||
pad_right : int
|
||||
Padding size on right.
|
||||
"""
|
||||
# compute the padding size
|
||||
if isinstance(padding, tuple | list):
|
||||
if len(padding) == len(kernel):
|
||||
pad_dimensions = [p * 2 for p in padding]
|
||||
elif len(padding) == len(kernel) * 2:
|
||||
return (
|
||||
[padding[i] for i in range(len(kernel))],
|
||||
[padding[len(kernel) + i] for i in range(len(kernel))],
|
||||
)
|
||||
else:
|
||||
raise ValueError("Size of padding can only be len(kernel) or len(kernel) * 2")
|
||||
elif isinstance(padding, int):
|
||||
pad_dimensions = [padding * 2 for _ in range(len(kernel))]
|
||||
elif padding == "VALID":
|
||||
pad_dimensions = [0 for _ in range(len(kernel))]
|
||||
elif padding == "SAME":
|
||||
pad_dimensions = [k - 1 for k in kernel]
|
||||
else:
|
||||
raise ValueError(f"Unknown padding option {padding}")
|
||||
pad_begin = [(p + 1) // 2 for p in pad_dimensions]
|
||||
return [pad_begin, [pd - pb for pb, pd in zip(pad_begin, pad_dimensions)]]
|
||||
|
||||
|
||||
def get_pad_tuple3d(padding, kernel):
|
||||
"""Common code to get the pad option
|
||||
|
||||
Parameters
|
||||
----------
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
kernel : tuple of int
|
||||
Conv kernel size
|
||||
|
||||
Returns
|
||||
-------
|
||||
pad_front : int
|
||||
Padding size on front.
|
||||
|
||||
pad_top : int
|
||||
Padding size on top
|
||||
|
||||
pad_left : int
|
||||
Padding size on left
|
||||
|
||||
pad_back : int
|
||||
Padding size on back.
|
||||
|
||||
pad_down : int
|
||||
Padding size on down.
|
||||
|
||||
pad_right : int
|
||||
Padding size on right.
|
||||
"""
|
||||
# compute the padding size
|
||||
if isinstance(padding, tuple | list):
|
||||
if len(padding) == 3:
|
||||
pad_d = padding[0] * 2
|
||||
pad_h = padding[1] * 2
|
||||
pad_w = padding[2] * 2
|
||||
elif len(padding) == 6:
|
||||
return padding[0], padding[1], padding[2], padding[3], padding[4], padding[5]
|
||||
else:
|
||||
raise ValueError("Size of padding can only be 3 or 6")
|
||||
elif isinstance(padding, int):
|
||||
pad_d = pad_w = pad_h = padding * 2
|
||||
elif padding == "VALID":
|
||||
pad_h = 0
|
||||
pad_w = 0
|
||||
pad_d = 0
|
||||
elif padding == "SAME":
|
||||
pad_d = kernel[0] - 1
|
||||
pad_h = kernel[1] - 1
|
||||
pad_w = kernel[2] - 1
|
||||
else:
|
||||
raise ValueError(f"Unknown padding option {padding}")
|
||||
pad_top = (pad_h + 1) // 2
|
||||
pad_left = (pad_w + 1) // 2
|
||||
pad_front = (pad_d + 1) // 2
|
||||
return pad_front, pad_top, pad_left, pad_d - pad_front, pad_h - pad_top, pad_w - pad_left
|
||||
|
||||
|
||||
def get_pad_tuple1d(padding, kernel):
|
||||
"""Common code to get the pad option
|
||||
|
||||
Parameters
|
||||
----------
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
kernel : tuple of int
|
||||
Conv kernel size
|
||||
|
||||
Returns
|
||||
-------
|
||||
pad_left : int
|
||||
Padding size on left
|
||||
|
||||
pad_right : int
|
||||
Padding size on right.
|
||||
"""
|
||||
# compute the padding size
|
||||
if isinstance(padding, tuple | list):
|
||||
if len(padding) == 1:
|
||||
pad_w = padding[0] * 2
|
||||
elif len(padding) == 2:
|
||||
return padding[0], padding[1]
|
||||
else:
|
||||
raise ValueError("Size of padding can only be 2 or 4")
|
||||
elif isinstance(padding, int):
|
||||
pad_w = padding * 2
|
||||
elif padding == "VALID":
|
||||
pad_w = 0
|
||||
elif padding == "SAME":
|
||||
pad_w = kernel[0] - 1
|
||||
else:
|
||||
raise ValueError(f"Unknown padding option {padding}")
|
||||
pad_left = (pad_w + 1) // 2
|
||||
return pad_left, pad_w - pad_left
|
||||
@@ -0,0 +1,181 @@
|
||||
# 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: E731
|
||||
"""Utility functions for implementing Winograd convolutions
|
||||
[*] Fast Algorithms for Convolutional Neural Networks
|
||||
Andrew Lavin, Scott Gray
|
||||
https://arxiv.org/abs/1509.09308
|
||||
https://github.com/andravin/wincnn
|
||||
"""
|
||||
|
||||
from functools import reduce
|
||||
from operator import mul
|
||||
|
||||
import numpy as np
|
||||
|
||||
from tvm.contrib.pickle_memoize import memoize
|
||||
|
||||
from ..utils import const_matrix
|
||||
|
||||
|
||||
# pylint: disable=invalid-name
|
||||
def _cook_toom_convolution(a, n, r):
|
||||
"""Compute Cook-Toom convolution A,B,G matrices"""
|
||||
|
||||
def _F_m(a, n):
|
||||
f = lambda j, i: reduce(mul, ((a[i] - a[k] if k != i else 1) for k in range(0, n - 1)), 1)
|
||||
F = np.fromfunction(np.vectorize(f), (1, n - 1), dtype=int)
|
||||
F = np.diagflat(F)
|
||||
F = np.append(F, np.zeros((n - 1, 1), dtype=int), axis=1)
|
||||
f = lambda i, j: 1 if j == (n - 1) else 0
|
||||
z = np.fromfunction(np.vectorize(f), (1, n), dtype=int)
|
||||
|
||||
return np.append(F, z, axis=0)
|
||||
|
||||
def _A_m(a, m, n):
|
||||
f = lambda i, j: a[i] ** j
|
||||
A = np.fromfunction(np.vectorize(f), (m - 1, n), dtype=int)
|
||||
f = lambda i, j: 1 if j == (n - 1) else 0
|
||||
z = np.fromfunction(np.vectorize(f), (1, n), dtype=int)
|
||||
|
||||
return np.append(A, z, axis=0)
|
||||
|
||||
def _B_m(a, n):
|
||||
f = lambda j, i: reduce(mul, ((a[i] - a[k] if k != i else 1) for k in range(0, n - 1)), 1)
|
||||
Ff = np.fromfunction(np.vectorize(f), (1, n - 1), dtype=int)
|
||||
f = lambda i, nth: (
|
||||
(
|
||||
reduce(mul, [(np.poly1d([1, -a[k]]) if k != i else 1) for k in range(0, n - 1)], 1)
|
||||
).coef[n - 1 - nth - 1]
|
||||
/ Ff[0, i]
|
||||
)
|
||||
F = np.fromfunction(np.vectorize(f), (n - 1, n - 1), dtype=int)
|
||||
f = lambda i, j: -(a[i] ** (n - 1))
|
||||
t = np.fromfunction(np.vectorize(f), (n - 1, 1), dtype=int)
|
||||
T = np.append(np.eye(n - 1), t, axis=1)
|
||||
|
||||
return np.append(F.T.dot(T), np.array([np.eye(n)[n - 1]]), axis=0)
|
||||
|
||||
alpha = n + r - 1
|
||||
|
||||
f = _F_m(a, alpha)
|
||||
|
||||
if f[0, 0] < 0:
|
||||
f[0, :] *= -1
|
||||
|
||||
A = _A_m(a, alpha, n)
|
||||
|
||||
G = _A_m(a, alpha, r).T
|
||||
G = G.dot(np.linalg.inv(f)).T
|
||||
|
||||
B = _B_m(a, alpha)
|
||||
B = B.dot(f.T)
|
||||
|
||||
return (A, B, G)
|
||||
|
||||
|
||||
def _interpolation_points(degree):
|
||||
"""Propose filter points"""
|
||||
|
||||
assert 2 < degree < 18
|
||||
|
||||
# Default interpolation lookup table
|
||||
#
|
||||
# [1] Error Analysis and Improving the Accuracy of Winograd Convolution for Deep Neural Networks
|
||||
# Barbara Barabasz, Andrew Anderson, Kirk M. Soodhalter, David Gregg
|
||||
# https://arxiv.org/abs/1803.10986
|
||||
#
|
||||
|
||||
# pylint: disable=bad-whitespace,line-too-long
|
||||
in_pts = [
|
||||
# {invalid}
|
||||
[],
|
||||
# 01 {E=4.63E-08 on conv2d [1]}
|
||||
[],
|
||||
# 02 {E=7.65E-08 on F( 2,3) [1]}
|
||||
[0, -1, 1],
|
||||
# 03 {E=2.35E-07 on F( 3,3) [1]}
|
||||
[0, -1, 1, 1 / 2],
|
||||
# 04 {E=3.29E-07 on F( 4,3) [1]}
|
||||
[0, -1, 1, 1 / 2, -2],
|
||||
# 05 {E=6.81E-07 on F( 5,3) [1]}
|
||||
[0, -1, 1, 1 / 2, -2, -1 / 2],
|
||||
# 06 {E=8.79E-07 on F( 6,3) [1]}
|
||||
[0, -1, 1, 1 / 2, -1 / 2, 2, -2],
|
||||
# 07 {E=3.71E-06 on F( 7,3) [1]}
|
||||
[0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4],
|
||||
# 08 {E=7.35E-06 on F( 8,3) [1]}
|
||||
[0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4],
|
||||
# 09 {E=2.20E-05 on F( 9,3) [1]}
|
||||
[0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 3 / 4, -4 / 3],
|
||||
# 10 {E=3.22E-05 on F(10,3) [1]}
|
||||
[0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 3 / 4, -4 / 3],
|
||||
# 11 {E=1.09E-04 on F(11,3) [1]}
|
||||
[0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 3 / 4, -4 / 3, 1 / 4],
|
||||
# 12 {E=1.99E-04 on F(12,3) [1]}
|
||||
[0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 1 / 4, -3 / 4, 4 / 3, -4],
|
||||
# 13 {E=5.54E-04 on F(13,3) [1]}
|
||||
[0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 1 / 4, -3 / 4, 4 / 3, 3 / 4, -4 / 3],
|
||||
# 14 {E=8.80E-04 on F(14,3) [1]}
|
||||
[0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 1 / 4, -3 / 4, 4 / 3, -4, 3 / 4, -4 / 3],
|
||||
# 15 {E=1.07E-02 on F(15,3) [1]}
|
||||
[0, -1, 1, 1 / 2, -1 / 2, 2, -2, -1 / 4, 4, 1 / 4, -3 / 4, 4 / 3, -4, 2 / 3, -3 / 2, 3 / 2],
|
||||
# 16 {E=1.93E-02 on F(16,3) [1]}
|
||||
[
|
||||
0,
|
||||
-1,
|
||||
1,
|
||||
1 / 2,
|
||||
-1 / 2,
|
||||
2,
|
||||
-2,
|
||||
-1 / 4,
|
||||
4,
|
||||
1 / 4,
|
||||
-3 / 4,
|
||||
4 / 3,
|
||||
-4,
|
||||
2 / 3,
|
||||
-3 / 2,
|
||||
-2 / 3,
|
||||
3 / 2,
|
||||
],
|
||||
] # pylint: enable=bad-whitespace,line-too-long
|
||||
|
||||
return np.array(in_pts[degree - 1], dtype=np.float64)
|
||||
|
||||
|
||||
@memoize("topi.nn.winograd_matrices", save_at_exit=False)
|
||||
def winograd_transform_matrices(tile_size, kernel_size, out_dtype):
|
||||
"""Compute the A, B, and G transform matrices for `tile_size` as a `tvm.Expr`."""
|
||||
if not 1 < tile_size < 9:
|
||||
raise ValueError(f"Unsupported tile size for Winograd: {tile_size}")
|
||||
if not 2 < kernel_size < 8:
|
||||
raise ValueError(f"Unsupported kernel size for Winograd: {kernel_size}")
|
||||
|
||||
degree = tile_size + kernel_size - 2
|
||||
|
||||
intp_pts = _interpolation_points(degree)
|
||||
A_data, B_data, G_data = _cook_toom_convolution(intp_pts, tile_size, kernel_size)
|
||||
|
||||
out_dtype = "uint16" if out_dtype == "bfloat16" else out_dtype
|
||||
return (
|
||||
const_matrix(A_data.astype(out_dtype), "A"),
|
||||
const_matrix(B_data.astype(out_dtype), "B"),
|
||||
const_matrix(G_data.astype(out_dtype), "G"),
|
||||
)
|
||||
Reference in New Issue
Block a user