[FORK][FEATURE] Added prelu as binary post op

Author: antonvor

include/oneapi/dnnl/dnnl.hpp

@@ -459,6 +459,8 @@ enum class algorithm {
     binary_eq = dnnl_binary_eq,
     /// Binary not equal
     binary_ne = dnnl_binary_ne,
+    /// Binary prelu
+    binary_prelu = dnnl_binary_prelu,
     /// Nearest Neighbor resampling method
     resampling_nearest = dnnl_resampling_nearest,
     /// Linear (Bilinear, Trilinear) resampling method

include/oneapi/dnnl/dnnl_types.h

@@ -2166,6 +2166,8 @@ typedef enum {
     dnnl_binary_eq = 0x1fffa,
     /// Binary not equal
     dnnl_binary_ne = 0x1fffb,
+    /// Binary prelu
+    dnnl_binary_prelu = 0x1fffc,
     /// Nearest Neighbor Resampling Method
     dnnl_resampling_nearest = 0x2fff0,
     /// Linear Resampling Method

src/common/binary.cpp

@@ -91,7 +91,7 @@ status_t dnnl_binary_primitive_desc_create(
     VCHECK_BINARY(
             one_of(alg_kind, binary_add, binary_mul, binary_max, binary_min,
                     binary_div, binary_sub, binary_ge, binary_gt, binary_le,
-                    binary_lt, binary_eq, binary_ne),
+                    binary_lt, binary_eq, binary_ne, binary_prelu),
             VERBOSE_BAD_ALGORITHM);
     // TODO - Add support for mutual or bi-directional broadcasts
     VCHECK_BINARY(!memory_desc_wrapper(src0_md).format_any(),

src/common/c_types_map.hpp

@@ -129,6 +129,7 @@ const alg_kind_t binary_le = dnnl_binary_le;
 const alg_kind_t binary_lt = dnnl_binary_lt;
 const alg_kind_t binary_eq = dnnl_binary_eq;
 const alg_kind_t binary_ne = dnnl_binary_ne;
+const alg_kind_t binary_prelu = dnnl_binary_prelu;
 const alg_kind_t resampling_nearest = dnnl_resampling_nearest;
 const alg_kind_t resampling_linear = dnnl_resampling_linear;
 const alg_kind_t reduction_max = dnnl_reduction_max;

src/common/dnnl_debug_autogenerated.cpp

@@ -1825,6 +1825,7 @@ const char *dnnl_alg_kind2str(dnnl_alg_kind_t v) {
     if (v == dnnl_binary_lt) return "binary_lt";
     if (v == dnnl_binary_eq) return "binary_eq";
     if (v == dnnl_binary_ne) return "binary_ne";
+    if (v == dnnl_binary_prelu) return "binary_prelu";
     if (v == dnnl_resampling_nearest) return "resampling_nearest";
     if (v == dnnl_resampling_linear) return "resampling_linear";
     if (v == dnnl_reduction_max) return "reduction_max";

src/common/primitive_attr.cpp

@@ -270,7 +270,7 @@ status_t post_ops_t::validate_binary(
     using namespace alg_kind;
     bool alg_ok = one_of(alg, binary_add, binary_mul, binary_max, binary_min,
             binary_div, binary_sub, binary_ge, binary_gt, binary_le, binary_lt,
-            binary_eq, binary_ne);
+            binary_eq, binary_ne, binary_prelu);
     if (!alg_ok) return invalid_arguments;
     if (!memory_desc_sanity_check(*user_src1_desc)) return invalid_arguments;
 

src/cpu/gemm_x8s8s32x_convolution.hpp

@@ -77,7 +77,8 @@ struct gemm_x8s8s32x_convolution_fwd_t : public primitive_t {
                                     | skip_mask_t::post_ops
                                     | skip_mask_t::sum_dt
                                     | primitive_attr_t::skip_mask_t::input_zero_points
-                                    | primitive_attr_t::skip_mask_t::output_compensations,
+                                    | primitive_attr_t::skip_mask_t::output_compensations
+                                    | primitive_attr_t::skip_mask_t::sum_dt,
                             dst_type),
                     VERBOSE_UNSUPPORTED_ATTR);
 

src/cpu/primitive_attr_postops.cpp

@@ -40,6 +40,7 @@ float compute_binary_scalar(alg_kind_t alg, float x, float y) {
         case binary_lt: return x < y;
         case binary_eq: return x == y;
         case binary_ne: return x != y;
+        case binary_prelu: return x >= 0 ? x : x * y;
         default: assert(!"not supported operation!"); return NAN;
     }
 }
@@ -139,7 +140,7 @@ ref_binary_scalar_t::ref_binary_scalar_t(alg_kind_t alg) : alg_(alg) {
             alg_kind::binary_min, alg_kind::binary_mul, alg_kind::binary_div,
             alg_kind::binary_sub, alg_kind::binary_ge, alg_kind::binary_gt,
             alg_kind::binary_le, alg_kind::binary_lt, alg_kind::binary_eq,
-            alg_kind::binary_ne));
+            alg_kind::binary_ne, alg_kind::binary_prelu));
 }
 
 ref_binary_scalar_t::ref_binary_scalar_t(
@@ -183,6 +184,8 @@ ref_post_ops_t::ref_post_ops_t(const post_ops_t &po, bool skip_sum)
             eltwise_po_.emplace_back(e.eltwise);
         } else if (po_.contain(primitive_kind::binary, idx)) {
             binary_po_.emplace_back(e.binary);
+        } else if (po_.contain(primitive_kind::depthwise, idx)) {
+            depthwise_po_.emplace_back(e.depthwise.alg);
         }
     }
 }
