api: add Swish eltwise support (oneapi-src#501)

doc/primitives/eltwise.md

@@ -28,6 +28,7 @@ The following operations are supported:
 | soft_relu    | #mkldnn_eltwise_soft_relu    | \f$ f(x) = \log_{e}(1+e^x) \f$
 | sqrt         | #mkldnn_eltwise_sqrt         | \f$ f(x) = \sqrt{x} \f$
 | square       | #mkldnn_eltwise_square       | \f$ f(x) = x^2 \f$
+| swish        | #mkldnn_eltwise_swish        | \f$ f(x) = x \frac{1}{1+e^{-\alpha x}} \f$
 | tanh         | #mkldnn_eltwise_tanh         | \f$ f(x) = \frac{e^z - e^{-z}}{e^z + e^{-z}} \f$
 
 #### Difference Between [Forward Training](#mkldnn_forward_training) and [Forward Inference](#mkldnn_forward_inference)
@@ -103,9 +104,11 @@ The eltwise primitive doesn't support any post-ops or attributes.
 @anchor dg_eltwise_impl_limits
 ## Implementation Limitations
 
-1. No primitive specific limitations. Refer to @ref dev_guide_data_types for
+1. Refer to @ref dev_guide_data_types for
    limitations related to data types support.
 
+2. **GPU**
+    - No support for swish (#mkldnn_eltwise_swish) operation
 
 ## Performance Tips
 

include/mkldnn.hpp

@@ -303,6 +303,8 @@ enum class algorithm {
     eltwise_abs = mkldnn_eltwise_abs,
     /// Eltwise: square root
     eltwise_sqrt = mkldnn_eltwise_sqrt,
+    /// Eltwise: x*sigmoid(a*x)
+    eltwise_swish = mkldnn_eltwise_swish,
     /// Eltwise: linear
     eltwise_linear = mkldnn_eltwise_linear,
     /// Eltwise: bounded_relu

include/mkldnn_types.h

@@ -681,6 +681,8 @@ typedef enum {
     /// @note Tanh approximation formula is used to approximate
     /// cumulative distribution function of a Gaussian
     mkldnn_eltwise_gelu = 0xcf,
+    /// Eltwise: swish
+    mkldnn_eltwise_swish = 0xdf,
     /// Max pooling
     mkldnn_pooling_max = 0x1ff,
     /// Average pooling include padding
@@ -1006,7 +1008,8 @@ typedef struct {
     /// #mkldnn_eltwise_tanh, #mkldnn_eltwise_elu, #mkldnn_eltwise_square,
     /// #mkldnn_eltwise_abs, #mkldnn_eltwise_sqrt, #mkldnn_eltwise_linear,
     /// #mkldnn_eltwise_bounded_relu, #mkldnn_eltwise_soft_relu,
-    /// #mkldnn_eltwise_logistic and #mkldnn_eltwise_exp.
+    /// #mkldnn_eltwise_swish, #mkldnn_eltwise_logistic and
+    /// #mkldnn_eltwise_exp.
     mkldnn_alg_kind_t alg_kind;
     /// Source and destination memory descriptor.
     mkldnn_memory_desc_t data_desc;
@@ -1021,6 +1024,7 @@ typedef struct {
     ///  - #mkldnn_eltwise_abs: @p alpha and @p beta ignored
     ///  - #mkldnn_eltwise_sqrt: @p alpha and @p beta ignored
     ///  - #mkldnn_eltwise_linear: @p alpha -- scale, @p beta -- shift
+    ///  - #mkldnn_eltwise_swish: @p alpha -- sigmoid arg scaling, @p beta ignored
     ///  - #mkldnn_eltwise_bounded_relu: @p alpha -- upper bound, @p beta ignored
     ///  - #mkldnn_eltwise_soft_relu: @p alpha and @p beta ignored
     ///  - #mkldnn_eltwise_logistic: @p alpha and @p beta ignored

src/common/c_types_map.hpp

@@ -68,6 +68,7 @@ namespace alg_kind {
     const alg_kind_t eltwise_square = mkldnn_eltwise_square;
     const alg_kind_t eltwise_abs = mkldnn_eltwise_abs;
     const alg_kind_t eltwise_sqrt = mkldnn_eltwise_sqrt;
+    const alg_kind_t eltwise_swish = mkldnn_eltwise_swish;
     const alg_kind_t eltwise_linear = mkldnn_eltwise_linear;
     const alg_kind_t eltwise_bounded_relu = mkldnn_eltwise_bounded_relu;
     const alg_kind_t eltwise_soft_relu = mkldnn_eltwise_soft_relu;

src/common/eltwise.cpp

@@ -39,7 +39,7 @@ status_t eltwise_desc_init(eltwise_desc_t *eltwise_desc, prop_kind_t prop_kind,
         && one_of(alg_kind, eltwise_relu, eltwise_tanh, eltwise_elu,
                   eltwise_square, eltwise_abs, eltwise_sqrt, eltwise_linear,
                   eltwise_bounded_relu, eltwise_soft_relu, eltwise_logistic,
-                  eltwise_exp, eltwise_gelu)
+                  eltwise_exp, eltwise_gelu, eltwise_swish)
         && IMPLICATION(prop_kind == backward_data, diff_data_desc != nullptr)
         && IMPLICATION(one_of(data_desc->data_type, mkldnn_s32, mkldnn_s8,
                     mkldnn_u8), alg_kind == eltwise_relu && alpha == 0);

src/common/math_utils.hpp

@@ -163,6 +163,19 @@ template <typename T, typename A,
     return (U)(dd * (s > 0 ? 1 : alpha * ::expf((float)s)));
 }
 
+template <typename T, typename A,
+        typename U = typename utils::remove_reference<T>::type>
+inline U swish_fwd(T s, A alpha) {
+    return (U)(s / (1 + ::expf(-alpha * (float)s)));
+}
+
+template <typename T, typename A,
+        typename U = typename utils::remove_reference<T>::type>
+inline U swish_bwd(T dd, T s, A alpha) {
+    U v = (U)(1 / (1 + ::expf((float)-s * alpha)));
+    return dd * (v + s * alpha * v * (1 - v));
+}
+
 template <typename T, typename U = typename utils::remove_reference<T>::type>
 inline U square_fwd(T s) {
     return s * s;

src/common/mkldnn_debug_autogenerated.cpp

@@ -394,6 +394,7 @@ const char *mkldnn_alg_kind2str(mkldnn_alg_kind_t v) {
     if (v == mkldnn_eltwise_sqrt) return "eltwise_sqrt";
     if (v == mkldnn_eltwise_linear) return "eltwise_linear";
     if (v == mkldnn_eltwise_bounded_relu) return "eltwise_bounded_relu";
+    if (v == mkldnn_eltwise_swish) return "eltwise_swish";
     if (v == mkldnn_eltwise_soft_relu) return "eltwise_soft_relu";
     if (v == mkldnn_eltwise_logistic) return "eltwise_logistic";
     if (v == mkldnn_eltwise_exp) return "eltwise_exp";

src/common/primitive_attr.cpp

@@ -70,7 +70,7 @@ status_t post_ops_t::append_eltwise(float scale, alg_kind_t alg, float alpha,
     bool known_alg = one_of(alg, eltwise_relu, eltwise_tanh, eltwise_elu,
             eltwise_square, eltwise_abs, eltwise_sqrt, eltwise_linear,
             eltwise_bounded_relu, eltwise_soft_relu, eltwise_logistic,
-            eltwise_exp, eltwise_gelu);
+            eltwise_exp, eltwise_gelu, eltwise_swish);
     if (!known_alg)
         return invalid_arguments;
 

src/cpu/jit_uni_eltwise.cpp

@@ -583,6 +583,23 @@ void jit_uni_eltwise_injector_f32<isa>::logistic_compute_vector(
     h->uni_vmovups(vmm_src, vmm_aux2);
 }
 
+template <cpu_isa_t isa>
+void jit_uni_eltwise_injector_f32<isa>::swish_compute_vector(
+        const Vmm &vmm_src) {
+    const int alpha_off = 25;
+    // Save src data on stack for later usage
+    h->sub(h->rsp, vlen);
+    h->uni_vmovups(h->ptr[h->rsp], vmm_src);
+    // x*alpha
+    h->uni_vmulps(vmm_src, vmm_src, table_val(alpha_off));
+    // sigmoid(x*alpha)
+    logistic_compute_vector(vmm_src);
+    // x*sigmoid(alpha*x)
+    h->uni_vmovups(vmm_aux0, h->ptr[h->rsp]);
+    h->add(h->rsp, vlen);
+    h->uni_vmulps(vmm_src, vmm_aux0);
+}
+
 template <cpu_isa_t isa>
 void jit_uni_eltwise_injector_f32<isa>::relu_prepare_table() {
     for (size_t d = 0; d < vlen / sizeof(float); ++d) h->dd(float2int(alpha_));
@@ -703,6 +720,7 @@ int jit_uni_eltwise_injector_f32<isa>::aux_vecs_count(alg_kind_t alg_) {
     case alg_kind::eltwise_square: return 0;
     case alg_kind::eltwise_abs: return 0;
     case alg_kind::eltwise_sqrt: return 2;
+    case alg_kind::eltwise_swish: return 4;
     case alg_kind::eltwise_linear: return 1;
     case alg_kind::eltwise_bounded_relu: return 0;
     case alg_kind::eltwise_soft_relu: return 4;
@@ -730,6 +748,7 @@ void jit_uni_eltwise_injector_f32<isa>::compute_body(size_t start_idx,
         case eltwise_square: square_compute_vector(Vmm(idx)); break;
         case eltwise_abs: abs_compute_vector(Vmm(idx)); break;
         case eltwise_sqrt: sqrt_compute_vector(Vmm(idx)); break;
+        case eltwise_swish: swish_compute_vector(Vmm(idx)); break;
         case eltwise_linear: linear_compute_vector(Vmm(idx)); break;
         case eltwise_bounded_relu: bounded_relu_compute_vector(Vmm(idx)); break;
         case eltwise_soft_relu: soft_relu_compute_vector(Vmm(idx)); break;
@@ -772,6 +791,7 @@ void jit_uni_eltwise_injector_f32<isa>::prepare_table(bool gen_table) {
         case eltwise_soft_relu: soft_relu_prepare_table(); break;
         case eltwise_abs: abs_prepare_table(); break;
         case eltwise_sqrt: sqrt_prepare_table(); break;
+        case eltwise_swish: elu_prepare_table(); break;
         case eltwise_linear: linear_prepare_table(); break;
         case eltwise_bounded_relu: bounded_relu_prepare_table(); break;
         case eltwise_square: break;
@@ -1136,7 +1156,7 @@ struct jit_uni_kernel_fwd : public jit_uni_eltwise_kernel,
         assert(utils::one_of(desc.alg_kind, eltwise_tanh, eltwise_elu,
                     eltwise_square, eltwise_abs, eltwise_sqrt, eltwise_linear,
                     eltwise_bounded_relu, eltwise_soft_relu, eltwise_logistic,
-                    eltwise_exp, eltwise_gelu));
+                    eltwise_exp, eltwise_gelu, eltwise_swish));
 
         preamble();
 
@@ -1271,7 +1291,7 @@ status_t jit_uni_eltwise_fwd_t<isa, d_type>::pd_t::init() {
         && utils::one_of(desc()->alg_kind, eltwise_relu, eltwise_tanh,
                 eltwise_elu, eltwise_square, eltwise_abs, eltwise_sqrt,
                 eltwise_linear, eltwise_bounded_relu, eltwise_soft_relu,
-                eltwise_logistic, eltwise_exp, eltwise_gelu)
+                eltwise_logistic, eltwise_exp, eltwise_gelu, eltwise_swish)
         && memory_desc_wrapper(src_md()).is_dense(true)
         && IMPLICATION(!memory_desc_wrapper(src_md()).is_dense(false),
                 math::eltwise_fwd_preserves_zero(desc()->alg_kind, true))

src/cpu/jit_uni_eltwise.hpp

@@ -48,7 +48,7 @@ struct jit_uni_eltwise_injector_f32 {
         assert(utils::one_of(alg_, eltwise_relu, eltwise_tanh, eltwise_elu,
                     eltwise_square, eltwise_abs, eltwise_sqrt, eltwise_linear,
                     eltwise_bounded_relu, eltwise_soft_relu, eltwise_logistic,
-                    eltwise_exp, eltwise_gelu));
+                    eltwise_exp, eltwise_gelu, eltwise_swish));
     }
 
     // note that eltwise.scale is ignored
@@ -120,6 +120,7 @@ struct jit_uni_eltwise_injector_f32 {
     void soft_relu_compute_vector(const Vmm &vmm_src);
     void logistic_compute_vector(const Vmm &vmm_src);
     void gelu_compute_vector(const Vmm &vmm_src);
+    void swish_compute_vector(const Vmm &vmm_src);
 
     void relu_prepare_table();
     void elu_prepare_table();

src/cpu/ref_eltwise.cpp

@@ -35,7 +35,7 @@ ref_eltwise_scalar_fwd_t::ref_eltwise_scalar_fwd_t(alg_kind_t alg, float alpha,
     assert(utils::one_of(alg_, eltwise_relu, eltwise_tanh, eltwise_elu,
                 eltwise_square, eltwise_abs, eltwise_sqrt, eltwise_linear,
                 eltwise_bounded_relu, eltwise_soft_relu, eltwise_logistic,
-                eltwise_exp, eltwise_gelu));
+                eltwise_exp, eltwise_gelu, eltwise_swish));
 }
 
 ref_eltwise_scalar_fwd_t::ref_eltwise_scalar_fwd_t(
@@ -56,6 +56,7 @@ float ref_eltwise_scalar_fwd_t::compute_scalar(float s) {
         case eltwise_logistic: return logistic_fwd(s);
         case eltwise_exp: return exp_fwd(s);
         case eltwise_gelu: return gelu_fwd(s);
+        case eltwise_swish: return swish_fwd(s, alpha_);
         default: assert(!"unknown eltwise alg_kind");
     }
 
@@ -149,6 +150,7 @@ void ref_eltwise_fwd_t<data_type>::execute_forward_generic(
             case eltwise_logistic: d = logistic_fwd(s); break;
             case eltwise_exp: d = exp_fwd(s); break;
             case eltwise_gelu: d = gelu_fwd(s); break;
+            case eltwise_swish: d = swish_fwd(s, alpha); break;
             default: assert(!"unknown eltwise alg_kind");
         }
     });
@@ -194,6 +196,7 @@ void ref_eltwise_fwd_t<data_type>::execute_forward_dense(
         case eltwise_logistic: d = logistic_fwd(s); break;
         case eltwise_exp: d = exp_fwd(s); break;
         case eltwise_gelu: d = gelu_fwd(s); break;
+        case eltwise_swish: d = swish_fwd(s, alpha); break;
         default: assert(!"unknown eltwise alg_kind");
         }
     });
@@ -247,6 +250,7 @@ void ref_eltwise_bwd_t<data_type>::execute_backward_generic(
             case eltwise_logistic: ds = logistic_bwd(dd, s); break;
             case eltwise_exp: ds = exp_bwd(dd, s); break;
             case eltwise_gelu: ds = gelu_bwd(dd, s); break;
+            case eltwise_swish: ds = swish_bwd(dd, s, alpha); break;
             default: assert(!"unknown eltwise alg_kind");
         }
     });
@@ -289,6 +293,7 @@ void ref_eltwise_bwd_t<data_type>::execute_backward_dense(
         case eltwise_logistic: ds = logistic_bwd(dd, s); break;
         case eltwise_exp: ds = exp_bwd(dd, s); break;
         case eltwise_gelu: ds = gelu_bwd(dd, s); break;
+        case eltwise_swish: ds = swish_bwd(dd, s, alpha); break;
         default: assert(!"unknown eltwise alg_kind");
         }
     });

