src: add a specialization for bf16 bwd ref impl

* Per element conversion in parallel_nd slows down the execution due to
* lack of vectorization and multiple function calls to cast operator.

Author: dzarukin

src/common/memory_tracking.hpp

@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2018-2019 Intel Corporation
+* Copyright 2018-2020 Intel Corporation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
@@ -173,6 +173,8 @@ enum {
     key_conv_wei_reduction,
     key_conv_wei_bia_reduction,
     key_conv_wei_bia_reduction_bctx,
+    key_eltwise_diff_dst,
+    key_eltwise_src,
     key_iprod_bias_bf16_convert_wsp,
     key_iprod_dst_bf16_convert_wsp,
     key_iprod_int_dat_in_acc_dt,

src/cpu/cpu_eltwise_list.cpp

@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2019 Intel Corporation
+* Copyright 2019-2020 Intel Corporation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
@@ -51,7 +51,6 @@ static const pd_create_f impl_list[] = {
         INSTANCE(ref_eltwise_fwd_t<s32>),
         INSTANCE(ref_eltwise_fwd_t<s8>),
         INSTANCE(ref_eltwise_fwd_t<u8>),
-        INSTANCE(ref_eltwise_bwd_t<s32>),
         /* eol */
         nullptr,
 };

src/cpu/ref_eltwise.cpp

@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2016-2019 Intel Corporation
+* Copyright 2016-2020 Intel Corporation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
@@ -229,8 +229,8 @@ void ref_eltwise_bwd_t<data_type>::execute_backward_generic(
             });
 }
 
-template <impl::data_type_t data_type>
-void ref_eltwise_bwd_t<data_type>::execute_backward_dense(
+template <>
+void ref_eltwise_bwd_t<data_type::f32>::execute_backward_dense(
         const exec_ctx_t &ctx) const {
     auto src = CTX_IN_MEM(const data_t *, DNNL_ARG_SRC);
     auto diff_dst = CTX_IN_MEM(const data_t *, DNNL_ARG_DIFF_DST);
@@ -239,7 +239,7 @@ void ref_eltwise_bwd_t<data_type>::execute_backward_dense(
     const memory_desc_wrapper data_d(pd()->src_md());
     const memory_desc_wrapper diff_data_d(pd()->diff_src_md());
 
-    const ptrdiff_t nelems = static_cast<ptrdiff_t>(data_d.nelems(true));
+    const auto nelems = data_d.nelems(true);
     const auto alg_kind = pd()->desc()->alg_kind;
     const float alpha = pd()->desc()->alpha;
     const float beta = pd()->desc()->beta;
@@ -248,11 +248,57 @@ void ref_eltwise_bwd_t<data_type>::execute_backward_dense(
     diff_dst += diff_data_d.offset0();
     diff_src += diff_data_d.offset0();
 
-    parallel_nd(nelems, [&](ptrdiff_t e) {
-        const data_t dd = diff_dst[e];
-        const data_t s = src[e];
-        data_t &ds = diff_src[e];
-        ds = compute_eltwise_scalar_bwd(alg_kind, dd, s, alpha, beta);
+    parallel(0, [&](const int ithr, const int nthr) {
+        dim_t start = 0, end = 0;
+        balance211(nelems, nthr, ithr, start, end);
+        if (start == end) return;
+
+        for (dim_t i = start; i < end; i++) {
+            diff_src[i] = compute_eltwise_scalar_bwd(
+                    alg_kind, diff_dst[i], src[i], alpha, beta);
+        }
+    });
+}
+
+template <>
+void ref_eltwise_bwd_t<data_type::bf16>::execute_backward_dense(
+        const exec_ctx_t &ctx) const {
+    using namespace memory_tracking::names;
+
+    auto src = CTX_IN_MEM(const data_t *, DNNL_ARG_SRC);
+    auto diff_dst = CTX_IN_MEM(const data_t *, DNNL_ARG_DIFF_DST);
+    auto diff_src = CTX_OUT_MEM(data_t *, DNNL_ARG_DIFF_SRC);
+
+    auto scratchpad = ctx.get_scratchpad_grantor();
+    auto s_f = scratchpad.template get<float>(key_eltwise_src);
+    auto dd_f = scratchpad.template get<float>(key_eltwise_diff_dst);
+
+    const memory_desc_wrapper data_d(pd()->src_md());
+    const memory_desc_wrapper diff_data_d(pd()->diff_src_md());
+
+    const auto nelems = data_d.nelems(true);
+    const auto alg_kind = pd()->desc()->alg_kind;
+    const float alpha = pd()->desc()->alpha;
+    const float beta = pd()->desc()->beta;
+
+    src += data_d.offset0();
+    diff_dst += diff_data_d.offset0();
+    diff_src += diff_data_d.offset0();
+
+    parallel(0, [&](const int ithr, const int nthr) {
+        dim_t start = 0, end = 0;
+        balance211(nelems, nthr, ithr, start, end);
+        if (start == end) return;
+
+        cvt_bfloat16_to_float(s_f + start, src + start, end - start);
+        cvt_bfloat16_to_float(dd_f + start, diff_dst + start, end - start);
+
+        for (dim_t i = start; i < end; i++) {
+            dd_f[i] = compute_eltwise_scalar_bwd(
+                    alg_kind, dd_f[i], s_f[i], alpha, beta);
+        }
+
+        cvt_float_to_bfloat16(diff_src + start, dd_f + start, end - start);
     });
 }
 
@@ -264,7 +310,6 @@ template struct ref_eltwise_fwd_t<data_type::u8>;
 
 template struct ref_eltwise_bwd_t<data_type::f32>;
 template struct ref_eltwise_bwd_t<data_type::bf16>;
-template struct ref_eltwise_bwd_t<data_type::s32>;
 
 } // namespace cpu
 } // namespace impl

src/cpu/ref_eltwise.hpp

@@ -1,5 +1,5 @@
 /*******************************************************************************
-* Copyright 2016-2019 Intel Corporation
+* Copyright 2016-2020 Intel Corporation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
@@ -136,10 +136,24 @@ struct ref_eltwise_bwd_t : public primitive_impl_t {
             if (use_generic && !one_of(diff_dst_d.ndims(), 4, 5))
                 return status::unimplemented;
 
+            if (data_type == data_type::bf16) init_scratchpad();
+
             return status::success;
         }
 
         bool use_dense_;
+
+    private:
+        void init_scratchpad() {
+            const memory_desc_wrapper data_d(src_md());
+            const memory_desc_wrapper diff_data_d(diff_dst_md());
+            using namespace memory_tracking::names;
+            auto scratchpad = scratchpad_registry().registrar();
+            const auto diff_dst_size = diff_data_d.nelems(true) * sizeof(float);
+            scratchpad.book(
+                    key_eltwise_src, data_d.nelems(true) * sizeof(float));
+            scratchpad.book(key_eltwise_diff_dst, diff_dst_size);
+        }
     };
 
     ref_eltwise_bwd_t(const pd_t *apd) : primitive_impl_t(apd) {}