gemm int8 support binary postops

Author: luo-cheng2021

src/cpu/gemm_x8s8s32x_convolution.hpp

@@ -97,7 +97,7 @@ struct gemm_x8s8s32x_convolution_fwd_t : public primitive_t {
                 bool ok = true;
 
                 for (int i = 0; i < po.len(); i++) {
-                    ok = ok && utils::one_of(po.entry_[i].kind, sum, eltwise, depthwise, quantization);
+                    ok = ok && utils::one_of(po.entry_[i].kind, sum, binary, eltwise, depthwise, quantization);
                 }
                 return ok;
             };

src/cpu/x64/gemm_bf16_convolution.cpp

@@ -111,7 +111,7 @@ gemm_bf16_convolution_fwd_t<dst_data_type>::pp_ker_t::pp_ker_t(const pd_t *pd)
         static constexpr size_t tail_size = 0;
         static constexpr bool use_exact_tail_scalar_bcast = false;
         const binary_injector::rhs_arg_static_params_t rhs_sp {
-            helper_vmm_idx, r13, r14, preserve_gpr,
+            helper_vmm_idx, reserved_eltwise_gpr, r13, r14, preserve_gpr,
             preserve_vmm, PARAM_OFF(post_ops_binary_rhs_arg_vec),
             PARAM_OFF(dst_orig), memory_desc_wrapper(pd->dst_md()),
             tail_size, kreg_rem_mask, use_exact_tail_scalar_bcast};

src/cpu/x64/jit_avx512_core_fork_bf16_dw_conv_kernel.cpp

