定点乘法#
定义一些辅助函数用于推理：
import set_env
import tvm
from tvm import relay
import numpy as np
def update_lib(lib, lib_name="lib.so"):
tmp_path = tvm.contrib.utils.tempdir()
lib_path = tmp_path.relpath(lib_name)
lib.export_library(lib_path, fcompile=False)
return tvm.runtime.load_module(lib_path)
def run_llvm(run_mod, params, input_dict, lib_name="lib.so"):
with tvm.transform.PassContext(opt_level=3, disabled_pass={"AlterOpLayout"}):
lib = relay.build(run_mod, target="llvm", params=params)
lib = update_lib(lib, lib_name=lib_name)
exe = tvm.contrib.graph_executor.GraphModule(lib["default"](tvm.cpu()))
exe.run(**input_dict)
return [
exe.get_output(k).asnumpy()
for k in range(exe.get_num_outputs())
]
relay 定点乘法#
ishape = (1, 2)
dtype = "int32"
a = relay.var("a", relay.TensorType(ishape, dtype))
for multiplier, shift, float_value in [
(1288490240, -2, 0.15),
(1395864320, 1, 1.3),
(1288490188, 0, 0.6),
]:
fpm = relay.fixed_point_multiply(a, multiplier, shift)
run_mod = tvm.IRModule.from_expr(fpm)
run_mod.show()
data_in = np.random.randint(0, 1000, size=ishape, dtype=dtype)
inputs = {"a": data_in}
expected_output = run_llvm(run_mod, {}, inputs, lib_name="lib.so")[0]
print(np.round(data_in*float_value), expected_output)
def @main(%a: Tensor[(1, 2), int32]) {
fixed_point_multiply(%a, multiplier=1288490240, shift=-2)
}
[[24. 9.]] [[24 9]]
[[120. 606.]] [[120 606]]
[[191. 480.]] [[191 480]]
def @main(%a: Tensor[(1, 2), int32]) {
fixed_point_multiply(%a, multiplier=1395864320, shift=1)
}
def @main(%a: Tensor[(1, 2), int32]) {
fixed_point_multiply(%a, multiplier=1288490188, shift=0)
}
relay 逐通道定点乘法#
定义用于 relay 的表达式：
from tvm.relay.op import _make
from tvm.relay.expr import Expr
def fixed_point_multiply_per_axis(
x: Expr,
y: Expr,
lshift: Expr,
rshift: Expr,
is_lshift_required : int,
is_rshift_required : int,
axes,
):
"""Fixed point multiplication between data and a fixed point constant expressed as
multiplier * 2^(-shift), where multiplier is a Q-number with 31 fractional bits
Parameters
----------
x : Expr
Input argument.
y : Expr
Multiplier of a fixed floating point number described as multiplier*2^(-shift).
lshift : Expr
Left shifts of a fixed floating point number described as multiplier*2^(-shift).
rshift : Expr
Right shifts of a fixed floating point number described as multiplier*2^(-shift).
is_lshift_required : int
Whether we need to do left shift or not.
is_rshift_required : int
Whether we need to do right shift or not.
Returns
-------
z : Expr
The result.
"""
return _make.fixed_point_multiply_per_axis(x, y, lshift, rshift, is_lshift_required, is_rshift_required, axes)
# from tvm.relay.testing.temp_op_attr import TempOpAttr
from tvm import te
a_shape = [2, 256, 16]
b_shape = [256]
shape = a_shape
# Verify accuracy
x_np = (
np.random.randint(-1000, 1000, size=np.prod(a_shape)).reshape(a_shape).astype("int32")
)
y_np = (
np.random.randint(-1000, 1000, size=np.prod(b_shape)).reshape(b_shape).astype("int32")
)
lsh_np = np.random.randint(0, 10, size=np.prod(b_shape)).reshape(b_shape).astype("int32")
rsh_np = np.random.randint(0, 10, size=np.prod(b_shape)).reshape(b_shape).astype("int32")
b_np = (
np.random.randint(-1000, 1000, size=np.prod(a_shape)).reshape(a_shape).astype("int32")
)
inputs = {"X":x_np, "Y": y_np, "l_shift": lsh_np, "l_shift": rsh_np}
shift_shape = [shape[1]]
x = relay.var("X", shape=shape, dtype="int32")
y = relay.var("Y", shape=shift_shape, dtype="int32")
l_shift = relay.var("l_shift", shape=shift_shape, dtype="int32")
r_shift = relay.var("r_shift", shape=shift_shape, dtype="int32")
out = fixed_point_multiply_per_axis(x, y, l_shift, r_shift, 31, 1, b_shape)
mod = tvm.IRModule.from_expr(out)
mod.show()
expected_output = run_llvm(run_mod, {}, inputs, lib_name="lib.so")[0]
print(expected_output)
def @main(%X: Tensor[(2, 256, 16), int32], %Y: Tensor[(256), int32], %l_shift: Tensor[(256), int32], %r_shift: Tensor[(256), int32]) {
fixed_point_multiply_per_axis(%X, %Y, %l_shift, %r_shift, is_lshift_required=True, is_rshift_required=True, axes=[256])
}
[[53496115 0]]
也可以使用 QNN：
axis = 1
ishape = [1, 128, 56, 56]
in_scale_const = (1.7, 0.6)
x = relay.var("data", shape=ishape, dtype="int32")
if isinstance(in_scale_const, tuple):
in_scale = list(in_scale_const) * (ishape[axis] // len(in_scale_const))
else:
in_scale = [in_scale_const] * ishape[axis]
assert len(in_scale) == ishape[axis]
iscale = relay.const(in_scale)
izero = relay.const(0)
oscale = relay.const(1.0)
ozero = relay.const(0)
op = relay.qnn.op.requantize(x, iscale, izero, oscale, ozero, axis=axis, out_dtype="int32")
mod = tvm.IRModule.from_expr(op)
mod = relay.transform.InferType()(mod)
mod.show()
x_np = (
np.random.randint(-1000, 1000, size=np.prod(ishape)).reshape(ishape).astype("int32")
)
inputs = {"data": x_np}
expected_output = run_llvm(mod, {}, inputs, lib_name="lib.so")[0]
def @main(%data: Tensor[(1, 128, 56, 56), int32] /* ty=Tensor[(1, 128, 56, 56), int32] */) -> Tensor[(1, 128, 56, 56), int32] {
qnn.requantize(%data, meta[relay.Constant][0] /* ty=Tensor[(128), float32] */, 0 /* ty=int32 */, 1f /* ty=float32 */, 0 /* ty=int32 */, axis=1, out_dtype="int32") /* ty=Tensor[(1, 128, 56, 56), int32] */
}
from tvm import te, topi, tir as T, relay
import tvm
from tvm.topi import tag
# from tvm.relay.op.op import register_compute, register_shape_func
# from tvm.relay.op.op import register_broadcast_schedule, register_injective_schedule
# from tvm.relay.op.op import register_pattern, OpPattern
@tvm.te.tag_scope(tag=tag.ELEMWISE)
def q_multiply_shift(
x: te.Tensor,
y: te.Tensor,
q: int,
left_shift: te.Tensor,
right_shift: te.Tensor,
is_left_shift_required: int):
# Only int32 types are supported (any number of lanes is allowed)
hp_dtype = "int64"
lp_dtype = "int32"
assert y.dtype == lp_dtype
assert left_shift.dtype == lp_dtype
assert right_shift.dtype == lp_dtype
one = T.const(1, hp_dtype)
def _compute(*indices):
# 0) 获取值
value = x(*indices)
multiplier = y(*indices)
ls = left_shift(*indices)
rs = right_shift(*indices)
# 1) Cast and Multiply the integer multiplier
value = value.astype(hp_dtype)
multiplier = multiplier.astype(hp_dtype)
value = T.Select(T.const(is_left_shift_required, "bool"),
value << ls, value)
# 2) Perform the multiplication in higher precision.
value = value * multiplier
# 3) Find the rounding scalar
total_right_shift = ls + q
pos_rounding_value = (one << (total_right_shift - 1))
value = value + pos_rounding_value
print(total_right_shift)
# 4) Simply right shift the result to get the final output.
value = value >> total_right_shift
# 5) The fixed point multiplication keeps the value in int32 range. Casting back to int32.
return value.astype(x.dtype)
return te.compute(x.shape, _compute)
shape = 1, 2
lp_dtype = "int32"
x = te.placeholder(shape, name="x", dtype="int32")
y = te.placeholder(shape, name="y", dtype=lp_dtype)
left_shift = te.placeholder(shape, name="left_shift", dtype=lp_dtype)
right_shift = te.placeholder(shape, name="right_shift", dtype=lp_dtype)
# multipliers_shifts = te.placeholder(shape, name="multipliers_shifts", dtype="int32")
q = 31 # int8
# q=8 -> uint8
z = q_multiply_shift(x, y, q, left_shift, right_shift, is_left_shift_required=1)
s = te.create_schedule(z.op)
f = tvm.build(s, [x, y, left_shift, right_shift, z], "llvm")
dev = tvm.cpu(0)
a_np = np.ones(shape).astype(x.dtype) * 125333333
multiplier_np = np.ones(shape).astype(lp_dtype) * 3650000
ls_np = np.ones(shape).astype(lp_dtype) * 1
rs_np = np.ones(shape).astype(lp_dtype) * -1
a = tvm.nd.array(a_np, dev)
multiplier = tvm.nd.array(multiplier_np, dev)
ls = tvm.nd.array(ls_np, dev)
rs = tvm.nd.array(rs_np, dev)
c = tvm.nd.array(np.zeros(shape, dtype=z.dtype), dev)
f(a, multiplier, ls, rs, c)
print(a, multiplier, c)
left_shift[i0, i1] + 31
[[125333333 125333333]] [[3650000 3650000]] [[213025 213025]]