[FORK][FEATURE] Binary networks support

Author: dmitry-gorokhov

include/oneapi/dnnl/dnnl.h

@@ -823,6 +823,9 @@ dnnl_status_t DNNL_API dnnl_post_ops_append_quantization(
         const void* input_scale, const void* input_shift,
         const void* output_scale, const void* output_shift);
 
+dnnl_status_t DNNL_API dnnl_post_ops_append_binarization(
+        dnnl_post_ops_t post_ops, dnnl_alg_kind_t alg, const float* weights_data, const float* output_mask);
+
 /// @} dnnl_api_attributes
 
 /// @} dnnl_api_primitives

include/oneapi/dnnl/dnnl.hpp

@@ -508,6 +508,7 @@ enum class algorithm {
 
     quantization_quantize_dequantize = dnnl_quantization_quantize_dequantize,
     quantization_quantize = dnnl_quantization_quantize,
+    binarization_depthwise = dnnl_binarization_depthwise,
 };
 
 /// Converts algorithm kind enum value from C++ API to C API type.
@@ -891,6 +892,8 @@ struct memory : public handle<dnnl_memory_t> {
         s4 = dnnl_s4,
         /// 4-bit unsigned integer.
         u4 = dnnl_u4,
+        /// 1-bit integer
+        bin = dnnl_bin
     };
 
     /// Returns size of data type in bytes.
@@ -1317,6 +1320,7 @@ struct memory : public handle<dnnl_memory_t> {
         aBCd4b4c = dnnl_aBCd4b4c,
         ABcd8a16b2a = dnnl_ABcd8a16b2a,
         ABcd8a8b = dnnl_ABcd8a8b,
+        ABcd8a32b = dnnl_ABcd8a32b,
         ABcd8a4b = dnnl_ABcd8a4b,
         ABcd8a2b = dnnl_ABcd8a2b,
         /// 4D tensor blocked by 2nd dimension with block size 8
@@ -1805,6 +1809,8 @@ struct memory : public handle<dnnl_memory_t> {
         OhwI8i8o = dnnl_OhwI8i8o,
         OIhw8o16i2o = dnnl_OIhw8o16i2o,
         OIhw8o8i = dnnl_OIhw8o8i,
+        OIhw8o32i = dnnl_OIhw8o32i,
+        OIhw16o32i = dnnl_OIhw16o32i,
         OIhw8o4i = dnnl_OIhw8o4i,
         OIhw2i8o4i = dnnl_OIhw2i8o4i,
         IOdhw16o16i = dnnl_IOdhw16o16i,
@@ -3956,6 +3962,11 @@ struct post_ops : public handle<dnnl_post_ops_t> {
                 input_scale, input_shift, output_scale, output_shift),
                           "could not append quantization");
     }
+
+    void append_binarization(algorithm alg, const float* weights_data, const float* output_mask) {
+        error::wrap_c_api(dnnl_post_ops_append_binarization(get(), convert_to_c(alg), weights_data, output_mask),
+                          "could not append binarization");
+    }
 };
 
 /// @cond DO_NOT_DOCUMENT_THIS

include/oneapi/dnnl/dnnl_common_types.h

