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

dmitrygo · dmitrygo · commit cc986b3180d3 · 2023-12-18T14:38:27.000+01:00
diff --git a/src/cpu/cpu_inner_product_list.cpp b/src/cpu/cpu_inner_product_list.cpp
@@ -90,10 +90,42 @@ const std::map<pk_dt_impl_key_t, std::vector<impl_list_item_t>> &impl_list_map()
         }},
         {{forward, bf16, u8, 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, u8, 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, bf16, nf4, 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, nf4, 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, bf16, s4, 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, s4, 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, bf16, u4, 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, u4, 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}, {
diff --git a/src/cpu/x64/brgemm/brgemm_utils.cpp b/src/cpu/x64/brgemm/brgemm_utils.cpp
@@ -51,7 +51,8 @@ void init_kernel_datatype(
     assert(dt_a != data_type::undef && dt_b != data_type::undef);
     brg->is_int8 = utils::one_of(dt_a, data_type::u8, data_type::s8)
             && utils::one_of(dt_b, data_type::u8, data_type::s8);
-    brg->is_bf16 = (dt_a == data_type::bf16) && (dt_b == data_type::bf16);
+    brg->is_bf16 = one_of(dt_a, data_type::bf16) &&
+                   one_of(dt_b, data_type::bf16, data_type::u8, data_type::nf4, data_type::s4, data_type::u4);
     brg->is_f32 = one_of(dt_a, data_type::f32) &&
                   one_of(dt_b, data_type::f32, data_type::u8, data_type::nf4, data_type::s4, data_type::u4);
     brg->is_f16 = utils::one_of(data_type::f16, dt_a, dt_b);
@@ -833,15 +834,15 @@ void init_brgemm_conf(brgemm_t *brg, cpu_isa_t isa, brgemm_batch_kind_t type,
     brg->bdb2_tail = 0;
 
     const bool is_b_in_vnni_format = !(brg->dt_b == data_type::f16 && brg->isa_impl == avx512_core_fp16) &&
-                                     !(brg->dt_a == data_type::f32 && one_of(brg->dt_b, data_type::u8));
+                                     !(one_of(brg->dt_a, data_type::f32, data_type::bf16) && one_of(brg->dt_b, data_type::u8));
     brg->ld_step
             = is_b_in_vnni_format ? data_type_vnni_granularity(brg->dt_b) : 1;
 
     const bool has_no_vnni_compute_instruction
             = (brg->is_f16
                       && one_of(brg->isa_impl, avx2_vnni_2, avx512_core_fp16))
             || (brg->is_bf16 && brg->isa_impl == avx2_vnni_2)
-            || (one_of(brg->dt_a, data_type::f32) &&
+            || (one_of(brg->dt_a, data_type::f32, data_type::bf16) &&
                 one_of(brg->dt_b, data_type::u8, data_type::nf4, data_type::s4, data_type::u4));
     brg->rd_step = has_no_vnni_compute_instruction
             ? 1
diff --git a/src/cpu/x64/brgemm/jit_brgemm_kernel.cpp b/src/cpu/x64/brgemm/jit_brgemm_kernel.cpp
@@ -2027,8 +2027,14 @@ void jit_brgemm_kernel_t<isa, Wmm>::gemm_microkernel(int bd_block2,
                         const auto bd_by_load_bytes
                                 = (bd >= bd_e - rows_by_load_bytes
                                         || brg.brgattr.wary_tail_read);
-                        broadcast(bcst(), A_offset(bd, rd),
-                                have_to_load_bytes && bd_by_load_bytes, brg.dt_a);
+                        if (brg.dt_a == data_type::bf16) {
+                            vpbroadcastw(bcst(), ptr[reg_aux_A + A_offset(bd, rd)]);
+                            uni_vpmovzxwd(bcst(), bcst());
+                            uni_vpslld(bcst(), bcst(), 16);
+                        } else {
+                            broadcast(bcst(), A_offset(bd, rd),
+                                    have_to_load_bytes && bd_by_load_bytes, brg.dt_a);
+                        }
                     }
                     if (prefetch_count_B < ld_block2) {
                         prefetcht0(ptr[reg_aux_B + B_offset(prefetch_count_B++, rd)
@@ -2044,10 +2050,9 @@ void jit_brgemm_kernel_t<isa, Wmm>::gemm_microkernel(int bd_block2,
                                 uni_vmulps(vmm, load(ld), bcst());
                         } else {
                             if (is_emdbd)
-                                uni_vfmadd231ps(vmm, load(ld),
-                                        ptr_b[reg_aux_A + A_offset(bd, rd)]);
+                                uni_vfmadd231ps(vmm, load(ld), ptr_b[reg_aux_A + A_offset(bd, rd)]);
                             else
-                                dot_product(vmm, load(ld), bcst());
+                                uni_vfmadd231ps(vmm, load(ld), bcst());
                         }
                     }
                 }
diff --git a/src/cpu/x64/jit_brgemm_inner_product.cpp b/src/cpu/x64/jit_brgemm_inner_product.cpp
@@ -307,7 +307,8 @@ status_t brgemm_inner_product_fwd_t<isa>::execute_forward(
                     const size_t decomp_buf_per_thr = jbgp.ic_block * jbgp.nb_ic_blocking * jbgp.oc_block * types::data_type_size(jbgp.wei_dt);
                     auto decomp_buf = decomp_buf_global + ithr * decomp_buf_per_thr + wei_ic_stride * b * ic_blocks_per_batch;
 
-                    const int ic_internal_block = is_amx || one_of(pd()->jbgp_.orig_wei_dt, data_type::nf4, data_type::s4, data_type::u4) ? 2 : 1;
+                    const int ic_internal_block = pd()->jbgp_.wei_dt == data_type::bf16 ||
+                                                  one_of(pd()->jbgp_.orig_wei_dt, data_type::nf4, data_type::s4, data_type::u4) ? 2 : 1;
                     auto wei_zero_points_ptr = wei_zero_points + wei_zero_points_oc_stride * oc;
                     auto wei_scales_ptr = wei_scales + wei_scales_oc_stride * oc;
 
@@ -408,7 +409,8 @@ status_t brgemm_inner_product_fwd_t<isa>::execute_forward(
                 const size_t decomp_buf_per_thr = jbgp.ic_block * jbgp.nb_ic_blocking * jbgp.oc_block * types::data_type_size(jbgp.wei_dt);
                 auto decomp_buf = decomp_buf_global + ithr * decomp_buf_per_thr;
 
-                const int ic_internal_block = is_amx || one_of(pd()->jbgp_.orig_wei_dt, data_type::nf4, data_type::s4, data_type::u4) ? 2 : 1;
+                const int ic_internal_block = pd()->jbgp_.wei_dt == data_type::bf16 ||
+                                              one_of(pd()->jbgp_.orig_wei_dt, data_type::nf4, data_type::s4, data_type::u4) ? 2 : 1;
                 auto wei_zero_points_ptr = wei_zero_points + wei_zero_points_oc_stride * oc;
                 auto wei_scales_ptr = wei_scales + wei_scales_oc_stride * oc;
 
diff --git a/src/cpu/x64/jit_brgemm_inner_product.hpp b/src/cpu/x64/jit_brgemm_inner_product.hpp
@@ -212,7 +212,8 @@ struct brgemm_inner_product_fwd_t : public primitive_t {
         if (pd()->jbgp_.weights_decompression && pd()->jbgp_.wei_decomp_algo == weights_decomp_kind_t::prepack) {
             weights_decompression_compile_params_t jcp = {};
             jcp.oc_size = pd()->jbgp_.oc_block;
-            jcp.ic_internal_size = pd()->jbgp_.is_amx || utils::one_of(pd()->jbgp_.orig_wei_dt, data_type::nf4, data_type::s4, data_type::u4) ? 2 : 1;
+            jcp.ic_internal_size = pd()->jbgp_.wei_dt == data_type::bf16 ||
+                                   utils::one_of(pd()->jbgp_.orig_wei_dt, data_type::nf4, data_type::s4, data_type::u4) ? 2 : 1;
             jcp.with_scales = !pd()->attr()->scales_.get(DNNL_ARG_WEIGHTS).has_default_values();
             jcp.broadcast_scales = pd()->attr()->scales_.get(DNNL_ARG_WEIGHTS).dims_[0] == 1;
             jcp.with_zero_points = !pd()->attr()->zero_points_.has_default_values(DNNL_ARG_WEIGHTS);
diff --git a/src/cpu/x64/jit_brgemm_inner_product_utils.cpp b/src/cpu/x64/jit_brgemm_inner_product_utils.cpp
@@ -151,7 +151,7 @@ jit_brgemm_ip_conf_t::get_desired_weights_tag() const {
     const bool is_xf16 = utils::one_of(jbgp.wei_dt, bf16, f16);
     const bool is_not_vnni_tag = (jbgp.wei_dt == f32
             || (jbgp.wei_dt == f16 && jbgp.isa == avx512_core_fp16)) && !jbgp.weights_decompression;
-    if (is_not_vnni_tag || (jbgp.weights_decompression && jbgp.orig_wei_dt == u8 && jbgp.wei_dt != bf16)) {
+    if (is_not_vnni_tag || (jbgp.weights_decompression && jbgp.orig_wei_dt == u8)) {
         if (is_superset(jbgp.isa, avx512_core))
             return {{64,
                             pick(n_sp_dims, OI16i64o, OIw16i64o, OIhw16i64o,
@@ -176,7 +176,7 @@ jit_brgemm_ip_conf_t::get_desired_weights_tag() const {
                             pick(n_sp_dims, OI8i16o, OIw8i16o, OIhw8i16o,
                                     OIdhw8i16o)},
                     {8, pick(n_sp_dims, OI8i8o, OIw8i8o, OIhw8i8o, OIdhw8i8o)}};
-    } else if (is_xf16 || (jbgp.weights_decompression && jbgp.orig_wei_dt == u8 && jbgp.wei_dt == bf16)) {
+    } else if (is_xf16) {
         if (jbgp.is_amx) {
             return {{64,
                             pick(n_sp_dims, OI16i64o2i, OIw16i64o2i,
@@ -374,10 +374,11 @@ int jit_brgemm_ip_conf_t::get_adjusted_oc_block() const {
     // time for weights reorder are key optimization points there.
     const size_t wei_size = static_cast<size_t>(jbgp.ic * jbgp.oc) * types::data_type_size(jbgp.wei_dt);
     // Use oc block to be 32 if weight size >= 8MB on amx bf16 to optimized memory consumption.
-    if (jbgp.is_amx && jbgp.wei_dt == bf16 && !jbgp.is_bf32 && wei_size >= 8 * (1 << 20))
+    if (jbgp.is_amx && jbgp.orig_wei_dt == bf16 && !jbgp.is_bf32 && wei_size >= 8 * (1 << 20))
         return 32;
     // Use oc block to be 64 if weight size >= 16MB on avx512 f32 to optimized memory consumption.
-    if (is_f32_compute_avx512 && wei_size >= 16 * (1 << 20))
+    if ((is_f32_compute_avx512 || (jbgp.is_amx && jbgp.orig_wei_dt != bf16 && !jbgp.is_bf32))
+        && wei_size >= 16 * (1 << 20))
         return 64;
     // Use oc block to be 24 if weight size >= 16MB on avx2 f32 to optimized memory consumption.
     if (is_f32_compute_avx2 && wei_size >= 16 * (1 << 20))
@@ -1339,8 +1340,7 @@ status_t jit_brgemm_ip_conf_t::init_conf_base(cpu_isa_t isa,
 
     jbgp.weights_decompression = one_of(jbgp.src_dt, f32, bf16) &&
                                  one_of(jbgp.wei_dt, u8, nf4, s4, u4);
-    jbgp.wei_decomp_algo = jbgp.is_amx ? weights_decomp_kind_t::prepack
-                                       : weights_decomp_kind_t::immediate;
+    jbgp.wei_decomp_algo = weights_decomp_kind_t::immediate;
     jbgp.orig_wei_dt = jbgp.wei_dt;
     jbgp.with_grouped_weights_decompression = false;
     if (jbgp.weights_decompression) {
@@ -1366,6 +1366,10 @@ status_t jit_brgemm_ip_conf_t::init_conf_base(cpu_isa_t isa,
         }
     }
 
+    // Current AMX implementation cannot provide perfromance benefit for immediate algorithm over avx512 version
+    if (jbgp.is_amx && jbgp.weights_decompression && jbgp.wei_decomp_algo == weights_decomp_kind_t::immediate)
+        return status::unimplemented;
+
     jbgp.bia_dt = jbgp.with_bias
             ? pick_by_prop_kind(jbgp.prop_kind, ipd.bias_desc.data_type,
                     data_type::undef, ipd.diff_bias_desc.data_type)
@@ -1382,7 +1386,8 @@ status_t jit_brgemm_ip_conf_t::init_conf_base(cpu_isa_t isa,
                     everyone_is(bf16, jbgp.wei_dt, jbgp.dst_dt)
                             && jbgp.src_dt == f32,
                     everyone_is(bf16, jbgp.src_dt, jbgp.dst_dt)
-                            && jbgp.wei_dt == f32);
+                            && jbgp.wei_dt == f32)
+            || (jbgp.weights_decompression && jbgp.src_dt == bf16 && one_of(jbgp.dst_dt, f32, bf16));
     const bool is_f16 = everyone_is(f16, jbgp.src_dt, jbgp.wei_dt, jbgp.dst_dt)
             || pick_by_prop_kind(jbgp.prop_kind,
                     everyone_is(f16, jbgp.src_dt, jbgp.wei_dt)
diff --git a/src/cpu/x64/jit_brgemm_weights_decompression_kernel.cpp b/src/cpu/x64/jit_brgemm_weights_decompression_kernel.cpp
diff --git a/src/cpu/x64/jit_brgemm_weights_decompression_kernel.hpp b/src/cpu/x64/jit_brgemm_weights_decompression_kernel.hpp
