[FORK][FEATURE] Asymmetric quntization for activations

Author: dmitry-gorokhov

include/oneapi/dnnl/dnnl.h

@@ -515,6 +515,24 @@ dnnl_status_t DNNL_API dnnl_primitive_attr_set_rounding(
 dnnl_status_t DNNL_API dnnl_primitive_attr_get_rounding(
         dnnl_primitive_attr_t attr, int arg, dnnl_rounding_mode_t *mode);
 
+dnnl_status_t DNNL_API dnnl_primitive_attr_get_output_compensations(
+        const_dnnl_primitive_attr_t attr, int *count, int *mask, const int32_t **compensations);
+
+dnnl_status_t DNNL_API dnnl_primitive_attr_set_output_compensations(
+        dnnl_primitive_attr_t attr, int count, int mask, const int32_t *compensations);
+
+dnnl_status_t DNNL_API dnnl_primitive_attr_get_input_zero_points(
+        const_dnnl_primitive_attr_t attr, int *count, int *mask, const uint8_t **zero_points);
+
+dnnl_status_t DNNL_API dnnl_primitive_attr_set_input_zero_points(
+        dnnl_primitive_attr_t attr, int count, int mask, const uint8_t *zero_points);
+
+dnnl_status_t DNNL_API dnnl_primitive_attr_get_weights_zero_points(
+        const_dnnl_primitive_attr_t attr, int *count, int *mask, const float **zero_points);
+
+dnnl_status_t DNNL_API dnnl_primitive_attr_set_weights_zero_points(
+        dnnl_primitive_attr_t attr, int count, int mask, const float *zero_points);
+
 /// Returns primitive attributes post-ops.
 ///
 /// @warning

include/oneapi/dnnl/dnnl.hpp

@@ -4204,6 +4204,69 @@ struct primitive_attr : public handle<dnnl_primitive_attr_t> {
                 "could not set zero points primitive attribute");
     }
 
+    void get_output_compensations(int &mask, std::vector<int32_t> &compensations) const
+    {
+        int count, c_mask;
+        const int32_t *c_compensations;
+        error::wrap_c_api(dnnl_primitive_attr_get_output_compensations(get(),
+                    &count, &c_mask, &c_compensations),
+                "could not get int output compensations");
+        compensations.resize(count);
+
+        mask = c_mask;
+        for (int c = 0; c < count; ++c)
+            compensations[c] = c_compensations[c];
+    }
+
+    void set_output_compensations(int mask, const std::vector<int32_t> &compensations)
+    {
+        error::wrap_c_api(dnnl_primitive_attr_set_output_compensations(get(),
+                    (int)compensations.size(), mask, &compensations[0]),
+                "could not set int output compensations");
+    }
+
+    void get_input_zero_points(int &mask, std::vector<uint8_t> &zero_points) const
+    {
+        int count, c_mask;
+        const uint8_t *c_zero_points;
+        error::wrap_c_api(dnnl_primitive_attr_get_input_zero_points(get(),
+                    &count, &c_mask, &c_zero_points),
+                "could not get int input zero_points");
+        zero_points.resize(count);
+
+        mask = c_mask;
+        for (int c = 0; c < count; ++c)
+            zero_points[c] = c_zero_points[c];
+    }
+
+    void set_input_zero_points(int mask, const std::vector<uint8_t> &zero_points)
+    {
+        error::wrap_c_api(dnnl_primitive_attr_set_input_zero_points(get(),
+                    (int)zero_points.size(), mask, &zero_points[0]),
+                "could not set int input zero_points");
+    }
+
+    void get_weights_zero_points(int &mask, std::vector<int8_t> &zero_points) const
+    {
+        int count, c_mask;
+        const float *c_zero_points;
+        error::wrap_c_api(dnnl_primitive_attr_get_weights_zero_points(get(),
+                    &count, &c_mask, &c_zero_points),
+                "could not get int weights zero_points");
+        zero_points.resize(count);
+
+        mask = c_mask;
+        for (int c = 0; c < count; ++c)
+            zero_points[c] = c_zero_points[c];
+    }
+
+    void set_weights_zero_points(int mask, const std::vector<float> &zero_points)
+    {
+        error::wrap_c_api(dnnl_primitive_attr_set_weights_zero_points(get(),
+                    (int)zero_points.size(), mask, &zero_points[0]),
+                "could not set int weights zero_points");
+    }
+
     /// Returns post-ops previously set via set_post_ops().
     ///
     /// @returns Post-ops.