@@ -105,6 +105,8 @@ typedef enum {
     /// [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,
 
+    /// 1-bit integer.
+    dnnl_bin = 14,
     /// 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,

include/oneapi/dnnl/dnnl_types.h

@@ -286,6 +286,7 @@ typedef enum {
     dnnl_ABcd8a16b2a,
     dnnl_ABcd2b8a4b,
     dnnl_ABcd8a8b,
+    dnnl_ABcd8a32b,
     dnnl_ABcd8a4b,
     /// 4D tensor blocked by 2nd dimension with block size 8
     dnnl_aBcd8b,
@@ -1391,6 +1392,8 @@ typedef enum {
     dnnl_OIhw2i8o4i = dnnl_ABcd2b8a4b,
     dnnl_IOhw8o16i2o = dnnl_BAcd8a16b2a,
     dnnl_OIhw8o8i = dnnl_ABcd8a8b,
+    dnnl_OIhw8o32i = dnnl_ABcd8a32b,
+    dnnl_OIhw16o32i = dnnl_ABcd16a32b,
     dnnl_OIhw8o4i = dnnl_ABcd8a4b,
     dnnl_Owhi16o = dnnl_Adcb16a,
     dnnl_OIhw8i32o = dnnl_ABcd8b32a,
@@ -2000,6 +2003,8 @@ typedef enum {
     dnnl_depthwise,
     /// A quantization primitive.
     dnnl_quantization,
+    /// A binatization primitive.
+    dnnl_binarization,
     /// An LRN primitive.
     dnnl_lrn,
     /// A batch normalization primitive.
@@ -2196,6 +2201,7 @@ typedef enum {
 
     dnnl_quantization_quantize_dequantize = 0x4fff0,
     dnnl_quantization_quantize = 0x4fff1,
+    dnnl_binarization_depthwise = 0x4fff2,
 } dnnl_alg_kind_t;
 
 /// Flags for normalization primitives.
@@ -2915,6 +2921,7 @@ typedef enum {
     dnnl_cpu_isa_avx10_1_512_amx_fp16 = 0x1fef,
     /// @copydoc dnnl_cpu_isa_avx10_1_512_amx_fp16
     dnnl_cpu_isa_avx512_core_amx_fp16 = dnnl_cpu_isa_avx10_1_512_amx_fp16,
+    dnnl_cpu_isa_avx512_vpopcnt = 0x3fef,
 } dnnl_cpu_isa_t;
 
 /// CPU ISA hints flags

src/common/c_types_map.hpp

@@ -148,6 +148,7 @@ const alg_kind_t depthwise_scale_shift = dnnl_depthwise_scale_shift;
 const alg_kind_t depthwise_prelu = dnnl_depthwise_prelu;
 const alg_kind_t quantization_quantize_dequantize = dnnl_quantization_quantize_dequantize;
 const alg_kind_t quantization_quantize = dnnl_quantization_quantize;
+const alg_kind_t binarization_depthwise = dnnl_binarization_depthwise;
 } // namespace alg_kind
 
 using data_type_t = dnnl_data_type_t;
@@ -171,6 +172,7 @@ const data_type_t data_type_max = dnnl_data_type_max;
 // Not exposed through API as all current uses are internal only
 const data_type_t tf32 = static_cast<data_type_t>(1 << 8);
 
+const data_type_t bin = dnnl_bin;
 } // namespace data_type
 
 using fpmath_mode_t = dnnl_fpmath_mode_t;
@@ -482,6 +484,7 @@ const format_tag_t aBCd4b4c = dnnl_aBCd4b4c;
 const format_tag_t ABcd8a16b2a = dnnl_ABcd8a16b2a;
 const format_tag_t BAcd8a16b2a = dnnl_BAcd8a16b2a;
 const format_tag_t ABcd8a8b = dnnl_ABcd8a8b;
+const format_tag_t ABcd8a32b = dnnl_ABcd8a32b;
 const format_tag_t ABcd8a4b = dnnl_ABcd8a4b;
 const format_tag_t ABcd8a2b = dnnl_ABcd8a2b;
 const format_tag_t aBcd8b = dnnl_aBcd8b;
@@ -1287,6 +1290,8 @@ const format_tag_t OhwI8i8o = dnnl_OhwI8i8o;
 const format_tag_t OIhw8o16i2o = dnnl_OIhw8o16i2o;
 const format_tag_t IOhw8o16i2o = dnnl_IOhw8o16i2o;
 const format_tag_t OIhw8o8i = dnnl_OIhw8o8i;
+const format_tag_t OIhw8o32i = dnnl_OIhw8o32i;
+const format_tag_t OIhw16o32i = dnnl_OIhw16o32i;
 const format_tag_t OIhw8o4i = dnnl_OIhw8o4i;
 const format_tag_t Owhi16o = dnnl_Owhi16o;
 const format_tag_t Odwhi16o = dnnl_Odwhi16o;
@@ -1964,6 +1969,7 @@ const primitive_kind_t layer_normalization = dnnl_layer_normalization;
 const primitive_kind_t group_normalization = dnnl_group_normalization;
 const primitive_kind_t depthwise = dnnl_depthwise;
 const primitive_kind_t quantization = dnnl_quantization;
+const primitive_kind_t binarization = dnnl_binarization;
 
 // Internal only primitive kinds.
 const primitive_kind_t internal_only_start = (primitive_kind_t)(1 << 12);

src/common/dnnl_debug_autogenerated.cpp

@@ -57,6 +57,7 @@ const char *dnnl_dt2str(dnnl_data_type_t v) {
     if (v == dnnl_s4) return "s4";
     if (v == dnnl_u4) return "u4";
     if (v == dnnl_e8m0) return "e8m0";
+    if (v == dnnl_bin) return "bin";
     if (v == dnnl_data_type_max) return "data_type_max";
     assert(!"unknown dt");
     return "unknown dt";
@@ -204,6 +205,8 @@ const char *dnnl_fmt_tag2str(dnnl_format_tag_t v) {
     if (v == dnnl_ABcd8a16b2a) return "ABcd8a16b2a";
     if (v == dnnl_ABcd2b8a4b) return "ABcd2b8a4b";
     if (v == dnnl_ABcd8a8b) return "ABcd8a8b";
+    if (v == dnnl_ABcd8a32b) return "ABcd8a32b";
+    if (v == dnnl_ABcd16a32b) return "ABcd16a32b";
     if (v == dnnl_ABcd8a4b) return "ABcd8a4b";
     if (v == dnnl_aBcd8b) return "aBcd8b";
     if (v == dnnl_aBCd4c8b2c) return "aBCd4c8b2c";
@@ -1177,6 +1180,8 @@ const char *dnnl_fmt_tag2str(dnnl_format_tag_t v) {
     if (v == dnnl_OIhw8o16i2o) return "OIhw8o16i2o";
     if (v == dnnl_OIhw2i8o4i) return "OIhw2i8o4i";
     if (v == dnnl_IOhw8o16i2o) return "IOhw8o16i2o";
+    if (v == dnnl_OIhw8o32i) return "OIhw8o23i";
+    if (v == dnnl_OIhw16o32i) return "OIhw16o23i";
     if (v == dnnl_OIhw8o8i) return "OIhw8o8i";
     if (v == dnnl_OIhw8o4i) return "OIhw8o4i";
     if (v == dnnl_Owhi16o) return "Owhi16o";
@@ -1843,6 +1848,7 @@ const char *dnnl_alg_kind2str(dnnl_alg_kind_t v) {
     if (v == dnnl_depthwise_prelu) return "depthwise_prelu";
     if (v == dnnl_quantization_quantize_dequantize) return "quantization_quantize_dequantize";
     if (v == dnnl_quantization_quantize) return "quantization_quantize";
+    if (v == dnnl_binarization_depthwise) return "binarization_depthwise";
     assert(!"unknown alg_kind");
     return "unknown alg_kind";
 }

src/common/dnnl_traits.hpp

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

src/common/memory_desc_wrapper.cpp

@@ -397,6 +397,7 @@ status_t memory_desc_wrapper::compute_blocking(
         C(ABcd8a16b2a, {0, 1, 2, 3}, {8, 16, 2}, {0, 1, 0});
         C(BAcd8a16b2a, {1, 0, 2, 3}, {8, 16, 2}, {0, 1, 0});
         C(ABcd8a8b, {0, 1, 2, 3}, {8, 8}, {0, 1});
+        C(ABcd8a32b, {0, 1, 2, 3}, {8, 32}, {0, 1});
         C(ABcd8a4b, {0, 1, 2, 3}, {8, 4}, {0, 1});
         C(ABcd8a2b, {0, 1, 2, 3}, {8, 2}, {0, 1});
         C(aBcd8b, {0, 1, 2, 3}, {8}, {1});

src/common/memory_zero_pad.cpp

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

src/common/primitive_attr.cpp

@@ -380,6 +380,23 @@ status_t post_ops_t::append_quantization(alg_kind_t alg,
     return success;
 }
 
+status_t post_ops_t::append_binarization(alg_kind_t alg, const float* weights_data, const float* output_mask_data) {
+    using namespace dnnl::impl::alg_kind;
+    if (len() == post_ops_limit) return out_of_memory;
+    bool known_alg = one_of(alg, binarization_depthwise);
+    if (!known_alg)
+        return invalid_arguments;
+
+    entry_.emplace_back();
+    auto &e = entry_.back();
+    e.kind = primitive_kind::binarization;
+    e.binarization.alg = alg;
+    e.binarization.weights_data = weights_data;
+    e.binarization.output_mask_data = output_mask_data;
+
+    return success;
+}
+
 bool post_ops_t::defined() const {
     for (int idx = 0; idx < len(); ++idx) {
         auto kind = entry_[idx].kind;
@@ -398,6 +415,8 @@ bool post_ops_t::defined() const {
             // depthwise is always defined
         } else if (kind == primitive_kind::quantization) {
             // quantization is always defined
+        } else if (kind == primitive_kind::binarization) {
+            // binarization is always defined
         } else {
             assert(!"unreachable");
         }
@@ -940,6 +959,14 @@ status_t dnnl_post_ops_append_quantization(post_ops_t *post_ops, alg_kind_t kind
     return post_ops->append_quantization(kind, crop_low, crop_high, input_scale, input_shift, output_scale, output_shift);
 }
 
+status_t dnnl_post_ops_append_binarization(post_ops_t *post_ops, alg_kind_t kind, const float* weights_data,
+                                           const float* output_mask_data) {
+    if (post_ops == nullptr)
+        return invalid_arguments;
+
+    return post_ops->append_binarization(kind, weights_data, output_mask_data);
+}
+
 status_t dnnl_primitive_attr_set_rnn_data_qparams(
         primitive_attr_t *attr, const float scale, const float shift) {
     if (attr == nullptr) return invalid_arguments;

src/common/primitive_attr.hpp

@@ -677,6 +677,12 @@ struct dnnl_post_ops : public dnnl::impl::c_compatible {
             const dnnl::impl::shifts_t<float>* output_shift_data;
         };
 
+        struct binarization_t {
+            dnnl::impl::alg_kind_t alg;
+            const float* weights_data;
+            const float* output_mask_data;
+        } ;
+
         dnnl::impl::primitive_kind_t kind
                 = dnnl::impl::primitive_kind::undefined;
         union {
@@ -687,6 +693,7 @@ struct dnnl_post_ops : public dnnl::impl::c_compatible {
             prelu_t prelu;
             depthwise_t depthwise;
             quantization_t quantization;
+            binarization_t binarization;
         };
 
         bool is_eltwise(bool require_scale_one = false) const {
@@ -735,6 +742,11 @@ struct dnnl_post_ops : public dnnl::impl::c_compatible {
             return kind == primitive_kind::quantization;
         }
 
+        bool is_binarization() const {
+            using namespace dnnl::impl;
+            return kind == primitive_kind::binarization;
+        }
+
         dnnl::impl::status_t set_depthwise_scales(const float *scales);
 
         bool operator==(const entry_t &rhs) const {
@@ -790,6 +802,11 @@ struct dnnl_post_ops : public dnnl::impl::c_compatible {
                           && quantization.output_scale_data == rhs.quantization.output_scale_data
                           && quantization.output_shift_data == rhs.quantization.output_shift_data;
                     break;
+                case primitive_kind::binarization:
+                    ret = depthwise.alg == rhs.depthwise.alg
+                          && binarization.weights_data == rhs.binarization.weights_data
+                          && binarization.output_mask_data == rhs.binarization.output_mask_data;
+                    break;
                 default: assert(!"unsupported post_op");
             }
             return ret;
@@ -820,6 +837,8 @@ struct dnnl_post_ops : public dnnl::impl::c_compatible {
             const void* crop_low, const void* crop_high,
             const void* input_scale, const void* input_shift,
             const void* output_scale, const void* output_shift);
+    dnnl::impl::status_t append_binarization(dnnl::impl::alg_kind_t alg, const float* weights_data,
+            const float* output_mask_data);
 
     dnnl::impl::status_t prepend_binary(dnnl::impl::alg_kind_t alg,
             const dnnl::impl::memory_desc_t *user_src1_desc);

src/common/tag_traits.hpp

@@ -89,6 +89,7 @@ enum class inner_blk_t {
     _8b16c,
     _8b24c,
     _8b32a,
+    _8a32b,
     _8b8c,
     _8c2b,
     _8c4b,
@@ -209,7 +210,7 @@ constexpr int AB_or_BC_blk_off(int x0, int x1) {
             utils::one_of(f, ib::_4a4b, ib::_4b4a, ib::_4b4c, ib::_4c4b,
                     ib::_8a2b, ib::_8a4b, ib::_8b2a, ib::_8b4a, ib::_8b2c,
                     ib::_8c2b, ib::_8c4b, ib::_8b4c, ib::_8a8b, ib::_8b8a,
-                    ib::_8b16a, ib::_8b24a, ib::_8b32a, ib::_8b8c, ib::_8c8b,
+                    ib::_8b16a, ib::_8b24a, ib::_8b32a, ib::_8a32b, ib::_8b8c, ib::_8c8b,
                     ib::_16a16b, ib::_16b64a, ib::_16b48a, ib::_16b32a,
                     ib::_16b16a, ib::_16b16c, ib::_16c16b, ib::_32a32b,
                     ib::_16a2b, ib::_16a4b, ib::_16b2a, ib::_16b4a, ib::_16b2c,
@@ -241,9 +242,10 @@ constexpr int AB_or_BC_blk_off(int x0, int x1) {
         : (utils::one_of(f, ib::_2a24b, ib::_2b24c, ib::_8a24b, ib::_8b24c)) ? 24 * x0 + x1
         : (f == ib::_4a4b || f == ib::_4b4c) ? 4 * x0 + x1
         : (f == ib::_4b4a || f == ib::_4c4b) ? 4 * x1 + x0
-        : (f == ib::_8a8b || f == ib::_8b8c) ? 8 * x0 + x1
+        : (f == ib::_8a8b || f == ib::_8a32b || f == ib::_8b8c) ? 8 * x0 + x1
         : (f == ib::_8b8a || f == ib::_8c8b) ? 8 * x1 + x0
         : (utils::one_of(f, ib::_16a16b, ib::_16b16c, ib::_8a16b, ib::_8b16c)) ? 16 * x0 + x1
+        : (f == ib::_16a16b || f == ib::_16a32b || f == ib::_16b16c) ? 16 * x0 + x1
         : (f == ib::_16b64a) ? 64 * x1 + x0
         : (f == ib::_16b48a) ? 48 * x1 + x0
         : (f == ib::_8b32a || f == ib::_16b32a) ? 32 * x1 + x0
@@ -668,6 +670,8 @@ DECL_TRAITS(ABcde16b48a2b, _AB, _16b48a2b, 5);
 DECL_TRAITS(ABcde16b64a2b, _AB, _16b64a2b, 5);
 DECL_TRAITS(ABcd8a16b2a, _AB, _8a16b2a, 4);
 DECL_TRAITS(ABcd8a8b, _AB, _8a8b, 4);
+DECL_TRAITS(ABcd8a32b, _AB, _8a32b, 4);
+DECL_TRAITS(ABcd16a32b, _AB, _16a32b, 4);
 DECL_TRAITS(aBcd8b, _B, _8b, 4);
 DECL_TRAITS(ABcd8b16a2b, _AB, _8b16a2b, 4);
 DECL_TRAITS(AcdB8b16a2b, _AB, _8b16a2b, 4);

src/common/type_helpers.hpp

@@ -96,6 +96,7 @@ inline size_t data_type_size(data_type_t data_type) {
         case s4: return sizeof(prec_traits<s4>::type);
         case u4: return sizeof(prec_traits<u4>::type);
         case boolean: return sizeof(prec_traits<boolean>::type);
+        case bin: return sizeof(prec_traits<u8>::type);
         case data_type::undef:
         default: assert(!"unknown data_type");
     }
@@ -1253,7 +1254,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, f8_e5m2, f8_e4m3, f16, bf16, f32, f64,
-                    s32, s8, u8, s4, u4);
+                    s32, s8, u8, s4, u4, bin);
     if (!ok) return false;
 
     bool has_runtime_dims = false;

src/cpu/reorder/cpu_reorder.cpp

@@ -36,6 +36,7 @@ regular_impl_list_map() {
             {{f32, u8, 0}, &regular_f32_u8_impl_list_map()},
             {{f8_e5m2, data_type::undef, 0}, &regular_fp8_impl_list_map()},
             {{f8_e4m3, data_type::undef, 0}, &regular_fp8_impl_list_map()},
+            {{f32, bin, 0}, &regular_f32_bin_impl_list_map()},
             {{bf16, data_type::undef, 0}, &regular_bf16_impl_list_map()},
             {{f16, data_type::undef, 0}, &regular_f16_impl_list_map()},
             {{s32, data_type::undef, 0}, &regular_s32_impl_list_map()},
@@ -45,6 +46,7 @@ regular_impl_list_map() {
             {{f32, u4, 0}, &regular_u4_impl_list_map()},
             {{s4, f32, 0}, &regular_s4_impl_list_map()},
             {{u4, f32, 0}, &regular_u4_impl_list_map()},
+            {{bin, data_type::undef, 0}, &regular_bin_impl_list_map()},
     };
     return the_map;
 }