inference-engine/thirdparty/mkl-dnn/src/cpu/jit_transpose_src_utils.cpp

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