Add support for quantized multiply to Relay (#4141)

This patch adds multiply operator for quantized tensors.
The details of the quantized multiplication are outlined
in the code.

This builds on pull request 3927 and includes the changes
Animesh mentions in the comments on that request.

Change-Id: I555715b53d0266a91d5c03dc3dfe8fc31e7ce4e1

python/tvm/relay/qnn/op/qnn.py

@@ -349,3 +349,45 @@ 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)

src/relay/qnn/op/mul.cc

+/*
+ * 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 <tvm/relay/analysis.h>
+#include <tvm/relay/op_attr_types.h>
+#include <tvm/relay/qnn/attrs.h>
+#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<Expr>& new_args,
+                        const Array<tvm::relay::Type>& 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<QnnBinaryOpAttrs>();
+  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<TensorTypeNode>();
+  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<FTVMLegalize>("FTVMQnnCanonicalize", QnnMulCanonicalize);
+
+}  // namespace qnn
+}  // namespace relay
+}  // namespace tvm

tests/python/relay/test_qnn_mul.py

+# 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()