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

3.5 squash list:

[FORK][FIX] Corrected brgemm rd_step for bf16 compressed weights

Author: dmitry-gorokhov

src/cpu/cpu_inner_product_list.cpp

@@ -88,6 +88,11 @@ 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, bf16, 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/x64/brgemm/brgemm_utils.cpp

@@ -55,7 +55,7 @@ void init_kernel_datatype(
     brg->is_bf16 = one_of(dt_a, data_type::bf16) &&
                    one_of(dt_b, data_type::bf16, data_type::u8, data_type::s8, 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::s8, data_type::nf4, data_type::s4, data_type::u4);
+                  one_of(dt_b, data_type::f32, data_type::f16, data_type::bf16, data_type::u8, data_type::s8, data_type::nf4, data_type::s4, data_type::u4);
     brg->is_f16 = utils::one_of(data_type::f16, dt_a, dt_b);
     brg->is_fp8 = one_of(dt_a, data_type::f8_e5m2, data_type::f8_e4m3)
             && one_of(dt_b, data_type::f8_e5m2, data_type::f8_e4m3);
@@ -905,10 +905,9 @@ void init_brgemm_conf(brgemm_desc_t *brg, cpu_isa_t isa,
     brg->bdb2 = 0;
     brg->bdb2_tail = 0;
 
-    const bool is_vcvtph2ps_kernel = (brg->dt_b == data_type::f16 && brg->dt_a == data_type::f32);
     const bool is_b_in_vnni_format = !(brg->dt_b == data_type::f16 && brg->isa_impl == avx512_core_fp16) &&
                                      !(one_of(brg->dt_a, data_type::f32, data_type::bf16) && one_of(brg->dt_b, data_type::u8, data_type::s8)) &&
-                                     !is_vcvtph2ps_kernel;
+                                     !(one_of(brg->dt_a, data_type::f32) && one_of(brg->dt_b, data_type::bf16, data_type::f16));
     brg->ld_step
             = is_b_in_vnni_format ? data_type_vnni_granularity(brg->dt_b) : 1;
     // const data_type_t ld_step_compute_dt
@@ -924,7 +923,8 @@ void init_brgemm_conf(brgemm_desc_t *brg, cpu_isa_t isa,
                       && 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, data_type::bf16) &&
-                one_of(brg->dt_b, data_type::u8, data_type::s8, data_type::nf4, data_type::s4, data_type::u4, data_type::f16));
+                one_of(brg->dt_b, data_type::u8, data_type::s8, data_type::nf4, data_type::s4, data_type::u4))
+            || (one_of(brg->dt_a, data_type::f32) && one_of(brg->dt_b, data_type::bf16, data_type::f16));
     brg->rd_step = has_no_vnni_compute_instruction
             ? 1
             : data_type_vnni_granularity(brg->dt_b);

src/cpu/x64/brgemm/jit_brgemm_kernel.cpp

@@ -2556,12 +2556,14 @@ void jit_brgemm_kernel_t<Wmm>::gemm_microkernel(int bd_block2, bool is_bdb_tail,
                         } else
                             vcvtph2ps(vmm_load, addr);
                     }
