1/*******************************************************************************
2* Copyright 2017-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 "c_types_map.hpp"
18#include "nstl.hpp"
19#include "type_helpers.hpp"
20#include "cpu_memory.hpp"
21
22#include "jit_sse42_conv_kernel_f32.hpp"
23
24#define GET_OFF(field) offsetof(jit_conv_call_s, field)
25
26namespace mkldnn {
27namespace impl {
28namespace cpu {
29
30using namespace mkldnn::impl::format_tag;
31using namespace mkldnn::impl::prop_kind;
32using namespace mkldnn::impl::utils;
33
34using namespace Xbyak;
35
36void jit_sse42_conv_fwd_kernel_f32::oh_step_unroll_kw(int ur_w,
37 int pad_l, int pad_r, int oc_blocks)
38{
39 int iw = jcp.iw;
40 int ih = jcp.ih;
41 int kw = jcp.kw;
42 int kh = jcp.kh;
43 int nb_ic = jcp.nb_ic;
44 int stride_w = jcp.stride_w;
45 int dilate_w = jcp.dilate_w + 1;
46 int ic_blk = jcp.ic_block;
47 int oc_blk = jcp.oc_block;
48
49 for (int ki = 0; ki < kw; ki++) {
50 int jj_start = nstl::max(0, div_up(pad_l - ki * dilate_w, stride_w));
51 int jj_end = ur_w
52 - nstl::max(0, div_up(ki*dilate_w + pad_r - (kw-1)*dilate_w, stride_w));
53 for (int ifm2 = 0; ifm2 < ic_blk; ifm2++) {
54 for (int jj = jj_start; jj < jj_end; jj++) {
55 int inp_off;
56 if (one_of(jcp.src_tag, ncw, nchw))
57 inp_off = ifm2*ih*iw + (ki*dilate_w + jj*stride_w - pad_l);
58 else
59 inp_off = (ki*dilate_w + jj*stride_w - pad_l)*ic_blk + ifm2;
60
61 movss(Xmm(oc_blocks * ur_w + jj + 1),
62 ptr[aux_reg_input + sizeof(float) * inp_off]);
63 shufps(Xmm(oc_blocks * ur_w + jj + 1),
64 Xmm(oc_blocks * ur_w + jj + 1), 0x0);
65 }
66
67 for (int ii = 0; ii < oc_blocks; ii++) {
68 int ker_off = ii * nb_ic * kh * kw * ic_blk * oc_blk
69 + ki * ic_blk * oc_blk + ifm2 * oc_blk;
70
71 for (int jj = jj_start; jj < jj_end; jj++)
72 {
73 movups(xmm0,
74 ptr[aux_reg_kernel + sizeof(float) * ker_off]);
75 mulps(xmm0, Xmm(oc_blocks * ur_w + jj + 1));
76 addps(Xmm(ur_w * ii + jj + 1), xmm0);
77 }
78 }
79 }
80 }
81}
82
83void jit_sse42_conv_fwd_kernel_f32::oh_step_nopad(int ur_w,
84 int pad_l, int pad_r, int oc_blocks)
85{
86 Label kw_loop;
87
88 int iw = jcp.iw;
89 int ih = jcp.ih;
90 int kw = jcp.kw;
91 int kh = jcp.kh;
92 int nb_ic = jcp.nb_ic;
93 int stride_w = jcp.stride_w;
94 int dilate_w = jcp.dilate_w + 1;
95 int ic_blk = jcp.ic_block;
96 int oc_blk = jcp.oc_block;
97
98 xor_(ki_iter, ki_iter);
99 L(kw_loop);
100 {
101 int jj_start = 0;
102 int jj_end = ur_w;
103 for (int ifm2 = 0; ifm2 < ic_blk; ifm2++) {
104 for (int jj = jj_start; jj < jj_end; jj++) {
105 int inp_off;
106 if (one_of(jcp.src_tag, ncw, nchw))
107 inp_off = ifm2 * ih * iw + (jj * stride_w - pad_l);
108 else
109 inp_off = (jj * stride_w - pad_l) * ic_blk + ifm2;
110
111 movss(Xmm(oc_blocks * ur_w + jj + 1),
112 ptr[aux_reg_input + sizeof(float) * inp_off]);
113 shufps(Xmm(oc_blocks * ur_w + jj + 1),
114 Xmm(oc_blocks * ur_w + jj + 1), 0x0);
115 }
116 for (int ii = 0; ii < oc_blocks; ii++) {
117 int aux_kernel_offset = ii * nb_ic * kh * kw * ic_blk * oc_blk
118 + ifm2 * oc_blk;
119 for (int jj = jj_start; jj < jj_end; jj++) {
120 movups(xmm0,
121 ptr[aux_reg_kernel + sizeof(float) * aux_kernel_offset]);
122 mulps(xmm0, Xmm(oc_blocks * ur_w + jj + 1));
123 addps(Xmm(ur_w * ii + jj + 1), xmm0);
124 }
125 }
126 }
127 add(aux_reg_kernel, sizeof(float) * oc_blk * ic_blk);
128 add(aux_reg_input, sizeof(float) * (one_of(jcp.src_tag, ncw, nchw) ?
129 dilate_w : ic_blk * dilate_w));
130
131 inc(ki_iter);
132 cmp(ki_iter, kw);
133 jl(kw_loop, T_NEAR);
134 }
135}
136
137void jit_sse42_conv_fwd_kernel_f32::width_blk_step(int ur_w,
138 int pad_l, int pad_r, int oc_blocks)
139{
140 int iw = jcp.iw;
141 int kw = jcp.kw;
142 int ow = jcp.ow;
143 int oh = jcp.oh;
144 int dilate_h = jcp.dilate_h + 1;
145 int dilate_w = jcp.dilate_w + 1;
146 int ic_blk = jcp.ic_block;
147 int oc_blk = jcp.oc_block;
148 const int inp_mult = one_of(jcp.src_tag, ncw, nchw)
149 ? dilate_h : ic_blk * dilate_h;
150 const int inp_off = one_of(jcp.src_tag, ncw, nchw)
151 ? dilate_w : ic_blk * dilate_w;
152
153 xor_(simd_iter, simd_iter);
154
155 mov(aux_reg_input, reg_input);
156 mov(aux_reg_kernel, reg_kernel);
157
158 Label init_simd_iter_loop;
159 Label init_done;
160 Label init_first;
161
162 L(init_simd_iter_loop);
163
164 if (!jcp.with_sum) {
165 test(reg_ci_flag, FLAG_IC_FIRST);
166 jne(init_first, T_NEAR);
167 }
168
169 for (int ii = 0; ii < oc_blocks; ii++)
170 for (int jj = 0; jj < ur_w; jj++)
171 movups(Xmm(ur_w * ii + jj + 1), xword[reg_output
172 + sizeof(float) * (ii * oh * ow + jj) * oc_blk]);
173
174 if (jcp.with_sum && jcp.with_bias) {
175 test(reg_ci_flag, FLAG_IC_FIRST);
176 je(init_done, T_NEAR);
177
178 for (int ii = 0; ii < oc_blocks; ii++)
179 for (int jj = 0; jj < ur_w; jj++)
180 addps(Xmm(ur_w * ii + jj + 1),
181 xword[reg_bias + sizeof(float) * ii * oc_blk]);
182 }
183
184 jmp(init_done);
185
186 L(init_first);
187 if (this->jcp.with_bias) {
188 for (int ii = 0; ii < oc_blocks; ii++)
189 for (int jj = 0; jj < ur_w; jj++)
190 movups(Xmm(ur_w * ii + jj + 1),
191 xword[reg_bias + sizeof(float) * ii * oc_blk]);
192 } else {
193 for (int ii = 0; ii < oc_blocks; ii++)
194 for (int jj = 0; jj < ur_w; jj++)
195 pxor(Xmm(ur_w * ii + jj + 1), Xmm(ur_w * ii + jj + 1));
196 }
197
198 L(init_done);
199
200 Label skip_kh_loop;
201 mov(kj, reg_kh);
202 if ((jcp.dilate_h >= jcp.ih)
203 || (jcp.kh - 1) * (jcp.dilate_h + 1) < nstl::max(jcp.t_pad, jcp.b_pad)) {
204 cmp(kj, 0);
205 je(skip_kh_loop, T_NEAR);
206 }
207 Label kh_loop;
208 L(kh_loop);
209 {
210 if (jcp.kw >= 5 && pad_l == 0 && pad_r == 0) {
211 oh_step_nopad(ur_w, pad_l, pad_r, oc_blocks);
212 sub(aux_reg_input, sizeof(float) * kw * inp_off);
213 add(aux_reg_input, sizeof(float) * iw * inp_mult);
214 } else {
215 oh_step_unroll_kw(ur_w, pad_l, pad_r, oc_blocks);
216 add(aux_reg_kernel, sizeof(float) * kw * oc_blk * ic_blk);
217 add(aux_reg_input, sizeof(float) * iw * inp_mult);
218 }
219
220 dec(kj);
221 cmp(kj, 0);
222 jg(kh_loop, T_NEAR);
223 }
224
225 L(skip_kh_loop);
226
227 if (jcp.with_eltwise) {
228 Label regular_store;
229 test(reg_ci_flag, FLAG_IC_LAST);
230 je(regular_store, T_NEAR);
231
232 eltwise_injector_->compute_vector_range(1, oc_blocks * ur_w + 1);
233
234 L(regular_store);
235 }
236
237 for (int ii = 0; ii < oc_blocks; ii++) {
238 for (int jj = 0; jj < ur_w; jj++) {
239 const size_t o_off = (ii * oh * ow + jj) * oc_blk;
240
241 Xmm reg_out = Xmm(ur_w * ii + jj + 1);
242 movups(xword[reg_output + sizeof(float) * o_off], reg_out);
243 }
244 }
245
246 mov(aux_reg_kernel, reg_kernel);
247 mov(aux_reg_input, reg_input);
248 add(aux_reg_kernel, sizeof(float) * 4);
249 add(reg_output, sizeof(float) * 4);
250 add(reg_bias, sizeof(float) * 4);
251
252 inc(simd_iter);
253 cmp(simd_iter, 2);
254 jl(init_simd_iter_loop, T_NEAR);
255
256 sub(reg_output, sizeof(float) * 8);
257 sub(reg_bias, sizeof(float) * 8);
258}
259
260inline void jit_sse42_conv_fwd_kernel_f32::solve_common(int oc_blocks)
261{
262 int ur_w = jcp.ur_w;
263 int ur_w_tail = jcp.ur_w_tail;
264 int n_oi = jcp.ow / ur_w;
265 int iw = jcp.iw;
266 int kw = jcp.kw;
267 int ic_blk = jcp.ic_block;
268 int oc_blk = jcp.oc_block;
269 int dilate_w = jcp.dilate_w + 1;
270 int str_w = jcp.stride_w;
271 const int inp_mult = one_of(jcp.src_tag, ncw, nchw) ? 1 : ic_blk;
272
273 int l_pad = jcp.l_pad;
274 int r_pad = nstl::max(0, (int(jcp.ow) - 1) * str_w + (kw - 1) * dilate_w
275 - (iw + l_pad - 1));
276 int r_pad1 = (ur_w * n_oi - 1) * str_w + (kw - 1) * dilate_w
277 - (iw + l_pad - 1);
278 if (r_pad1 > 0) n_oi--;
279
280 if (l_pad > 0) {
281 n_oi--;
282 if (n_oi < 0 && r_pad1 > 0)
283 width_blk_step(ur_w, l_pad, r_pad1, oc_blocks); // "lrpad"
284 else
285 width_blk_step(ur_w, l_pad, 0, oc_blocks); // "lpad"
286 add(reg_input, sizeof(float) * (ur_w * str_w - l_pad) * inp_mult);
287 add(reg_output, sizeof(float) * ur_w * oc_blk);
288 }
289
290 Label ow_loop;
291 xor_(oi_iter, oi_iter);
292
293 if (n_oi > 0) {
294 L(ow_loop);
295
296 width_blk_step(ur_w, 0, 0, oc_blocks); // "middle"
297 add(reg_input, sizeof(float) * ur_w * str_w * inp_mult);
298 add(reg_output, sizeof(float) * ur_w * oc_blk);
299
300 inc(oi_iter);
301 cmp(oi_iter, n_oi);
302 jl(ow_loop, T_NEAR);
303 }
304
305 if (r_pad1 > 0 && n_oi >=0) {
306 width_blk_step(ur_w, 0, r_pad1, oc_blocks); // "rpad"
307 add(reg_input, sizeof(float) * ur_w * str_w * inp_mult);
308 add(reg_output, sizeof(float) * ur_w * oc_blk);
309 }
310
311 if (ur_w_tail != 0)
312 width_blk_step(ur_w_tail, 0, r_pad, oc_blocks); // "tail"
313}
314
315void jit_sse42_conv_fwd_kernel_f32::generate()
316{
317 this->preamble();
318
319 mov(reg_input, ptr[this->param1 + GET_OFF(src)]);
320 mov(reg_output, ptr[this->param1 + GET_OFF(dst)]);
321 mov(reg_kernel, ptr[this->param1 + GET_OFF(filt)]);
322 if (jcp.with_bias)
323 mov(reg_bias, ptr[this->param1 + GET_OFF(bias)]);
324 mov(reg_kh, ptr[this->param1 + GET_OFF(kh_padding)]);
325 mov(reg_ci_flag, ptr[this->param1 + GET_OFF(flags)]);
326 mov(reg_oc_blocks, ptr[this->param1 + GET_OFF(oc_blocks)]);
327
328 int nb_oc_tail = jcp.nb_oc % jcp.nb_oc_blocking;
329 Label tail, exit;
330
331 cmp(reg_oc_blocks, jcp.nb_oc_blocking);
332 jne(nb_oc_tail ? tail : exit, T_NEAR);
333
334 solve_common(jcp.nb_oc_blocking);
335 jmp(exit, T_NEAR);
336
337 if (nb_oc_tail) {
338 L(tail);
339 cmp(reg_oc_blocks, nb_oc_tail);
340 jne(exit, T_NEAR);
341 solve_common(nb_oc_tail);
342 }
343
344 L(exit);
345
346 this->postamble();
347
348 if (jcp.with_eltwise)
349 eltwise_injector_->prepare_table();
350}
351
352bool jit_sse42_conv_fwd_kernel_f32::post_ops_ok(
353 jit_conv_conf_t &jcp, const primitive_attr_t &attr) {
354 const auto &p = attr.post_ops_;
355
356 auto is_eltwise = [&](int idx) { return p.entry_[idx].is_eltwise(); };
357 auto is_sum = [&](int idx) { return p.entry_[idx].is_sum(); };
358
359 switch (p.len_) {
360 case 0: return true; // no post_ops
361 case 1: return is_eltwise(0) || is_sum(0); // sum OR eltwise
362 case 2: return is_sum(0) && is_eltwise(1); // sum -> eltwise
363 default: return false;
364 }
365
366 return false;
367}
368
369status_t jit_sse42_conv_fwd_kernel_f32::init_conf(jit_conv_conf_t &jcp,
370 const convolution_desc_t &cd, const memory_desc_wrapper &src_d,
371 const memory_desc_wrapper &weights_d, const memory_desc_wrapper &dst_d,
372 const primitive_attr_t &attr)
373{
374 if (!mayiuse(sse42)) return status::unimplemented;
375
376 jcp.prop_kind = cd.prop_kind;
377
378 const bool with_groups = weights_d.ndims() == src_d.ndims() + 1;
379 const int ndims = src_d.ndims();
380 jcp.ndims = ndims;
381
382 jcp.ngroups = with_groups ? weights_d.dims()[0] : 1;
383 jcp.mb = src_d.dims()[0];
384
385 jcp.oc = dst_d.dims()[1] / jcp.ngroups;
386 jcp.ic = src_d.dims()[1] / jcp.ngroups;
387
388 jcp.ih = (ndims == 3) ? 1 : src_d.dims()[2];
389 jcp.iw = src_d.dims()[ndims - 1];
390 jcp.oh = (ndims == 3) ? 1 : dst_d.dims()[2];
391 jcp.ow = dst_d.dims()[ndims - 1];
392
393 jcp.kh = (ndims == 3) ? 1 : weights_d.dims()[with_groups + 2];
394 jcp.kw = weights_d.dims()[with_groups + ndims - 1];
395
396 jcp.t_pad = (ndims == 3) ? 0 : cd.padding[0][0];
397 jcp.l_pad = cd.padding[0][ndims - 3];
398
399 jcp.stride_h = (ndims == 3) ? 1 : cd.strides[0];
400 jcp.stride_w = cd.strides[ndims - 3];
401
402 jcp.dilate_h = (ndims == 3) ? 0 : cd.dilates[0];
403 jcp.dilate_w = cd.dilates[ndims - 3];
404 jcp.b_pad = (jcp.oh - 1) * jcp.stride_h + (jcp.kh - 1) * (jcp.dilate_h + 1)
405 - (jcp.ih + jcp.t_pad - 1);
406
407 if (ndims == 3) {
408 jcp.src_tag = src_d.matches_one_of_tag(ncw, nwc, nCw8c);
409 jcp.wei_tag = weights_d.matches_one_of_tag(
410 Owi8o, gOwi8o, OIw8i8o, gOIw8i8o);
411 jcp.dst_tag = dst_d.matches_one_of_tag(nCw8c);
412 } else if (ndims == 4) {
413 jcp.src_tag = src_d.matches_one_of_tag(nchw, nhwc, nChw8c);
414 jcp.wei_tag = weights_d.matches_one_of_tag(
415 Ohwi8o, gOhwi8o, OIhw8i8o, gOIhw8i8o);
416 jcp.dst_tag = dst_d.matches_one_of_tag(nChw8c);
417 }
418 jcp.with_bias = cd.bias_desc.format_kind != format_kind::undef;
419
420 if (!post_ops_ok(jcp, attr))
421 return status::unimplemented;
422
423 const auto &p = attr.post_ops_;
424 jcp.with_sum = p.find(primitive_kind::sum) != -1;
425 const int eltwise_ind = p.find(primitive_kind::eltwise);
426 jcp.with_eltwise = eltwise_ind != -1;
427 if (jcp.with_eltwise)
428 jcp.eltwise = p.entry_[eltwise_ind].eltwise;
429
430 const bool flat = jcp.ic == 3;
431 const bool mimo = !flat;
432
433 bool args_ok = true
434 && IMPLICATION(flat, one_of(jcp.src_tag, ncw, nwc, nchw, nhwc)
435 && one_of(jcp.wei_tag, Owi8o, gOwi8o, Ohwi8o, gOhwi8o))
436 && IMPLICATION(mimo, one_of(jcp.src_tag, nCw8c, nChw8c)
437 && one_of(jcp.wei_tag, OIw8i8o, gOIw8i8o, OIhw8i8o, gOIhw8i8o))
438 && one_of(jcp.dst_tag, nCw8c, nChw8c);
439 if (!args_ok) return status::unimplemented;
440
441 const int simd_w = 8; // 2 SSE vectors processing at once
442
443 jcp.ur_h = 1; /* no code-unrolling by h so far */
444 jcp.ur_w = 3;
445 if (jcp.ow < jcp.ur_w) jcp.ur_w = jcp.ow;
446 jcp.ur_w_tail = jcp.ow % jcp.ur_w;
447
448 jcp.nb_oc_blocking = 4; /* the optimal value for the kernel */
449
450 args_ok = true
451 && jcp.oc % simd_w == 0
452 && jcp.l_pad <= jcp.ur_w
453 && IMPLICATION(jcp.kw > 7, (jcp.t_pad == 0 && jcp.l_pad == 0)
454 || (jcp.stride_w == 1 && jcp.stride_h == 1))
455 && IMPLICATION(mimo, jcp.ic % simd_w == 0);
456 if (!args_ok) return status::unimplemented;
457
458 int r_pad_no_tail = nstl::max(0, (jcp.ow - jcp.ur_w_tail - 1) * jcp.stride_w
459 + (jcp.kw - 1) * (jcp.dilate_w + 1) - (jcp.iw + jcp.l_pad - 1));
460
461 // kernel needs 1 temporary YMM register
462 const int num_avail_regs = 15;
463 if (r_pad_no_tail > jcp.ur_w * jcp.stride_w && jcp.ow / jcp.ur_w > 1) {
464 /* recalculate ur_w, nb_oc_blocking and ur_w_tail */
465 jcp.ur_w = nstl::min(r_pad_no_tail / jcp.stride_w + jcp.ur_w_tail,
466 nstl::min(jcp.ow, num_avail_regs / 2));
467 jcp.nb_oc_blocking = (num_avail_regs - jcp.ur_w) / jcp.ur_w;
468 jcp.ur_w_tail = jcp.ow % jcp.ur_w;
469 /* check again ... */
470 r_pad_no_tail = nstl::max(0, (jcp.ow - jcp.ur_w_tail - 1) * jcp.stride_w
471 + (jcp.kw - 1) * (jcp.dilate_w + 1) - (jcp.iw + jcp.l_pad - 1));
472 if (jcp.ur_w < nstl::max(jcp.l_pad, r_pad_no_tail))
473 return status::unimplemented;
474 }
475 assert(jcp.nb_oc_blocking > 0);
476 assert(jcp.ur_w * (jcp.nb_oc_blocking + 1) <= num_avail_regs);
477
478 jcp.ic_block = (jcp.ic % simd_w != 0) ? jcp.ic : simd_w;
479 jcp.nb_ic = jcp.ic / jcp.ic_block;
480
481 jcp.oc_block = simd_w;
482 jcp.nb_oc = jcp.oc / jcp.oc_block;
483
484 if (one_of(jcp.prop_kind, forward_training, forward_inference)) {
485 jcp.nb_ic_blocking = 12;
486 jcp.nb_ic_blocking_max = 16;
487 } else {
488 jcp.nb_ic_blocking = 1;
489 jcp.nb_ic_blocking_max = jcp.nb_ic_blocking;
490 }
491
492 return status::success;
493}
494
495}
496}
497}
498