[RFC][Quantization] Designing and lowering of quantized ops

include/tvm/relay/qnn/attrs.h

@@ -0,0 +1,66 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/relay/qnn/attrs.h
+ * \brief Auxiliary attributes for quantized nn operators.
+ */
+#ifndef TVM_RELAY_QNN_ATTRS_H_
+#define TVM_RELAY_QNN_ATTRS_H_
+
+#include <tvm/attrs.h>
+#include <string>
+
+namespace tvm {
+namespace relay {
+
+struct QuantizeAttrs : public tvm::AttrsNode<QuantizeAttrs> {
+  int32_t output_zero_point;
+  double output_scale;
+  DataType out_dtype;
+
+  TVM_DECLARE_ATTRS(QuantizeAttrs, "relay.attrs.QuantizeAttrs") {
+    TVM_ATTR_FIELD(out_dtype)
+        .describe("Output data type, can be one of [int8 or uint8].");
+
+    TVM_ATTR_FIELD(output_zero_point)
+        .describe("The zero_point for the activation of this op.");
+
+    TVM_ATTR_FIELD(output_scale)
+        .describe("The scale for the activation of this op.");
+  }
+};
+
+struct DequantizeAttrs : public tvm::AttrsNode<DequantizeAttrs> {
+  int32_t input_zero_point;
+  double input_scale;
+
+  TVM_DECLARE_ATTRS(QuantizeAttrs, "relay.attrs.QuantizeAttrs") {
+    TVM_ATTR_FIELD(input_zero_point)
+        .describe("The zero_point for the input tensor of this op.");
+
+    TVM_ATTR_FIELD(input_scale)
+        .describe("The scale for the input tensor of this op.");
+  }
+};
+
+}  // namespace relay
+}  // namespace tvm
+
+#endif  // TVM_RELAY_QNN_ATTRS_H_

python/tvm/relay/__init__.py

@@ -49,6 +49,9 @@
 from . import backend
 from . import quantize
 
+# Dialects
+from . import qnn
+
 from .scope_builder import ScopeBuilder
 
 # Span

python/tvm/relay/qnn/__init__.py

@@ -0,0 +1,21 @@
+# 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.
+
+"""Neural network related operators."""
+from __future__ import absolute_import as _abs
+from . import op
+from . import ir_pass

python/tvm/relay/qnn/_qnn.py

@@ -0,0 +1,23 @@
+# 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.
+
+"""Internal module for quantization."""
+
+from __future__ import absolute_import
+from tvm._ffi.function import _init_api
+
+_init_api("relay._qnn", __name__)

python/tvm/relay/qnn/ir_pass.py

@@ -0,0 +1,37 @@
+# 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.
+
+"""Automatic quantization toolkit."""
+from __future__ import absolute_import
+
+from . import _qnn
+
+def rewrite(expr):
+    """
+    Rewrites the high-level quantized ops into low-level exisiting Relay ops.
+
+    Parameters
+    ----------
+    expr : tvm.relay.Expr
+        The input expression.
+
+    Returns
+    -------
+    expr : tvm.relay.Expr
+        The output expression.
+    """
+    return _qnn.rewrite(expr)

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

@@ -0,0 +1,21 @@
+# 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.
+
+"""Neural network related operators."""
+
+from __future__ import absolute_import as _abs
+from .qnn import *

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

@@ -0,0 +1,22 @@
+# 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.
+
+"""Constructor APIs"""
+
+from ...._ffi.function import _init_api
+
+_init_api("relay.op.qnn._make", __name__)

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

@@ -0,0 +1,66 @@
+# 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.
+
+"""Neural network operations."""
+
+from __future__ import absolute_import as _abs
+from . import _make
+
+def quantize(input_data, output_zero_point, output_scale, out_dtype='int8'):
+    r""" Quantize op
+     This operator takes float32 as input and produces quantized int8 or unit8 as output.
+     The input tensor can be of any shape. The output shape is the same as input shape.
+     ..math::
+            \mbox{out}[x] =
+                \mbox{clamp(round(input_tensor/output_scale) + output_zero_point);
+                 out_dtype::min, out_dtype::max}
+     Parameters
+    ----------
+    input_data : tvm.relay.Expr
+        The input tensor to be quantized. Can be of type float32.
+    output_zero_point :
+        The output zero_point.
+    output_scale:
+        The output scale.
+    input_dtype:
+        The data type of the input tensor. Can be [int8, uint8]
+    Returns
+    -------
+    result : tvm.relay.Expr
+        The computed result.
+    """
+    return _make.quantize(input_data, output_zero_point, output_scale, out_dtype)
+
+def dequantize(input_data, input_zero_point, input_scale):
+    r""" Dequantize op
+     This operator takes quantized int8 and unit8 as input and produces
+    dequantized float32 as output. The output shape is the same as input shape. The input
+    tensor can be of any shape.
+     Parameters
+    ----------
+    input_data : tvm.relay.Expr
+        The input tensor to be dequantized. Can be of type [int8, uint8].
+    input_zero_point :
+        The output zero_point.
+    input_scale:
+        The output scale.
+    Returns
+    -------
+    result : tvm.relay.Expr
+        The computed result.
+    """
+    return _make.dequantize(input_data, input_zero_point, input_scale)

