Quantization

include/tvm/relay/dataflow_matcher.h

@@ -87,7 +87,8 @@ bool MatchPattern(DFPattern pattern, Expr expr);
  * \return Return An Expr with every match of the pattern inside the callbacks rewritten by the
  * functions inside the callbacks
  */
-Expr RewritePatterns(Array<DFPatternCallback> callbacks, Expr expr, IRModule mod = IRModule());
+Expr RewritePatterns(Array<DFPatternCallback> callbacks, Expr expr, IRModule mod = IRModule(),
+                     int allow_overlapping_groups = 0);
 
 /*!
  * \brief Partition all matches of a DFPattern inside an Expr into separate Function calls

include/tvm/relay/qnn/attrs.h

@@ -78,13 +78,18 @@ struct QuantizeAttrs : public tvm::AttrsNode<QuantizeAttrs> {
 /*! \brief Attribute for dequantize operator */
 struct DequantizeAttrs : public tvm::AttrsNode<DequantizeAttrs> {
   int axis;
+  DataType out_dtype;
 
   TVM_DECLARE_ATTRS(DequantizeAttrs, "relay.attrs.DequantizeAttrs") {
     TVM_ATTR_FIELD(axis)
         .describe(
             "The channel axis for channel wise dequantization. Default value is -1,"
             "which corresponds to the last axis.")
         .set_default(-1);
+    TVM_ATTR_FIELD(out_dtype)
+        .describe(
+            "The datatype we are dequantizing to (float32 or int32). Defaults to float32.")
+        .set_default(DataType::Float(32));
   }
 };
 

python/tvm/data/__init__.py

@@ -0,0 +1,18 @@
+# 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.
+# pylint: disable=wildcard-import, redefined-builtin, invalid-name
+from ._dataset_manager import DatasetManager, TFDatasetManager, RandomDatasetManager

python/tvm/data/_dataset_manager.py

@@ -0,0 +1,136 @@
+# 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.
+
+"""Wrapper classes to expose datasets during quantization."""
+
+import numpy as np
+
+class DatasetManager:
+    """Simple wrapper class to expose datasets in quantization."""
+
+    def get_next_batch(self):
+        """Returns the next batch of data.
+
+        Returns
+        -------
+        inputs : List
+            The inputs to be provided to the graph.
+            The list is of the form [batched_input_1, batched_input_2, ..., batched_input_n]
+
+        labels: List
+            The expected outputs of the graph.
+            The length of labels should be equal to the batch size.
+        """
+        raise NotImplementedError
+
+    def batch_size(self):
+        """Returns the size of each batch the dataset manager has.
+
+        Returns
+        -------
+        batch_size : int
+            The number of inputs in each batch.
+        """
+    def num_batches(self):
+        """Returns the number of batches the dataset manager has.
+
+        Returns
+        ------
+        num_batches : int
+            The number of batches the dataset manager contains.
+        """
+        raise NotImplementedError
+
+    def is_empty(self):
+        """Checks whether the dataset manager has gone through
+        all its batches.
+        Returns
+        -------
+        is_empty : bool
+            True if there are batches left, False if there are no more
+            batches.
+        """
+        raise NotImplementedError
+
+    def reset(self):
+        """Resets the counter in the dataset manager to the beginning."""
+        raise NotImplementedError
+
+
+class TFDatasetManager(DatasetManager):
+    """DatasetManager wrapping a tensorflow dataset."""
+
+    def __init__(self, tf_dataset, batch_size, total_batches):
+        self.idx = 0
+        self.total_batches = total_batches
+        self.batch_sz = batch_size
+        self.tf_dataset = tf_dataset
+        self.tf_iter = iter(self.tf_dataset)
+
+    def get_next_batch(self):
+        if self.is_empty():
+            raise IndexError
+        self.idx += 1
+
+        data, label = next(self.tf_iter)
+
+        return [data.numpy()], label.numpy()
+
+    def num_batches(self):
+        return self.total_batches
+
+    def batch_size(self):
+        return self.batch_sz
+
+    def is_empty(self):
+        return self.idx >= self.total_batches
+
+    def reset(self):
+        self.tf_iter = iter(self.tf_dataset)
+        self.idx = 0
+
+
+class RandomDatasetManager(DatasetManager):
+    """DatasetManager that creates a random input of a specific shape.
+    This class is mostly used for testing, and as an example of how to
+    implement a DatasetManager.
+    """
+
+    def __init__(self, data_shape, dtype, batch_size, total_batches):
+        self.idx = 0
+        self.data_shape = data_shape
+        self.dtype = dtype
+        self.batch_sz = batch_size
+        self.total_batches = total_batches
+
+    def get_next_batch(self):
+        if self.is_empty():
+            raise IndexError
+        self.idx += 1
+        return [np.random.randn(*self.data_shape).astype(self.dtype)], [None]
+
+    def batch_size(self):
+        return self.batch_sz
+
+    def num_batches(self):
+        return self.total_batches
+
+    def is_empty(self):
+        return self.idx >= self.total_batches
+
+    def reset(self):
+        self.idx = 0