@@ -264,12 +267,13 @@ status_t ref_post_ops_t::init(const memory_desc_t *dst_md) {
     return status::success;
 }
 
-void ref_post_ops_t::execute(float &res, const args_t &args) const {
+void ref_post_ops_t::execute(float &res, const args_t &args, const size_t oc) const {
     if (po_.len() == 0) return;
 
     auto it_eltwise_po = eltwise_po_.begin();
     auto it_binary_po = binary_po_.begin();
     auto it_prelu_md = prelu_md_.begin();
+    auto it_depthwise_po = depthwise_po_.begin();
     for (auto idx = 0; idx < po_.len(); ++idx) {
         const auto &e = po_.entry_[idx];
         switch (e.kind) {
@@ -330,6 +334,40 @@ void ref_post_ops_t::execute(float &res, const args_t &args) const {
                 res = weights_value * res;
                 ++it_prelu_md;
             } break;
+            case primitive_kind::depthwise: {
+                auto depthwise_weights = e.depthwise.weights_data;
+                auto depthwise_bias = e.depthwise.biases_data;
+                res = it_depthwise_po->compute_scalar(res, depthwise_weights + oc, depthwise_bias + oc);
+                ++it_depthwise_po;
+            } break;
+            case primitive_kind::quantization: {
+                bool do_dequantization = e.quantization.alg == alg_kind::quantization_quantize_dequantize;
+                bool do_rounding = do_dequantization || args.dst_md->data_type == dnnl_f32 || idx != po_.len() - 1;
+
+                auto quant = e.quantization;
+                auto pcl = quant.crop_low_data->shifts_;
+                auto pch = quant.crop_high_data->shifts_;
+                auto pisc = quant.input_scale_data->scales_;
+                auto pish = quant.input_shift_data->shifts_;
+                auto posc = quant.output_scale_data->scales_;
+                auto posh = quant.output_shift_data->shifts_;
+
+                int cl_idx = quant.crop_low_data->count_ == 1 ? 0 : oc;
+                int ch_idx = quant.crop_high_data->count_ == 1 ? 0 : oc;
+                int isc_idx = quant.input_scale_data->count_ == 1 ? 0 : oc;
+                int ish_idx = quant.input_shift_data->count_ == 1 ? 0 : oc;
+                int osc_idx = quant.output_scale_data->count_ == 1 ? 0 : oc;
+                int osh_idx = quant.output_shift_data->count_ == 1 ? 0 : oc;
+
+                res = nstl::min(pch[ch_idx], nstl::max(pcl[cl_idx], res));
+                res = res * pisc[isc_idx] + pish[ish_idx];
+
+                if (do_rounding)
+                    res = roundf(res);
+
+                if (do_dequantization)
+                    res = res * posc[osc_idx] + posh[osh_idx];
+            } break;
             default: assert(!"unsupported post op primitive kind!");
         }
     }

src/cpu/primitive_attr_postops.hpp

@@ -22,6 +22,8 @@
 #include "common/primitive.hpp"
 #include "common/primitive_attr.hpp"
 
+#include "ref_depthwise_injector.hpp"
+
 namespace dnnl {
 namespace impl {
 namespace cpu {
@@ -71,7 +73,7 @@ struct ref_post_ops_t {
 
     status_t init(const memory_desc_t *dst_md);
 
-    void execute(float &res, const args_t &args = args_t()) const;
+    void execute(float &res, const args_t &args = args_t(), const size_t oc = 0) const;
 
     static bool primitive_kind_ok(const post_ops_t &po) {
         using namespace primitive_kind;
@@ -86,6 +88,7 @@ struct ref_post_ops_t {
 
     std::vector<ref_eltwise_scalar_fwd_t> eltwise_po_;
     std::vector<ref_binary_scalar_t> binary_po_;
+    std::vector<ref_depthwise_scalar_fwd_t> depthwise_po_;
     std::vector<memory_desc_t> prelu_md_;
 };
 

src/cpu/ref_depthwise_injector.cpp

@@ -64,7 +64,7 @@ ref_depthwise_scalar_fwd_t::ref_depthwise_scalar_fwd_t(const alg_kind_t alg_)
     assert(utils::one_of(alg, depthwise_scale_shift, depthwise_prelu));
 }
 
-float ref_depthwise_scalar_fwd_t::compute_scalar(float s, const float* weights, const float* bias) {
+float ref_depthwise_scalar_fwd_t::compute_scalar(float s, const float* weights, const float* bias) const {
     switch (alg) {
         case depthwise_scale_shift: return scale_shift_fwd(s, *weights, *bias);
         case depthwise_prelu: return prelu_fwd(s, *weights);

src/cpu/ref_depthwise_injector.hpp

@@ -27,7 +27,7 @@ namespace cpu {
 struct ref_depthwise_scalar_fwd_t {
 public:
     explicit ref_depthwise_scalar_fwd_t(alg_kind_t alg);
-    float compute_scalar(float s, const float* weights, const float* bias);
+    float compute_scalar(float s, const float* weights, const float* bias) const;
 
 private:
     alg_kind_t alg;

src/cpu/x64/injectors/jit_uni_binary_injector.cpp

@@ -199,12 +199,14 @@ rhs_arg_static_params_t::rhs_arg_static_params_t(
         bool preserve_vmm_helper, std::size_t abi_param_offset,
         std::size_t dst_orig_offset, const memory_desc_wrapper &dst_d,
         std::size_t tail_size, const Xbyak::Opmask &tail_opmask,
-        bool use_exact_tail_scalar_bcast)
+        bool use_exact_tail_scalar_bcast, std::size_t rhs_prelu_helper_vmm_idx)
     : rhs_arg_static_params_t(rhs_dt_helper_vmm_idx, rhs_addr_reg,
             rhs_helper_reg, rhs_addr_cache_reg, preserve_gpr_helpers,
             preserve_vmm_helper, abi_param_offset, dst_orig_offset, dst_d,
             tail_size, tail_opmask, use_exact_tail_scalar_bcast, rhs_helper_reg,
-            true /*is_opmask_set*/) {}
+            true /*is_opmask_set*/) {
+    this->rhs_prelu_helper_vmm_idx = rhs_prelu_helper_vmm_idx;
+}
 
 rhs_arg_static_params_t::rhs_arg_static_params_t(
         std::size_t rhs_dt_helper_vmm_idx, const Xbyak::Reg64 &rhs_addr_reg,
@@ -213,12 +215,14 @@ rhs_arg_static_params_t::rhs_arg_static_params_t(
         bool preserve_vmm_helper, std::size_t abi_param_offset,
         std::size_t dst_orig_offset, const memory_desc_wrapper &dst_d,
         std::size_t tail_size, const Xbyak::Opmask &tail_opmask,
-        const Xbyak::Reg64 &reg_tail_size, bool use_exact_tail_scalar_bcast)
+        const Xbyak::Reg64 &reg_tail_size, bool use_exact_tail_scalar_bcast, std::size_t rhs_prelu_helper_vmm_idx)
     : rhs_arg_static_params_t(rhs_dt_helper_vmm_idx, rhs_addr_reg,
             rhs_helper_reg, rhs_addr_cache_reg, preserve_gpr_helpers,
             preserve_vmm_helper, abi_param_offset, dst_orig_offset, dst_d,
             tail_size, tail_opmask, use_exact_tail_scalar_bcast, reg_tail_size,
-            true /*is_opmask_set*/) {}
+            true /*is_opmask_set*/) {
+    this->rhs_prelu_helper_vmm_idx = rhs_prelu_helper_vmm_idx;
+}
 
 rhs_arg_static_params_t::rhs_arg_static_params_t(
         std::size_t rhs_dt_helper_vmm_idx, const Xbyak::Reg64 &rhs_addr_reg,
@@ -2295,7 +2299,7 @@ void jit_uni_binary_injector_t<isa, Vmm>::inject_binary(
             = rhs_arg_data_type != data_type::f32 || (scalar_f32 && !is_avx512_)
             || with_tail_not_fusable_to_binary_op
             || !binary_op_with_unaligned_mem_operand_allowed_
-            || (cmp_op && !is_avx512_);
+            || ((cmp_op || alg == alg_kind::binary_prelu) && !is_avx512_);
 
     if (process_rhs_arg_using_tmp_vmm) {
 
@@ -3192,6 +3196,23 @@ jit_uni_binary_injector_t<isa, Vmm>::execute_cmp_binary(const Vmm &dst,
     pop_opmask(host_, cmp_mask);
 }
 
+template <cpu_isa_t isa, typename Vmm>
+template <typename T>
+typename std::enable_if<std::is_same<T, Xbyak::Zmm>::value
+        || std::is_same<T, Xbyak::Address>::value>::type
+jit_uni_binary_injector_t<isa, Vmm>::execute_prelu_binary(const Vmm &dst, const Vmm &lhs, const T &rhs) const {
+    const auto &cmp_mask = rhs_arg_static_params_.tail_opmask;
+    const Xbyak::Zmm zmm_aux0
+            = Xbyak::Zmm(rhs_arg_static_params_.rhs_prelu_helper_vmm_idx);
+
+    push_opmask(host_, cmp_mask);
+    host_->uni_vpxor(zmm_aux0, zmm_aux0, zmm_aux0);
+    host_->vcmpps(cmp_mask, lhs, zmm_aux0, jit_generator::_cmp_lt_os);
+    host_->uni_vmulps(dst | cmp_mask, lhs, rhs);
+    pop_opmask(host_, cmp_mask);
+}
+
+
 // SSE4.1., AVX and AVX2 implementation
 template <cpu_isa_t isa, typename Vmm>
 template <typename T>
@@ -3211,6 +3232,23 @@ jit_uni_binary_injector_t<isa, Vmm>::execute_cmp_binary(const Vmm &dst,
     host_->uni_vminps(dst, dst, vreg_one);
 }
 
+// todo: [antonvor] check sse41 path
+template <cpu_isa_t isa, typename Vmm>
+template <typename T>
+typename std::enable_if<!(std::is_same<T, Xbyak::Zmm>::value
+                          || std::is_same<T, Xbyak::Address>::value)>::type
+jit_uni_binary_injector_t<isa, Vmm>::execute_prelu_binary(const Vmm &dst,
+                                                        const Vmm &lhs, const T &rhs) const {
+    const Vmm vmm_aux0 = Vmm(rhs_arg_static_params_.rhs_prelu_helper_vmm_idx);
+
+    push_vmm(host_, vmm_aux0);
+    host_->uni_vmulps(rhs, rhs, lhs);
+    host_->vpxor(vmm_aux0, vmm_aux0, vmm_aux0);
+    host_->vcmpltps(vmm_aux0, lhs, vmm_aux0);
+    host_->uni_vblendvps(dst, lhs, rhs, vmm_aux0);
+    pop_vmm(host_, vmm_aux0);
+}
+
 template <cpu_isa_t isa, typename Vmm>
 template <typename T>
 void jit_uni_binary_injector_t<isa, Vmm>::execute_binary(alg_kind_t binary_alg,
@@ -3240,6 +3278,9 @@ void jit_uni_binary_injector_t<isa, Vmm>::execute_binary(alg_kind_t binary_alg,
         case alg_kind::binary_ne:
             execute_cmp_binary(dst, lhs, rhs, jit_generator::_cmp_neq_uq);
             break;
+        case alg_kind::binary_prelu:
+            execute_prelu_binary(dst, lhs, rhs);
+            break;
         default: assert(!"unsupported algorithm");
     }
 }

src/cpu/x64/injectors/jit_uni_binary_injector.hpp

@@ -110,7 +110,7 @@ struct rhs_arg_static_params_t {
             bool preserve_vmm_helper, std::size_t abi_param_offset,
             std::size_t dst_orig_offset, const memory_desc_wrapper &dst_d,
             std::size_t tail_size, const Xbyak::Opmask &tail_opmask,
-            bool use_exact_tail_scalar_bcast);
+            bool use_exact_tail_scalar_bcast, std::size_t rhs_prelu_helper_vmm_idx = 0);
     rhs_arg_static_params_t(std::size_t rhs_dt_helper_vmm_idx,
             const Xbyak::Reg64 &rhs_addr_reg,
             const Xbyak::Reg64 &rhs_helper_reg,
@@ -119,7 +119,7 @@ struct rhs_arg_static_params_t {
             std::size_t dst_orig_offset, const memory_desc_wrapper &dst_d,
             std::size_t tail_size, const Xbyak::Opmask &tail_opmask,
             const Xbyak::Reg64 &reg_tail_size,
-            bool use_exact_tail_scalar_bcast);
+            bool use_exact_tail_scalar_bcast, std::size_t rhs_prelu_helper_vmm_idx = 0);
 
     bool is_opmask_set() const noexcept { return is_opmask_set_; }
 
@@ -138,6 +138,8 @@ struct rhs_arg_static_params_t {
     Xbyak::Reg64 reg_tail_size;
     bool is_tail;
 
+    mutable std::size_t rhs_prelu_helper_vmm_idx;
+
 private:
     rhs_arg_static_params_t(std::size_t rhs_dt_helper_vmm_idx,
             const Xbyak::Reg64 &rhs_addr_reg,
@@ -509,11 +511,19 @@ class jit_uni_binary_injector_t {
     execute_cmp_binary(const Vmm &dst, const Vmm &lhs, const T &rhs,
             const unsigned int cmp_predicate) const;
     template <typename T>
+    typename std::enable_if<std::is_same<T, Xbyak::Zmm>::value
+            || std::is_same<T, Xbyak::Address>::value>::type
+    execute_prelu_binary(const Vmm &dst, const Vmm &lhs, const T &rhs) const;
+    template <typename T>
     typename std::enable_if<!(std::is_same<T, Xbyak::Zmm>::value
             || std::is_same<T, Xbyak::Address>::value)>::type
     execute_cmp_binary(const Vmm &dst, const Vmm &lhs, const T &rhs,
             const unsigned int cmp_predicate) const;
     template <typename T>
+    typename std::enable_if<!(std::is_same<T, Xbyak::Zmm>::value
+            || std::is_same<T, Xbyak::Address>::value)>::type
+    execute_prelu_binary(const Vmm &dst, const Vmm &lhs, const T &rhs) const;
+    template <typename T>
     void execute_binary(alg_kind_t binary_alg, const Vmm &dst, const Vmm &lhs,
             const T &rhs) const;
     void execute_prelu(const Vmm &dst, const Xbyak::Operand &rhs) const;