[Relay][Quantization] Extend FakeQuantizationToInteger to more ops (a…

…pache#8241)

* support scalars in quantize and requantize

* Add affine type support for ops with multipe output, use it in concat, move to header

* support new ops, refactor tests

* add more binary ops

fix pylint

fix black

black broke pylint

oops on black

* fix a typo in a branch and add a test that hits it

* improve comments

include/tvm/ir/affine_type.h

@@ -0,0 +1,146 @@
+/*
+ * 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/ir/affine_type.h
+ * \brief Quantized Tensor Types.
+ */
+#ifndef TVM_IR_AFFINE_TYPE_H_
+#define TVM_IR_AFFINE_TYPE_H_
+
+#include <tvm/ir/expr.h>
+#include <tvm/ir/type.h>
+
+namespace tvm {
+
+/*!
+ * \brief AffineType representation
+ * \sa AffineType
+ */
+class AffineTypeNode : public Object {
+ public:
+  /*!
+   * \brief Span that points to the original source code.
+   *        Reserved debug information.
+   */
+  mutable Span span;
+
+  static constexpr const char* _type_key = "AffineType";
+  static constexpr const bool _type_has_method_sequal_reduce = true;
+  static constexpr const bool _type_has_method_shash_reduce = true;
+  TVM_DECLARE_BASE_OBJECT_INFO(AffineTypeNode, Object);
+};
+
+/*!
+ * \brief Managed reference to AffineTypeNode.
+ * \sa AffineTypeNode
+ */
+class AffineType : public ObjectRef {
+ public:
+  TVM_DEFINE_OBJECT_REF_METHODS(AffineType, ObjectRef, AffineTypeNode);
+};
+
+/*!
+ * \brief TensorAffineType representation
+ * \sa TensorAffineType
+ *
+ *  This Type represents a quantized integer tensor that can be converted
+ *  back to real space via the x_real = scale * (x_quant - zero_point)
+ */
+class TensorAffineTypeNode : public AffineTypeNode {
+ public:
+  /*! \brief The scale of this type */
+  RelayExpr scale;
+  /*! \brief The zero point of this type */
+  RelayExpr zero_point;
+  /*! \brief The data type of this type */
+  DataType dtype;
+
+  void VisitAttrs(tvm::AttrVisitor* v) {
+    v->Visit("scale", &scale);
+    v->Visit("zero_point", &zero_point);
+    v->Visit("dtype", &dtype);
+  }
+
+  bool SEqualReduce(const TensorAffineTypeNode* other, SEqualReducer equal) const {
+    equal->MarkGraphNode();
+    return equal(scale, other->scale) && equal(zero_point, other->zero_point) &&
+           equal(dtype, other->dtype);
+  }
+
+  void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce->MarkGraphNode();
+    hash_reduce(scale);
+    hash_reduce(zero_point);
+    hash_reduce(dtype);
+  }
+
+  static constexpr const char* _type_key = "TensorAffineType";
+  TVM_DECLARE_BASE_OBJECT_INFO(TensorAffineTypeNode, AffineTypeNode);
+};
+
+/*!
+ * \brief Managed reference to AffineTypes.
+ * \sa AffineTypeNode
+ */
+class TensorAffineType : public AffineType {
+ public:
+  TVM_DLL TensorAffineType(RelayExpr scale, RelayExpr zero_point, DataType dtype);
+
+  TVM_DEFINE_OBJECT_REF_METHODS(TensorAffineType, AffineType, TensorAffineTypeNode);
+};
+
+/*!
+ * \brief TupleAffineType representation
+ * \sa TupleAffineType
+ */
+class TupleAffineTypeNode : public AffineTypeNode {
+ public:
+  /*! \brief The types of this tuple*/
+  Array<TensorAffineType> types;
+
+  void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("types", &types); }
+
+  bool SEqualReduce(const TupleAffineTypeNode* other, SEqualReducer equal) const {
+    equal->MarkGraphNode();
+    return equal(types, other->types);
+  }
+
+  void SHashReduce(SHashReducer hash_reduce) const {
+    hash_reduce->MarkGraphNode();
+    hash_reduce(types);
+  }
+
+  static constexpr const char* _type_key = "TupleAffineType";
+  TVM_DECLARE_BASE_OBJECT_INFO(TupleAffineTypeNode, AffineTypeNode);
+};
+
+/*!
+ * \brief Managed reference to TupleAffineTypes.
+ * \sa TupleAffineType
+ */
+class TupleAffineType : public AffineType {
+ public:
+  TVM_DLL TupleAffineType(Array<TensorAffineType> types);
+
+  TVM_DEFINE_OBJECT_REF_METHODS(TupleAffineType, AffineType, TupleAffineTypeNode);
+};
+
+}  // namespace tvm
+#endif  // TVM_IR_AFFINE_TYPE_H_

python/tvm/ir/__init__.py

@@ -21,6 +21,7 @@
 from .type import Type, TypeKind, PrimType, PointerType, TypeVar, GlobalTypeVar, TupleType
 from .type import TypeConstraint, FuncType, IncompleteType, RelayRefType
 from .tensor_type import TensorType
+from .affine_type import TensorAffineType, TupleAffineType
 from .type_relation import TypeCall, TypeRelation
 from .expr import BaseExpr, PrimExpr, RelayExpr, GlobalVar, Range
 from .op import Op, register_op_attr, register_intrin_lowering

python/tvm/ir/affine_type.py

@@ -0,0 +1,69 @@
+# 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.
+"""Types for quantized Tensors."""
+import tvm._ffi
+
+from .base import Node
+from . import _ffi_api
+
+
+class AffineType(Node):
+    """The base class of Affine Types."""
+
+    def __eq__(self, other):
+        """Compare two types for structural equivalence."""
+        return bool(tvm.ir.structural_equal(self, other))
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+
+@tvm._ffi.register_object("TensorAffineType")
+class TensorAffineType(AffineType):
+    """The quantized type of a tensor, with scale, zero point, and datatype
+
+    The real space value is calculated as x = x_q * scale + zero_point
+
+    Parameters
+    ----------
+    scale: Expr
+        The scale
+
+    zero_point: Expr
+        The zero_point
+
+    dtype : str
+        The content data type.
+    """
+
+    def __init__(self, scale, zero_point, dtype):
+        self.__init_handle_by_constructor__(_ffi_api.TensorAffineType, scale, zero_point, dtype)
+
+
+@tvm._ffi.register_object("TupleAffineType")
+class TupleAffineType(AffineType):
+    """Affine types of a node with multiple outputs
+
+    Parameters
+    ----------
+    types : List[TensorAffineType]
+        The shape of the Tensor
+
+    """
+
+    def __init__(self, types):
+        self.__init_handle_by_constructor__(_ffi_api.TupleAffineType, types)

python/tvm/relay/frontend/onnx.py

@@ -1008,7 +1008,7 @@ def _impl_v11(cls, inputs, attr, params):
         if len(inputs) == 3:
             value = fold_constant(_op.take(inputs[2], _op.const(0)))
         else:
-            value = 0
+            value = 0.0
 
         pad_width_expr = fold_constant(_op.transpose(_op.reshape(pads, (2, -1))))
         pad_mode = attr.get("mode", b"constant").decode("utf-8")