[FORK][FEATURE] InnerProduct primitive: squashed weight decompression

U8,i8,u4,i4,fp16 and dynamic quantization support on U8,i8,u4,i4.

ONEDNN3.5 migration has squashed these changes into one:

    [FORK][FEATURE] InnerProduct primitive: 8bit weights decompression support

    [FORK][FEATURE] InnerProduct primitive: 8bit weights decompression support on AMX

    [FORK][FEATURE] InnerProduct primitive: 4bit weights decompression support

    [FORK][FEATURE] Enable prepack algorithm for 4bit weights decompression

    [FORK][FEATURE] InnerProduct primitive: 4bit weights decompression support on SPR

    [FORK][FEATURE] InnerProduct primitive: src dynamic quantization

    [FORK][FIX] Fixed behavior for unaligned src and weights ic groups
    [FORK][FEATURE] InnerProduct primitive: src dynamic quantization

    [FORK][FEATURE] Support (f32,fp16,f32) inner-product

    [FORK][FEATURE] Support s8 for weight-compresseion ip

    [FORK][FIX] Squash FC weight compression fix when migrating 3.5.

		[FORK][FIX]fix ld_step

		[FORK][FEATURE]Remove sub_byte_data_type_multiplier

		[Fork][Fix]Fix offset for dynamic quantization

		[FORK][FIX]Fix zp check and u4/s4 reorder

		[FORK][FIX]Fix zp set API

		[Fork][Fix]Skip runtime scale & zp check with weight compression

   [Fork][Fix] fix bug in brgemm introduced by (f32,fp16,f32) inner-product

Co-Authored-By: Tingqian Li <Tingqian.Li@intel.com>, Yi Zhang<Yi3.Zhang@intel.com>

Author: 2 people

include/oneapi/dnnl/dnnl.h

@@ -421,6 +421,8 @@ dnnl_status_t DNNL_API dnnl_primitive_attr_set_scratchpad_mode(
 ///     otherwise.
 dnnl_status_t DNNL_API dnnl_primitive_attr_set_scales_mask(
         dnnl_primitive_attr_t attr, int arg, int mask);
+dnnl_status_t DNNL_API dnnl_primitive_attr_set_scales_dims(
+        dnnl_primitive_attr_t attr, int arg, const dnnl_dims_t dims, int ndims);
 
 /// Sets primitive attributes scaling factors for primitive operations for a
 /// given memory argument. The scaling factors must be passed at execution time
@@ -468,6 +470,8 @@ dnnl_status_t DNNL_API dnnl_primitive_attr_set_scales(
 ///     otherwise.
 dnnl_status_t DNNL_API dnnl_primitive_attr_set_zero_points_mask(
         dnnl_primitive_attr_t attr, int arg, int mask);
+dnnl_status_t DNNL_API dnnl_primitive_attr_set_zero_points_dims(
+        dnnl_primitive_attr_t attr, int arg, const dnnl_dims_t dims, int ndims, dnnl_data_type_t data_type);
 
 /// Sets primitive attributes zero points for primitive operations for a given
 /// memory argument. The zero points must be passed at execution time
@@ -2703,6 +2707,10 @@ dnnl_status_t DNNL_API dnnl_primitive_attr_get_rnn_weights_projection_qparams(
         const_dnnl_primitive_attr_t attr, dnnl_dim_t *count, int *mask,
         const float **scales);
 
+dnnl_status_t DNNL_API dnnl_primitive_attr_set_src_dyn_quant_params(
+        dnnl_primitive_attr_t attr, uint64_t group_size);
+dnnl_status_t DNNL_API  dnnl_primitive_attr_get_src_dyn_quant_params(
+        dnnl_primitive_attr_t attr, uint64_t* group_size);
 /// @} dnnl_api_attributes
 
 /// @addtogroup dnnl_api_rnn

include/oneapi/dnnl/dnnl.hpp

@@ -900,7 +900,9 @@ struct memory : public handle<dnnl_memory_t> {
         /// 4-bit unsigned integer.
         u4 = dnnl_u4,
         /// 1-bit integer
-        bin = dnnl_bin
+        bin = dnnl_bin,
+        /// 4-bit normalized float.
+        nf4 = dnnl_nf4,
     };
 
     /// Returns size of data type in bytes.
@@ -4237,6 +4239,10 @@ struct primitive_attr : public handle<dnnl_primitive_attr_t> {
         error::wrap_c_api(dnnl_primitive_attr_set_scales_mask(get(), arg, mask),
                 "could not set scales primitive attribute");
     }
+    void set_scales_dims(int arg, const memory::dims& dims) {
+        error::wrap_c_api(dnnl_primitive_attr_set_scales_dims(get(), arg, dims.data(), dims.size()),
+                "could not set scales primitive attribute");
+    }
 
     /// Sets scaling factors for primitive operations for a given memory
     /// argument. The scaling factors must be passed at execution time
@@ -4282,6 +4288,11 @@ struct primitive_attr : public handle<dnnl_primitive_attr_t> {
                 dnnl_primitive_attr_set_zero_points_mask(get(), arg, mask),
                 "could not set zero points primitive attribute");
     }
+    void set_zero_points_dims(int arg, const memory::dims& dims, memory::data_type dt) {
+        error::wrap_c_api(
+                dnnl_primitive_attr_set_zero_points_dims(get(), arg, dims.data(), dims.size(), memory::convert_to_c(dt)),
+                "could not set zero points primitive attribute");
+    }
 
     /// Sets zero points for primitive operations for a given memory argument.
     /// The zero points must be passed at execution time as an argument with
@@ -4550,6 +4561,16 @@ struct primitive_attr : public handle<dnnl_primitive_attr_t> {
         for (dnnl_dim_t c = 0; c < count; c++)
             scales[c] = c_scales[c];
     }
+
+    void set_src_dyn_quant_params(uint64_t group_size) {
+        error::wrap_c_api(dnnl_primitive_attr_set_src_dyn_quant_params(get(), group_size),
+                "could not set src dynamic quantization parameters primitive attribute");
+    }
+
+    void get_src_dyn_quant_params(uint64_t& group_size) const {
+        error::wrap_c_api(dnnl_primitive_attr_get_src_dyn_quant_params(get(), &group_size),
+                "could not get src dynamic quantization parameters primitive attribute");
+    }
 };
 
 /// @} dnnl_api_attributes

