1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements.  See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership.  The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License.  You may obtain a copy of the License at
9  *
10  *   http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied.  See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
20 /*!
21  * \file src/relay/qnn/op/mul.cc
22  * \brief QNN mul operator.
23  */
24 #include <tvm/relay/analysis.h>
25 #include <tvm/relay/op_attr_types.h>
26 #include <tvm/relay/qnn/attrs.h>
27 #include "../../pass/pattern_util.h"
28 #include "../util.h"
29 #include "op_common.h"
30 
31 namespace tvm {
32 namespace relay {
33 namespace qnn {
34 
35 /*
36  * \brief Canonicalizes the QNN mul op.
37  * \param attrs The QNN concatenate attrs.
38  * \param new_args The new mutated args to the call node.
39  * \param arg_types The types of input and output.
40  * \return The sequence of Relay ops for mul op.
41  */
QnnMulCanonicalize(const Attrs & attrs,const Array<Expr> & new_args,const Array<tvm::relay::Type> & arg_types)42 Expr QnnMulCanonicalize(const Attrs& attrs, const Array<Expr>& new_args,
43                         const Array<tvm::relay::Type>& arg_types) {
44   // Get the attrs.
45   CHECK_EQ(new_args.size(), 2);
46   auto& lhs = new_args[0];
47   auto& rhs = new_args[1];
48   const auto* binary_op_attrs = attrs.as<QnnBinaryOpAttrs>();
49   CHECK(binary_op_attrs != nullptr);
50   auto lhs_scale = binary_op_attrs->lhs_scale;
51   auto lhs_zero_point = binary_op_attrs->lhs_zero_point;
52   auto rhs_scale = binary_op_attrs->rhs_scale;
53   auto rhs_zero_point = binary_op_attrs->rhs_zero_point;
54   auto output_scale = binary_op_attrs->output_scale;
55   auto output_zero_point = binary_op_attrs->output_zero_point;
56 
57   // Get the input dtype and shape.
58   CHECK_EQ(arg_types.size(), 3);
59   auto tensor_type = arg_types[0].as<TensorTypeNode>();
60   auto input_dtype = tensor_type->dtype;
61   auto input_shape = tensor_type->shape;
62 
63   /*
64   A tensor multiplication c = a * b can be written in terms of respective
65   quantized tensors, scales and zero points as
66   S_c * (Q_c - zp_c) = S_a * (Q_a - zp_a) * S_b * (Q_b - zp_b).
67 
68   We can consider the product (Q_a - zp_a) * (Q_b - zp_b) as a different
69   quantized tensor of c, Q', with corresponding scale S' = S_a * S_b and zp' =
70   0. The quantized multiplication then becomes
71   Q_c = S'/S_c Q' + z_c,
72   which is essentially a requantization of tensor Q' into tensor Q_c.
73   */
74 
75   auto lhs_shifted = Cast(lhs, Int(32));
76   auto rhs_shifted = Cast(rhs, Int(32));
77 
78   if (lhs_zero_point != 0) {
79     auto lhs_zp = MakeConstantScalar(Int(32), lhs_zero_point);
80     lhs_shifted = Subtract(lhs_shifted, lhs_zp);
81   }
82 
83   if (rhs_zero_point != 0) {
84     auto rhs_zp = MakeConstantScalar(Int(32), rhs_zero_point);
85     rhs_shifted = Subtract(rhs_shifted, rhs_zp);
86   }
87 
88   // Create a new tensor Q'
89   auto output = Multiply(lhs_shifted, rhs_shifted);
90 
91   auto scale_new = rhs_scale * lhs_scale;
92 
93   // Requantize to get Q_c
94   output = Requantize(output, input_shape, scale_new, 0, output_scale,
95     output_zero_point, input_dtype);
96 
97   return output;
98 }
99 
100 // QNN Multiplication operator.
101 QNN_REGISTER_BINARY_OP("mul")
102 .describe("Elementwise mul with with broadcasting for quantized tensors.")
103 .set_support_level(11)
104 .set_attr<FTVMLegalize>("FTVMQnnCanonicalize", QnnMulCanonicalize);
105 
106 }  // namespace qnn
107 }  // namespace relay
108 }  // namespace tvm
109