add op AscendQuant AscendDequant

mindspore/ccsrc/transform/convert.cc

@@ -205,6 +205,8 @@ const char kNameL2Loss[] = "L2Loss";
 const char kNameCTCLoss[] = "CTCLoss";
 const char kNameRange[] = "Range";
 const char kNameSquareSumAll[] = "SquareSumAll";
+const char kNameAscendQuant[] = "AscendQuant";
+const char kNameAscendDequant[] = "AscendDequant";
 
 // -----------------OpAdapter initialization--------------
 std::unordered_map<std::string, OpAdapterDescPtr> &DfGraphConvertor::get_adpt_map() {
@@ -408,7 +410,9 @@ std::unordered_map<std::string, OpAdapterDescPtr> &DfGraphConvertor::get_adpt_ma
     {string(kNameL2Loss), ADPT_DESC(L2Loss)},
     {string(kNameCTCLoss), ADPT_DESC(CTCLoss)},
     {string(kNameRange), ADPT_DESC(RangeD)},
-    {string(kNameSquareSumAll), ADPT_DESC(SquareSumAll)}};
+    {string(kNameSquareSumAll), ADPT_DESC(SquareSumAll)},
+    {string(kNameAscendQuant), ADPT_DESC(AscendQuant)},
+    {string(kNameAscendDequant), ADPT_DESC(AscendDequant)}};
 #ifdef ENABLE_GE
   adpt_map[string(kNamePrint)] = ADPT_DESC(Print);
   adpt_map[string(kNameApplyAdam)] = ADPT_DESC(ApplyAdamD);

mindspore/ccsrc/transform/op_declare.cc

@@ -1277,6 +1277,19 @@ ATTR_MAP(CTCLoss) = {
   {"ignore_longer_outputs_than_inputs", ATTR_DESC(ignore_longer_outputs_than_inputs, AnyTraits<bool>())}};
 OUTPUT_MAP(CTCLoss) = {{0, OUTPUT_DESC(loss)}, {1, OUTPUT_DESC(gradient)}};
 
+// AscendQuant
+INPUT_MAP(AscendQuant) = {{1, INPUT_DESC(x)}};
+ATTR_MAP(AscendQuant) = {{"scale", ATTR_DESC(scale, AnyTraits<float>())},
+                         {"offset", ATTR_DESC(offset, AnyTraits<float>())},
+                         {"sqrt_mode", ATTR_DESC(sqrt_mode, AnyTraits<bool>())},
+                         {"round_mode", ATTR_DESC(round_mode, AnyTraits<std::string>())}};
+OUTPUT_MAP(AscendQuant) = {{0, OUTPUT_DESC(y)}};
+
+// AscendDequant
+INPUT_MAP(AscendDequant) = {{1, INPUT_DESC(x)}, {2, INPUT_DESC(deq_scale)}};
+ATTR_MAP(AscendDequant) = {{"sqrt_mode", ATTR_DESC(sqrt_mode, AnyTraits<bool>())},
+                           {"relu_flag", ATTR_DESC(relu_flag, AnyTraits<bool>())}};
+OUTPUT_MAP(AscendDequant) = {{0, OUTPUT_DESC(y)}};
 #ifdef ENABLE_GE
 // Print
 INPUT_MAP(Print) = EMPTY_INPUT_MAP;

mindspore/ccsrc/transform/op_declare.h

@@ -481,6 +481,10 @@ DECLARE_OP_ADAPTER(L2Loss)
 DECLARE_OP_USE_OUTPUT(L2Loss)
 DECLARE_OP_ADAPTER(CTCLoss)
 DECLARE_OP_USE_OUTPUT(CTCLoss)
+DECLARE_OP_ADAPTER(AscendQuant)
+DECLARE_OP_USE_OUTPUT(AscendQuant)
+DECLARE_OP_ADAPTER(AscendDequant)
+DECLARE_OP_USE_OUTPUT(AscendDequant)
 #ifdef ENABLE_GE
 DECLARE_OP_ADAPTER(Print)
 DECLARE_OP_USE_DYN_INPUT(Print)

mindspore/ops/operations/_quant_ops.py

@@ -39,6 +39,8 @@ __all__ = ["FakeQuantPerLayer",
            "BatchNormFold2_D",
            "BatchNormFold2GradD",
            "BatchNormFold2GradReduce",
+           "AscendQuant",
+           "AscendDequant",
            ]
 
 
@@ -975,3 +977,104 @@ class FakeQuantMinMaxPerChannelUpdate(PrimitiveWithInfer):
         validator.check_tensor_type_same(
             {"max": max_type}, valid_types, self.name)
         return min_type, max_type