python/tvm/relay/dataflow_pattern/__init__.py

@@ -799,7 +799,7 @@ def __init__(self, require_type=False):
         self.pattern = None
         self.require_type = require_type
 
-    def rewrite(self, expr: Expr) -> Expr:
+    def rewrite(self, expr: Expr, allow_overlapping_groups: bool = False) -> Expr:
         """
         Rewrite expression with this callback
 
@@ -813,7 +813,7 @@ def rewrite(self, expr: Expr) -> Expr:
         result : tvm.relay.Expr
             The Expression with matched subgraphs rewritten by the callbacks.
         """
-        return rewrite(self, expr)
+        return rewrite(self, expr, allow_overlapping_groups = allow_overlapping_groups)
 
     def callback(self, pre: Expr, post: Expr, node_map: tvm.ir.container.Map) -> Expr:
         """
@@ -843,7 +843,8 @@ def __init__(self, pattern, callback, require_type):
         self.__init_handle_by_constructor__(ffi.DFPatternCallback, pattern, callback, require_type)
 
 
-def rewrite(callbacks, expr: Expr, mod: Optional[_ir.IRModule] = None) -> Expr:
+def rewrite(callbacks, expr: Expr, mod: Optional[_ir.IRModule] = None,
+            allow_overlapping_groups: bool = False) -> Expr:
     """
     Rewrite expression with the given callbacks.
 
@@ -868,8 +869,7 @@ def rewrite(callbacks, expr: Expr, mod: Optional[_ir.IRModule] = None) -> Expr:
     for callback in callbacks:
         assert callback.pattern is not None
         tmp.append(_DFPatternCallback(callback.pattern, callback.callback, callback.require_type))
-
-    return ffi.rewrite(tmp, expr, mod)
+    return ffi.rewrite(tmp, expr, mod, allow_overlapping_groups)
 
 
 def partition(

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

@@ -20,10 +20,9 @@
 from __future__ import absolute_import as _abs
 from tvm.relay.expr import Tuple, TupleWrapper
 from tvm.relay.op.nn.utils import get_pad_tuple2d
-from . import _make
 from ... import op as reg
 from ...op import OpPattern
-
+from . import _make
 
 def requantize(
     data,
@@ -118,7 +117,7 @@ def quantize(data, output_scale, output_zero_point, axis=-1, out_dtype="int8"):
     return _make.quantize(data, output_scale, output_zero_point, axis, out_dtype)
 
 
-def dequantize(data, input_scale, input_zero_point, axis=-1):
+def dequantize(data, input_scale, input_zero_point, axis=-1, out_dtype="float32"):
     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
@@ -134,13 +133,15 @@ def dequantize(data, input_scale, input_zero_point, axis=-1):
         The input scale.
     axis : int
         The channel axis for quantization. Default value is -1 which corresponds to the last axis.
+    out_dtype : str, optional
+        The output type to dequantize to. Can be either float32 or int32.
     Returns
     -------
     result : tvm.relay.Expr
         The computed result.
     """
 
-    return _make.dequantize(data, input_scale, input_zero_point, axis)
+    return _make.dequantize(data, input_scale, input_zero_point, axis, out_dtype)
 
 
 def concatenate(data, input_scales, input_zero_points, output_scale, output_zero_point, axis):
@@ -611,7 +612,6 @@ def subtract(
         output_zero_point,
     )
 
-
 # register fuse pattern for qnn ops
 reg.register_pattern("qnn.quantize", OpPattern.OPAQUE)
 reg.register_pattern("qnn.dequantize", OpPattern.OPAQUE)

python/tvm/relay/transform/quantize/__init__.py

@@ -0,0 +1,47 @@
+# 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.
+# pylint: disable=wildcard-import, redefined-builtin, invalid-name
+"""The namespace containing quantization and calibration passes"""
+from ._calibration_callback import (
+    CalibrationCallback,
+    GlobalCalibrationCallback,
+    AverageMaxCalibrationCallback,
+)
+from ._quantizer_patterns import (
+    QuantizerPattern,
+    Conv2DBiasAddPattern,
+    Conv2DPattern,
+    DensePattern,
+    DenseBiasAddPattern,
+    AddPattern,
+    MultiplyPattern,
+    PerChannelPattern
+)
+from ._average_max_channel_patterns import (
+    AverageMaxPerChannelConv2DBiasAddPattern,
+    AverageMaxPerChannelConv2DPattern,
+    AverageMaxPerChannelDenseBiasAddPattern,
+    AverageMaxPerChannelDensePattern
+)
+
+from ._quantizer_pattern_utils import all_patterns, average_max_per_channel_patterns
+
+from ._quantizer import Quantizer
+from ._calibrator import QuantizationCalibrator
+from ._requantizer import Requantizer
+
+from . import _ffi as ffi