src/common/primitive_attr.cpp

@@ -177,6 +177,9 @@ bool primitive_attr_t::has_default_values(dnnl_primitive_attr::skip_mask_t mask,
     CHECK_ARG(
             IMPLICATION((bool)(~mask & smask_t::zero_points_runtime_data_type),
                     zero_points_.has_default_data_type()));
+    CHECK_MASK(smask_t::input_zero_points, input_zero_points_);
+    CHECK_MASK(smask_t::weights_zero_points, weights_zero_points_);
+    CHECK_MASK(smask_t::output_compensations, output_compensations_);
     CHECK_MASK(smask_t::post_ops, post_ops_);
     CHECK_MASK(smask_t::rnn_data_qparams, rnn_data_qparams_);
     CHECK_MASK(smask_t::rnn_weights_qparams, rnn_weights_qparams_);
@@ -205,14 +208,6 @@ bool primitive_attr_t::has_default_values(dnnl_primitive_attr::skip_mask_t mask,
 #undef CHECK_ARG
 }
 
-bool primitive_attr_t::has_asymmetric_quantization() const {
-    return true
-           && output_scales_.has_default_values()
-           && rnn_data_qparams_.has_default_values()
-           && rnn_weights_qparams_.has_default_values()
-           && (!input_zero_points_.has_default_values() || !weights_zero_points_.has_default_values());
-}
-
 bool primitive_attr_t::defined(dnnl_primitive_attr::skip_mask_t mask) const {
     using smask_t = skip_mask_t;
     bool ok = true;
@@ -700,6 +695,69 @@ status_t dnnl_primitive_attr_set_rounding(
     return attr->rounding_mode_.set(arg, mode);
 }
 
+status_t dnnl_primitive_attr_get_output_compensations(const primitive_attr_t *attr,
+        int *count, int *mask, const int32_t **compensations) {
+    if (any_null(attr, count, mask, compensations))
+        return invalid_arguments;
+
+    *count = attr->output_compensations_.count_;
+    *mask = attr->output_compensations_.mask_;
+    *compensations = attr->output_compensations_.shifts_;
+
+    return success;
+}
+
+status_t dnnl_primitive_attr_set_output_compensations(primitive_attr_t *attr,
+        int count, int mask, const int32_t *compensations) {
+    bool ok = !any_null(attr, compensations) && count > 0 && mask >= 0;
+    if (!ok)
+        return invalid_arguments;
+
+    return attr->output_compensations_.set(count, mask, compensations);
+}
+
+status_t dnnl_primitive_attr_get_input_zero_points(const primitive_attr_t *attr,
+        int *count, int *mask, const uint8_t **zero_points) {
+    if (any_null(attr, count, mask, zero_points))
+        return invalid_arguments;
+
+    *count = attr->input_zero_points_.count_;
+    *mask = attr->input_zero_points_.mask_;
+    *zero_points = attr->input_zero_points_.shifts_;
+
+    return success;
+}
+
+status_t dnnl_primitive_attr_set_input_zero_points(primitive_attr_t *attr,
+        int count, int mask, const uint8_t *zero_points) {
+    bool ok = !any_null(attr, zero_points) && count > 0 && mask >= 0;
+    if (!ok)
+        return invalid_arguments;
+
+    return attr->input_zero_points_.set(count, mask, zero_points);
+}
+
+status_t dnnl_primitive_attr_get_weights_zero_points(const primitive_attr_t *attr,
+        int *count, int *mask, const float **zero_points) {
+    if (any_null(attr, count, mask, zero_points))
+        return invalid_arguments;
+
+    *count = attr->weights_zero_points_.count_;
+    *mask = attr->weights_zero_points_.mask_;
+    *zero_points = attr->weights_zero_points_.shifts_;
+
+    return success;
+}
+
+status_t dnnl_primitive_attr_set_weights_zero_points(primitive_attr_t *attr,
+        int count, int mask, const float *zero_points) {
+    bool ok = !any_null(attr, zero_points) && count > 0 && mask >= 0;
+    if (!ok)
+        return invalid_arguments;
+
+    return attr->weights_zero_points_.set(count, mask, zero_points);
+}
+
 status_t dnnl_primitive_attr_get_post_ops(
         const primitive_attr_t *attr, const post_ops_t **post_ops) {
     if (any_null(attr, post_ops)) return invalid_arguments;

src/common/primitive_attr.hpp

@@ -190,18 +190,15 @@ struct shifts_t: public c_compatible {
     shifts_t(): count_(1), mask_(0), shifts_(shifts_buf_)
     { set(0); }
 
-    shifts_t(const shifts_t &rhs): shifts_t()
-    { set(rhs.count_, rhs.mask_, rhs.shifts_); }
-
     ~shifts_t() { cleanup(); }
 
-    shifts_t &operator=(const shifts_t &rhs) {
-        if (&rhs == this)
-            return *this;
-        status_t status = set(rhs.count_, rhs.mask_, rhs.shifts_);
-        assert(status == status::success);
-        (void)status;
-        return *this;
+    bool operator==(const shifts_t<T> &rhs) const {
+        bool ret = count_ == rhs.count_ && mask_ == rhs.mask_
+                   && !utils::any_null(shifts_, rhs.shifts_)
+                   && defined() == rhs.defined()
+                   && IMPLICATION(defined(),
+                                  utils::array_cmp(shifts_, rhs.shifts_, count_));
+        return ret;
     }
 
     bool has_default_values() const {
@@ -211,10 +208,16 @@ struct shifts_t: public c_compatible {
         return true;
     }
 
+    bool defined() const { return !is_runtime_value(shifts_[0]); }
+
     status_t set(int count, int mask, const T *zero_points);
     status_t set(T single_zero_point) { return this->set(1, 0, &single_zero_point); }
 
-    int count_;
+    status_t copy_from(const shifts_t &other) {
+        return set(other.count_, other.mask_, other.shifts_);
+    }
+
+    dim_t count_;
     int mask_;
     T *shifts_;
 
@@ -230,6 +233,8 @@ struct shifts_t: public c_compatible {
         mask_ = 0;
         shifts_ = shifts_buf_;
     }
+
+    DNNL_DISALLOW_COPY_AND_ASSIGN(shifts_t);
 };
 
 struct runtime_scales_t : public c_compatible {
@@ -948,6 +953,9 @@ struct dnnl_primitive_attr : public dnnl::impl::c_compatible {
         CHECK(rnn_tparams_.copy_from(other.rnn_tparams_));
         if (other.gpu_attr_) gpu_attr_ = other.gpu_attr_->clone();
         dropout_ = other.dropout_;
+        input_zero_points_.copy_from(other.input_zero_points_);
+        weights_zero_points_.copy_from(other.weights_zero_points_);
+        output_compensations_.copy_from(other.output_compensations_);
 
         return status::success;
     }
@@ -978,6 +986,9 @@ struct dnnl_primitive_attr : public dnnl::impl::c_compatible {
         = (unsigned)zero_points_runtime | (1u << 18),
         dropout = 1u << 19,
         rounding_mode = 1u << 20,
+        input_zero_points = 1 << 21,
+        weights_zero_points = 1 << 22,
+        output_compensations = 1 << 23,
     };
 
     /** Returns true if the attributes have default values.
@@ -986,8 +997,6 @@ struct dnnl_primitive_attr : public dnnl::impl::c_compatible {
     bool has_default_values(skip_mask_t mask = skip_mask_t::none,
             dnnl::impl::data_type_t dst_dt = dnnl_data_type_undef) const;
 
-    bool has_asymmetric_quantization() const;
-
     /** Returns true if the attributes are fully defined. */
     bool defined(skip_mask_t mask = skip_mask_t::none) const;
 
@@ -1007,7 +1016,10 @@ struct dnnl_primitive_attr : public dnnl::impl::c_compatible {
                             && gpu_attr_->is_equal(*rhs.gpu_attr_))
                         || (!gpu_attr_ && !rhs.gpu_attr_))
                 && dropout_ == rhs.dropout_
-                && rounding_mode_ == rhs.rounding_mode_;
+                && rounding_mode_ == rhs.rounding_mode_
+                && input_zero_points_ == rhs.input_zero_points_
+                && weights_zero_points_ == rhs.weights_zero_points_
+                && output_compensations_ == rhs.output_compensations_;
         return ret;
     }