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 | #ifndef BATCH_NORMALIZATION_PD_HPP |
18 | #define BATCH_NORMALIZATION_PD_HPP |
19 | |
20 | #include "mkldnn.h" |
21 | |
22 | #include "c_types_map.hpp" |
23 | #include "primitive_desc.hpp" |
24 | #include "utils.hpp" |
25 | |
26 | namespace mkldnn { |
27 | namespace impl { |
28 | |
29 | struct batch_normalization_fwd_pd_t; |
30 | |
31 | struct batch_normalization_pd_t: public primitive_desc_t { |
32 | static constexpr auto base_pkind = primitive_kind::batch_normalization; |
33 | |
34 | batch_normalization_pd_t(engine_t *engine, |
35 | const batch_normalization_desc_t *adesc, |
36 | const primitive_attr_t *attr, |
37 | const batch_normalization_fwd_pd_t *hint_fwd_pd) |
38 | : primitive_desc_t(engine, attr, base_pkind) |
39 | , desc_(*adesc) |
40 | , hint_fwd_pd_(hint_fwd_pd) |
41 | , data_md_(desc_.data_desc) |
42 | , stat_md_(desc_.mean_desc) |
43 | , scaleshift_md_(desc_.data_scaleshift_desc) |
44 | , ws_md_() |
45 | {} |
46 | |
47 | const batch_normalization_desc_t *desc() const { return &desc_; } |
48 | virtual const op_desc_t *op_desc() const override |
49 | { return reinterpret_cast<const op_desc_t *>(this->desc()); } |
50 | virtual void init_info() override { impl::init_info(this, this->info_); } |
51 | |
52 | virtual status_t query(query_t what, int idx, void *result) const override { |
53 | switch (what) { |
54 | case query::batch_normalization_d: |
55 | *(const batch_normalization_desc_t**)result = desc(); break; |
56 | default: return primitive_desc_t::query(what, idx, result); |
57 | } |
58 | return status::success; |
59 | } |
60 | |
61 | /* common batch_normalization aux functions */ |
62 | |
63 | dim_t MB() const { return data_desc().dims[0]; } |
64 | dim_t C() const { return data_desc().dims[1]; } |
65 | dim_t D() const { return ndims() >= 5 ? data_desc().dims[ndims() - 3] : 1; } |
66 | dim_t H() const { return ndims() >= 4 ? data_desc().dims[ndims() - 2] : 1; } |
67 | dim_t W() const { return ndims() >= 3 ? data_desc().dims[ndims() - 1] : 1; } |
68 | |
69 | int ndims() const { return desc_.data_desc.ndims; } |
70 | |
71 | bool stats_is_src() const { return desc_.flags & mkldnn_use_global_stats; } |
72 | bool use_scaleshift() const { return desc_.flags & mkldnn_use_scaleshift; } |
73 | bool use_global_stats() const |
74 | { return desc_.flags & mkldnn_use_global_stats; } |
75 | bool fuse_bn_relu() const { return desc_.flags & mkldnn_fuse_bn_relu; } |
76 | bool with_relu_post_op() const { |
77 | const auto &p = this->attr()->post_ops_; |
78 | return p.len_ == 1 && p.entry_[0].is_relu(true, true); |
79 | } |
80 | |
81 | bool is_fwd() const { |
82 | return utils::one_of(desc_.prop_kind, prop_kind::forward_training, |
83 | prop_kind::forward_inference); |
84 | } |
85 | bool is_bwd() const { return !this->is_fwd(); } |
86 | bool is_training() const |
87 | { return desc_.prop_kind == prop_kind::forward_training; } |
88 | |
89 | bool has_zero_dim_memory() const |
90 | { return memory_desc_wrapper(desc_.data_desc).has_zero_dim(); } |
91 | |
92 | protected: |
93 | batch_normalization_desc_t desc_; |
94 | const batch_normalization_fwd_pd_t *hint_fwd_pd_; |
95 | |
96 | memory_desc_t data_md_; |
97 | memory_desc_t stat_md_; |
98 | memory_desc_t scaleshift_md_; |
99 | |
100 | memory_desc_t ws_md_; |
101 | |
102 | void init_default_ws(size_t bits_per_element) { |
103 | const auto data_mdw = memory_desc_wrapper(data_md_); |
104 | |
105 | const dim_t data_nelems = data_mdw.nelems(true); |
106 | const dim_t bits_per_byte = 8; |
107 | const dims_t ws_sz = { (dim_t)utils::div_up( |
108 | data_nelems * bits_per_element, bits_per_byte) }; |
109 | mkldnn_memory_desc_init_by_tag(&ws_md_, 1, ws_sz, impl::data_type::u8, |
110 | format_tag::x); |
111 | } |
112 | |
113 | private: |
114 | const memory_desc_t &data_desc() const { return desc_.data_desc; } |
115 | }; |
116 | |
117 | struct batch_normalization_fwd_pd_t: public batch_normalization_pd_t { |
118 | typedef batch_normalization_fwd_pd_t base_class; |
119 | typedef batch_normalization_fwd_pd_t hint_class; |
120 | |
121 | batch_normalization_fwd_pd_t(engine_t *engine, |
122 | const batch_normalization_desc_t *adesc, |
123 | const primitive_attr_t *attr, |
124 | const batch_normalization_fwd_pd_t *hint_fwd_pd) |
125 | : batch_normalization_pd_t(engine, adesc, attr, hint_fwd_pd) |
126 | {} |
127 | |
128 | virtual arg_usage_t arg_usage(primitive_arg_index_t arg) const override { |
129 | if (arg == MKLDNN_ARG_SRC) return arg_usage_t::input; |
130 | if (arg == MKLDNN_ARG_DST) return arg_usage_t::output; |
131 | |
132 | if (utils::one_of(arg, MKLDNN_ARG_MEAN, MKLDNN_ARG_VARIANCE)) { |
133 | if (stats_is_src()) return arg_usage_t::input; |
134 | if (!stats_is_src() && is_training()) return arg_usage_t::output; |
135 | return arg_usage_t::unused; |
136 | } |
137 | |
138 | if (arg == MKLDNN_ARG_SCALE_SHIFT && use_scaleshift()) |
139 | return arg_usage_t::input; |
140 | |
141 | if (arg == MKLDNN_ARG_WORKSPACE && is_training() && fuse_bn_relu()) |
142 | return arg_usage_t::output; |
143 | |
144 | return primitive_desc_t::arg_usage(arg); |
145 | } |
146 | |
147 | virtual const memory_desc_t *src_md(int index = 0) const override { |
148 | if (index == 0) return &data_md_; |
149 | if (stats_is_src() && (index == 1 || index == 2)) return &stat_md_; |
150 | return nullptr; |
151 | } |
152 | |
153 | virtual const memory_desc_t *dst_md(int index = 0) const override { |
154 | if (index == 0) return &data_md_; |
155 | if (!stats_is_src() && is_training() && (index == 1 || index == 2)) |
156 | return &stat_md_; |
157 | return nullptr; |
158 | } |
159 | |
160 | virtual const memory_desc_t *weights_md(int index = 0) const override |
161 | { return index == 0 ? &scaleshift_md_ : nullptr; } |
162 | |
163 | virtual const memory_desc_t *workspace_md(int index = 0) const override |
164 | { return index == 0 && is_training() && fuse_bn_relu() ? &ws_md_ : nullptr; } |
165 | |
166 | const memory_desc_t *stat_md() const |
167 | { return stats_is_src() ? src_md(1) : dst_md(1); } |
168 | |
169 | virtual int n_inputs() const override |
170 | { return 1 + 2 * stats_is_src() + use_scaleshift(); } |
171 | virtual int n_outputs() const override |
172 | { return 1 + (fuse_bn_relu() + 2 * (!stats_is_src())) * is_training(); } |
173 | }; |
174 | |
175 | struct batch_normalization_bwd_pd_t: public batch_normalization_pd_t { |
176 | typedef batch_normalization_bwd_pd_t base_class; |
177 | typedef batch_normalization_fwd_pd_t hint_class; |
178 | |
179 | batch_normalization_bwd_pd_t(engine_t *engine, |
180 | const batch_normalization_desc_t *adesc, |
181 | const primitive_attr_t *attr, |
182 | const batch_normalization_fwd_pd_t *hint_fwd_pd) |
183 | : batch_normalization_pd_t(engine, adesc, attr, hint_fwd_pd) |
184 | , diff_data_md_(desc_.diff_data_desc) |
185 | , diff_scaleshift_md_(desc_.diff_data_scaleshift_desc) |
186 | {} |
187 | |
188 | virtual arg_usage_t arg_usage(primitive_arg_index_t arg) const override { |
189 | if (utils::one_of(arg, MKLDNN_ARG_SRC, MKLDNN_ARG_MEAN, |
190 | MKLDNN_ARG_VARIANCE, MKLDNN_ARG_DIFF_DST)) |
191 | return arg_usage_t::input; |
192 | |
193 | if (arg == MKLDNN_ARG_SCALE_SHIFT && use_scaleshift()) |
194 | return arg_usage_t::input; |
195 | |
196 | if (arg == MKLDNN_ARG_WORKSPACE && fuse_bn_relu()) |
197 | return arg_usage_t::input; |
198 | |
199 | if (arg == MKLDNN_ARG_DIFF_SRC) |
200 | return arg_usage_t::output; |
201 | |
202 | if (arg == MKLDNN_ARG_DIFF_SCALE_SHIFT && use_scaleshift()) |
203 | return arg_usage_t::output; |
204 | |
205 | return primitive_desc_t::arg_usage(arg); |
206 | } |
207 | |
208 | virtual const memory_desc_t *src_md(int index = 0) const override |
209 | { return index == 0 ? &data_md_ : index <= 2 ? &stat_md_ : nullptr; } |
210 | virtual const memory_desc_t *diff_dst_md(int index = 0) const override |
211 | { return index == 0 ? &diff_data_md_ : nullptr; } |
212 | virtual const memory_desc_t *diff_src_md(int index = 0) const override |
213 | { return index == 0 ? &diff_data_md_ : nullptr; } |
214 | |
215 | virtual const memory_desc_t *weights_md(int index = 0) const override |
216 | { return index == 0 ? &scaleshift_md_ : nullptr; } |
217 | virtual const memory_desc_t *diff_weights_md(int index = 0) const override |
218 | { return index == 0 ? &diff_scaleshift_md_ : nullptr; } |
219 | |
220 | virtual const memory_desc_t *workspace_md(int index = 0) const override |
221 | { return index == 0 && fuse_bn_relu() ? &ws_md_ : nullptr; } |
222 | |
223 | const memory_desc_t *stat_md() const { return src_md(1); } |
224 | |
225 | virtual int n_inputs() const override |
226 | { return 4 + use_scaleshift() + fuse_bn_relu(); } |
227 | virtual int n_outputs() const override |
228 | { return 1 + (desc_.prop_kind == prop_kind::backward); } |
229 | |
230 | protected: |
231 | memory_desc_t diff_data_md_; |
232 | memory_desc_t diff_scaleshift_md_; |
233 | }; |
234 | |
235 | } |
236 | } |
237 | |
238 | #endif |
239 | |
240 | // vim: et ts=4 sw=4 cindent cino^=l0,\:0,N-s |
241 | |