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