wehub-resource-sync
277 lines
8.9 KiB
Python
277 lines
8.9 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# pylint: disable=invalid-name, 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
|