1 /*******************************************************************************
2 * Copyright 2018 Intel Corporation
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
8 * http://www.apache.org/licenses/LICENSE-2.0
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 *******************************************************************************/
22 #include "c_types_map.hpp"
23 #include "memory_pd.hpp"
24 #include "primitive_desc.hpp"
29 // struct rnn_fwd_pd_t;
31 struct rnn_pd_t : public primitive_desc_t {
32 static constexpr auto base_pkind = primitive_kind::rnn;
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)
38 , hint_pd_(hint_pd) {}
39 virtual ~rnn_pd_t() {}
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());
45 virtual void init_info() override { init_info_rnn(this, this->info_); }
47 virtual status_t query(query_t what, int idx, void *result) const override {
49 case query::rnn_d: *(const rnn_desc_t **)result = desc(); break;
50 default: return primitive_desc_t::query(what, idx, result);
52 return status::success;
55 inline bool is_training() const {
56 return utils::one_of(desc_.prop_kind, prop_kind::forward_training,
60 inline bool is_fwd() const {
61 return utils::one_of(desc_.prop_kind, prop_kind::forward_training,
62 prop_kind::forward_inference);
65 int T() const { return desc_.src_layer_desc.dims[0]; }
66 int MB() const { return desc_.src_layer_desc.dims[1]; }
68 int L() const { return desc_.weights_layer_desc.dims[0]; }
69 int D() const { return desc_.weights_layer_desc.dims[1]; }
71 int SIC() const { return desc_.weights_iter_desc.dims[2]; }
73 int SLC() const { return desc_.weights_layer_desc.dims[2]; }
74 int G() const { return desc_.weights_layer_desc.dims[3]; }
75 int DIC() const { return desc_.weights_layer_desc.dims[4]; }
77 int DLC() const { return desc_.dst_layer_desc.dims[2]; }
79 bool with_bias() const {
80 return !memory_desc_wrapper(desc_.bias_desc).is_zero();
83 bool with_src_iter() const {
84 return !(memory_desc_wrapper(desc_.src_iter_desc).is_zero());
87 bool with_dst_iter() const {
88 return !memory_desc_wrapper(desc_.dst_iter_desc).is_zero();
91 mkldnn::impl::alg_kind_t cell_kind() const {
92 return desc_.cell_desc.cell_kind;
94 mkldnn::impl::alg_kind_t activation_kind() const {
95 return desc_.cell_desc.activation_kind;
99 return cell_kind() == mkldnn_gru_linear_before_reset;
102 mkldnn_rnn_direction_t direction() const { return desc_.direction; }
106 const rnn_pd_t *hint_pd_;
109 struct rnn_fwd_pd_t : public rnn_pd_t {
110 typedef rnn_fwd_pd_t base_class;
111 typedef rnn_fwd_pd_t hint_class;
113 using rnn_pd_t::rnn_pd_t;
114 virtual ~rnn_fwd_pd_t() {}
116 virtual const memory_pd_t *input_pd(int index = 0) const override {
117 if (index == 0) return src_pd(0);
118 if (with_src_iter() && index == 1) return src_pd(1);
119 index = index - 1 - with_src_iter();
121 if (index < 3) return weights_pd(index);
126 virtual const memory_pd_t *output_pd(int index = 0) const override {
127 if (index == 0) return dst_pd(0);
128 if (with_dst_iter() && index == 1) return dst_pd(1);
129 index = index - 1 - with_dst_iter();
131 if (is_training() && index == 0) return workspace_pd();
136 virtual int n_inputs() const override {
137 return 3 + with_bias() + with_src_iter();
140 virtual int n_outputs() const override {
141 return 1 + with_dst_iter() + is_training();
144 int ws_idx() const { return 1 + with_dst_iter(); }
147 struct rnn_bwd_pd_t : public rnn_pd_t {
148 typedef rnn_bwd_pd_t base_class;
149 typedef rnn_fwd_pd_t hint_class;
151 using rnn_pd_t::rnn_pd_t;
152 virtual ~rnn_bwd_pd_t() {}
154 virtual const memory_pd_t *input_pd(int index = 0) const override {
155 if (index == 0) return src_pd(0);
156 if (with_src_iter() && index == 1) return src_pd(1);
157 index = index - 1 - with_src_iter();
159 if (index < 2) return weights_pd(index);
160 if (with_bias() && index == 2) return weights_pd(2);
161 index = index - 2 - with_bias();
163 if (index == 0) return dst_pd(0);
164 if (with_dst_iter() && index == 1) return dst_pd(1);
165 index = index - 1 - with_dst_iter();
167 if (index == 0) return diff_dst_pd(0);
168 if (with_dst_iter() && index == 1) return diff_dst_pd(1);
169 index = index - 1 - with_dst_iter();
171 if (index == 0) return workspace_pd();
176 virtual const memory_pd_t *output_pd(int index = 0) const override {
177 if (index == 0) return diff_src_pd(0);
178 if (with_src_iter() && index == 1) return diff_src_pd(1);
179 index = index - 1 - with_src_iter();
181 if (index < 3) return diff_weights_pd(index);
186 virtual int n_inputs() const override {
187 return 6 + with_src_iter() + with_bias() + 2 * with_dst_iter();
189 virtual int n_outputs() const override {
190 return 3 + with_src_iter() + with_bias();
194 return 5 + with_src_iter() + with_bias() + 2 * with_dst_iter();