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

Author: dmitrygo

include/oneapi/dnnl/dnnl.h

@@ -334,6 +334,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 zero points for primitive operations for a given
 /// memory argument. The zero points must be passed at execution time
@@ -354,6 +356,8 @@ dnnl_status_t DNNL_API dnnl_primitive_attr_set_scales_mask(
 ///     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_status_t DNNL_API dnnl_primitive_attr_set_output_compensations(
         dnnl_primitive_attr_t attr, int count, int mask);

include/oneapi/dnnl/dnnl.hpp

@@ -869,7 +869,13 @@ struct memory : public handle<dnnl_memory_t> {
         /// 8-bit unsigned integer.
         u8 = dnnl_u8,
         /// 1-bit integer
-        bin = dnnl_bin
+        bin = dnnl_bin,
+        /// 4-bit normalized float.
+        nf4 = dnnl_nf4,
+        /// 4-bit signed integer.
+        s4 = dnnl_s4,
+        /// 4-bit unsigned integer.
+        u4 = dnnl_u4
     };
 
     /// Returns size of data type in bytes.
@@ -3874,6 +3880,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 zero points for primitive operations for a given memory argument.
     /// The zero points must be passed at execution time as an argument with
@@ -3893,6 +3903,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) {
+        error::wrap_c_api(
+                dnnl_primitive_attr_set_zero_points_dims(get(), arg, dims.data(), dims.size()),
+                "could not set zero points primitive attribute");
+    }
 
     void set_output_compensations(dnnl_dim_t count, int mask)
     {

include/oneapi/dnnl/dnnl_common_types.h

@@ -94,6 +94,12 @@ typedef enum {
     dnnl_boolean = 8,
     /// 1-bit integer.
     dnnl_bin = 9,
+    /// 4-bit normalized float.
+    dnnl_nf4 = 10,
+    /// 4-bit signed integer.
+    dnnl_s4 = 11,
+    /// 4-bit unsigned integer.
+    dnnl_u4 = 12,
     /// 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

@@ -168,6 +168,9 @@ const data_type_t boolean = dnnl_boolean;
 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;
+const data_type_t s4 = dnnl_s4;
+const data_type_t u4 = dnnl_u4;
 } // namespace data_type
 
 using fpmath_mode_t = dnnl_fpmath_mode_t;

src/common/dnnl_debug_autogenerated.cpp

@@ -51,6 +51,9 @@ const char *dnnl_dt2str(dnnl_data_type_t v) {
     if (v == dnnl_f64) return "f64";
     if (v == dnnl_boolean) return "boolean";
     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

@@ -79,6 +79,18 @@ template <> struct prec_traits<data_type::bin> {
     typedef uint8_t type;
 };
 
+template <> struct prec_traits<data_type::nf4> {
+    typedef uint8_t type;
+};
+
+template <> struct prec_traits<data_type::s4> {
+    typedef uint8_t type;
+};
+
+template <> struct prec_traits<data_type::u4> {
+    typedef uint8_t type;
+};
+
 template <>
 struct data_traits<float16_t> {
     static constexpr data_type_t data_type = data_type::f16;

src/common/memory_zero_pad.cpp

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

src/common/primitive_attr.cpp

@@ -20,6 +20,7 @@
 #include "primitive_attr.hpp"
 #include "type_helpers.hpp"
 #include "utils.hpp"
+#include <iostream>
 
 using namespace dnnl::impl;
 using namespace dnnl::impl::status;
@@ -119,6 +120,22 @@ status_t zero_points_t::set(int arg, int mask) {
     return status::success;
 }
 
+status_t zero_points_t::set(int arg, const dims_t dims, int ndims) {
+    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);
+            break;
+    }
+    return status::success;
+}
+
 } // namespace impl
 } // namespace dnnl
 