+
+
+class AscendQuant(PrimitiveWithInfer):
+    r"""
+    Returns the quantized value of input_x.
+
+    If `sqrt_mode` is False:
+
+    .. math::
+        y = round(scale * x + offset)
+    If `sqrt_mode` is True:
+
+    .. math::
+        y = round(scale * x * scale + offset)
+
+    Note:
+        This operation only support Ascend 310 inference environment.
+
+    Args:
+        scale (float) : Specifies the scaling ratio.
+        offset (float): Specifies the offset.
+        sqrt_mode (bool) : Specifies whether to perform square root on `scale`. Default: False.
+        round_mode (str): Specifies the way to round. Should be one of ["Round", "Floor", "Ceil", "Trunc"].
+          Default: "Round".
+
+    Inputs:
+        - **input_x** (Tensor) : Input tensor. Its data type should be mindspore.float16 or mindspore.float32.
+
+    Outputs:
+        - Tensor: The quantized output tensor of type mindspore.int8.
+
+    Examples:
+        >>> input_x = Tensor([100.0, 150.0], mstype.float32)
+        >>> quant = P.AscendQuant(80.0, 0.0, False, "Round")
+        >>> y = quant(input_x)
+    """
+
+    @prim_attr_register
+    def __init__(self, scale, offset, sqrt_mode=False, round_mode="Round"):
+        self.scale = validator.check_value_type("scale", scale, [float], self.name)
+        self.offset = validator.check_value_type("offset", offset, [float], self.name)
+        self.sqrt_mode = validator.check_value_type("sqrt_mode", sqrt_mode, [bool], self.name)
+        self.round_mode = validator.check_string("round_mode", round_mode,
+                                                 ["Round", "Floor", "Ceil", "Trunc"], self.name)
+
+    def infer_shape(self, x_shape):
+        return x_shape
+
+    def infer_dtype(self, x_type):
+        validator.check_subclass("input_x", x_type, mstype.tensor, self.name)
+        validator.check_type_name("input_x", x_type, [mstype.float16, mstype.float32], self.name)
+        return mstype.int8
+
+
+class AscendDequant(PrimitiveWithInfer):
+    r"""
+    Returns the dequantized value of input_x.
+    This operation will do ReLU to the dequantized value if `relu_flag` is True.
+
+    If `sqrt_mode` is False:
+
+    .. math::
+        y = x * deq\_scale
+    If `sqrt_mode` is True:
+
+    .. math::
+        y = x * deq\_scale * deq\_scale
+
+    Note:
+        This operation only support Ascend 310 inference environment.
+
+    Args:
+        sqrt_mode (bool) : Specifies whether to perform square root on `scale`. Default: False.
+        relu_flag (bool): Specifies whether to perform ReLU. Default: False.
+
+    Inputs:
+        - **input_x** (Tensor) : Input tensor. Should be mindspore.int32.
+        - **deq_scale** (Tensor) : Specifies the scaling ratio.
+          Data type should be mindspore.float16 or mindspore.uint64
+
+    Outputs:
+        - Tensor: The quantized output tensor of type mindspore.float16.
+
+    Examples:
+        >>> input_x = Tensor([100.0, 150.0], mstype.float32)
+        >>> dequant = P.AscendDequant(False, False)
+        >>> y = dequant(input_x)
+    """
+    @prim_attr_register
+    def __init__(self, sqrt_mode=False, relu_flag=False):
+        self.sqrt_mode = validator.check_value_type("sqrt_mode", sqrt_mode, [bool], self.name)
+        self.relu_flag = validator.check_value_type("relu_flag", relu_flag, [bool], self.name)
+
+    def infer_shape(self, x_shape, deq_scale_shape):
+        return x_shape
+
+    def infer_dtype(self, x_type, deq_scale_type):
+        validator.check_subclass("x", x_type, mstype.tensor, self.name)
+        validator.check_type_name("x", x_type, [mstype.int32], self.name)
+        validator.check_type_name("deq_scale", deq_scale_type, [mstype.float16, mstype.uint64], self.name)
+        return mstype.float16

tests/ut/python/ops/test_ops.py

@@ -1614,7 +1614,43 @@ test_case_other_ops = [
 
 ]
 
-test_case_lists = [test_case_nn_ops, test_case_math_ops, test_case_array_ops, test_case_other_ops]
+
+test_case_quant_ops = [
+    ('AscendQuant_1', {
+        'block': P.AscendQuant(0.5, 0.0, False, "Round"),
+        'desc_inputs': [Tensor(np.random.rand(1,2,4,4), mstype.float32)],
+        'skip': ['backward']}),
+    ('AscendQuant_2', {
+        'block': P.AscendQuant(80.0, 10.0, True, "Round"),
+        'desc_inputs': [Tensor([100.0, 200.0], mstype.float32)],
+        'skip': ['backward']}),
+    ('AscendQuant_3', {
+        'block': P.AscendQuant(80.0, 0.0, False, "Floor"),
+        'desc_inputs': [Tensor([100.0, 200.0], mstype.float32)],
+        'skip': ['backward']}),
+    ('AscendQuant_4', {
+        'block': P.AscendQuant(80.0, 0.0, False, "Ceil"),
+        'desc_inputs': [Tensor([100.0, 200.0], mstype.float32)],
+        'skip': ['backward']}),
+    ('AscendQuant_5', {
+        'block': P.AscendQuant(80.0, 0.0, False, "Trunc"),
+        'desc_inputs': [Tensor([100.0, 200.0], mstype.float32)],
+        'skip': ['backward']}),
+    ('AscendQuant_6', {
+        'block': P.AscendQuant(-80.0, 10.0, False, "Round"),
+        'desc_inputs': [Tensor([100.0, 200.0], mstype.float32)],
+        'skip': ['backward']}),
+    ('AscendQuant_7', {
+        'block': P.AscendQuant(80.0, -10.0, False, "Round"),
+        'desc_inputs': [Tensor([100.0, 200.0], mstype.float32)],
+        'skip': ['backward']}),
+    ('AscendQuant_8', {
+        'block': P.AscendQuant(80.0, 10.0, False, "Round"),
+        'desc_inputs': [Tensor([100.0, 200.0], mstype.float16)],
+        'skip': ['backward']}),
+]
+
+test_case_lists = [test_case_nn_ops, test_case_math_ops, test_case_array_ops, test_case_other_ops, test_case_quant_ops]
 test_case = functools.reduce(lambda x, y: x + y, test_case_lists)
 # use -k to select certain testcast
 # pytest tests/python/ops/test_ops.py::test_backward -k LayerNorm