tests/benchdnn/dnn_types.cpp

@@ -222,6 +222,7 @@ attr_t::post_ops_t::kind_t attr_t::post_ops_t::str2kind(const char *str) {
     CASE(LOGISTIC);
     CASE(EXP);
     CASE(GELU);
+    CASE(SWISH);
 #undef CASE
     assert(!"unknown attr::post_ops::kind");
     return KIND_TOTAL;
@@ -242,6 +243,7 @@ const char *attr_t::post_ops_t::kind2str(attr_t::post_ops_t::kind_t kind) {
     CASE(LOGISTIC, "logistic");
     CASE(EXP, "exp");
     CASE(GELU, "gelu");
+    CASE(SWISH, "swish");
 #undef CASE
     assert(!"unknown attr::post_ops::kind");
     return "unknown attr::post_ops::kind";
@@ -262,6 +264,7 @@ mkldnn_alg_kind_t attr_t::post_ops_t::kind2mkldnn_kind(
     CASE(LOGISTIC, mkldnn_eltwise_logistic);
     CASE(EXP, mkldnn_eltwise_exp);
     CASE(GELU, mkldnn_eltwise_gelu);
+    CASE(SWISH, mkldnn_eltwise_swish);
 #undef CASE
     assert(!"unknown attr::post_ops::kind");
     return mkldnn_alg_kind_undef;
@@ -399,6 +402,7 @@ std::ostream &operator<<(std::ostream &s, const attr_t::post_ops_t &post_ops) {
         case pk::LOGISTIC:
         case pk::EXP:
         case pk::GELU:
+        case pk::SWISH:
             s << kind2str(e.kind) << ":" << e.eltwise.alpha;
             if (e.eltwise.beta != 0.f || e.eltwise.scale != 1.f)
                 s << ":" << e.eltwise.beta << ":" << e.eltwise.scale;
@@ -502,6 +506,7 @@ mkldnn_primitive_attr_t create_mkldnn_attr(const attr_t &attr,
             case attr_t::post_ops_t::LOGISTIC:
             case attr_t::post_ops_t::EXP:
             case attr_t::post_ops_t::GELU:
+            case attr_t::post_ops_t::SWISH:
                 DNN_SAFE_V(mkldnn_post_ops_append_eltwise(ops, e.eltwise.scale,
                             e.eltwise.alg, e.eltwise.alpha, e.eltwise.beta));
                 break;
@@ -564,6 +569,7 @@ float compute_eltwise_fwd(attr_t::post_ops_t::kind_t kind, float src,
     case pk::LOGISTIC: return scale * logistic_fwd(src);
     case pk::EXP: return scale * exp_fwd(src);
     case pk::GELU: return scale * gelu_fwd(src);
+    case pk::SWISH: return scale * swish_fwd(src, alpha);
     default: assert(!"unknown attr::post_ops::kind");
     };
     return NAN;
@@ -587,6 +593,7 @@ float compute_eltwise_bwd(attr_t::post_ops_t::kind_t kind, float d_dst,
     case pk::LOGISTIC: return logistic_bwd(d_dst, src);
     case pk::EXP: return exp_bwd(d_dst, src);
     case pk::GELU: return gelu_bwd(d_dst, src);
+    case pk::SWISH: return swish_bwd(d_dst, src, alpha);
     default: assert(!"unknown attr::post_ops::kind");
     }
     return NAN;

tests/benchdnn/dnn_types.hpp

@@ -100,6 +100,7 @@ struct attr_t {
             LOGISTIC,
             EXP,
             GELU,
+            SWISH,
             KIND_TOTAL
         };
         static kind_t str2kind(const char *str);

tests/benchdnn/eltwise/bench_eltwise.cpp

@@ -86,6 +86,7 @@ void check_correctness() {
         case pk::BRELU:
         case pk::ELU:
         case pk::RELU:
+        case pk::SWISH:
             // Test several alpha values but single beta
             if (i_beta != 0) continue;
 

tests/benchdnn/inputs/eltwise/test_eltwise_all

@@ -6,7 +6,7 @@
 --dir=FWD_D,BWD_D
 --dt=f32
 --tag=nchw,nhwc,nChw8c,nChw16c
---alg=relu,tanh,elu,square,abs,sqrt,linear,brelu,srelu,logistic,exp,gelu
+--alg=relu,tanh,elu,square,abs,sqrt,linear,brelu,srelu,logistic,exp,gelu,swish
 4x8x3x3 3x7x4x5 2x16x6x2 3x19x1x2
 
 --dir=FWD_I
@@ -18,7 +18,7 @@
 --dir=FWD_D,BWD_D
 --dt=f32
 --tag=ncdhw,ndhwc,nCdhw8c,nCdhw16c
---alg=relu,tanh,elu,square,abs,sqrt,linear,brelu,srelu,logistic,exp,gelu
+--alg=relu,tanh,elu,square,abs,sqrt,linear,brelu,srelu,logistic,exp,gelu,swish
 2x16x6x2x8 3x15x5x2x3
 
 --dir=FWD_I

tests/benchdnn/inputs/eltwise/test_eltwise_bfloat16

@@ -6,12 +6,12 @@
 --dir=FWD_D,BWD_D
 --dt=bf16
 --tag=nchw,nhwc,nChw16c
---alg=relu,tanh,elu,square,abs,sqrt,linear,brelu,logistic,exp,gelu,srelu
+--alg=relu,tanh,elu,square,abs,sqrt,linear,brelu,logistic,exp,gelu,srelu,swish
 4x8x3x3 3x7x4x5 2x16x6x2 3x19x1x2
 
 --dir=FWD_D,BWD_D
 --dt=bf16
 --tag=ncdhw,ndhwc,nCdhw16c
---alg=relu,tanh,elu,square,abs,sqrt,linear,brelu,logistic,exp,gelu,srelu
+--alg=relu,tanh,elu,square,abs,sqrt,linear,brelu,logistic,exp,gelu,srelu,swish
 2x16x6x2x8 3x15x5x2x3
 

tests/gtests/test_convolution_eltwise_forward_common.hpp

@@ -91,6 +91,7 @@ void compute_ref_conv_eltwise_fwd(const test_convolution_sizes_t &c,
             case algorithm::eltwise_soft_relu: d = soft_relu_fwd(d); break;
             case algorithm::eltwise_logistic: d = logistic_fwd(d); break;
             case algorithm::eltwise_exp: d = exp_fwd(d); break;
+            case algorithm::eltwise_swish: d = swish_fwd(d, elt_alpha); break;
             default: assert(!"unknown alg_kind");
             }
         }

tests/gtests/test_convolution_eltwise_forward_f32.cpp

@@ -61,10 +61,10 @@ TEST_P(convolution_test, TestConvolutionEltwise)
     EXPAND_ARGS(PARAMS_CONV(algorithm::eltwise_linear, __VA_ARGS__)), \
     EXPAND_ARGS(PARAMS_CONV(algorithm::eltwise_bounded_relu, __VA_ARGS__)), \
     EXPAND_ARGS(PARAMS_CONV(algorithm::eltwise_soft_relu, __VA_ARGS__)), \
-    EXPAND_ARGS(PARAMS_CONV(algorithm::eltwise_logistic, __VA_ARGS__))
+    EXPAND_ARGS(PARAMS_CONV(algorithm::eltwise_logistic, __VA_ARGS__)), \
+    EXPAND_ARGS(PARAMS_CONV(algorithm::eltwise_swish, __VA_ARGS__))
 //  Not testing due to not scaled output
 //  EXPAND_ARGS(PARAMS_CONV(algorithm::eltwise_exp, __VA_ARGS__))
-
 #define ELTWISE_ALPHA 0.5f
 #define ELTWISE_BETA 1.5f
 

tests/gtests/test_convolution_eltwise_forward_x8s8f32s32.cpp

@@ -60,7 +60,8 @@ using convolution_test_s8s8s32f32 =
     EXPAND_ARGS(PARAMS_CONV(algorithm::eltwise_linear, __VA_ARGS__)), \
     EXPAND_ARGS(PARAMS_CONV(algorithm::eltwise_bounded_relu, __VA_ARGS__)), \
     EXPAND_ARGS(PARAMS_CONV(algorithm::eltwise_soft_relu, __VA_ARGS__)), \
-    EXPAND_ARGS(PARAMS_CONV(algorithm::eltwise_logistic, __VA_ARGS__))
+    EXPAND_ARGS(PARAMS_CONV(algorithm::eltwise_logistic, __VA_ARGS__)), \
+    EXPAND_ARGS(PARAMS_CONV(algorithm::eltwise_swish, __VA_ARGS__))
 //  Not testing due to not scaled output
 //  EXPAND_ARGS(PARAMS_CONV(algorithm::eltwise_exp, __VA_ARGS__))