@@ -548,6 +565,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_zero_points_mask(
         primitive_attr_t *attr, int arg, int mask) {
@@ -556,6 +580,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) {
+    bool ok = attr && ndims > 0;
+    if (!ok) return invalid_arguments;
+
+    return attr->zero_points_.set(arg, dims, ndims);
+}
 
 status_t dnnl_primitive_attr_set_output_compensations(primitive_attr_t *attr,
         int count, int mask) {

src/common/primitive_attr.hpp

@@ -242,8 +242,17 @@ struct runtime_scales_t : public c_compatible {
         return status::success;
     }
 
+    status_t set(const dims_t dims, int ndims) {
+        is_set_ = true;
+        ndims_ = ndims;
+        mask_ = 1;
+        utils::array_copy(dims_, dims, ndims_);
+        return status::success;
+    }
+
     bool operator==(const runtime_scales_t &rhs) const {
-        return mask_ == rhs.mask_ && is_set_ == rhs.is_set_;
+        return mask_ == rhs.mask_ && is_set_ == rhs.is_set_ &&
+               ndims_ == rhs.ndims_ && utils::array_cmp(dims_, rhs.dims_, ndims_);
     }
 
     bool has_default_values() const { return !is_set_; }
@@ -259,6 +268,9 @@ struct runtime_scales_t : public c_compatible {
     // Hide `mask_` under `private:` to force interface usage.
     int mask_ = 0;
     bool is_set_ = false;
+
+    int ndims_ = 0;
+    dnnl::impl::dims_t dims_;
 };
 
 struct arg_scales_t : public c_compatible {
@@ -295,6 +307,10 @@ struct arg_scales_t : public c_compatible {
         if (!check_arg(arg)) return status::invalid_arguments;
         return scales_[arg].set(mask);
     }
+    status_t set(int arg, const dims_t dims, int ndims) {
+        if (!check_arg(arg)) return status::invalid_arguments;
+        return scales_[arg].set(dims, ndims);
+    }
 
     status_t get(int arg, int *mask, bool *is_set) const {
         if (!check_arg(arg)) return status::invalid_arguments;
@@ -354,7 +370,8 @@ struct zero_points_t : public c_compatible {
     bool operator==(const zero_points_t &rhs) const {
         return mask_src == rhs.mask_src && mask_wei == rhs.mask_wei
                 && mask_dst == rhs.mask_dst && is_set_src == rhs.is_set_src
-                && is_set_wei == rhs.is_set_wei && is_set_dst == rhs.is_set_dst;
+                && is_set_wei == rhs.is_set_wei && is_set_dst == rhs.is_set_dst
+                && IMPLICATION(ndims_wei > 0, ndims_wei == rhs.ndims_wei && utils::array_cmp(dims_wei, rhs.dims_wei, ndims_wei));
     }
 
     // arg-specific checks
@@ -373,12 +390,26 @@ struct zero_points_t : public c_compatible {
     int get(int arg) const; // Returns 0 if dimension is unset
 
     status_t set(int arg, int mask);
+    status_t set(int arg, const dims_t dims, int ndims);
     status_t set(int arg) { return set(arg, 0); }
 
+    const dims_t & get_dims(int /*arg*/) const {
+        return dims_wei;
+    }
+    int get_ndims(int arg) const {
+        switch (arg) {
+            case DNNL_ARG_WEIGHTS: return ndims_wei; break;
+            default: return 0;
+        }
+    }
+
 private:
     bool is_set_src = false, is_set_wei = false, is_set_dst = false;
     int mask_src = 0, mask_wei = 0, mask_dst = 0;
 
+    int ndims_wei = 0;
+    dnnl::impl::dims_t dims_wei;
+
     int get_mask(int arg) const {
         int mask = 0;
         switch (arg) {

src/common/type_helpers.hpp

@@ -93,6 +93,9 @@ inline size_t data_type_size(data_type_t data_type) {
         case u8: return sizeof(prec_traits<u8>::type);
         case boolean: return sizeof(prec_traits<boolean>::type);
         case bin: return sizeof(prec_traits<u8>::type);
+        case nf4: return sizeof(prec_traits<u8>::type);
+        case u4: return sizeof(prec_traits<u8>::type);
+        case s4: return sizeof(prec_traits<u8>::type);
         case data_type::undef:
         default: assert(!"unknown data_type");
     }
@@ -284,7 +287,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) && wei_dt == u8) return f32;
+    if (one_of(src_dt, f32, bf16) && one_of(wei_dt, u8, nf4, s4, u4)) return f32;
     if (everyone_is(f64, src_dt, wei_dt)) return f64;
 
     if (one_of(prop_kind, forward_training, forward_inference)) {
@@ -949,7 +952,7 @@ inline bool memory_desc_sanity_check(int ndims, const dims_t dims,
     if (ndims == 0) return true;
 
     bool ok = dims != nullptr && 0 < ndims && ndims <= DNNL_MAX_NDIMS
-            && utils::one_of(data_type, f16, bf16, f32, f64, s32, s8, u8, bin);
+            && utils::one_of(data_type, f16, bf16, f32, f64, s32, s8, u8, bin, nf4, s4, u4);
     if (!ok) return false;
 
     bool has_runtime_dims = false;

src/cpu/cpu_inner_product_list.cpp

@@ -57,6 +57,21 @@ 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, nf4, f32}, {
+            CPU_INSTANCE_AVX512(brgemm_inner_product_fwd_t, avx512_core)
+            CPU_INSTANCE_AVX2(brgemm_inner_product_fwd_t, avx2)
+            nullptr,
+        }},
+        {{forward, f32, s4, f32}, {
+            CPU_INSTANCE_AVX512(brgemm_inner_product_fwd_t, avx512_core)
+            CPU_INSTANCE_AVX2(brgemm_inner_product_fwd_t, avx2)
+            nullptr,
+        }},
+        {{forward, f32, u4, f32}, {
+            CPU_INSTANCE_AVX512(brgemm_inner_product_fwd_t, avx512_core)
+            CPU_INSTANCE_AVX2(brgemm_inner_product_fwd_t, avx2)
+            nullptr,
+        }},
         {{forward, bf16, bf16, f32}, {
             CPU_INSTANCE_AMX(brgemm_inner_product_fwd_t, avx512_core_amx)
             CPU_INSTANCE_AVX512(brgemm_inner_product_fwd_t, avx512_core_bf16)

src/cpu/cpu_primitive.hpp

@@ -67,7 +67,7 @@
             if (scales == nullptr) return status::invalid_arguments; \
             const auto scales_d = ctx.memory_mdw(DNNL_ARG_ATTR_SCALES | arg); \
             bool ok = scales_d.data_type() == data_type::f32 \
-                    && scales_d.ndims() == 1; \
+                    && (scales_d.ndims() == 1 || scales_d.ndims() == 2); \
             if (!ok) return status::invalid_arguments; \
             if (scales_d.dims()[0] == 1) { \
                 if (utils::one_of(arg, DNNL_ARG_DST, \

src/cpu/reorder/cpu_reorder.cpp

@@ -38,6 +38,9 @@ regular_impl_list_map() {
             {{s8, data_type::undef, 0}, &regular_s8_impl_list_map()},
             {{u8, data_type::undef, 0}, &regular_u8_impl_list_map()},
             {{bin, data_type::undef, 0}, &regular_bin_impl_list_map()},
+            {{nf4, data_type::undef, 0}, &regular_nf4_impl_list_map()},
+            {{s4, data_type::undef, 0}, &regular_s4_impl_list_map()},
+            {{u4, data_type::undef, 0}, &regular_u4_impl_list_map()},
     };
     return the_map;
 }

src/cpu/reorder/cpu_reorder.hpp

@@ -56,7 +56,7 @@ struct reorder_impl_key_t {
     }
 
 private:
-    enum { MAX_DT_NUM = 10 };
+    enum { MAX_DT_NUM = 12 };
     size_t value() const {
         return ((size_t)ndims * MAX_DT_NUM + (size_t)src_dt) * MAX_DT_NUM
                 + (size_t)dst_dt;
@@ -80,6 +80,9 @@ extern const impl_list_map_t &regular_s32_impl_list_map();
 extern const impl_list_map_t &regular_s8_impl_list_map();
 extern const impl_list_map_t &regular_u8_impl_list_map();
 extern const impl_list_map_t &regular_bin_impl_list_map();
+extern const impl_list_map_t &regular_nf4_impl_list_map();
+extern const impl_list_map_t &regular_s4_impl_list_map();
+extern const impl_list_map_t &regular_u4_impl_list_map();
 
 /* conv reorders w/ compensation */
 extern const impl_list_map_t &comp_f32_s8_impl_list_map();

src/cpu/reorder/cpu_reorder_regular_nf4.cpp

@@ -0,0 +1,44 @@
+/*******************************************************************************
+* Copyright 2021 Intel Corporation
+*
+* Licensed 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.
+*******************************************************************************/
+
+#include "cpu/reorder/cpu_reorder.hpp"
+
+namespace dnnl {
+namespace impl {
+namespace cpu {
+
+// clang-format off
+
+const impl_list_map_t &regular_nf4_impl_list_map() {
+    static const impl_list_map_t the_map = REG_REORDER_P({
+        // nf4 ->
+        {{nf4, data_type::undef, 0}, {
+            REG_SR(nf4, any, nf4, OI8i8o2i, fmt_order_keep)
+            REG_SR(nf4, any, nf4, OI8i16o2i, fmt_order_keep)
+            REG_SR(nf4, any, nf4, OI8i24o2i, fmt_order_keep)
+            REG_SR(nf4, any, nf4, OI8i32o2i, fmt_order_keep)
+            REG_SR(nf4, any, nf4, OI8i64o2i, fmt_order_keep)
+            nullptr,
+        }},
+    });
+    return the_map;
+}
+
+// clang-format on
+
+} // namespace cpu
+} // namespace impl
+} // namespace dnnl