-                } else if (brg.dt_b == data_type::bf16
-                        && brg.isa_impl == avx2_vnni_2) {
+                } else if (brg.dt_b == data_type::bf16 && brg.isa_impl == avx2_vnni_2) {
                     if (rd % 2 == 0)
                         vcvtneebf162ps(vmm_load, addr);
                     else
                         vcvtneobf162ps(vmm_load, addr);
+                } else if (brg.dt_b == data_type::bf16 && brg.is_f32) {
+                    vpmovzxwd(vmm_load, addr);
+                    uni_vpslld(vmm_load, vmm_load, 16);
                 } else if (is_ld_tail) {
                     if (is_superset(brg.isa_impl, avx512_core)) {
                         uni_vmovups(vmm_load, addr);
@@ -2852,12 +2854,14 @@ void jit_brgemm_kernel_t<Wmm>::gemm_microkernel(int bd_block2, bool is_bdb_tail,
                         } else
                             vcvtph2ps(vmm_load, addr);
                     }
-                } else if (brg.dt_b == data_type::bf16
-                        && brg.isa_impl == avx2_vnni_2) {
+                } else if (brg.dt_b == data_type::bf16 && brg.isa_impl == avx2_vnni_2) {
                     if (rd % 2 == 0)
                         vcvtneebf162ps(vmm_load, addr);
                     else
                         vcvtneobf162ps(vmm_load, addr);
+                } else if (brg.dt_b == data_type::bf16 && brg.is_f32) {
+                    vpmovzxwd(vmm_load, addr);
+                    uni_vpslld(vmm_load, vmm_load, 16);
                 } else if (is_ld_tail) {
                     if (is_superset(brg.isa_impl, avx512_core)) {
                         uni_vmovups(vmm_load, addr);

src/cpu/x64/jit_brgemm_inner_product.hpp

@@ -63,8 +63,8 @@ struct brgemm_inner_product_fwd_t : public primitive_t {
             auto dst_dt = invariant_dst_md()->data_type;
             auto wei_dt = invariant_wei_md()->data_type;
             const bool is_int8 = one_of(src_dt, u8, s8);
-            const bool is_wei_decomp = one_of(src_dt, f32, bf16) &&
-                                       one_of(wei_dt, u8, s8, nf4, s4, u4, f16);
+            const bool is_wei_decomp = (one_of(src_dt, f32, bf16) && one_of(wei_dt, u8, s8, nf4, s4, u4)) ||
+                                       (one_of(src_dt, f32) && one_of(wei_dt, f16, bf16));
 
             using skip_mask_t = primitive_attr_t::skip_mask_t;
             auto skip_mask = skip_mask_t::post_ops | skip_mask_t::sum_dt
@@ -120,8 +120,8 @@ struct brgemm_inner_product_fwd_t : public primitive_t {
             const float beta = 1.0;
             const float beta_init = 0.0;
 
-            // f16 weight decompression doesn't need scales/zero-points which handles by normal brgemm kernel
-            bool brgemm_with_wei_decomp = is_wei_decomp && jbgp_.wei_decomp_algo == weights_decomp_kind_t::immediate && wei_dt != f16;
+            // f16/bf16 weights decompression doesn't need scales/zero-points which is handled by normal brgemm kernel
+            bool brgemm_with_wei_decomp = is_wei_decomp && jbgp_.wei_decomp_algo == weights_decomp_kind_t::immediate && !one_of(wei_dt, f16, bf16);
 
             for_(int i_bs = 0; i_bs < 2; i_bs++)
             for_(int i_init = 0; i_init < 2; i_init++)

src/cpu/x64/jit_brgemm_inner_product_utils.cpp

@@ -172,8 +172,8 @@ 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;
-    const bool is_vcvtph2ps_kernel = (jbgp.orig_wei_dt == f16 && jbgp.src_dt == f32);
-    if (is_not_vnni_tag || (jbgp.weights_decompression && (jbgp.orig_wei_dt == u8 || jbgp.orig_wei_dt == s8 || is_vcvtph2ps_kernel) && !jbgp.with_src_dynamic_quant)) {
+    const bool is_half_prc_weights = (one_of(jbgp.orig_wei_dt, f16, bf16) && jbgp.src_dt == f32);
+    if (is_not_vnni_tag || (jbgp.weights_decompression && (one_of(jbgp.orig_wei_dt, u8, s8) || is_half_prc_weights) && !jbgp.with_src_dynamic_quant)) {
         if (is_superset(jbgp.isa, avx512_core))
             return {{64,
                             pick(n_sp_dims, OI16i64o, OwI16i64o, OhwI16i64o,
@@ -1419,8 +1419,8 @@ status_t jit_brgemm_ip_conf_t::init_conf_base(cpu_isa_t isa,
     jbgp.dst_dt = dst_d.data_type();
     jbgp.wei_dt = weights_d.data_type();
 
-    jbgp.weights_decompression = one_of(jbgp.src_dt, f32, bf16) &&
-                                 one_of(jbgp.wei_dt, u8, s8, nf4, s4, u4, f16);
+    jbgp.weights_decompression = (one_of(jbgp.src_dt, f32, bf16) && one_of(jbgp.wei_dt, u8, s8, nf4, s4, u4)) ||
+                                 (one_of(jbgp.src_dt, f32) && one_of(jbgp.wei_dt, f16, bf16));
     jbgp.wei_decomp_algo = weights_decomp_kind_t::immediate;
     jbgp.orig_wei_dt = jbgp.wei_dt;
     jbgp.with_grouped_weights_decompression = false;

src/cpu/x64/jit_brgemm_weights_decompression_kernel.cpp

@@ -132,6 +132,11 @@ void jit_brgemm_weights_decompression_kernel_t<isa>::load_weights(Vmm vmm_load,
             vcvtph2ps(vmm_load, addr);
             break;
         }
+        case data_type::bf16: {
+            vpmovzxwd(vmm_load, addr);
+            uni_vpslld(vmm_load, vmm_load, 16);
+            break;
+        }
         default: assert(!"unsupported data type");
     }
 }