include/oneapi/dnnl/dnnl_common_types.h

@@ -104,9 +104,11 @@ typedef enum {
     dnnl_u4 = 12,
     /// [MX-compliant 8-bit compliant scale data type](https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf) with 8-bit exponent.
     dnnl_e8m0 = 13,
-
+    /// 4-bit normalized float.
+    dnnl_nf4 = 14,
     /// 1-bit integer.
-    dnnl_bin = 14,
+    dnnl_bin = 15,
+
     /// Parameter to allow internal only data_types without undefined behavior.
     /// This parameter is chosen to be valid for so long as sizeof(int) >= 2.
     dnnl_data_type_max = 0x7fff,

src/common/c_types_map.hpp

@@ -174,6 +174,7 @@ const data_type_t data_type_max = dnnl_data_type_max;
 const data_type_t tf32 = static_cast<data_type_t>(1 << 8);
 
 const data_type_t bin = dnnl_bin;
+const data_type_t nf4 = dnnl_nf4;
 } // namespace data_type
 
 using fpmath_mode_t = dnnl_fpmath_mode_t;

src/common/dnnl_debug_autogenerated.cpp

@@ -58,6 +58,9 @@ const char *dnnl_dt2str(dnnl_data_type_t v) {
     if (v == dnnl_u4) return "u4";
     if (v == dnnl_e8m0) return "e8m0";
     if (v == dnnl_bin) return "bin";
+    if (v == dnnl_nf4) return "nf4";
+    if (v == dnnl_s4) return "s4";
+    if (v == dnnl_u4) return "u4";
     if (v == dnnl_data_type_max) return "data_type_max";
     assert(!"unknown dt");
     return "unknown dt";

src/common/dnnl_traits.hpp

@@ -100,6 +100,10 @@ template <> struct prec_traits<data_type::bin> {
     typedef uint8_t type;
 };
 
+template <> struct prec_traits<data_type::nf4> {
+    typedef uint8_t type;
+};
+
 template <>
 struct data_traits<float8_e5m2_t> {
     static constexpr data_type_t data_type = data_type::f8_e5m2;

src/common/inner_product.cpp

@@ -109,13 +109,27 @@ status_t ip_attr_check(const inner_product_desc_t &desc, const engine_t *engine,
     using smask_t = primitive_attr_t::skip_mask_t;
 
     if (attr == nullptr) return status::success;
-    if (attr->has_default_values()) return status::success;
+    const data_type_t src_dt = desc.src_desc.data_type;
+    const data_type_t wei_dt = desc.weights_desc.data_type;
+    bool is_weight_compression = (one_of(src_dt, data_type::f32, data_type::bf16) && one_of(wei_dt, data_type::u8, data_type::s8, data_type::nf4, data_type::s4, data_type::u4)) ||
+                                 (one_of(src_dt, data_type::f32) && one_of(wei_dt, data_type::f16, data_type::bf16));
+    auto attr_mask = smask_t::none;
+    // From oneDNN 3.5, those checks must be skipped if wei_decomp is enabled
+    // reference from src/plugins/intel_cpu/thirdparty/onednn/src/common/matmul.cpp:L62
+    if (is_weight_compression) {
+        attr_mask |= smask_t::zero_points_runtime_data_type;
+        attr_mask |= smask_t::zero_points_runtime_groups;
+        attr_mask |= smask_t::scales_runtime_data_type;
+        attr_mask |= smask_t::scales_runtime_groups;
+    }
+    if (attr->has_default_values(attr_mask)) return status::success;
 
     // Check attributes
     if (utils::one_of(desc.prop_kind, prop_kind::forward_inference,
                 prop_kind::forward_training)) {
         const data_type_t src_dt = desc.src_desc.data_type;
         const data_type_t dst_dt = desc.dst_desc.data_type;
+        const data_type_t wei_dt = desc.weights_desc.data_type;
 
         auto fwd_attr_mask
                 = smask_t::post_ops | smask_t::sum_dt | smask_t::fpmath_mode;
@@ -125,8 +139,12 @@ status_t ip_attr_check(const inner_product_desc_t &desc, const engine_t *engine,
             is_int8 = is_int8
                     || utils::one_of(dst_dt, data_type::s8, data_type::u8,
                             data_type::s32);
-        if (is_int8) fwd_attr_mask |= smask_t::scales_runtime | smask_t::zero_points_runtime;
-
+        if (engine->kind() == engine_kind::cpu)
+            is_int8 |= one_of(wei_dt, data_type::u8, data_type::s8, data_type::nf4, data_type::s4, data_type::u4);
+        if (is_int8) fwd_attr_mask |= smask_t::scales_runtime | smask_t::zero_points_runtime | smask_t::src_dyn_quant_params;
+        if (is_weight_compression) {
+            fwd_attr_mask |= attr_mask;
+        }
         VCHECK_IP_UNIMPL(attr->has_default_values(fwd_attr_mask, dst_dt),
                 VERBOSE_UNSUPPORTED_ATTR);
 

src/common/memory_desc_wrapper.hpp

@@ -298,8 +298,7 @@ struct memory_desc_wrapper : public c_compatible {
                 max_size = utils::array_product(bd.inner_blks, bd.inner_nblks);
             }
 
-            size_t data_size = max_size * data_type_size()
-                    / sub_byte_data_type_multiplier();
+            size_t data_size = max_size * data_type_size() / sub_byte_data_type_multiplier();
             if (is_additional_buffer()) {
                 // The additional buffers, typically of data type int32_t, float
                 // are stored at the end of data. Pad the data, so that the

src/common/memory_tracking.hpp

@@ -316,6 +316,10 @@ enum {
     key_wino_V,
     key_wino_M,
     key_wino_workspace,
+    key_decompression_scales,
+    key_decompression_zero_points,
+    key_src_quantized,
+    key_src_dequantized_scales,
     // These two keys should always be the last ones,
     // even though they are not in alphabetical order
     key_nested,

src/common/memory_zero_pad.cpp

@@ -294,6 +294,7 @@ static status_t zero_pad(const memory_t *memory, const exec_ctx_t &ctx) {
         case s4: return typed_zero_pad<s8>(memory, ctx);
         case u4: return typed_zero_pad<u8>(memory, ctx);
         case bin: return typed_zero_pad<u8>(memory, ctx);
+        case nf4: return typed_zero_pad<u8>(memory, ctx);
         default: assert(!"memory is undefined"); return unimplemented;
     }
     return unimplemented;

src/common/primitive_attr.cpp

@@ -123,6 +123,8 @@ status_t zero_points_t::set(int arg, int mask, int ndims, const dims_t groups,
             data_type_wei = data_type;
             group_ndims_wei = ndims;
             utils::array_copy(group_dims_wei, groups, group_ndims_wei);
+            utils::array_copy(dims_wei, groups, ndims);
+            ndims_wei = ndims;
             break;
         case DNNL_ARG_DST:
             is_set_dst = true;
@@ -132,6 +134,23 @@ status_t zero_points_t::set(int arg, int mask, int ndims, const dims_t groups,
     return status::success;
 }
 
+status_t zero_points_t::set(int arg, const dims_t dims, int ndims, data_type_t data_type) {
+   const bool supported_arg
+           = utils::one_of(arg, DNNL_ARG_WEIGHTS);
+   if (!supported_arg) return status::unimplemented;
+
+   switch (arg) {
+       case DNNL_ARG_WEIGHTS:
+           is_set_wei = true;
+           ndims_wei = ndims;
+           mask_wei = 1;
+           utils::array_copy(dims_wei, dims, ndims);
+           data_type_wei = data_type;
+           break;
+   }
+   return status::success;
+}
+
 status_t dropout_t::set_default_formats(const memory_desc_t *dst_md) {
     auto is_any_or_undef = [](format_kind_t kind) {
         return one_of(kind, dnnl_format_kind_any, dnnl_format_kind_undef);
@@ -177,6 +196,8 @@ 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_ARG(IMPLICATION((bool)(~mask & smask_t::zero_points),
+            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_);
@@ -185,6 +206,7 @@ bool primitive_attr_t::has_default_values(dnnl_primitive_attr::skip_mask_t mask,
     CHECK_MASK(smask_t::rnn_weights_qparams, rnn_weights_qparams_);
     CHECK_MASK(smask_t::rnn_weights_projection_qparams,
             rnn_weights_projection_qparams_);
+    CHECK_MASK(smask_t::src_dyn_quant_params, src_dyn_quant_params_);
     CHECK_ARG(IMPLICATION((bool)(~mask & smask_t::sum_dt),
             post_ops_.sum_with_default_dt(dst_dt)));
     bool gpu_attr_ok = IMPLICATION((bool)(~mask & smask_t::gpu_attr),
@@ -222,6 +244,7 @@ bool primitive_attr_t::defined(dnnl_primitive_attr::skip_mask_t mask) const {
     CHECK_MASK(smask_t::rnn_weights_qparams, rnn_weights_qparams_);
     CHECK_MASK(smask_t::rnn_weights_projection_qparams,
             rnn_weights_projection_qparams_);
+    CHECK_MASK(smask_t::src_dyn_quant_params, src_dyn_quant_params_);
     return ok;
 #undef CHECK_MASK
 #undef CHECK_ARG
@@ -673,6 +696,13 @@ status_t dnnl_primitive_attr_set_scales_mask(
     if (!ok) return invalid_arguments;
     return attr->scales_.set(arg, mask);
 }
+status_t dnnl_primitive_attr_set_scales_dims(
+        primitive_attr_t *attr, int arg, const dims_t dims, int ndims) {
+    bool ok = attr && arg >= 0 && ndims > 0
+            && attr->output_scales_.has_default_values();
+    if (!ok) return invalid_arguments;
+    return attr->scales_.set(arg, dims, ndims);
+}
 
 status_t dnnl_primitive_attr_set_scales(primitive_attr_t *attr, int arg,
         int mask, int ndims, const dims_t group_dims, data_type_t data_type) {
@@ -700,6 +730,13 @@ status_t dnnl_primitive_attr_set_zero_points_mask(
 
     return attr->zero_points_.set(arg, mask);
 }
+status_t dnnl_primitive_attr_set_zero_points_dims(
+        primitive_attr_t *attr, int arg, const dims_t dims, int ndims, dnnl_data_type_t data_type) {
+    bool ok = attr && ndims > 0;
+    if (!ok) return invalid_arguments;
+
+    return attr->zero_points_.set(arg, dims, ndims, data_type);
+}
 
 dnnl_status_t DNNL_API dnnl_primitive_attr_set_zero_points(
         dnnl_primitive_attr_t attr, int arg, int mask, int ndims,
@@ -1028,6 +1065,21 @@ status_t DNNL_API dnnl_primitive_attr_set_rnn_tparams(
     return attr->rnn_tparams_.set(mode, ngates, scales, cscale);
 }
 
+status_t dnnl_primitive_attr_set_src_dyn_quant_params(
+        primitive_attr_t *attr, const uint64_t group_size) {
+    if (attr == nullptr) return invalid_arguments;
+
+    return attr->src_dyn_quant_params_.set(group_size);
+}
+
+status_t dnnl_primitive_attr_get_src_dyn_quant_params(
+        primitive_attr_t *attr, uint64_t* group_size) {
+    if (attr == nullptr) return invalid_arguments;
+
+    if (group_size) *group_size = attr->src_dyn_quant_params_.get();
+    return success;
+}
+
 template struct dnnl::impl::shifts_t<uint8_t>;
 template struct dnnl::impl::shifts_t<int32_t>;
 template struct dnnl::impl::shifts_t<float>;