wehub-resource-sync
111 lines
3.7 KiB
Python
111 lines
3.7 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
|
|
"""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:])
|