diff --git a/python/tvm/relay/qnn/op/qnn.py b/python/tvm/relay/qnn/op/qnn.py index c8ebfc00a21b0..eb82c640807e6 100644 --- a/python/tvm/relay/qnn/op/qnn.py +++ b/python/tvm/relay/qnn/op/qnn.py @@ -349,3 +349,44 @@ def dense(data, input_zero_point, kernel_zero_point, out_dtype) + + +def mul(lhs, rhs, lhs_scale, lhs_zero_point, rhs_scale, rhs_zero_point, + output_scale, output_zero_point): + """Quantized multiplication with numpy-style broadcasting. + + Parameters + ---------- + lhs : relay.Expr + The left hand side quantized input data. + + rhs : relay.Expr + The right hand side quantized input data. + + lhs_scale: float + The scale of the lhs quantized expr. + lhs_zero_point: int + The zero point of lhs quantized expr. + + rhs_scale: float + The scale of the rhs quantized expr. + + rhs_zero_point: int + The zero point of rhs quantized expr. + + output_scale: float + The scale of the output quantized expr. + + output_zero_point: int + The zero point of output quantized expr. + + Returns + ------- + result : relay.Expr + The computed result. + + """ + return _make.mul(lhs, rhs, + lhs_scale, lhs_zero_point, + rhs_scale, rhs_zero_point, + output_scale, output_zero_point) diff --git a/src/relay/qnn/op/mul.cc b/src/relay/qnn/op/mul.cc new file mode 100644 index 0000000000000..958aed8dca249 --- /dev/null +++ b/src/relay/qnn/op/mul.cc @@ -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. + */ + +/*! + * Copyright (c) 2019 by Contributors + * \file src/relay/qnn/op/mul.cc + * \brief QNN mul operator. + */ +#include +#include +#include +#include "../../pass/pattern_util.h" +#include "../util.h" +#include "op_common.h" + +namespace tvm { +namespace relay { +namespace qnn { + +/* + * \brief Canonicalizes the QNN mul op. + * \param attrs The QNN concatenate attrs. + * \param new_args The new mutated args to the call node. + * \param arg_types The types of input and output. + * \return The sequence of Relay ops for mul op. + */ +Expr QnnMulCanonicalize(const Attrs& attrs, const Array& new_args, + const Array& arg_types) { + // Get the attrs. + CHECK_EQ(new_args.size(), 2); + auto& lhs = new_args[0]; + auto& rhs = new_args[1]; + const auto* binary_op_attrs = attrs.as(); + CHECK(binary_op_attrs != nullptr); + auto lhs_scale = binary_op_attrs->lhs_scale; + auto lhs_zero_point = binary_op_attrs->lhs_zero_point; + auto rhs_scale = binary_op_attrs->rhs_scale; + auto rhs_zero_point = binary_op_attrs->rhs_zero_point; + auto output_scale = binary_op_attrs->output_scale; + auto output_zero_point = binary_op_attrs->output_zero_point; + + // Get the input dtype and shape. + CHECK_EQ(arg_types.size(), 3); + auto tensor_type = arg_types[0].as(); + auto input_dtype = tensor_type->dtype; + auto input_shape = tensor_type->shape; + + +/* +A tensor multiplication c = a * b can be written in terms of respective +quantized tensors, scales and zero points as +S_c * (Q_c - zp_c) = S_a * (Q_a - zp_a) * S_b * (Q_b - zp_b). + +We can consider the product (Q_a - zp_a) * (Q_b - zp_b) as a different +quantized tensor of c, Q', with corresponding scale S' = S_a * S_b and zp' = +0. The quantized multiplication then becomes +Q_c = S'/S_c Q' + z_c, +which is essentially a requantization of tensor Q' into tensor Q_c. +*/ + + auto lhs_shifted = Cast(lhs, Int(32)); + auto rhs_shifted = Cast(rhs, Int(32)); + + if (lhs_zero_point != 0) { + auto lhs_zp = MakeConstantScalar(Int(32), lhs_zero_point); + lhs_shifted = Subtract(lhs_shifted, lhs_zp); + } + + if (rhs_zero_point != 0) { + auto rhs_zp = MakeConstantScalar(Int(32), rhs_zero_point); + rhs_shifted = Subtract(rhs_shifted, rhs_zp); + } + + // Create a new tensor Q' + auto output = Multiply(lhs_shifted, rhs_shifted); + + auto scale_new = rhs_scale * lhs_scale; + + // Requantize to get Q_c + output = Requantize(output, input_shape, scale_new, 0, output_scale, + output_zero_point, input_dtype); + + return output; +} + +// QNN Multiplication operator. +QNN_REGISTER_BINARY_OP("mul") +.describe("Elementwise mul with with broadcasting for quantized tensors.") +.set_support_level(11) +.set_attr("FTVMQnnCanonicalize", QnnMulCanonicalize); + +} // namespace qnn +} // namespace relay +} // namespace tvm diff --git a/tests/python/relay/test_qnn_mul.py b/tests/python/relay/test_qnn_mul.py new file mode 100644 index 0000000000000..3bcb48632cfa0 --- /dev/null +++ b/tests/python/relay/test_qnn_mul.py @@ -0,0 +1,251 @@ +# 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. + +import tvm +import numpy as np +from tvm import relay +from tvm.contrib import graph_runtime +import topi.testing + +# "unquantize" a quantized tensor +def recover(data, scale, zp): + return scale * (np.asarray(data) - zp) + + +def generate_golden_output(x_recovered, y_recovered, scale, zp): + mul = x_recovered * y_recovered + output = np.around(mul / scale + zp) + + q_min = np.iinfo(np.uint8).min + q_max = np.iinfo(np.uint8).max + return np.clip(output, q_min, q_max) + + +def test_tflite_same_io_qnn_params(): + data_dtype = "uint8" + + lhs_scale = rhs_scale = output_scale = 0.00784314 + lhs_zero_point = rhs_zero_point = output_zero_point = 127 + + x = relay.var("x", shape=(1, 4), dtype=data_dtype) + y = relay.var("y", shape=(1, 4), dtype=data_dtype) + z = relay.qnn.op.mul( + lhs=x, + rhs=y, + lhs_scale=lhs_scale, + lhs_zero_point=lhs_zero_point, + rhs_scale=rhs_scale, + rhs_zero_point=rhs_zero_point, + output_scale=output_scale, + output_zero_point=output_zero_point, + ) + + func = relay.Function([x, y], z) + mod = relay.Module.from_expr(func) + mod = relay.qnn.transform.CanonicalizeOps()(mod) + func = mod["main"] + + x_datas = [ + np.array((1, 153, 2, 178)).reshape((1, 4)), + np.array((25, 1, 178, 216)).reshape((1, 4)), + np.array((25, 153, 1, 165)).reshape((1, 4)), + ] + y_datas = [ + np.array((204, 178, 1, 8)).reshape((1, 4)), + np.array((204, 178, 191, 1)).reshape((1, 4)), + np.array((204, 178, 1, 191)).reshape((1, 4)), + ] + + for i in range(0, 3): + x_data = x_datas[i] + y_data = y_datas[i] + + x_rec = recover(x_data, lhs_scale, lhs_zero_point) + y_rec = recover(y_data, rhs_scale, rhs_zero_point) + golden = generate_golden_output(x_rec, y_rec, output_scale, + output_zero_point) + + intrp = relay.create_executor("graph", ctx=tvm.cpu(0), target="llvm") + op_res = intrp.evaluate(func)(x_data, y_data) + + np.testing.assert_equal(op_res.asnumpy(), np.uint8(golden)) + + +def test_tflite_different_io_qnn_params(): + data_dtype = "uint8" + + lhs_scale = 0.0156863 + lhs_zero_point = 127 + rhs_scale = 0.0117647 + rhs_zero_point = 85 + output_scale = 0.0235294 + output_zero_point = 128 + + x = relay.var("x", shape=(1, 4), dtype=data_dtype) + y = relay.var("y", shape=(1, 4), dtype=data_dtype) + z = relay.qnn.op.mul( + lhs=x, + rhs=y, + lhs_scale=lhs_scale, + lhs_zero_point=lhs_zero_point, + rhs_scale=rhs_scale, + rhs_zero_point=rhs_zero_point, + output_scale=output_scale, + output_zero_point=output_zero_point, + ) + + func = relay.Function([x, y], z) + mod = relay.Module.from_expr(func) + mod = relay.qnn.transform.CanonicalizeOps()(mod) + func = mod["main"] + + x_datas = [ + np.array((76, 140, 153, 172)).reshape((1, 4)), + np.array((133, 140, 146, 153)).reshape((1, 4)), + np.array((76, 140, 172, 146)).reshape((1, 4)), + ] + y_datas = [ + np.array((136, 119, 128, 17)).reshape((1, 4)), + np.array((136, 119, 111, 94)).reshape((1, 4)), + np.array((136, 119, 17, 128)).reshape((1, 4)), + ] + + for i in range(0, 3): + x_data = x_datas[i] + y_data = y_datas[i] + + x_rec = recover(x_data, lhs_scale, lhs_zero_point) + y_rec = recover(y_data, rhs_scale, rhs_zero_point) + golden = generate_golden_output(x_rec, y_rec, output_scale, + output_zero_point) + + intrp = relay.create_executor("graph", ctx=tvm.cpu(0), target="llvm") + op_res = intrp.evaluate(func)(x_data, y_data) + np.testing.assert_equal(op_res.asnumpy(), np.uint8(golden)) + + +def test_saturation(): + # Same params + data_dtype = "uint8" + lhs_scale = rhs_scale = output_scale = 0.125 + lhs_zero_point = rhs_zero_point = output_zero_point = 0 + + x = relay.var("x", shape=(1, 4), dtype=data_dtype) + y = relay.var("y", shape=(1, 4), dtype=data_dtype) + z = relay.qnn.op.mul( + lhs=x, + rhs=y, + lhs_scale=lhs_scale, + lhs_zero_point=lhs_zero_point, + rhs_scale=rhs_scale, + rhs_zero_point=rhs_zero_point, + output_scale=output_scale, + output_zero_point=output_zero_point, + ) + + func = relay.Function([x, y], z) + mod = relay.Module.from_expr(func) + mod = relay.qnn.transform.CanonicalizeOps()(mod) + func = mod["main"] + + x_data = np.array((255, 1, 1, 0)).reshape((1, 4)) + y_data = np.array((255, 255, 128, 0)).reshape((1, 4)) + + x_rec = recover(x_data, lhs_scale, lhs_zero_point) + y_rec = recover(y_data, rhs_scale, rhs_zero_point) + + golden = generate_golden_output(x_rec, y_rec, output_scale, + output_zero_point) + + intrp = relay.create_executor("graph", ctx=tvm.cpu(0), target="llvm") + op_res = intrp.evaluate(func)(x_data, y_data) + np.testing.assert_equal(op_res.asnumpy(), np.uint8(golden)) + + # Same params, different scale + + lhs_scale = rhs_scale = 0.125 + output_scale = 0.25 + + z = relay.qnn.op.mul( + lhs=x, + rhs=y, + lhs_scale=lhs_scale, + lhs_zero_point=lhs_zero_point, + rhs_scale=rhs_scale, + rhs_zero_point=rhs_zero_point, + output_scale=output_scale, + output_zero_point=output_zero_point, + ) + + func = relay.Function([x, y], z) + mod = relay.Module.from_expr(func) + mod = relay.qnn.transform.CanonicalizeOps()(mod) + func = mod["main"] + + x_data = np.array((255, 1, 1, 0)).reshape((1, 4)) + y_data = np.array((255, 255, 127, 0)).reshape((1, 4)) + + x_rec = recover(x_data, lhs_scale, lhs_zero_point) + y_rec = recover(y_data, rhs_scale, rhs_zero_point) + + golden = generate_golden_output(x_rec, y_rec, output_scale, + output_zero_point) + + intrp = relay.create_executor("graph", ctx=tvm.cpu(0), target="llvm") + op_res = intrp.evaluate(func)(x_data, y_data) + np.testing.assert_equal(op_res.asnumpy(), np.uint8(golden)) + + # All params different + + lhs_scale = 0.5 + rhs_scale = 0.25 + output_scale = 0.125 + + z = relay.qnn.op.mul( + lhs=x, + rhs=y, + lhs_scale=lhs_scale, + lhs_zero_point=lhs_zero_point, + rhs_scale=rhs_scale, + rhs_zero_point=rhs_zero_point, + output_scale=output_scale, + output_zero_point=output_zero_point, + ) + + func = relay.Function([x, y], z) + mod = relay.Module.from_expr(func) + mod = relay.qnn.transform.CanonicalizeOps()(mod) + func = mod["main"] + + x_data = np.array((255, 0, 1, 0)).reshape((1, 4)) + y_data = np.array((0, 128, 64, 0)).reshape((1, 4)) + + x_rec = recover(x_data, lhs_scale, lhs_zero_point) + y_rec = recover(y_data, rhs_scale, rhs_zero_point) + + golden = generate_golden_output(x_rec, y_rec, output_scale, + output_zero_point) + + intrp = relay.create_executor("graph", ctx=tvm.cpu(0), target="llvm") + op_res = intrp.evaluate(func)(x_data, y_data) + np.testing.assert_equal(op_res.asnumpy(), np.uint8(golden)) + + +if __name__ == "__main__": + test_tflite_same_io_qnn_params() + test_tflite_different_io_qnn_params() + test_saturation()