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 "type_helpers.hpp"
19#include "nstl.hpp"
20#include "utils.hpp"
21#include "jit_generator.hpp"
22#include "cpu_barrier.hpp"
23
24#include "jit_transpose_src_utils.hpp"
25
26namespace mkldnn {
27namespace impl {
28namespace cpu {
29
30using namespace Xbyak;
31
32#define GET_OFF(x) offsetof(ctx_t, x)
33
34struct jit_trans_iw_ic_t: public jit_trans_src_t, public jit_generator {
35 DECLARE_CPU_JIT_AUX_FUNCTIONS(jit_trans_iw_ic_t)
36
37 jit_trans_iw_ic_t(const jit_conv_conf_t *conf): jit_trans_src_t(conf) {
38 generate();
39 ker_ = (decltype(ker_))this->getCode();
40 }
41
42private:
43 using reg64_t = const Xbyak::Reg64;
44 using reg32_t = const Xbyak::Reg32;
45 using opmask_t = const Xbyak::Opmask;
46
47 enum { typesize = sizeof(float), transpose_size = 16, small_spatial = 14 };
48 int src_stride, tr_src_stride;
49 int tail;
50 bool enable_prefetch;
51
52 opmask_t k3333 = k1;
53 opmask_t k5555 = k2;
54 opmask_t kAAAA = k3;
55 opmask_t kCCCC = k4;
56 opmask_t k0F0F = k5;
57 opmask_t kF0F0 = k6;
58 opmask_t kTail = k7;
59
60 reg64_t reg_src = r8;
61 reg64_t reg_tr_src = r9;
62 reg64_t reg_src_prf = r10;
63 reg64_t reg_tr_src_prf = r11;
64 reg64_t reg_loop = r12;
65 reg64_t reg_tr_src_tmp = r13;
66 reg32_t regw_tmp = r14d;
67
68 void transpose(int nrows, int l_pad, int r_pad, bool nontemporal_stores);
69 void generate();
70};
71
72void jit_trans_iw_ic_t::transpose(int nrows, int l_pad, int r_pad,
73 bool nontemporal_stores) {
74 assert(nrows >= 0 && nrows <= transpose_size);
75 static_assert(transpose_size == 16, "Unsupported transpose size");
76 if (!nrows)
77 return;
78
79 auto pf_src_t0 = [=](int i) {
80 if(enable_prefetch) prefetcht0(EVEX_compress_addr(reg_src,
81 (transpose_size + i) * src_stride));
82 };
83
84 auto pf_tr_src_t0 = [=](int i) {
85 int offset = (transpose_size) * typesize + i * tr_src_stride;
86 if(enable_prefetch) prefetcht0(EVEX_compress_addr(reg_tr_src, offset));
87 if(enable_prefetch) prefetcht0(EVEX_compress_addr(reg_tr_src,
88 offset + 64));
89 };
90
91 auto pf_src_t1 = [=](int i) {
92 if(enable_prefetch) prefetcht1(EVEX_compress_addr(reg_src_prf,
93 i * src_stride));
94 };
95
96 auto pf_tr_src_t1 = [=](int i) {
97 if(enable_prefetch) prefetchwt1(EVEX_compress_addr(reg_tr_src_prf,
98 i * tr_src_stride));
99 };
100
101 auto src_zmm = [=](int i) {
102 assert(i >= 0 && i < 16);
103 return Zmm(i);
104 };
105
106 auto tmp_zmm = [=](int i) {
107 assert(i >= 0 && i < 16);
108 return Zmm(16 + i);
109 };
110
111 auto load = [=](int i) {
112 vmovups(src_zmm(i), EVEX_compress_addr(reg_src, i * src_stride));
113 };
114
115 auto store = [=](Zmm r, int i) {
116 auto kmovw = [=](Opmask k, unsigned w) {
117 mov(regw_tmp, w);
118 jit_generator::kmovw(k, regw_tmp);
119 };
120
121 auto padding = [=] (Reg64 reg, int pad) {
122 kmovw(kTail, (1 << pad) - 1);
123 auto k = kTail;
124 auto base = reg;
125 base.setOpmaskIdx(k.getIdx(), true);
126
127 auto zmm_zero = r;
128 vpxord(zmm_zero, zmm_zero, zmm_zero);
129 auto addr = EVEX_compress_addr(base, i * tr_src_stride);
130 vmovups(addr, zmm_zero);
131 };
132
133 mov(reg_tr_src_tmp, reg_tr_src);
134 if (l_pad > 0)
135 add(reg_tr_src_tmp, l_pad * typesize);
136
137 if (tail != transpose_size)
138 kmovw(kTail, (1 << tail) - 1);
139
140 // Xbyak does not allow k0 to be specified explicitly via the '|'
141 // operator, so we have to do this via a method call (implicitly
142 // EVEX encoding uses k0 to mean 'no mask')
143 bool partial_store = nrows < 16;
144 auto k = partial_store ? kTail : k0;
145 auto base = reg_tr_src_tmp;
146 base.setOpmaskIdx(k.getIdx(), true);
147
148 auto addr = EVEX_compress_addr(base, i * tr_src_stride);
149 if (nontemporal_stores && !partial_store)
150 vmovntps(addr, r);
151 else
152 vmovups(addr, r);
153
154 if (r_pad > 0) {
155 add(reg_tr_src_tmp, tail * typesize);
156 padding(reg_tr_src_tmp, r_pad);
157 }
158
159 if (l_pad > 0) {
160 padding(reg_tr_src, l_pad);
161 }
162 };
163
164 auto transpose16x8 = [=](int base_idx) {
165 assert(base_idx == 0 || base_idx == 8);
166
167 // swap 1
168 for (int i = 0; i < 4; i++) {
169 int src_idx0 = base_idx + i * 2;
170 int src_idx1 = src_idx0 + 1;
171
172 int next_src_idx0 = src_idx0 + 2;
173 int next_src_idx1 = src_idx1 + 2;
174 bool load_next = base_idx == 0 || i < 3;
175
176 if (base_idx == 0 && i == 0) {
177 load(src_idx0);
178 load(src_idx1);
179 }
180
181 auto tmp0 = tmp_zmm(src_idx0);
182 auto tmp1 = tmp_zmm(src_idx1);
183 auto src0 = src_zmm(src_idx0);
184 auto src1 = src_zmm(src_idx1);
185
186 if (next_src_idx0 < nrows && load_next)
187 load(next_src_idx0);
188 valignd(tmp0, src0, src0, 0x1);
189 pf_src_t1(base_idx + i);
190
191 if (next_src_idx1 < nrows && load_next)
192 load(next_src_idx1);
193 valignd(tmp1, src1, src1, 0xf);
194 pf_src_t0(base_idx + i);
195
196 vmovaps(src0 | kAAAA, tmp1);
197 vmovaps(src1 | k5555, tmp0);
198 }
199 // swap 2
200 for (int i = 0; i < 4; i++) {
201 int select_half = (i < 2) ? 0 : 2;
202 int src_idx0 = base_idx + i + select_half + 0;
203 int src_idx2 = src_idx0 + 2;
204
205 auto tmp0 = tmp_zmm(src_idx0);
206 auto tmp1 = tmp_zmm(src_idx2);
207 auto src0 = src_zmm(src_idx0);
208 auto src2 = src_zmm(src_idx2);
209
210 valignd(tmp0, src0, src0, 0x2);
211 pf_src_t1(base_idx + 4 + i);
212 valignd(tmp1, src2, src2, 0xe);
213 pf_src_t0(base_idx + 4 + i);
214 vmovaps(src2 | k3333, tmp0);
215 vmovaps(src0 | kCCCC, tmp1);
216 }
217
218 // swap 4
219 for (int i = 0; i < 4; i++) {
220 int src_idx0 = base_idx + i;
221 int src_idx4 = src_idx0 + 4;
222
223 auto tmp0 = tmp_zmm(src_idx0);
224 auto src0 = src_zmm(src_idx0);
225 auto src4 = src_zmm(src_idx4);
226
227 vmovaps(tmp0, src0);
228 vshuff32x4(src0 | kF0F0, src4, src4, 0xb1);
229 pf_tr_src_t1(base_idx / 2 + i);
230 vshuff32x4(src4 | k0F0F, tmp0, tmp0, 0xb1);
231 pf_tr_src_t0(base_idx / 2 + i);
232 }
233 };
234
235 auto fixup16x16 = [=]() {
236 // swap 8
237 for (int i = 0; i < 8; i++) {
238 auto tmp = tmp_zmm(i);
239 auto src0 = src_zmm(i);
240 auto src8 = src_zmm(8 + i);
241 vshuff64x2(tmp, src0, src8, 0x44);
242 store(tmp, i);
243 if (i % 2 == 0) {
244 pf_tr_src_t1(8 + i / 2);
245 pf_tr_src_t0(8 + i / 2);
246 }
247 }
248
249 for (int i = 0; i < 8; i++) {
250 auto tmp = tmp_zmm(8 + i);
251 auto src0 = src_zmm(i);
252 auto src8 = src_zmm(8 + i);
253 vshuff64x2(tmp, src0, src8, 0xee);
254 store(tmp, 8 + i);
255 if (i % 2 == 0) {
256 pf_tr_src_t1(12 + i / 2);
257 pf_tr_src_t0(12 + i / 2);
258 }
259 }
260 };
261
262 transpose16x8(0);
263 transpose16x8(8);
264 fixup16x16();
265}
266
267void jit_trans_iw_ic_t::generate() {
268 preamble();
269
270 const int ic_block = conf_->ic_block;
271 const int iw = conf_->iw;
272 const int tr_iw = conf_->tr_iw;
273 const int transposes = utils::div_up(iw, transpose_size);
274 int loop_iters = nstl::max(0, transposes - 1);
275 tail = iw - loop_iters * transpose_size;
276
277 src_stride = ic_block * typesize;
278 assert(src_stride == 64);
279 tr_src_stride = tr_iw * typesize;
280
281 bool nontemporal_stores = false;
282 enable_prefetch = iw > small_spatial ? 1 : 0;
283
284 assert(transpose_size == ic_block);
285 const int src_step = ic_block * transpose_size * typesize;
286 const int tr_src_step = ic_block * typesize;
287
288 const int left_pad = conf_->l_pad;
289 const int right_pad = tr_iw - iw - left_pad;
290
291 mov(reg_src, ptr [param1 + GET_OFF(src)]);
292 mov(reg_tr_src, ptr [param1 + GET_OFF(tr_src)]);
293 mov(reg_src_prf, ptr [param1 + GET_OFF(src_prf)]);
294 mov(reg_tr_src_prf, ptr [param1 + GET_OFF(tr_src_prf)]);
295
296 auto kmovw = [=](Opmask k, unsigned w) {
297 mov(regw_tmp, w);
298 jit_generator::kmovw(k, regw_tmp);
299 };
300
301 kmovw(k3333, 0x3333); // 0011001100110011
302 kmovw(k5555, 0x5555); // 0101010101010101
303 kmovw(kAAAA, 0xaaaa); // 1010101010101010
304 kmovw(kCCCC, 0xcccc); // 1100110011001100
305 kmovw(k0F0F, 0x0f0f); // 0000111100001111
306 kmovw(kF0F0, 0xf0f0); // 1111000011110000
307
308 if (left_pad > 0 && loop_iters > 0) {
309 loop_iters--;
310 transpose(transpose_size, left_pad, 0, nontemporal_stores);
311 add(reg_src, src_step);
312 add(reg_tr_src, tr_src_step + left_pad * typesize);
313 add(reg_src_prf, src_step);
314 add(reg_tr_src_prf, tr_src_step + left_pad * typesize);
315 }
316
317 if (loop_iters) {
318 mov(reg_loop, loop_iters);
319 Label loop;
320 L(loop); {
321 transpose(transpose_size, 0, 0, nontemporal_stores);
322 add(reg_src, src_step);
323 add(reg_tr_src, tr_src_step);
324 add(reg_src_prf, src_step);
325 add(reg_tr_src_prf, tr_src_step);
326 sub(reg_loop, 1);
327 jnz(loop);
328 }
329 }
330 if (transposes > 1)
331 transpose(tail, 0, right_pad, nontemporal_stores);
332 else
333 transpose(tail, left_pad, right_pad, nontemporal_stores);
334
335 postamble();
336}
337
338struct jit_trans_iw_ic_int16_t: public jit_trans_src_t, public jit_generator {
339 DECLARE_CPU_JIT_AUX_FUNCTIONS(jit_trans_iw_ic_int16_t)
340 jit_trans_iw_ic_int16_t(const jit_conv_conf_t *conf):
341 jit_trans_src_t(conf) {
342 generate();
343 ker_ = (decltype(ker_))this->getCode();
344 }
345
346private:
347 using reg64_t = const Xbyak::Reg64;
348 using reg32_t = const Xbyak::Reg32;
349 using opmask_t = const Xbyak::Opmask;
350
351 enum { typesize = sizeof(int16_t), transpose_size = 16, small_spatial = 14 };
352 int src_stride, tr_src_stride;
353 int tail;
354 bool enable_prefetch;
355
356 opmask_t kFFFF = k1;
357 opmask_t k5555 = k2;
358 opmask_t kAAAA = k3;
359 opmask_t kAA = k4;
360 opmask_t k55 = k5;
361 opmask_t kCC = k6;
362 opmask_t k33 = k7;
363 opmask_t kTail = k1;
364
365 reg64_t reg_src = r8;
366 reg64_t reg_tr_src = r9;
367 reg64_t reg_src_prf = r10;
368 reg64_t reg_tr_src_prf = r11;
369 reg64_t reg_loop = r12;
370 reg64_t reg_tr_src_tmp = r13;
371 reg32_t regw_tmp = r14d;
372 reg64_t imm_addr64 = rbx;
373
374 Xbyak::Zmm vidx1 = zmm31;
375 Xbyak::Zmm vidx2 = zmm30;
376 Xbyak::Zmm vidx3 = zmm29;
377 Xbyak::Zmm vidx4 = zmm28;
378 Xbyak::Zmm vidx5 = zmm27;
379 Xbyak::Zmm zmm_tmp = zmm26;
380
381
382 void transpose(int nrows, int l_pad, int r_pad, bool nontemporal_stores);
383 void generate();
384};
385
386void jit_trans_iw_ic_int16_t::transpose(int nrows, int l_pad, int r_pad,
387 bool nontemporal_stores) {
388 assert(nrows >= 0 && nrows <= transpose_size);
389 static_assert(transpose_size == 16, "Unsupported transpose size");
390 if (!nrows)
391 return;
392
393 auto src_zmm = [=](int i) {
394 return Zmm(i);
395 };
396
397 auto src_ymm = [=](int i) {
398 assert(i >= 0 && i < 16);
399 return Ymm(i);
400 };
401
402 auto load_ymm = [=](int i) {
403 vmovups(src_ymm(i), EVEX_compress_addr(reg_src, i * src_stride));
404 };
405
406 auto kmovw = [=](Opmask k, unsigned w) {
407 mov(regw_tmp, w);
408 jit_generator::kmovw(k, regw_tmp);
409 };
410
411 auto store = [=](Zmm r, int i) {
412
413 auto padding = [=] (Reg64 reg, int pad) {
414 kmovw(kTail, (1 << pad) - 1);
415 auto k = kTail;
416 auto base = reg;
417 base.setOpmaskIdx(k.getIdx(), true);
418
419 auto zmm_zero = zmm_tmp;
420 vpxord(zmm_zero, zmm_zero, zmm_zero);
421 auto addr = EVEX_compress_addr(base, i * tr_src_stride);
422 vmovups(addr, zmm_zero);
423 };
424
425 int store_tail = (nrows%2) ? nrows+1 : nrows;
426
427 int store_pad = (l_pad%2) ? l_pad/2 + 1 : l_pad/2;
428 mov(reg_tr_src_tmp, reg_tr_src);
429 if (l_pad > 0) {
430 padding(reg_tr_src, store_pad);
431 add(reg_tr_src_tmp, l_pad * typesize);
432 }
433 if (r_pad > 0) {
434 store_pad = (r_pad%2) ? r_pad/2 + 1 : r_pad/2;
435 int addr_shift = (r_pad%2) ? 1 : 0;
436 add(reg_tr_src_tmp, (nrows - addr_shift) * typesize);
437 padding(reg_tr_src_tmp, store_pad);
438 }
439
440 mov(reg_tr_src_tmp, reg_tr_src);
441 add(reg_tr_src_tmp, l_pad * typesize);
442
443 kmovw(kTail, (1 << store_tail/2) - 1);
444 auto k = kTail;
445 auto base = reg_tr_src_tmp;
446 base.setOpmaskIdx(k.getIdx(), true);
447
448 auto addr = EVEX_compress_addr(base, i * tr_src_stride);
449 vmovups(addr, r);
450
451 };
452
453 kmovw(kFFFF, 0xffff);
454 //all loads
455 for (int i=0; i<16; i++){
456 vpxord(src_zmm(i), src_zmm(i), src_zmm(i));
457 }
458
459 for (int i = 0; i < nrows/2; i++) {
460 auto src0 = src_ymm(2*i);
461 auto src1 = src_ymm(2*i+1);
462 auto zmm_src0 = src_zmm(2*i);
463 load_ymm(2*i);
464
465 vpunpcklwd(src1, src0,
466 EVEX_compress_addr(reg_src, (2*i+1) * src_stride));
467 vpunpckhwd(src0, src0,
468 EVEX_compress_addr(reg_src, (2*i+1) * src_stride));
469 vinserti64x4(zmm_src0, zmm_src0, src1, 1);
470 vpermps(zmm_src0 | kFFFF, vidx4, zmm_src0);
471 }
472
473 // for odd numbers we need to mix row with zeroes
474 if (nrows%2) {
475 int i = nrows-1;
476 auto src0 = src_ymm(i);
477 auto src1 = src_ymm(i+1); //zero
478
479 auto zmm_src0 = src_zmm(i);
480 vpxor(src1, src1, src1);
481
482 load_ymm(i);
483 vpunpckhwd(src0, src0, src1);
484 vinserti64x4(zmm_tmp, zmm_tmp, src0, 0);
485 vpxor(src0, src0, src0);
486 load_ymm(i);
487 vpunpcklwd(src1, src0, src1);
488 vinserti64x4(zmm_tmp, zmm_tmp, src1, 1);
489 vpxord(zmm_src0, zmm_src0, zmm_src0);
490 vmovups(zmm_src0, zmm_tmp);
491 vpermps(zmm_src0 | kFFFF, vidx4, zmm_src0);
492 }
493
494 // swap 1
495 for (int i=0; i<4; i++) {
496 auto zmm0 = src_zmm(4*i);
497 auto zmm1 = src_zmm(4*i+2);
498 auto tmp0 = src_zmm(4*i+1);
499 auto tmp1 = src_zmm(4*i+3);
500
501 vmovups(tmp0, zmm0);
502 vmovups(tmp1, zmm1);
503
504 vpermps(tmp0 | kAAAA, vidx3, zmm1);
505 vpermps(tmp1 | k5555, vidx3, zmm0);
506 }
507 // swap 2
508 int base_idx;
509 base_idx=0;
510 for (int i=0; i<2; i++) {
511 auto zmm0 = src_zmm(base_idx+2*i+1);
512 auto zmm1 = src_zmm(base_idx+2*i+5);
513
514 auto tmp0 = src_zmm(base_idx+2*i);
515 auto tmp1 = src_zmm(base_idx+2*i+4);
516
517 vmovupd(tmp0, zmm0);
518 vmovupd(tmp1, zmm1);
519
520 vpermpd(tmp0 | kAA, vidx2, zmm1);
521 vpermpd(tmp1 | k55, vidx2, zmm0);
522 }
523 base_idx=8;
524 for (int i=0; i<2; i++) {
525 auto zmm0 = src_zmm(base_idx+2*i+1);
526 auto zmm1 = src_zmm(base_idx+2*i+5);
527
528 auto tmp0 = src_zmm(base_idx+2*i);
529 auto tmp1 = src_zmm(base_idx+2*i+4);
530
531 vmovupd(tmp0, zmm0);
532 vmovupd(tmp1, zmm1);
533
534 vpermpd(tmp0 | kAA, vidx2, zmm1);
535 vpermpd(tmp1 | k55, vidx2, zmm0);
536 }
537
538 // swap 3
539 for (int i=0; i<4; i++) {
540 auto zmm0 = src_zmm(2*i);
541 auto zmm1 = src_zmm(2*i+8);
542
543 auto tmp0 = src_zmm(2*i+1);
544 auto tmp1 = src_zmm(2*i+9);
545
546 vmovupd(tmp0, zmm0);
547 vmovupd(tmp1, zmm1);
548
549 vpermpd(tmp0 | kCC, vidx1, zmm1);
550 vpermpd(tmp1 | k33, vidx1, zmm0);
551 }
552
553 // all stores
554 for (int i=0; i<8; i++)
555 vextracti64x4(src_ymm(2*i), src_zmm(2*i+1), 1);
556
557 store(src_zmm(1), 0);
558 store(src_zmm(0), 1);
559 store(src_zmm(3), 2);
560 store(src_zmm(2), 3);
561 store(src_zmm(9), 4);
562 store(src_zmm(8), 5);
563 store(src_zmm(11), 6);
564 store(src_zmm(10), 7);
565 store(src_zmm(5), 8);
566 store(src_zmm(4), 9);
567 store(src_zmm(7), 10);
568 store(src_zmm(6), 11);
569 store(src_zmm(13), 12);
570 store(src_zmm(12), 13);
571 store(src_zmm(15), 14);
572 store(src_zmm(14), 15);
573
574}
575
576void jit_trans_iw_ic_int16_t::generate() {
577 preamble();
578
579 alignas(64) static constexpr const int64_t idx1[8]
580 = { 2, 3, 0, 1, 6, 7, 4, 5 };
581 alignas(64) static constexpr const int64_t idx2[8]
582 = { 1, 0, 3, 2, 5, 4, 7, 6 };
583 alignas(64) static constexpr const int32_t idx3[16]
584 = { 1, 0, 3, 2, 5, 4, 7, 6, 9, 8, 11, 10, 13, 12, 15, 14 };
585 alignas(64) static constexpr const int32_t idx4[16]
586 = { 8, 10, 12, 14, 0, 2, 4, 6, 9, 11, 13, 15, 1, 3, 5, 7 };
587 alignas(64) static constexpr const int32_t idx5[16]
588 = { 8, 10, 12, 14, 0, 2, 4, 6, 9, 11, 13, 15, 1, 3, 5, 7 };
589
590 const int ic_block = conf_->ic_block;
591 const int iw = conf_->iw;
592 const int tr_iw = conf_->tr_iw;
593 const int transposes = utils::div_up(iw, transpose_size);
594 int loop_iters = nstl::max(0, transposes - 1);
595 tail = iw - loop_iters * transpose_size;
596
597 src_stride = ic_block * typesize;
598 tr_src_stride = tr_iw * typesize;
599
600 bool nontemporal_stores = false;
601 enable_prefetch = iw > small_spatial ? 1 : 0;
602
603 assert(transpose_size == ic_block);
604 const int src_step = ic_block * transpose_size * typesize;
605 const int tr_src_step = ic_block * typesize;
606
607 const int left_pad = conf_->l_pad;
608 const int right_pad = tr_iw - iw - left_pad;
609
610 mov(reg_src, ptr [param1 + GET_OFF(src)]);
611 mov(reg_tr_src, ptr [param1 + GET_OFF(tr_src)]);
612 mov(reg_src_prf, ptr [param1 + GET_OFF(src_prf)]);
613 mov(reg_tr_src_prf, ptr [param1 + GET_OFF(tr_src_prf)]);
614
615 auto kmovw = [=](Opmask k, unsigned w) {
616 mov(regw_tmp, w);
617 jit_generator::kmovw(k, regw_tmp);
618 };
619
620 kmovw(kFFFF, 0xffff);
621 kmovw(k5555, 0x5555);
622 kmovw(kAAAA, 0xaaaa);
623 kmovw(kAA, 0xaa);
624 kmovw(k55, 0x55);
625 kmovw(kCC, 0xcc);
626 kmovw(k33, 0x33);
627
628 auto vmovdqa64 = [=](Zmm z, const int64_t *addr) {
629 mov(imm_addr64, reinterpret_cast<size_t>(addr));
630 jit_generator::vmovdqa64(z, ptr[imm_addr64]);
631 };
632
633 auto vmovdqa32 = [=](Zmm z, const int32_t *addr) {
634 mov(imm_addr64, reinterpret_cast<size_t>(addr));
635 jit_generator::vmovdqa32(z, ptr[imm_addr64]);
636 };
637
638 vmovdqa64(vidx1, idx1);
639 vmovdqa64(vidx2, idx2);
640 vmovdqa32(vidx3, idx3);
641 vmovdqa32(vidx4, idx4);
642 vmovdqa32(vidx5, idx5);
643
644 if (left_pad > 0 && loop_iters > 0) {
645 loop_iters--;
646 transpose(transpose_size, left_pad, 0, nontemporal_stores);
647 add(reg_src, src_step);
648 add(reg_tr_src, tr_src_step + left_pad * typesize);
649 add(reg_src_prf, src_step);
650 add(reg_tr_src_prf, tr_src_step + left_pad * typesize);
651 }
652
653 if (loop_iters) {
654 mov(reg_loop, loop_iters);
655 Label loop;
656 L(loop); {
657 transpose(transpose_size, 0, 0, nontemporal_stores);
658 add(reg_src, src_step);
659 add(reg_tr_src, tr_src_step);
660 add(reg_src_prf, src_step);
661 add(reg_tr_src_prf, tr_src_step);
662 sub(reg_loop, 1);
663 jnz(loop);
664 }
665 }
666 if (transposes > 1)
667 transpose(tail, 0, right_pad, nontemporal_stores);
668 else
669 transpose(tail, left_pad, right_pad, nontemporal_stores);
670
671 postamble();
672
673}
674
675struct jit_trans_ow_oc_t: public jit_trans_dst_t, public jit_generator {
676 DECLARE_CPU_JIT_AUX_FUNCTIONS(jit_trans_ow_oc_t)
677 jit_trans_ow_oc_t(const jit_conv_conf_t *conf): jit_trans_dst_t(conf) {
678 generate();
679 ker_ = (decltype(ker_))this->getCode();
680 }
681
682private:
683 using reg64_t = const Xbyak::Reg64;
684 using reg32_t = const Xbyak::Reg32;
685 using opmask_t = const Xbyak::Opmask;
686 using zmm = const Xbyak::Zmm;
687
688 enum { typesize = sizeof(int16_t), transpose_size = 16, small_spatial = 14 };
689 int src_stride, tr_src_stride;
690 int tail;
691 bool enable_prefetch;
692
693 opmask_t kFF = k1;
694
695 zmm vidx1 = zmm31;
696
697 reg64_t reg_src = r8;
698 reg64_t reg_tr_src = r9;
699 reg64_t reg_src_prf = r10;
700 reg64_t reg_tr_src_prf = r11;
701 reg64_t reg_loop = r12;
702 reg64_t reg_tr_src_tmp = r13;
703 reg32_t regw_tmp = r14d;
704 reg64_t imm_addr64 = rbx;
705
706 void transpose(int nrows, int l_pad, int r_pad, bool nontemporal_stores);
707 void generate();
708};
709
710void jit_trans_ow_oc_t::transpose(int nrows, int l_pad, int r_pad,
711 bool nontemporal_stores) {
712 assert(nrows >= 0 && nrows <= transpose_size);
713 static_assert(transpose_size == 16, "Unsupported transpose size");
714 if (!nrows)
715 return;
716
717 auto src_zmm = [=](int i) {
718 return Zmm(i);
719 };
720
721 auto src_ymm = [=](int i) {
722 assert(i >= 0 && i < 16);
723 return Ymm(i);
724 };
725
726 auto load_ymm = [=](int i) {
727 vmovups(src_ymm(i), EVEX_compress_addr(reg_src, i * src_stride));
728 };
729
730
731 auto store = [=](Zmm r, int i) {
732 auto addr = EVEX_compress_addr(reg_tr_src, i * tr_src_stride);
733 if (nontemporal_stores)
734 vmovntps(addr, r);
735 else
736 vmovups(addr, r);
737 };
738
739 for (int i = 0; i < nrows/2; i++) {
740 auto src0 = src_ymm(2*i);
741 auto src1 = src_ymm(2*i+1);
742 auto zmm_src0 = src_zmm(2*i);
743 load_ymm(2*i);
744 vpunpcklwd(src1, src0,
745 EVEX_compress_addr(reg_src, (2*i+1) * src_stride));
746 vpunpckhwd(src0, src0,
747 EVEX_compress_addr(reg_src, (2*i+1) * src_stride));
748 vinserti64x4(zmm_src0, zmm_src0, src1, 1);
749 vpermpd(zmm_src0 | kFF, vidx1, zmm_src0);
750 store(zmm_src0, 2*i);
751 }
752 if (r_pad > 0) {
753 auto src0 = src_ymm(nrows-1);
754 auto src1 = src_ymm(nrows);
755 auto zmm_src0 = src_zmm(30);
756 load_ymm(nrows-1);
757
758 vpxor(src1, src1, src1);
759 vpunpckhwd(src1, src0, src1);
760 vinserti64x4(zmm_src0, zmm_src0, src1, 0);
761 vpxor(src1, src1, src1);
762 vpunpcklwd(src0, src0, src1);
763 vinserti64x4(zmm_src0, zmm_src0, src0, 1);
764 vpermpd(zmm_src0 | kFF, vidx1, zmm_src0);
765 store(zmm_src0, nrows-1);
766 }
767}
768
769void jit_trans_ow_oc_t::generate() {
770 preamble();
771
772 alignas(64) static constexpr const int64_t idx1[8]
773 = { 4, 5, 0, 1, 6, 7, 2, 3 };
774
775 const int oc_block = conf_->oc_block;
776 const int ow = conf_->ow;
777 const int transposes = utils::div_up(ow, transpose_size);
778 int loop_iters = nstl::max(0, transposes - 1);
779 tail = ow - loop_iters * transpose_size;
780
781 src_stride = oc_block * typesize;
782 tr_src_stride = oc_block * typesize;
783
784 bool nontemporal_stores = false;
785 enable_prefetch = ow > small_spatial ? 1 : 0;
786
787 const int src_step = oc_block * transpose_size * typesize;
788 const int tr_src_step = oc_block * transpose_size * typesize;
789 const int right_pad = ow % 2;
790
791 mov(reg_src, ptr [param1 + GET_OFF(src)]);
792 mov(reg_tr_src, ptr [param1 + GET_OFF(tr_src)]);
793 mov(reg_src_prf, ptr [param1 + GET_OFF(src_prf)]);
794 mov(reg_tr_src_prf, ptr [param1 + GET_OFF(tr_src_prf)]);
795
796 auto kmovw = [=](Opmask k, unsigned w) {
797 mov(regw_tmp, w);
798 jit_generator::kmovw(k, regw_tmp);
799 };
800
801 kmovw(kFF, 0xFF);
802
803 auto vmovdqa64 = [=](Zmm z, const int64_t *addr) {
804 mov(imm_addr64, reinterpret_cast<size_t>(addr));
805 jit_generator::vmovdqa64(z, ptr[imm_addr64]);
806 };
807
808 vmovdqa64(vidx1, idx1);
809 if (loop_iters) {
810 mov(reg_loop, loop_iters);
811 Label loop;
812 L(loop); {
813 transpose(transpose_size, 0, 0, nontemporal_stores);
814 add(reg_src, src_step);
815 add(reg_tr_src, tr_src_step);
816 add(reg_src_prf, src_step);
817 add(reg_tr_src_prf, tr_src_step);
818 sub(reg_loop, 1);
819 jnz(loop);
820 }
821 }
822 transpose(tail, 0, right_pad, nontemporal_stores);
823
824 postamble();
825}
826
827struct jit_trans_iw_x4_4x_t: public jit_trans_src_t, public jit_generator {
828 DECLARE_CPU_JIT_AUX_FUNCTIONS(jit_trans_iw_x4_4x_t)
829
830 jit_trans_iw_x4_4x_t(const jit_conv_conf_t *conf): jit_trans_src_t(conf) {
831 generate();
832 ker_ = (decltype(ker_))this->getCode();
833 }
834
835 void generate();
836 enum { typesize = (int)sizeof(float) };
837};
838
839/** @brief transposition of the form [:][iw/4][4] -> [:][4][iw/4]
840 * required for 1st 4fma backward by weights convolution */
841void jit_trans_iw_x4_4x_t::generate() {
842 using namespace utils;
843
844 /* TODO: put into code */
845 static int mask[16] = {
846 0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15, };
847
848 const auto &c = *conf_;
849 const int simd_w = cpu_isa_traits<avx512_common>::vlen / typesize;
850 const int niters = c.tr_ld / simd_w;
851
852 assert(niters <= 4); /* [bwd_w:tr_src:r1] */
853
854 Reg64 reg_ptr_src = r8;
855 Reg64 reg_ptr_tr_src = r9;
856
857 Reg64 reg_ih = rax;
858 Reg64 reg_ih_end = rbx;
859
860 Reg64 reg_nthr_oc_b = rsi;
861 Reg64 reg_ptr_tr_src_bctx = abi_not_param1;
862
863 Reg64 reg_tmp = rdx;
864
865 Zmm vmsk = Zmm(31);
866 Opmask kmsk = k7;
867
868 auto emit_tr_sync = [&]() {
869 simple_barrier::generate(*this, reg_ptr_tr_src_bctx, reg_nthr_oc_b);
870 };
871
872 auto emit_tr_iw = [&]() {
873 auto vreg = [](int iter, int i) {
874 assert(4 * iter + i < 24);
875 return Zmm(4 * iter + i);
876 };
877 auto vtmp = [](int i) { return Zmm(24 + i); };
878
879 auto emit_load = [&](int iter) {
880 for (int i = 0; i < 4; ++i) {
881 auto v = vreg(iter, i);
882 const int off = (iter * 4 + i) * simd_w;
883
884 if (off + simd_w <= c.iw)
885 vmovups(v, ptr[reg_ptr_src + off * typesize]);
886 else if (off < c.iw)
887 vmovups(v | kmsk | T_z, ptr[reg_ptr_src + off * typesize]);
888 else
889 vpxord(v, v, v);
890 }
891 };
892
893 auto emit_tr = [&](int iter) {
894 for (int i = 0; i < 4; ++i)
895 vpermps(vreg(iter, i), vmsk, vreg(iter, i));
896
897 vshuff32x4(vtmp(0), vreg(iter, 0), vreg(iter, 1), 0x88);
898 vshuff32x4(vtmp(1), vreg(iter, 0), vreg(iter, 1), 0xdd);
899 vshuff32x4(vtmp(2), vreg(iter, 2), vreg(iter, 3), 0x88);
900 vshuff32x4(vtmp(3), vreg(iter, 2), vreg(iter, 3), 0xdd);
901
902 vshuff32x4(vreg(iter, 0), vtmp(0), vtmp(2), 0x88);
903 vshuff32x4(vreg(iter, 2), vtmp(0), vtmp(2), 0xdd);
904 vshuff32x4(vreg(iter, 1), vtmp(1), vtmp(3), 0x88);
905 vshuff32x4(vreg(iter, 3), vtmp(1), vtmp(3), 0xdd);
906 };
907
908 auto emit_store = [&]() {
909 for (int i = 0; i < 4; ++i) {
910 for (int iter = 0; iter < niters; ++iter) {
911 const size_t off = i * c.tr_ld + iter * simd_w;
912 vmovups(ptr[reg_ptr_tr_src + off * typesize], vreg(iter, i));
913 }
914 }
915 };
916
917 for (int iter = 0; iter < niters; ++iter)
918 emit_load(iter);
919
920 for (int iter = 0; iter < niters; ++iter)
921 emit_tr(iter);
922
923 emit_store();
924 };
925
926 preamble();
927
928 mov(reg_ptr_src, ptr[abi_param1 + GET_OFF(src)]);
929 mov(reg_ptr_tr_src, ptr[abi_param1 + GET_OFF(tr_src)]);
930
931 mov(reg_nthr_oc_b.cvt32(), ptr[abi_param1 + GET_OFF(nthr_oc_b)]);
932 mov(reg_ih.cvt32(), ptr[abi_param1 + GET_OFF(tr_src_ih_start)]);
933 mov(reg_ih_end.cvt32(), ptr[abi_param1 + GET_OFF(tr_src_ih_end)]);
934 mov(reg_ptr_tr_src_bctx, ptr[abi_param1 + GET_OFF(tr_src_bctx)]);
935
936 emit_tr_sync();
937
938 Label l_ih_loop, l_tr_done;
939 cmp(reg_ih, reg_ih_end);
940 je(l_tr_done, T_NEAR);
941
942 mov(reg_tmp, (size_t)&mask[0]);
943 vmovups(vmsk, ptr[reg_tmp]);
944
945 if (c.iw % simd_w) {
946 const char load_mask = (1 << (c.iw % simd_w)) - 1;
947 mov(reg_tmp, load_mask);
948 kmovw(kmsk, reg_tmp.cvt32());
949 }
950
951 /* src += ih_start * c.iw; */
952 imul(reg_tmp, reg_ih, c.iw * typesize);
953 add(reg_ptr_src, reg_tmp);
954 /* tr_src += ih_start * c.stride_w * c.tr_ld; */
955 imul(reg_tmp, reg_ih, c.stride_w * c.tr_ld * typesize);
956 add(reg_ptr_tr_src, reg_tmp);
957
958 L(l_ih_loop); {
959 emit_tr_iw();
960
961 add(reg_ptr_src, c.iw * typesize);
962 add(reg_ptr_tr_src, c.stride_w * c.tr_ld * typesize);
963
964 inc(reg_ih);
965 cmp(reg_ih, reg_ih_end);
966 jl(l_ih_loop, T_NEAR);
967 }
968
969 L(l_tr_done);
970
971 emit_tr_sync();
972
973 postamble();
974}
975
976/*
977// -------------------------------------------------
978// jit_transpose4x16_src
979// -------------------------------------------------
980*/
981
982void jit_transpose4x16_src::transpose(int nrows)
983{
984 assert(nrows >= 0 && nrows <= transpose_size);
985 static_assert(transpose_size == 4, "Unsupported transpose size");
986 if (!nrows)
987 return;
988
989 auto pf_src_t0 = [=](int i) {
990 if (tparams->src_pf0_distance)
991 prefetcht0(EVEX_compress_addr(
992 reg_src, (tparams->src_pf0_distance + i) * src_stride));
993 };
994
995 auto pf_tr_src_t0 = [=](int i) {
996 if (tparams->tr_src_pf0_distance)
997 prefetcht0(EVEX_compress_addr(reg_tr_src,
998 (tparams->tr_src_pf0_distance + i) * src_stride));
999 };
1000
1001 auto pf_src_t1 = [=](int i) {
1002 if (tparams->src_pf1)
1003 prefetcht1(EVEX_compress_addr(reg_src_prf, i * src_stride));
1004 };
1005
1006 auto pf_tr_src_t1 = [=](int i) {
1007 if (tparams->tr_src_pf1)
1008 prefetchwt1(EVEX_compress_addr(reg_tr_src_prf, i * tr_src_stride));
1009 };
1010
1011 auto src_zmm = [=](int i) {
1012 assert(i >= 0 && i < 4);
1013 return Zmm(i);
1014 };
1015
1016 auto tmp_zmm = [=](int i) {
1017 assert(i >= 0 && i < 4);
1018 return Zmm(4 + i);
1019 };
1020
1021 auto load = [=](int i) {
1022 vmovups(src_zmm(i), EVEX_compress_addr(reg_src, i * src_stride));
1023 };
1024
1025 auto store = [=](Zmm r, int i) {
1026 vmovups(EVEX_compress_addr(reg_tr_src, i * tr_src_stride), r);
1027 };
1028
1029 auto tmp0 = tmp_zmm(0);
1030 auto tmp1 = tmp_zmm(1);
1031 auto tmp2 = tmp_zmm(2);
1032 auto tmp3 = tmp_zmm(3);
1033
1034 auto src0 = src_zmm(0);
1035 auto src1 = src_zmm(1);
1036 auto src2 = src_zmm(2);
1037 auto src3 = src_zmm(3);
1038 for (int i = 0; i < nrows; i++) {
1039 load(i);
1040 }
1041
1042 for (size_t i = nrows; i < 4; i++) {
1043 vpxord(src_zmm(i), src_zmm(i), src_zmm(i));
1044 }
1045
1046 vmovupd(tmp0, src0);
1047 vmovupd(tmp1, src1);
1048 pf_src_t0(0);
1049 vpermpd(tmp0 | kF0, vidx01, src2);
1050 vpermpd(tmp1 | kF0, vidx01, src3);
1051
1052 valignd(src0, src0, src0, 8);
1053 valignd(src1, src1, src1, 8);
1054 pf_src_t0(1);
1055 vmovupd(tmp2, src0);
1056 vmovupd(tmp3, src1);
1057 pf_src_t0(2);
1058 vpermpd(tmp2 | kF0, vidx10, src2);
1059 vpermpd(tmp3 | kF0, vidx10, src3);
1060 pf_src_t0(3);
1061
1062 vmovupd(src0, tmp0);
1063 pf_src_t1(0);
1064 vmovupd(src1, tmp2);
1065 pf_src_t1(1);
1066 vmovupd(src2, tmp1);
1067 pf_src_t1(2);
1068 vmovupd(src3, tmp3);
1069 pf_src_t1(3);
1070 vpermpd(src0 | kCC, vidx1, tmp1);
1071 vpermpd(src1 | kCC, vidx1, tmp3);
1072 pf_tr_src_t0(0);
1073 vpermpd(src2 | k33, vidx1, tmp0);
1074 vpermpd(src3 | k33, vidx1, tmp2);
1075 pf_tr_src_t0(1);
1076
1077 vmovupd(tmp0, src0);
1078 vmovupd(tmp1, src2);
1079 pf_tr_src_t0(2);
1080 vmovupd(tmp2, src1);
1081 vmovupd(tmp3, src3);
1082 pf_tr_src_t0(3);
1083 vpermps(tmp0 | kFFFF, vidxP, src0);
1084 pf_tr_src_t1(0);
1085 vpermps(tmp1 | kFFFF, vidxP, src2);
1086 pf_tr_src_t1(1);
1087 vpermps(tmp2 | kFFFF, vidxP, src1);
1088 pf_tr_src_t1(3);
1089 vpermps(tmp3 | kFFFF, vidxP, src3);
1090 pf_tr_src_t1(4);
1091
1092 store(tmp0, 0);
1093 store(tmp1, 1);
1094 store(tmp2, 2);
1095 store(tmp3, 3);
1096}
1097
1098alignas(64) static constexpr const int64_t idx01[8]
1099 = { 0, 0, 0, 0, 0, 1, 2, 3 };
1100alignas(64) static constexpr const int64_t idx10[8]
1101 = { 0, 0, 0, 0, 4, 5, 6, 7 };
1102alignas(64) static constexpr const int64_t idx1[8] = { 2, 3, 0, 1, 6, 7, 4, 5 };
1103alignas(64) static constexpr const int32_t idxP[16]
1104 = { 0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15 };
1105
1106void jit_transpose4x16_src::generate()
1107{
1108 preamble();
1109
1110 const int ic_block = params->ic_block;
1111 const int is = params->is;
1112 int tail = is % transpose_size;
1113
1114 src_stride = ic_block * typesize;
1115 assert(src_stride == 64);
1116 tr_src_stride = ic_block * typesize;
1117
1118 const int src_step = ic_block * transpose_size * typesize;
1119 const int tr_src_step = ic_block * transpose_size * typesize;
1120
1121#define GET_TR_OFF(x) offsetof(jit_src_transpose_s, x)
1122 mov(reg_loop, ptr[param1 + GET_TR_OFF(size)]);
1123 mov(reg_src, ptr[param1 + GET_TR_OFF(src)]);
1124 mov(reg_tr_src, ptr[param1 + GET_TR_OFF(tr_src)]);
1125 mov(reg_src_prf, ptr[param1 + GET_TR_OFF(src_prf)]);
1126 mov(reg_tr_src_prf, ptr[param1 + GET_TR_OFF(tr_src_prf)]);
1127#undef GET_TR_OFF
1128
1129 auto kmovw = [=](Opmask k, unsigned w) {
1130 mov(regw_tmp, w);
1131 jit_generator::kmovw(k, regw_tmp);
1132 };
1133
1134 auto vmovdqa64 = [=](Zmm z, const int64_t *addr) {
1135 mov(imm_addr64, reinterpret_cast<size_t>(addr));
1136 jit_generator::vmovdqa64(z, ptr[imm_addr64]);
1137 };
1138
1139 auto vmovdqa32 = [=](Zmm z, const int32_t *addr) {
1140 mov(imm_addr64, reinterpret_cast<size_t>(addr));
1141 jit_generator::vmovdqa32(z, ptr[imm_addr64]);
1142 };
1143
1144 kmovw(kF0, 0xf0); // 11110000
1145 kmovw(kCC, 0xcc); // 11001100
1146 kmovw(k33, 0x33); // 00110011
1147 kmovw(kFFFF, 0xffff); // 1111111111111111
1148
1149 vmovdqa64(vidx01, idx01);
1150 vmovdqa64(vidx10, idx10);
1151 vmovdqa64(vidx1, idx1);
1152 vmovdqa32(vidxP, idxP);
1153
1154 Label loop_label;
1155 Label tail_label;
1156
1157 cmp(reg_loop, transpose_size);
1158 jl(tail_label, T_NEAR);
1159
1160 L(loop_label);
1161 {
1162 transpose(transpose_size);
1163 add(reg_src, src_step);
1164 add(reg_tr_src, tr_src_step);
1165 add(reg_src_prf, src_step);
1166 add(reg_tr_src_prf, tr_src_step);
1167 sub(reg_loop, transpose_size);
1168 cmp(reg_loop, transpose_size);
1169 jge(loop_label, T_NEAR);
1170 }
1171 L(tail_label);
1172 transpose(tail);
1173
1174 postamble();
1175}
1176
1177jit_trans_src_t *create_trans_src(const jit_conv_conf_t *conf) {
1178 if (conf->ver == ver_4fma && !conf->is_1stconv)
1179 return new jit_trans_iw_ic_t(conf);
1180 if (conf->ver == ver_4fma && conf->is_1stconv)
1181 return new jit_trans_iw_x4_4x_t(conf);
1182 assert(!"unsupported configuration");
1183 return nullptr;
1184}
1185
1186jit_trans_dst_t *create_trans_dst(const jit_conv_conf_t *conf) {
1187 assert(!"unsupported configuration");
1188 return nullptr;
1189}
1190}
1191}
1192}
1193