[FORK][CPU][FEATURE] InnerProduct primitive: u2 weights decompression

Author: xuchen-intel

include/oneapi/dnnl/dnnl.hpp

@@ -907,6 +907,8 @@ struct memory : public handle<dnnl_memory_t> {
         s4 = dnnl_s4,
         /// 4-bit unsigned integer.
         u4 = dnnl_u4,
+        /// 2-bit unsigned integer.
+        u2 = dnnl_u2,
         /// 1-bit integer
         bin = dnnl_bin,
         /// 4-bit normalized float.

include/oneapi/dnnl/dnnl_common_types.h

@@ -112,6 +112,8 @@ typedef enum {
     dnnl_nf4 = 16,
     /// 1-bit integer.
     dnnl_bin = 17,
+    /// 2-bit unsigned integer.
+    dnnl_u2 = 18,
 
     /// Parameter to allow internal only data_types without undefined behavior.
     /// This parameter is chosen to be valid for so long as sizeof(int) >= 2.

src/common/c_types_map.hpp

@@ -176,6 +176,7 @@ const data_type_t s8 = dnnl_s8;
 const data_type_t u8 = dnnl_u8;
 const data_type_t s4 = dnnl_s4;
 const data_type_t u4 = dnnl_u4;
+const data_type_t u2 = dnnl_u2;
 const data_type_t boolean = dnnl_boolean;
 const data_type_t data_type_max = dnnl_data_type_max;
 

src/common/dnnl_traits.hpp

@@ -107,6 +107,10 @@ template <> struct prec_traits_t<data_type::nf4> {
     using type = uint8_t;
 };
 
+template <> struct prec_traits_t<data_type::u2> {
+    using type = uint8_t;
+};
+
 template <>
 struct data_traits_t<float4_e3m0_t> {
     static constexpr data_type_t data_type = data_type::f4_e3m0;

src/common/inner_product.cpp

@@ -112,7 +112,7 @@ status_t ip_attr_check(const inner_product_desc_t &desc, const engine_t *engine,
     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, data_type::f4_e2m1)) ||
+                                  one_of(wei_dt, data_type::u8, data_type::s8, data_type::nf4, data_type::s4, data_type::u4, data_type::f4_e2m1, data_type::u2)) ||
                                  (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
@@ -142,7 +142,7 @@ status_t ip_attr_check(const inner_product_desc_t &desc, const engine_t *engine,
                             data_type::s32);
 
         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, data_type::f4_e2m1);
+            is_int8 |= one_of(wei_dt, data_type::u8, data_type::s8, data_type::nf4, data_type::s4, data_type::u4, data_type::f4_e2m1, data_type::u2);
         if (is_int8) fwd_attr_mask |= smask_t::scales | smask_t::zero_points | smask_t::src_dyn_quant_params;
 
         if (is_weight_compression) {

src/common/memory_desc_wrapper.hpp

@@ -157,6 +157,7 @@ struct memory_desc_wrapper : public c_compatible {
      * For the rest data types returns 1. */
     size_t sub_byte_data_type_multiplier() const {
         if (utils::one_of(data_type(), data_type::s4, data_type::u4)) return 2;
+        if (data_type() == data_type::u2) return 4;
         return 1;
     }
 

src/common/type_helpers.hpp

@@ -108,6 +108,7 @@ inline size_t data_type_size(data_type_t data_type) {
         case u8: return sizeof(prec_traits_t<u8>::type);
         case s4: return sizeof(prec_traits_t<s4>::type);
         case u4: return sizeof(prec_traits_t<u4>::type);
+        case u2: return sizeof(prec_traits_t<u2>::type);
         case boolean: return sizeof(prec_traits_t<boolean>::type);
         case bin: return sizeof(prec_traits_t<u8>::type);
         case nf4: return sizeof(prec_traits_t<u8>::type);
@@ -124,6 +125,7 @@ inline size_t elements_to_bytes(data_type_t data_type, size_t count) {
         case f4_e3m0:
         case s4:
         case u4: return (count + 1) >> 1;
+        case u2: return (count + 3) >> 2;
         default: return data_type_size(data_type) * count;
     }
 }
@@ -135,6 +137,7 @@ inline size_t bytes_to_elements(data_type_t data_type, size_t bytes) {
         case f4_e3m0:
         case s4:
         case u4: return bytes * 2;
+        case u2: return bytes * 4;
         default: return utils::div_up(bytes, data_type_size(data_type));
     }
 }
@@ -453,7 +456,7 @@ inline data_type_t default_accum_data_type(data_type_t src_dt,
 
     /* prop_kind doesn't matter */
     if (everyone_is(f32, src_dt, wei_dt)) return f32;
-    if (one_of(src_dt, f32, bf16) && one_of(wei_dt, u8, s8, nf4, s4, u4, f4_e2m1)) return f32;
+    if (one_of(src_dt, f32, bf16) && one_of(wei_dt, u8, s8, nf4, s4, u4, f4_e2m1, u2)) return f32;
     if (everyone_is(f64, src_dt, wei_dt)) return f64;
 
     if (one_of(prop_kind, forward_training, forward_inference)) {
@@ -1301,7 +1304,7 @@ inline bool memory_desc_sanity_check(int ndims, const dims_t dims,
 
     bool ok = dims != nullptr && 0 < ndims && ndims <= DNNL_MAX_NDIMS
             && utils::one_of(data_type, f4_e3m0, f4_e2m1, e8m0, f8_e5m2,
-                    f8_e4m3, f16, bf16, f32, f64, s32, s8, u8, s4, u4, bin, nf4);
+                    f8_e4m3, f16, bf16, f32, f64, s32, s8, u8, s4, u4, u2, bin, nf4);
     if (!ok) return false;
 
     bool has_runtime_dims = false;

src/cpu/cpu_inner_product_list.cpp

@@ -99,6 +99,13 @@ const std::map<pk_dt_impl_key_t, std::vector<impl_list_item_t>> &impl_list_map()
             CPU_INSTANCE_AVX2(brgemm_inner_product_fwd_t, avx2)
             nullptr,
         }},
+        {{forward, f32, u2, f32}, {
+            CPU_INSTANCE_AVX512(brgemm_inner_product_fwd_t, avx512_core_vnni)
+            CPU_INSTANCE_AVX512(brgemm_inner_product_fwd_t, avx512_core)
+            CPU_INSTANCE_AVX2(brgemm_inner_product_fwd_t, avx2_vnni)
+            CPU_INSTANCE_AVX2(brgemm_inner_product_fwd_t, avx2)
+            nullptr,
+        }},
         {{forward, f32, f16, f32}, {
             CPU_INSTANCE_AVX512(brgemm_inner_product_fwd_t, avx512_core)
             CPU_INSTANCE_AVX2(brgemm_inner_product_fwd_t, avx2)
@@ -188,6 +195,16 @@ const std::map<pk_dt_impl_key_t, std::vector<impl_list_item_t>> &impl_list_map()
             CPU_INSTANCE_AVX512(brgemm_inner_product_fwd_t, avx512_core_bf16)
             nullptr,
         }},
+        {{forward, bf16, u2, f32}, {
+            CPU_INSTANCE_AMX(brgemm_inner_product_fwd_t, avx512_core_amx)
+            CPU_INSTANCE_AVX512(brgemm_inner_product_fwd_t, avx512_core_bf16)
+            nullptr,
+        }},
+        {{forward, bf16, u2, bf16}, {
+            CPU_INSTANCE_AMX(brgemm_inner_product_fwd_t, avx512_core_amx)
+            CPU_INSTANCE_AVX512(brgemm_inner_product_fwd_t, avx512_core_bf16)
+            nullptr,
+        }},
         {{forward, f16, f16, f32}, {
             //CPU_INSTANCE_X64(matmul_inner_product_fwd_t)
             CPU_INSTANCE_AMX(brgemm_inner_product_fwd_t,avx512_core_amx_fp16)

src/cpu/cpu_primitive.hpp

@@ -84,7 +84,7 @@
                     = ctx.memory_mdw(DNNL_ARG_ATTR_ZERO_POINTS | (arg)); \
             VCHECK_ATTR(utils::one_of(zero_points_d.data_type(), \
                                 data_type::s32, data_type::s8, data_type::u8, \
-                                data_type::s4, data_type::u4, data_type::f32), \
+                                data_type::s4, data_type::u4, data_type::u2, data_type::f32), \
                     VERBOSE_INVALID_DATATYPE, "zero points"); \
         } \
     } \

src/cpu/reorder/cpu_reorder.cpp

@@ -51,6 +51,7 @@ regular_impl_list_map() {
             {{f32, u4, 0}, &regular_u4_impl_list_map()},
             {{s4, data_type::undef, 0}, &regular_s4_impl_list_map()},
             {{u4, data_type::undef, 0}, &regular_u4_impl_list_map()},
+            {{u2, data_type::undef, 0}, &regular_u2_impl_list_map()},
             {{bin, data_type::undef, 0}, &regular_bin_impl_list_map()},
             {{nf4, data_type::undef, 0}, &regular_nf4_impl_list_map()},
             {{s4, f32, 0}, &regular_s4_impl_list_map()},