inference-engine/thirdparty/mkl-dnn/src/common/rnn_pd.hpp

   1 /*******************************************************************************
   2 * Copyright 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 #ifndef RNN_PD_HPP
  18 #define RNN_PD_HPP
  19
  20 #include "mkldnn.h"
  21
  22 #include "c_types_map.hpp"
  23 #include "memory_pd.hpp"
  24 #include "primitive_desc.hpp"
  25
  26 namespace mkldnn {
  27 namespace impl {
  28
  29 // struct rnn_fwd_pd_t;
  30
  31 struct rnn_pd_t : public primitive_desc_t {
  32     static constexpr auto base_pkind = primitive_kind::rnn;
  33
  34     rnn_pd_t(mkldnn::impl::engine_t *engine, const rnn_desc_t *adesc,
  35             const primitive_attr_t *attr, const rnn_pd_t *hint_pd)
  36         : primitive_desc_t(engine, attr, primitive_kind::rnn)
  37         , desc_(*adesc)
  38         , hint_pd_(hint_pd) {}
  39     virtual ~rnn_pd_t() {}
  40
  41     const rnn_desc_t *desc() const { return &desc_; }
  42     virtual const op_desc_t *op_desc() const override {
  43         return reinterpret_cast<const op_desc_t *>(this->desc());
  44     }
  45     virtual void init_info() override { init_info_rnn(this, this->info_); }
  46
  47     virtual status_t query(query_t what, int idx, void *result) const override {
  48         switch (what) {
  49         case query::rnn_d: *(const rnn_desc_t **)result = desc(); break;
  50         default: return primitive_desc_t::query(what, idx, result);
  51         }
  52         return status::success;
  53     }
  54
  55     inline bool is_training() const {
  56         return utils::one_of(desc_.prop_kind, prop_kind::forward_training,
  57                 prop_kind::backward);
  58     }
  59
  60     inline size_t ws_states_size() {
  61         int wic = nstl::max(SLC(), nstl::max(SIC(), DIC()));
  62         return (size_t)(L() + 1) * D() * (T() + 1) * S() * MB() * wic;
  63     }
  64
  65     inline size_t ws_diff_states_size() {
  66         int wic = nstl::max(SLC(), nstl::max(SIC(), DIC()));
  67         return (size_t)(L() + 1) * D() * (T() + 1) * (S() + 1) * MB() * wic;
  68     }
  69
  70     inline size_t ws_gates_size() {
  71         int n_layer = L();
  72         int n_direction = D();
  73         int n_iter = T();
  74         int n_gates = G();
  75         int batch = MB();
  76         int s_size = DIC();
  77
  78         return (size_t)n_layer * n_direction * n_iter * batch * n_gates
  79                 * s_size;
  80     }
  81
  82     inline size_t ws_cell_comp_size() {
  83         int n_gates = G();
  84         int batch = MB();
  85         int s_size = DIC();
  86         return (size_t)is_lbr() * n_gates * batch * s_size;
  87     }
  88
  89     inline size_t ws_grid_comp_size() {
  90         int n_layer = L();
  91         int n_direction = D();
  92         int n_iter = T();
  93         int batch = MB();
  94         int s_size = DIC();
  95         return (size_t)is_lbr() * is_training() * n_layer * n_direction * n_iter
  96                 * batch * s_size;
  97     }
  98
  99     inline int ws_per_cell() {
 100         int batch = MB();
 101         int s_size = DIC();
 102         return is_lbr() * is_training() * batch * s_size;
 103     }
 104
 105     inline void set_offsets(size_t &ws_gates_offset, size_t &ws_states_offset,
 106             size_t &ws_diff_states_offset, size_t &ws_grid_comp_offset,
 107             size_t &ws_cell_comp_offset) {
 108         const size_t page_size = 4096; // 2097152;
 109         ws_gates_offset
 110                 = 0; // assumes the workspace base pointer is page aligned
 111         ws_states_offset = utils::rnd_up(ws_gates_size(), page_size);
 112         ws_diff_states_offset
 113                 = utils::rnd_up(ws_states_offset + ws_states_size(), page_size);
 114         ws_grid_comp_offset = utils::rnd_up(ws_diff_states_offset
 115                 + ws_diff_states_size(), page_size);
 116
 117         ws_cell_comp_offset = utils::rnd_up(ws_grid_comp_offset
 118                 + ws_grid_comp_size(), page_size);
 119     }
 120
 121     inline size_t get_ws_size() {
 122         size_t ws_gates_offset, ws_states_offset, ws_diff_states_offset,
 123             ws_grid_comp_offset, ws_cell_comp_offset;
 124         set_offsets(
 125                 ws_gates_offset, ws_states_offset, ws_diff_states_offset,
 126                 ws_grid_comp_offset, ws_cell_comp_offset);
 127         return ws_grid_comp_offset + ws_grid_comp_size();
 128     }
 129
 130     inline size_t get_scratchpad_size() {
 131         size_t ws_gates_offset, ws_states_offset, ws_diff_states_offset,
 132             ws_grid_comp_offset, ws_cell_comp_offset;
 133         set_offsets(
 134                 ws_gates_offset, ws_states_offset, ws_diff_states_offset,
 135                 ws_grid_comp_offset, ws_cell_comp_offset);
 136         if (desc_.prop_kind == prop_kind::forward_inference)
 137             return ws_cell_comp_offset + ws_cell_comp_size();
 138         else
 139             return ws_cell_comp_size();
 140     }
 141
 142     int T() const { return desc_.src_layer_desc.dims[0]; }
 143     int MB() const { return desc_.src_layer_desc.dims[1]; }
 144
 145     int L() const { return desc_.weights_layer_desc.dims[0]; }
 146     int D() const { return desc_.weights_layer_desc.dims[1]; }
 147
 148     int SIC() const { return desc_.weights_iter_desc.dims[2]; }
 149
 150     int SLC() const { return desc_.weights_layer_desc.dims[2]; }
 151     int G() const { return desc_.weights_layer_desc.dims[3]; }
 152     int DIC() const { return desc_.weights_layer_desc.dims[4]; }
 153
 154     int DLC() const { return desc_.dst_layer_desc.dims[2]; }
 155
 156     int S() const { return mkldnn_rnn_cell_get_states_count(&desc_.cell_desc); }
 157
 158     bool with_bias() const {
 159         return !memory_desc_wrapper(desc_.bias_desc).is_zero();
 160     }
 161
 162     bool with_src_iter() const {
 163         return !(memory_desc_wrapper(desc_.src_iter_desc).is_zero());
 164     }
 165
 166     bool with_dst_iter() const {
 167         return !memory_desc_wrapper(desc_.dst_iter_desc).is_zero();
 168     }
 169
 170     mkldnn::impl::alg_kind_t cell_kind() const {
 171         return desc_.cell_desc.cell_kind;
 172     }
 173     mkldnn::impl::alg_kind_t activation_kind() const {
 174         return desc_.cell_desc.activation_kind;
 175     }
 176
 177     bool is_lbr() const {
 178         return cell_kind() == mkldnn_gru_linear_before_reset;
 179     }
 180
 181     mkldnn_rnn_direction_t direction() const { return desc_.direction; }
 182
 183 protected:
 184     rnn_desc_t desc_;
 185     const rnn_pd_t *hint_pd_;
 186 };
 187
 188 struct rnn_fwd_pd_t : public rnn_pd_t {
 189     typedef rnn_fwd_pd_t base_class;
 190     typedef rnn_fwd_pd_t hint_class;
 191
 192     using rnn_pd_t::rnn_pd_t;
 193     virtual ~rnn_fwd_pd_t() {}
 194
 195     virtual const memory_pd_t *input_pd(int index = 0) const override {
 196         switch (index) {
 197         case 0: return src_pd(0);
 198         case 1: return src_pd(1);
 199         case 2: return weights_pd(0);
 200         case 3: return weights_pd(1);
 201         case 4: return weights_pd(2);
 202         default: return nullptr;
 203         }
 204     }
 205
 206     virtual const memory_pd_t *output_pd(int index = 0) const override {
 207         switch (index) {
 208         case 0: return dst_pd(0);
 209         case 1: return dst_pd(1);
 210         case 2: return workspace_pd();
 211         default: return nullptr;
 212         }
 213     }
 214
 215     virtual int n_inputs() const override {
 216         return 3 + with_bias() + with_src_iter();
 217     }
 218
 219     virtual int n_outputs() const override {
 220         return 1 + with_dst_iter() + is_training();
 221     }
 222
 223     int ws_idx() const { return 1 + with_dst_iter(); }
 224 };
 225
 226 struct rnn_bwd_pd_t : public rnn_pd_t {
 227     typedef rnn_bwd_pd_t base_class;
 228     typedef rnn_bwd_pd_t hint_class;
 229
 230     using rnn_pd_t::rnn_pd_t;
 231     virtual ~rnn_bwd_pd_t() {}
 232
 233     virtual const memory_pd_t *input_pd(int index = 0) const override {
 234         switch (index) {
 235         case 0: return src_pd(0);
 236         case 1: return src_pd(1);
 237         case 2: return weights_pd(0);
 238         case 3: return weights_pd(1);
 239         case 4: return weights_pd(2);
 240         case 5: return dst_pd(0);
 241         case 6: return dst_pd(1);
 242         case 7: return diff_dst_pd(0);
 243         case 8: return diff_dst_pd(1);
 244         case 9: return workspace_pd();
 245         default: return nullptr;
 246         }
 247     }
 248
 249     virtual const memory_pd_t *output_pd(int index = 0) const override {
 250         switch (index) {
 251         case 0: return diff_src_pd(0);
 252         case 1: return diff_src_pd(1);
 253         case 2: return diff_weights_pd(0);
 254         case 3: return diff_weights_pd(1);
 255         case 4: return diff_weights_pd(2);
 256         default: return nullptr;
 257         }
 258     }
 259
 260     virtual int n_inputs() const override {
 261         return 6 + with_src_iter() + with_bias() + 2 * with_dst_iter();
 262     }
 263     virtual int n_outputs() const override {
 264         return 3 + with_src_iter() + with_bias();
 265     }
 266
 267     int ws_idx() const {
 268         return 5 + with_src_iter() + with_bias() + 2 * with_dst_iter();
 269     }
 270 };
 271 }
 272 }
 273
 274 #endif