@@ -591,7 +591,7 @@ void jit_avx512_fork_dw_conv_fwd_kernel_bf16::generate() {
                 % (cpu_isa_traits<avx512_core>::vlen / sizeof(float));
         static constexpr bool use_exact_tail_scalar_bcast = false;
         const binary_injector::rhs_arg_static_params_t rhs_sp {
-            helper_vmm_idx, r10, r11, preserve_gpr,
+            helper_vmm_idx, reserved_eltwise_gpr, r10, r11, preserve_gpr,
             preserve_vmm, GET_OFF(post_ops_binary_rhs_arg_vec),
             GET_OFF(dst_orig), memory_desc_wrapper(&dst_md_),
             tail_size, k_oc_tail_mask, use_exact_tail_scalar_bcast};

src/cpu/x64/jit_avx512_core_fork_bf16_dw_conv_kernel.hpp

@@ -95,6 +95,8 @@ struct jit_avx512_fork_dw_conv_fwd_kernel_bf16 : public jit_generator {
     mask_t ktail_mask = k_oc_tail_mask;
     mask_t k_ch_tail_mask_extended = Xbyak::Opmask(3);
 
+    Xbyak::Reg64 reserved_eltwise_gpr = r10;
+
     Xbyak::Zmm zmm_ker_reg = Xbyak::Zmm(0);
     Xbyak::Zmm zmm_src_reg = Xbyak::Zmm(1);
     Xbyak::Zmm zmm_prev_dst = Xbyak::Zmm(31);

src/cpu/x64/jit_gemm_convolution_utils.cpp

@@ -66,15 +66,15 @@ struct jit_pp_kernel_t : pp_kernel_t, public jit_generator {
             static constexpr size_t tail_size = 0;
             static constexpr bool use_exact_tail_scalar_bcast = false;
             const binary_injector::rhs_arg_static_params_t rhs_sp {
-                helper_vmm_idx, r13, r14, preserve_gpr,
+                helper_vmm_idx, reserved_eltwise_gpr, r13, r14, preserve_gpr,
                 preserve_vmm, PARAM_OFF(post_ops_binary_rhs_arg_vec),
                 PARAM_OFF(dst_orig), memory_desc_wrapper(pd->dst_md()),
                 tail_size, kreg_rem_mask, use_exact_tail_scalar_bcast};
 #undef PARAM_OFF
             const binary_injector::static_params_t bsp {this->reg_abi_bak, rhs_sp};
             jit_binary_injector_ = utils::make_unique<
                     binary_injector::jit_uni_binary_injector_t<isa>>(
-                    this, bsp);            
+                    this, bsp);
         }
         if (post_ops_.len() > 0 && !only_eltwise) {
             vreg_d_weights = Vmm(idx_compute_vreg_max_--);
@@ -155,6 +155,7 @@ struct jit_pp_kernel_t : pp_kernel_t, public jit_generator {
     Xbyak::Reg64 reg_d_bias = r15;
     Xbyak::Reg64 reg_post_ops_data = rax;
     Vmm vreg_d_weights, vreg_d_bias;
+    Xbyak::Reg64 reserved_eltwise_gpr = r10;
 
     int idx_compute_vreg_start_;
     int idx_compute_vreg_max_;

src/cpu/x64/jit_gemm_x8s8s32x_convolution_utils.cpp

@@ -18,6 +18,7 @@
 #include <functional>
 
 #include "cpu/x64/injectors/jit_uni_postops_injector.hpp"
+#include "cpu/x64/injectors/jit_uni_binary_injector.hpp"
 #include "cpu/x64/jit_gemm_x8s8s32x_conv_zp_src_pad_comp.hpp"
 #include "cpu/x64/jit_gemm_x8s8s32x_convolution_utils.hpp"
 #include "cpu/x64/jit_generator.hpp"
@@ -60,10 +61,14 @@ struct jit_pp_ker_t : pp_ker_t, public jit_generator {
             vreg_zero = Vmm(idx_compute_vreg_start_++);
         }
         bool only_eltwise_or_sum = true;
+        bool with_binary = false;
         for (int idx = 0; idx < post_ops_.len(); ++idx) {
             const auto &e = post_ops_.entry_[idx];
             if (e.is_eltwise(true)) {
                 do_eltwise_ = true;
+            } else if (e.is_binary()) {
+                with_binary = true;
+                only_eltwise_or_sum = false;
             } else if (e.is_sum()) {
                 do_sum_ = true;
                 sum_scale_ = e.sum.scale;
@@ -72,6 +77,24 @@ struct jit_pp_ker_t : pp_ker_t, public jit_generator {
                 only_eltwise_or_sum = false;
             }
         }
+        if (with_binary) {
+#define PARAM_OFF(field) offsetof(ker_args_t, field)
+            static constexpr bool preserve_gpr = true;
+            static constexpr bool preserve_vmm = true;
+            static constexpr size_t helper_vmm_idx = 1;
+            static constexpr size_t tail_size = 0;
+            static constexpr bool use_exact_tail_scalar_bcast = false;
+            const binary_injector::rhs_arg_static_params_t rhs_sp {
+                helper_vmm_idx, r13, r14, r15, preserve_gpr,
+                preserve_vmm, PARAM_OFF(post_ops_binary_rhs_arg_vec),
+                PARAM_OFF(dst_orig), memory_desc_wrapper(pd->dst_md()),
+                tail_size, kreg_rem_mask_short, use_exact_tail_scalar_bcast};
+#undef PARAM_OFF
+            const binary_injector::static_params_t bsp {this->reg_param_bak, rhs_sp};
+            jit_binary_injector_ = utils::make_unique<
+                    binary_injector::jit_uni_binary_injector_t<isa>>(
+                    this, bsp);
+        }
         if (post_ops_.len() > 0 && !only_eltwise_or_sum) {
             vreg_d_weights = Vmm(idx_compute_vreg_max_--);
             vreg_d_bias = Vmm(idx_compute_vreg_max_--);
@@ -133,6 +156,7 @@ struct jit_pp_ker_t : pp_ker_t, public jit_generator {
         args.dst = dst
                    + (os_offset * dst_os_stride_ + oc_offset)
                      * dst_data_type_size_;
+        args.dst_orig = dst_orig;
         args.bias = bias + (g * jcp_.oc + oc_offset) * bias_data_type_size_;
         args.scales = scales + scale_idx_mult_ * (g * jcp_.oc + oc_offset);
         args.sum_scale = sum_scale_;
@@ -149,6 +173,7 @@ struct jit_pp_ker_t : pp_ker_t, public jit_generator {
 
     struct ker_args_t {
         char *dst;
+        const void* dst_orig;
         const acc_data_t *acc;
         const char *bias;
         const float *scales;
@@ -163,11 +188,14 @@ struct jit_pp_ker_t : pp_ker_t, public jit_generator {
 
     nstl::vector<jit_uni_eltwise_injector_f32<isa> *> jit_eltwise_injectors_;
     nstl::vector<jit_uni_depthwise_injector_f32<isa> *> jit_depthwise_injectors_;
+    std::unique_ptr<binary_injector::jit_uni_binary_injector_t<isa>>
+            jit_binary_injector_;
 
     using Vmm = typename cpu_isa_traits<isa>::Vmm;
     static const size_t vlen = cpu_isa_traits<isa>::vlen / sizeof(float);
 
     Xbyak::Reg64 reg_param = abi_param1;
+    Xbyak::Reg64 reg_param_bak = r11;
     Xbyak::Reg64 reg_dst = rdx;
     Xbyak::Reg64 reg_acc = rax;
     Xbyak::Reg64 reg_bias = rbx;
@@ -264,6 +292,7 @@ void jit_pp_ker_t<isa>::generate() {
         }
     }
 
+    mov(reg_param_bak, reg_param);
     mov(reg_dst, ptr[reg_param + PARAM_OFF(dst)]);
     mov(reg_acc, ptr[reg_param + PARAM_OFF(acc)]);
     mov(reg_bias, ptr[reg_param + PARAM_OFF(bias)]);
@@ -285,10 +314,11 @@ void jit_pp_ker_t<isa>::generate() {
     if (utils::one_of(isa, avx2, sse41))
         mov(reg_table, l_table);
 
-    auto apply_post_ops = [&](size_t offset, int idx) {
+    auto apply_post_ops = [&](size_t offset, int idx, bool apply_mask) {
         std::size_t post_ops_data_offset = 0;
         int eltwise_inj_idx = 0;
         int depthwise_inj_idx = 0;
+        int binary_inj_idx = 0;
         for (int i = 0; i < post_ops_.len(); i++) {
             auto &post_op = post_ops_.entry_[i];
             if (post_op.is_sum()) {
@@ -312,6 +342,18 @@ void jit_pp_ker_t<isa>::generate() {
                     uni_vcvtdq2ps(vreg_prev_dst(idx), vreg_prev_dst(idx));
 
                 uni_vfmadd231ps(vreg_dst(idx), vreg_prev_dst(idx), vreg_sum_scale);
+            } else if (post_op.is_binary()) {
+                binary_injector::rhs_arg_dynamic_params_t rhs_arg_params;
+                auto dst_addr = ptr[reg_dst + offset * dst_data_type_size_];
+                rhs_arg_params.vmm_idx_to_out_addr.emplace(idx, dst_addr);
+                rhs_arg_params.vmm_idx_to_out_elem_off_val.emplace(
+                        idx, 0 * sizeof(float));
+                if (mayiuse(avx512_core) && apply_mask) 
+                    rhs_arg_params.vmm_tail_idx_.emplace(idx);
+                jit_binary_injector_->compute_vector(
+                        idx, binary_inj_idx, post_op, rhs_arg_params);
+
+                binary_inj_idx++;
             } else if (post_op.is_eltwise()) {
                 jit_eltwise_injectors_[eltwise_inj_idx]->compute_vector_range(vreg_dst(idx).getIdx(),
                                                                               vreg_dst(idx).getIdx() + 1);
@@ -329,6 +371,7 @@ void jit_pp_ker_t<isa>::generate() {
 
                 post_ops_data_offset += jit_depthwise_injectors_[depthwise_inj_idx]->memoryStep();
                 depthwise_inj_idx++;
+                binary_inj_idx++;
             } else if (post_op.is_quantization()) {
                 add(reg_oc_offset, reg_g_offset);
                 bool do_dequantization = post_op.quantization.alg == alg_kind::quantization_quantize_dequantize;
@@ -403,6 +446,7 @@ void jit_pp_ker_t<isa>::generate() {
                 sub(reg_oc_offset, reg_g_offset);
 
                 post_ops_data_offset += sizeof(float*);
+                binary_inj_idx++;
             }
         }
     };
@@ -488,7 +532,7 @@ void jit_pp_ker_t<isa>::generate() {
         if (do_scale_)
             uni_vmulps(vreg_dst(idx), vreg_dst(idx), vreg_scale);
 
-        apply_post_ops(offset, idx);
+        apply_post_ops(offset, idx, apply_mask);
 
         if (dst_data_type_ != data_type::f32) {
             if (isa == avx512_core) {

src/cpu/x64/jit_uni_fork_dw_conv_kernel_f32.cpp

@@ -772,7 +772,7 @@ void jit_uni_fork_dw_conv_fwd_kernel_f32<isa>::generate() {
                 % (cpu_isa_traits<isa>::vlen / sizeof(float));
         static constexpr bool use_exact_tail_scalar_bcast = false;
         const binary_injector::rhs_arg_static_params_t rhs_sp {
-            helper_vmm_idx, r10, r11, preserve_gpr,
+            helper_vmm_idx, reserved_eltwise_gpr, r10, r11, preserve_gpr,
             preserve_vmm, GET_OFF(post_ops_binary_rhs_arg_vec),
             GET_OFF(dst_orig), memory_desc_wrapper(&dst_md_),
             tail_size, k_oc_tail_mask, use_exact_tail_scalar_bcast};

src/cpu/x64/jit_uni_fork_dw_conv_kernel_f32.hpp

@@ -87,6 +87,8 @@ struct jit_uni_fork_dw_conv_fwd_kernel_f32 : public jit_generator {
     reg64_t aux_reg_ch_blocks = reg_ur_w;
     reg64_t aux_reg_blocks_offset = abi_not_param1;
 
+    Xbyak::Reg64 reserved_eltwise_gpr = r10;
+
     reg64_t reg_d_weights = imm_addr64;
     reg64_t reg_d_bias = iter_kh;
     int base_post_ops_data_offset = 0;