python/tvm/relay/quantize/__init__.py

@@ -19,4 +19,5 @@
 from __future__ import absolute_import as _abs
 
 from .quantize import *
+from .rewrite import *
 from ._annotate import register_annotate_function

src/relay/pass/pattern_util.h

@@ -34,7 +34,7 @@
 #include <tvm/relay/attrs/nn.h>
 #include <tvm/relay/attrs/transform.h>
 #include <string>
-
+#include <utility>
 
 namespace tvm {
 namespace relay {
@@ -373,6 +373,26 @@ inline Expr Copy(Expr data) {
 }
 
 
+inline Expr Where(const Expr& condition, const Expr& x, const Expr& y) {
+  static const Op& op = Op::Get("where");
+  return CallNode::make(op, {condition, x, y});
+}
+
+inline Expr GreaterEqual(const Expr& lhs, const Expr& rhs) {
+  static const Op& op = Op::Get("greater_equal");
+  return CallNode::make(op, {lhs, rhs}, Attrs(), {});
+}
+
+inline Expr Full(Expr fill_value,
+              Array<IndexExpr> shape,
+              DataType dtype) {
+  auto attrs = make_node<InitOpAttrs>();
+  attrs->shape = std::move(shape);
+  attrs->dtype = std::move(dtype);
+  static const Op& op = Op::Get("full");
+  return CallNode::make(op, {fill_value}, Attrs(attrs), {});
+}
+
 Expr MakeConcatenate(Expr data, int axis);
 
 Expr MakeStridedSlice(Expr data, Array<Integer> begin, Array<Integer> end, Array<Integer> strides);

src/relay/qnn/op/dequantize.cc

@@ -0,0 +1,76 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file src/relay/qnn/op/dequantize.cc
+ * \brief Dequantize operator that converts from quantized domain to
+ * unquantized domain.
+ */
+
+#include <tvm/relay/op.h>
+#include <tvm/relay/qnn/attrs.h>
+#include "../util.h"
+
+namespace tvm {
+namespace relay {
+
+TVM_REGISTER_NODE_TYPE(DequantizeAttrs);
+
+bool DequantizeRel(const Array<Type>& types,
+                 int num_inputs,
+                 const Attrs& attrs,
+                 const TypeReporter& reporter) {
+  CHECK_EQ(types.size(), 2);
+  const auto* data = types[0].as<TensorTypeNode>();
+  const auto input_dtype = data->dtype;
+  CHECK(IsValidOpInputType(QuantizeOpType::Dequantize, input_dtype))
+    << "Input type should be one of the quantized types [unit8, int8] but was " <<  input_dtype;
+  const Array<tvm::Expr> oshape = data->shape;
+  // assign output type
+  reporter->Assign(types[1], TensorTypeNode::make(oshape, Float(32)));
+  return true;
+}
+
+Expr MakeDequantize(Expr data,
+                  int32_t input_zero_point,
+                  double input_scale) {
+  auto attrs = make_node<DequantizeAttrs>();
+  attrs->input_scale = input_scale;
+  attrs->input_zero_point = input_zero_point;
+  static const Op& op = Op::Get("qnn.dequantize");
+  return CallNode::make(op, {data}, Attrs(attrs), {});
+}
+
+RELAY_REGISTER_OP("qnn.dequantize")
+    .describe(R"code(Quantizes the input and produces quantized output.
+
+The input is always quantized (int8, uint8) and will be converted to float32 given input scale and zero_point.
+- **data**: Quantized tensor of any shape to dequantize. The input data can be of floating point
+)code" TVM_ADD_FILELINE)
+.set_attrs_type_key("relay.attrs.DequantizeAttrs")
+.set_num_inputs(1)
+.add_argument("data", "Tensor", "The tensor to dequantize.")
+.set_support_level(10)
+.add_type_rel("Dequantize", DequantizeRel);
+
+TVM_REGISTER_API("relay.op.qnn._make.dequantize")
+.set_body_typed(MakeDequantize);
+
+}  // namespace relay
+}  // namespace tvm