inference-engine/thirdparty/mkl-dnn/src/cpu/ref_softmax.cpp

   1 /*******************************************************************************
   2 * Copyright 2016-2018 Intel Corporation
   3 *
   4 * Licensed under the Apache License, Version 2.0 (the "License");
   5 * you may not use this file except in compliance with the License.
   6 * You may obtain a copy of the License at
   7 *
   8 *     http://www.apache.org/licenses/LICENSE-2.0
   9 *
  10 * Unless required by applicable law or agreed to in writing, software
  11 * distributed under the License is distributed on an "AS IS" BASIS,
  12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 * See the License for the specific language governing permissions and
  14 * limitations under the License.
  15 *******************************************************************************/
  16
  17 #include <assert.h>
  18 #include <float.h>
  19 #include <math.h>
  20
  21 #include "c_types_map.hpp"
  22 #include "type_helpers.hpp"
  23 #include "mkldnn_thread.hpp"
  24
  25 #include "ref_softmax.hpp"
  26 #include "gemm/os_blas.hpp"
  27
  28 #ifdef USE_MKL
  29 #include "mkl_vml_functions.h"
  30 #endif
  31
  32 namespace mkldnn {
  33 namespace impl {
  34 namespace cpu {
  35
  36 template <impl::data_type_t data_type>
  37 void ref_softmax_fwd_t<data_type>::execute_forward_dense() const {
  38     auto src = reinterpret_cast<const data_t *>(this->input_memory(0));
  39     auto dst = reinterpret_cast<data_t *>(this->memory(0));
  40
  41     int outer_size_ = utils::array_product(pd()->src_pd()->desc()->dims, pd()->desc()->softmax_axis);
  42
  43     if (outer_size_ == 1) {
  44         for (int ou = 0; ou < outer_size_; ou++) {
  45             const data_t *src_data = src + ou * channels_;
  46             data_t *dst_data = dst + ou * channels_;
  47             data_t scalar = 0;
  48
  49             _max(channels_, src_data, &scalar);
  50             _sub(channels_, scalar, src_data, dst_data);
  51             _exp_parallel(channels_, dst_data, dst_data);
  52             _sum(channels_, dst_data, &scalar);
  53             _scal_parallel(channels_, data_t(1) / scalar, dst_data);
  54         }
  55     } else {
  56         parallel_nd(outer_size_, [&](int ou) {
  57             const data_t *src_data = src + ou * channels_;
  58             data_t *dst_data = dst + ou * channels_;
  59             data_t scalar = 0;
  60
  61             _max(channels_, src_data, &scalar);
  62             _sub(channels_, scalar, src_data, dst_data);
  63             _exp(channels_, dst_data, dst_data);
  64             _sum(channels_, dst_data, &scalar);
  65             _scal(channels_, data_t(1) / scalar, dst_data);
  66         });
  67     }
  68 }
  69
  70 template <impl::data_type_t data_type>
  71 void ref_softmax_fwd_t<data_type>::execute_forward_generic() const {
  72     auto src = reinterpret_cast<const data_t *>(this->input_memory(0));
  73     auto dst = reinterpret_cast<data_t *>(this->memory(0));
  74
  75     data_t space_max_val = 0, space_denom_val = 0;
  76     data_t *space_max = &space_max_val, *space_denom = &space_denom_val;
  77     if (inner_size_ > 1) {
  78         using namespace memory_tracking::names;
  79         space_max = scratchpad().template get<data_t>(key_softmax_reduction);
  80         space_denom = space_max + inner_size_;
  81     }
  82
  83     const memory_desc_wrapper data_d(pd()->src_pd());
  84     const size_t dim = channels_ * inner_size_;
  85
  86     int outer_size_ = utils::array_product(pd()->src_pd()->desc()->dims, pd()->desc()->softmax_axis);
  87
  88     for (int ou = 0; ou < outer_size_; ou++) {
  89         utils::array_set(space_max, -FLT_MAX, inner_size_);
  90         utils::array_set(space_denom, 0, inner_size_);
  91
  92         for (int c = 0; c < channels_; c++) {
  93             for(int in = 0; in < inner_size_; in++) {
  94                 size_t off = data_d.off_l(ou * dim + c * inner_size_ + in);
  95                 space_max[in] = nstl::max(space_max[in], src[off]);
  96             }
  97         }
  98
  99         for (int c = 0; c < channels_; c++) {
 100             for(int in = 0; in < inner_size_; in++) {
 101                 size_t off = data_d.off_l(ou * dim + c * inner_size_ + in);
 102                 space_denom[in] += dst[off] = exp(src[off] - space_max[in]);
 103             }
 104         }
 105
 106         for (int c = 0; c < channels_; c++) {
 107             for (int in = 0; in < inner_size_; in++) {
 108                 size_t off = data_d.off_l(ou * dim + c * inner_size_ + in);
 109                 dst[off] /= space_denom[in];
 110             }
 111         }
 112     }
 113 }
 114
 115 template <impl::data_type_t data_type>
 116 void ref_softmax_fwd_t<data_type>::_max(int n, const data_t *x,
 117         data_t *max_data) const {
 118 // Intel(R) C++ Compiler generates the maxps + shuffle pattern
 119 // for the max search which works faster
 120 #if !defined(__INTEL_COMPILER)
 121     // The code below makes a compiler to generate maxps instruction
 122     // rather than maxss, which is generated for the 'else' code path
 123     auto max_wrapper = [](data_t a, data_t b) { return nstl::max(a, b); };
 124     auto min_wrapper = [](int a, int b) { return nstl::min(a, b); };
 125
 126     constexpr int unroll_factor = 32;
 127     data_t max_values[unroll_factor];
 128
 129     if (n < unroll_factor) {
 130         data_t max_val = x[0];
 131         for (int i = 1; i < n; i++) {
 132             max_val = max_wrapper(max_val, x[i]);
 133         }
 134         max_data[0] = max_val;
 135         return;
 136     }
 137     for (int i = 0; i < unroll_factor; i++) {
 138         max_values[i] = x[i];
 139     }
 140     for (int i = unroll_factor; i < n; i += unroll_factor) {
 141         int offset = min_wrapper(i, n - unroll_factor);
 142         for (int j = 0; j < unroll_factor; j++) {
 143             max_values[j] = max_wrapper(max_values[j], x[offset + j]);
 144         }
 145     }
 146     data_t max_val = max_values[0];
 147     for (int i = 1; i < unroll_factor; i++) {
 148         max_val = max_wrapper(max_val, max_values[i]);
 149     }
 150     max_data[0] = max_val;
 151 #else
 152     max_data[0] = x[0];
 153     for (int c = 1; c < n; ++c)
 154         max_data[0] = nstl::max(max_data[0], x[c]);
 155 #endif
 156 }
 157
 158 template <impl::data_type_t data_type>
 159 void ref_softmax_fwd_t<data_type>::_sub(int n, data_t alpha, const data_t *x,
 160         data_t *y) const {
 161     constexpr int unroll_factor = 32;
 162     int tail = n % unroll_factor;
 163     for (int i = 0; i < n - tail; i += unroll_factor) {
 164         PRAGMA_OMP_SIMD()
 165         for (int j = 0; j < unroll_factor; j++) {
 166             y[i + j] = x[i + j] - alpha;
 167         }
 168     }
 169     PRAGMA_OMP_SIMD()
 170     for (int i = n - tail; i < n; i++) {
 171         y[i] = x[i] - alpha;
 172     }
 173 }
 174
 175 template <impl::data_type_t data_type>
 176 void ref_softmax_fwd_t<data_type>::_exp_parallel(int n, const data_t *a, data_t *r) const {
 177 #ifdef USE_MKL
 178     if (data_type == data_type::f32) {
 179         vsExp(n, a, r);
 180         return;
 181     }
 182 #endif
 183     parallel_nd(n, [&](int c) { r[c] = expf(a[c]); });
 184 }
 185
 186 template <impl::data_type_t data_type>
 187 void ref_softmax_fwd_t<data_type>::_exp(int n, const data_t *a, data_t *r) const {
 188     for (int c = 0; c < n; c++)
 189         r[c] = expf(a[c]);
 190 }
 191
 192 template <impl::data_type_t data_type>
 193 void ref_softmax_fwd_t<data_type>::_sum(int n, const data_t *x,
 194         data_t *sum_data) const {
 195 #ifdef USE_CBLAS
 196     // Here we are summing x's eg. e^z , which are positives
 197     // so we can use BLAS ASUM
 198     if (data_type == data_type::f32) {
 199         sum_data[0] = cblas_sasum(n, x, 1);
 200         return;
 201     }
 202 #endif
 203     data_t tsum = static_cast<data_t>(0);
 204     PRAGMA_OMP_SIMD(reduction(+ : tsum))
 205     for (int c = 0; c < n; ++c)
 206         tsum += x[c];
 207     sum_data[0] = tsum;
 208 }
 209
 210 template <impl::data_type_t data_type>
 211 void ref_softmax_fwd_t<data_type>::_scal_parallel(int n, data_t alpha, data_t *x) const {
 212 #ifdef USE_CBLAS
 213     if (data_type == data_type::f32) {
 214         cblas_sscal(n, alpha, x, 1);
 215         return;
 216     }
 217 #endif
 218     parallel_nd(n, [&](int c) { x[c] *= alpha; });
 219 }
 220
 221 template <impl::data_type_t data_type>
 222 void ref_softmax_fwd_t<data_type>::_scal(int n, data_t alpha, data_t *x) const {
 223     for (int c = 0; c < n; c++)
 224         x[c] *= alpha;
 225 }
 226
 227 template struct ref_softmax_fwd_t<data_type::f32>;
 228
 229
 230 // NC/NCHW softmax for along final axe (1 for NC, 3 for NCHW)
 231 template <impl::data_type_t data_type>
 232 void ref_softmax_bwd_t<data_type>::execute_backward_dense() const {
 233     auto data = reinterpret_cast<const data_t *>(this->input_memory(0));
 234     auto diff_dst = reinterpret_cast<const data_t *>(this->input_memory(1));
 235     auto diff_src = reinterpret_cast<data_t *>(this->memory(0));
 236
 237     parallel_nd(outer_size_, [&](int ou) {
 238         data_t sbr = 0;
 239         size_t off = channels_*ou;
 240         for (int c = 0; c < channels_; c++) {
 241             size_t loff = off + c;
 242             data_t ldata = data[loff];
 243             sbr += diff_dst[loff]*ldata;
 244             diff_src[loff] = ldata;
 245         }
 246
 247         for(int c=0; c < channels_ ; ++c) {
 248           size_t loff = off + c;
 249           diff_src[loff] *= (diff_dst[loff] - sbr);
 250         }
 251     });
 252 }
 253
 254 template <impl::data_type_t data_type>
 255 void ref_softmax_bwd_t<data_type>::execute_backward_generic() const {
 256     const size_t dim = channels_ * inner_size_;
 257     auto data = reinterpret_cast<const data_t *>(this->input_memory(0));
 258     auto diff_dst = reinterpret_cast<const data_t *>(this->input_memory(1));
 259     auto diff_src = reinterpret_cast<data_t *>(this->memory(0));
 260     const memory_desc_wrapper diff_d(pd()->diff_src_pd());
 261     const memory_desc_wrapper data_d(pd()->dst_pd());
 262
 263     parallel_nd(outer_size_, [&](int ou) {
 264         for (int in = 0; in < inner_size_; in++) {
 265             data_t sbr = 0;
 266             for (int c = 0; c < channels_; c++) {
 267                 size_t off_diff = diff_d.off_l(ou * dim + c * inner_size_ + in);
 268                 size_t off_data = diff_d.off_l(ou * dim + c * inner_size_ + in);
 269                 sbr += diff_dst[off_diff]*data[off_data];
 270             }
 271
 272             for(int c=0; c < channels_ ; ++c) {
 273               size_t off_diff = diff_d.off_l(ou * dim + c * inner_size_ + in);
 274               size_t off_data = data_d.off_l(ou * dim + c * inner_size_ + in);
 275               diff_src[off_diff] = data[off_data]*(diff_dst[off_diff] - sbr);
 276             }
 277         }
 278     });
 279 }
 280
 281 template struct ref_softmax_bwd_t<data_type::f32>;
 282
 283 }
 284 }
 285 }
 286
 287 // vim: et ts=4 sw=4 cindent cino^=l0,\:0,N-s