OpenI
/
MegEngine
Not watched
1
0
Fork 0
Code
Releases 31 Wiki Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: c85631aa77
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'c85631aa77'
${ noResults }
MegEngine/dnn/src/common/local/local_def.inl

450 lines
21 kB
Raw Blame History 

  1. /**

  2.  * \file dnn/src/common/local/local_def.inl

  3.  * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")

  4.  *

  5.  * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.

  6.  *

  7.  * Unless required by applicable law or agreed to in writing,

  8.  * software distributed under the License is distributed on an

  9.  * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

  10.  */

  11. // simd_macro/*_helper.h should be included before including this file.

  12. //

  13. // The following functions would be defined in this file:

  14. //

  15. // void local_xcorr_MEGDNN_SIMD_NAME(const LocalKParam &kparam);

  16. // void local_conv_MEGDNN_SIMD_NAME(const LocalKParam &kparam);

  17. //

  18. 

  19. #include "src/common/local/local_decl.inl"

  20. 

  21. #include "src/common/macro_helper.h"

  22. #include "src/common/utils.h"

  23. 

  24. namespace {

  25. 

  26. using namespace megdnn;

  27. 

  28. template <int N, int OC>

  29. void local_xcorr_tpl(const LocalKParam& kparam) MEGDNN_SIMD_ATTRIBUTE_TARGET;

  30. template <int N, int OC>

  31. void local_xcorr_tpl(const LocalKParam& kparam) {

  32.     const float* src = static_cast<const float*>(kparam.src.get_ptr());

  33.     const float* filter = static_cast<const float*>(kparam.filter.get_ptr());

  34.     float* dst = static_cast<float*>(kparam.dst.get_ptr());

  35.     float* workspace = static_cast<float*>(kparam.workspace);

  36.     const int IC = kparam.ic, IH = kparam.ih, IW = kparam.iw, OH = kparam.oh,

  37.               OW = kparam.ow, FH = kparam.fh, FW = kparam.fw;

  38.     const uint32_t PH = kparam.ph, PW = kparam.pw, SH = kparam.sh, SW = kparam.sw;

  39.     const ptrdiff_t INP_BS = kparam.inp_bs, OUT_BS = kparam.out_bs;

  40. 

  41.     float* dst2 = workspace;

  42.     const int width = MEGDNN_SIMD_WIDTH;

  43.     // dst2 is (H, W, N, C)

  44.     memset(dst2, 0, sizeof(float) * OH * OW * N * OC);

  45.     float* dst2_hwnc = dst2;

  46.     rep(oh, OH) rep(ow, OW) {

  47.         const float* src_bak = src;

  48.         rep(ic, IC) {

  49.             rep(fh, FH) for (int fw = 0; fw < FW; ++fw, filter += OC) {

  50.                 int ih = -PH + oh * SH + fh;

  51.                 int iw = -PW + ow * SW + fw;

  52.                 if (ih < 0 || ih >= IH || iw < 0 || iw >= IW)

  53.                     continue;

  54.                 float* dst2_bak = dst2;

  55.                 rep(n, N) {

  56.                     float s = src[n * INP_BS + ih * IW + iw];

  57.                     const float* filter_bak = filter;

  58.                     MEGDNN_SIMD_TYPE vs = MEGDNN_SIMD_SET1(s);

  59.                     int oc = 0;

  60.                     for (; oc + 4 * width <= OC; oc += 4 * width, filter += 4 * width) {

  61.                         MEGDNN_SIMD_TYPE vf0 = MEGDNN_SIMD_LOADU(filter + 0 * width);

  62.                         MEGDNN_SIMD_TYPE vf1 = MEGDNN_SIMD_LOADU(filter + 1 * width);

  63.                         MEGDNN_SIMD_TYPE vf2 = MEGDNN_SIMD_LOADU(filter + 2 * width);

  64.                         MEGDNN_SIMD_TYPE vf3 = MEGDNN_SIMD_LOADU(filter + 3 * width);

  65.                         MEGDNN_SIMD_TYPE vd0 = MEGDNN_SIMD_LOADU(dst2 + oc + 0 * width);

  66.                         MEGDNN_SIMD_TYPE vd1 = MEGDNN_SIMD_LOADU(dst2 + oc + 1 * width);

  67.                         MEGDNN_SIMD_TYPE vd2 = MEGDNN_SIMD_LOADU(dst2 + oc + 2 * width);

  68.                         MEGDNN_SIMD_TYPE vd3 = MEGDNN_SIMD_LOADU(dst2 + oc + 3 * width);

  69.                         vd0 = MEGDNN_SIMD_FMADD(vf0, vs, vd0);

  70.                         vd1 = MEGDNN_SIMD_FMADD(vf1, vs, vd1);

  71.                         vd2 = MEGDNN_SIMD_FMADD(vf2, vs, vd2);

  72.                         vd3 = MEGDNN_SIMD_FMADD(vf3, vs, vd3);

  73.                         MEGDNN_SIMD_STOREU(dst2 + oc + 0 * width, vd0);

  74.                         MEGDNN_SIMD_STOREU(dst2 + oc + 1 * width, vd1);

  75.                         MEGDNN_SIMD_STOREU(dst2 + oc + 2 * width, vd2);

  76.                         MEGDNN_SIMD_STOREU(dst2 + oc + 3 * width, vd3);

  77.                     }

  78.                     if (oc + 2 * width <= OC) {

  79.                         MEGDNN_SIMD_TYPE vf0 = MEGDNN_SIMD_LOADU(filter + 0 * width);

  80.                         MEGDNN_SIMD_TYPE vf1 = MEGDNN_SIMD_LOADU(filter + 1 * width);

  81.                         MEGDNN_SIMD_TYPE vd0 = MEGDNN_SIMD_LOADU(dst2 + oc + 0 * width);

  82.                         MEGDNN_SIMD_TYPE vd1 = MEGDNN_SIMD_LOADU(dst2 + oc + 1 * width);

  83.                         vd0 = MEGDNN_SIMD_FMADD(vf0, vs, vd0);

  84.                         vd1 = MEGDNN_SIMD_FMADD(vf1, vs, vd1);

  85.                         MEGDNN_SIMD_STOREU(dst2 + oc + 0 * width, vd0);

  86.                         MEGDNN_SIMD_STOREU(dst2 + oc + 1 * width, vd1);

  87.                         oc += 2 * width;

  88.                         filter += 2 * width;

  89.                     }

  90.                     if (oc + 1 * width <= OC) {

  91.                         MEGDNN_SIMD_TYPE vf0 = MEGDNN_SIMD_LOADU(filter + 0 * width);

  92.                         MEGDNN_SIMD_TYPE vd0 = MEGDNN_SIMD_LOADU(dst2 + oc + 0 * width);

  93.                         vd0 = MEGDNN_SIMD_FMADD(vf0, vs, vd0);

  94.                         MEGDNN_SIMD_STOREU(dst2 + oc + 0 * width, vd0);

  95.                         oc += 1 * width;

  96.                         filter += 1 * width;

  97.                     }

  98.                     for (; oc < OC; ++oc, ++filter) {

  99.                         dst2[oc] += s * (*filter);

  100.                     }

  101.                     filter = filter_bak;

  102.                     dst2 += OC;

  103.                 }

  104.                 dst2 = dst2_bak;

  105.             }

  106.             src += IH * IW;

  107.         }

  108.         src = src_bak;

  109.         dst2 += N * OC;

  110.     }

  111.     transpose_knc2nsck(dst2_hwnc, dst, OH * OW, N, OC, OUT_BS);

  112. }

  113. void local_xcorr_generic(const LocalKParam& kparam) MEGDNN_SIMD_ATTRIBUTE_TARGET;

  114. void local_xcorr_generic(const LocalKParam& kparam) {

  115.     UNPACK_LOCAL_FLOAT_NONCONTIG_BATCH_KERN_PARAM(kparam, float);

  116. 

  117.     float* dst2 = workspace;

  118.     const int width = MEGDNN_SIMD_WIDTH;

  119.     // dst2 is (H, W, N, C)

  120.     memset(dst2, 0, sizeof(float) * OH * OW * N * OC);

  121.     float* dst2_hwnc = dst2;

  122.     rep(oh, OH) rep(ow, OW) {

  123.         const float* src_bak = src;

  124.         rep(ic, IC) {

  125.             rep(fh, FH) for (int fw = 0; fw < FW; ++fw, filter += OC) {

  126.                 int ih = -PH + oh * SH + fh;

  127.                 int iw = -PW + ow * SW + fw;

  128.                 if (ih < 0 || ih >= IH || iw < 0 || iw >= IW)

  129.                     continue;

  130.                 float* dst2_bak = dst2;

  131.                 rep(n, N) {

  132.                     float s = src[n * INP_BS + ih * IW + iw];

  133.                     const float* filter_bak = filter;

  134.                     MEGDNN_SIMD_TYPE vs = MEGDNN_SIMD_SET1(s);

  135.                     int oc = 0;

  136.                     for (; oc + 4 * width <= OC; oc += 4 * width, filter += 4 * width) {

  137.                         MEGDNN_SIMD_TYPE vf0 = MEGDNN_SIMD_LOADU(filter + 0 * width);

  138.                         MEGDNN_SIMD_TYPE vf1 = MEGDNN_SIMD_LOADU(filter + 1 * width);

  139.                         MEGDNN_SIMD_TYPE vf2 = MEGDNN_SIMD_LOADU(filter + 2 * width);

  140.                         MEGDNN_SIMD_TYPE vf3 = MEGDNN_SIMD_LOADU(filter + 3 * width);

  141.                         MEGDNN_SIMD_TYPE vd0 = MEGDNN_SIMD_LOADU(dst2 + oc + 0 * width);

  142.                         MEGDNN_SIMD_TYPE vd1 = MEGDNN_SIMD_LOADU(dst2 + oc + 1 * width);

  143.                         MEGDNN_SIMD_TYPE vd2 = MEGDNN_SIMD_LOADU(dst2 + oc + 2 * width);

  144.                         MEGDNN_SIMD_TYPE vd3 = MEGDNN_SIMD_LOADU(dst2 + oc + 3 * width);

  145.                         vd0 = MEGDNN_SIMD_FMADD(vf0, vs, vd0);

  146.                         vd1 = MEGDNN_SIMD_FMADD(vf1, vs, vd1);

  147.                         vd2 = MEGDNN_SIMD_FMADD(vf2, vs, vd2);

  148.                         vd3 = MEGDNN_SIMD_FMADD(vf3, vs, vd3);

  149.                         MEGDNN_SIMD_STOREU(dst2 + oc + 0 * width, vd0);

  150.                         MEGDNN_SIMD_STOREU(dst2 + oc + 1 * width, vd1);

  151.                         MEGDNN_SIMD_STOREU(dst2 + oc + 2 * width, vd2);

  152.                         MEGDNN_SIMD_STOREU(dst2 + oc + 3 * width, vd3);

  153.                     }

  154.                     if (oc + 2 * width <= OC) {

  155.                         MEGDNN_SIMD_TYPE vf0 = MEGDNN_SIMD_LOADU(filter + 0 * width);

  156.                         MEGDNN_SIMD_TYPE vf1 = MEGDNN_SIMD_LOADU(filter + 1 * width);

  157.                         MEGDNN_SIMD_TYPE vd0 = MEGDNN_SIMD_LOADU(dst2 + oc + 0 * width);

  158.                         MEGDNN_SIMD_TYPE vd1 = MEGDNN_SIMD_LOADU(dst2 + oc + 1 * width);

  159.                         vd0 = MEGDNN_SIMD_FMADD(vf0, vs, vd0);

  160.                         vd1 = MEGDNN_SIMD_FMADD(vf1, vs, vd1);

  161.                         MEGDNN_SIMD_STOREU(dst2 + oc + 0 * width, vd0);

  162.                         MEGDNN_SIMD_STOREU(dst2 + oc + 1 * width, vd1);

  163.                         oc += 2 * width;

  164.                         filter += 2 * width;

  165.                     }

  166.                     if (oc + 1 * width <= OC) {

  167.                         MEGDNN_SIMD_TYPE vf0 = MEGDNN_SIMD_LOADU(filter + 0 * width);

  168.                         MEGDNN_SIMD_TYPE vd0 = MEGDNN_SIMD_LOADU(dst2 + oc + 0 * width);

  169.                         vd0 = MEGDNN_SIMD_FMADD(vf0, vs, vd0);

  170.                         MEGDNN_SIMD_STOREU(dst2 + oc + 0 * width, vd0);

  171.                         oc += 1 * width;

  172.                         filter += 1 * width;

  173.                     }

  174.                     for (; oc < OC; ++oc, ++filter) {

  175.                         dst2[oc] += s * (*filter);

  176.                     }

  177.                     filter = filter_bak;

  178.                     dst2 += OC;

  179.                 }

  180.                 dst2 = dst2_bak;

  181.             }

  182.             src += IH * IW;

  183.         }

  184.         src = src_bak;

  185.         dst2 += N * OC;

  186.     }

  187.     transpose_knc2nsck(dst2_hwnc, dst, OH * OW, N, OC, OUT_BS);

  188. }

  189. 

  190. template <int N, int OC>

  191. void local_conv_tpl(const LocalKParam& kparam) MEGDNN_SIMD_ATTRIBUTE_TARGET;

  192. template <int N, int OC>

  193. void local_conv_tpl(const LocalKParam& kparam) {

  194.     const float* src = static_cast<const float*>(kparam.src.get_ptr());

  195.     const float* filter = static_cast<const float*>(kparam.filter.get_ptr());

  196.     float* dst = static_cast<float*>(kparam.dst.get_ptr());

  197.     float* workspace = static_cast<float*>(kparam.workspace);

  198.     const int IC = kparam.ic, IH = kparam.ih, IW = kparam.iw, OH = kparam.oh,

  199.               OW = kparam.ow, FH = kparam.fh, FW = kparam.fw;

  200.     const uint32_t PH = kparam.ph, PW = kparam.pw, SH = kparam.sh, SW = kparam.sw;

  201.     const ptrdiff_t INP_BS = kparam.inp_bs, OUT_BS = kparam.out_bs;

  202. 

  203.     float* dst2 = workspace;

  204.     const int width = MEGDNN_SIMD_WIDTH;

  205.     // dst2 is (H, W, N, C)

  206.     memset(dst2, 0, sizeof(float) * OH * OW * N * OC);

  207.     float* dst2_hwnc = dst2;

  208.     rep(oh, OH) rep(ow, OW) {

  209.         const float* src_bak = src;

  210.         rep(ic, IC) {

  211.             rep(fh, FH) for (int fw = 0; fw < FW; ++fw, filter += OC) {

  212.                 int ih = -PH + oh * SH + (FH - fh - 1);

  213.                 int iw = -PW + ow * SW + (FW - fw - 1);

  214.                 if (ih < 0 || ih >= IH || iw < 0 || iw >= IW)

  215.                     continue;

  216.                 float* dst2_bak = dst2;

  217.                 rep(n, N) {

  218.                     float s = src[n * INP_BS + ih * IW + iw];

  219.                     const float* filter_bak = filter;

  220.                     MEGDNN_SIMD_TYPE vs = MEGDNN_SIMD_SET1(s);

  221.                     int oc = 0;

  222.                     for (; oc + 4 * width <= OC; oc += 4 * width, filter += 4 * width) {

  223.                         MEGDNN_SIMD_TYPE vf0 = MEGDNN_SIMD_LOADU(filter + 0 * width);

  224.                         MEGDNN_SIMD_TYPE vf1 = MEGDNN_SIMD_LOADU(filter + 1 * width);

  225.                         MEGDNN_SIMD_TYPE vf2 = MEGDNN_SIMD_LOADU(filter + 2 * width);

  226.                         MEGDNN_SIMD_TYPE vf3 = MEGDNN_SIMD_LOADU(filter + 3 * width);

  227.                         MEGDNN_SIMD_TYPE vd0 = MEGDNN_SIMD_LOADU(dst2 + oc + 0 * width);

  228.                         MEGDNN_SIMD_TYPE vd1 = MEGDNN_SIMD_LOADU(dst2 + oc + 1 * width);

  229.                         MEGDNN_SIMD_TYPE vd2 = MEGDNN_SIMD_LOADU(dst2 + oc + 2 * width);

  230.                         MEGDNN_SIMD_TYPE vd3 = MEGDNN_SIMD_LOADU(dst2 + oc + 3 * width);

  231.                         vd0 = MEGDNN_SIMD_FMADD(vf0, vs, vd0);

  232.                         vd1 = MEGDNN_SIMD_FMADD(vf1, vs, vd1);

  233.                         vd2 = MEGDNN_SIMD_FMADD(vf2, vs, vd2);

  234.                         vd3 = MEGDNN_SIMD_FMADD(vf3, vs, vd3);

  235.                         MEGDNN_SIMD_STOREU(dst2 + oc + 0 * width, vd0);

  236.                         MEGDNN_SIMD_STOREU(dst2 + oc + 1 * width, vd1);

  237.                         MEGDNN_SIMD_STOREU(dst2 + oc + 2 * width, vd2);

  238.                         MEGDNN_SIMD_STOREU(dst2 + oc + 3 * width, vd3);

  239.                     }

  240.                     if (oc + 2 * width <= OC) {

  241.                         MEGDNN_SIMD_TYPE vf0 = MEGDNN_SIMD_LOADU(filter + 0 * width);

  242.                         MEGDNN_SIMD_TYPE vf1 = MEGDNN_SIMD_LOADU(filter + 1 * width);

  243.                         MEGDNN_SIMD_TYPE vd0 = MEGDNN_SIMD_LOADU(dst2 + oc + 0 * width);

  244.                         MEGDNN_SIMD_TYPE vd1 = MEGDNN_SIMD_LOADU(dst2 + oc + 1 * width);

  245.                         vd0 = MEGDNN_SIMD_FMADD(vf0, vs, vd0);

  246.                         vd1 = MEGDNN_SIMD_FMADD(vf1, vs, vd1);

  247.                         MEGDNN_SIMD_STOREU(dst2 + oc + 0 * width, vd0);

  248.                         MEGDNN_SIMD_STOREU(dst2 + oc + 1 * width, vd1);

  249.                         oc += 2 * width;

  250.                         filter += 2 * width;

  251.                     }

  252.                     if (oc + 1 * width <= OC) {

  253.                         MEGDNN_SIMD_TYPE vf0 = MEGDNN_SIMD_LOADU(filter + 0 * width);

  254.                         MEGDNN_SIMD_TYPE vd0 = MEGDNN_SIMD_LOADU(dst2 + oc + 0 * width);

  255.                         vd0 = MEGDNN_SIMD_FMADD(vf0, vs, vd0);

  256.                         MEGDNN_SIMD_STOREU(dst2 + oc + 0 * width, vd0);

  257.                         oc += 1 * width;

  258.                         filter += 1 * width;

  259.                     }

  260.                     for (; oc < OC; ++oc, ++filter) {

  261.                         dst2[oc] += s * (*filter);

  262.                     }

  263.                     filter = filter_bak;

  264.                     dst2 += OC;

  265.                 }

  266.                 dst2 = dst2_bak;

  267.             }

  268.             src += IH * IW;

  269.         }

  270.         src = src_bak;

  271.         dst2 += N * OC;

  272.     }

  273.     transpose_knc2nsck(dst2_hwnc, dst, OH * OW, N, OC, OUT_BS);

  274. }

  275. 

  276. void local_conv_generic(const LocalKParam& kparam) MEGDNN_SIMD_ATTRIBUTE_TARGET;

  277. void local_conv_generic(const LocalKParam& kparam) {

  278.     UNPACK_LOCAL_FLOAT_NONCONTIG_BATCH_KERN_PARAM(kparam, float);

  279. 

  280.     float* dst2 = workspace;

  281.     const int width = MEGDNN_SIMD_WIDTH;

  282.     // dst2 is (H, W, N, C)

  283.     memset(dst2, 0, sizeof(float) * OH * OW * N * OC);

  284.     float* dst2_hwnc = dst2;

  285.     rep(oh, OH) rep(ow, OW) {

  286.         const float* src_bak = src;

  287.         rep(ic, IC) {

  288.             rep(fh, FH) for (int fw = 0; fw < FW; ++fw, filter += OC) {

  289.                 int ih = -PH + oh * SH + (FH - fh - 1);

  290.                 int iw = -PW + ow * SW + (FW - fw - 1);

  291.                 if (ih < 0 || ih >= IH || iw < 0 || iw >= IW)

  292.                     continue;

  293.                 float* dst2_bak = dst2;

  294.                 rep(n, N) {

  295.                     float s = src[n * INP_BS + ih * IW + iw];

  296.                     const float* filter_bak = filter;

  297.                     MEGDNN_SIMD_TYPE vs = MEGDNN_SIMD_SET1(s);

  298.                     int oc = 0;

  299.                     for (; oc + 4 * width <= OC; oc += 4 * width, filter += 4 * width) {

  300.                         MEGDNN_SIMD_TYPE vf0 = MEGDNN_SIMD_LOADU(filter + 0 * width);

  301.                         MEGDNN_SIMD_TYPE vf1 = MEGDNN_SIMD_LOADU(filter + 1 * width);

  302.                         MEGDNN_SIMD_TYPE vf2 = MEGDNN_SIMD_LOADU(filter + 2 * width);

  303.                         MEGDNN_SIMD_TYPE vf3 = MEGDNN_SIMD_LOADU(filter + 3 * width);

  304.                         MEGDNN_SIMD_TYPE vd0 = MEGDNN_SIMD_LOADU(dst2 + oc + 0 * width);

  305.                         MEGDNN_SIMD_TYPE vd1 = MEGDNN_SIMD_LOADU(dst2 + oc + 1 * width);

  306.                         MEGDNN_SIMD_TYPE vd2 = MEGDNN_SIMD_LOADU(dst2 + oc + 2 * width);

  307.                         MEGDNN_SIMD_TYPE vd3 = MEGDNN_SIMD_LOADU(dst2 + oc + 3 * width);

  308.                         vd0 = MEGDNN_SIMD_FMADD(vf0, vs, vd0);

  309.                         vd1 = MEGDNN_SIMD_FMADD(vf1, vs, vd1);

  310.                         vd2 = MEGDNN_SIMD_FMADD(vf2, vs, vd2);

  311.                         vd3 = MEGDNN_SIMD_FMADD(vf3, vs, vd3);

  312.                         MEGDNN_SIMD_STOREU(dst2 + oc + 0 * width, vd0);

  313.                         MEGDNN_SIMD_STOREU(dst2 + oc + 1 * width, vd1);

  314.                         MEGDNN_SIMD_STOREU(dst2 + oc + 2 * width, vd2);

  315.                         MEGDNN_SIMD_STOREU(dst2 + oc + 3 * width, vd3);

  316.                     }

  317.                     if (oc + 2 * width <= OC) {

  318.                         MEGDNN_SIMD_TYPE vf0 = MEGDNN_SIMD_LOADU(filter + 0 * width);

  319.                         MEGDNN_SIMD_TYPE vf1 = MEGDNN_SIMD_LOADU(filter + 1 * width);

  320.                         MEGDNN_SIMD_TYPE vd0 = MEGDNN_SIMD_LOADU(dst2 + oc + 0 * width);

  321.                         MEGDNN_SIMD_TYPE vd1 = MEGDNN_SIMD_LOADU(dst2 + oc + 1 * width);

  322.                         vd0 = MEGDNN_SIMD_FMADD(vf0, vs, vd0);

  323.                         vd1 = MEGDNN_SIMD_FMADD(vf1, vs, vd1);

  324.                         MEGDNN_SIMD_STOREU(dst2 + oc + 0 * width, vd0);

  325.                         MEGDNN_SIMD_STOREU(dst2 + oc + 1 * width, vd1);

  326.                         oc += 2 * width;

  327.                         filter += 2 * width;

  328.                     }

  329.                     if (oc + 1 * width <= OC) {

  330.                         MEGDNN_SIMD_TYPE vf0 = MEGDNN_SIMD_LOADU(filter + 0 * width);

  331.                         MEGDNN_SIMD_TYPE vd0 = MEGDNN_SIMD_LOADU(dst2 + oc + 0 * width);

  332.                         vd0 = MEGDNN_SIMD_FMADD(vf0, vs, vd0);

  333.                         MEGDNN_SIMD_STOREU(dst2 + oc + 0 * width, vd0);

  334.                         oc += 1 * width;

  335.                         filter += 1 * width;

  336.                     }

  337.                     for (; oc < OC; ++oc, ++filter) {

  338.                         dst2[oc] += s * (*filter);

  339.                     }

  340.                     filter = filter_bak;

  341.                     dst2 += OC;

  342.                 }

  343.                 dst2 = dst2_bak;

  344.             }

  345.             src += IH * IW;

  346.         }

  347.         src = src_bak;

  348.         dst2 += N * OC;

  349.     }

  350.     transpose_knc2nsck(dst2_hwnc, dst, OH * OW, N, OC, OUT_BS);

  351. }

  352. 

  353. }  // anonymous namespace

  354. 

  355. namespace megdnn {

  356. 

  357. #define FUNC_NAME CONCAT_STR(local_xcorr_, MEGDNN_SIMD_NAME)

  358. 

  359. void FUNC_NAME(const LocalKParam& kparam) {

  360.     auto N = kparam.n, OC = kparam.oc;

  361. #define DISPATCH_WITH_N_OC(N, OC)       \

  362.     do {                                \

  363.         local_xcorr_tpl<N, OC>(kparam); \

  364.         return;                         \

  365.     } while (0)

  366. 

  367. #define DISPATCH_WITH_N(N)             \

  368.     switch (OC) {                      \

  369.         case 16:                       \

  370.             DISPATCH_WITH_N_OC(N, 16); \

  371.             break;                     \

  372.         case 32:                       \

  373.             DISPATCH_WITH_N_OC(N, 32); \

  374.             break;                     \

  375.         case 48:                       \

  376.             DISPATCH_WITH_N_OC(N, 48); \

  377.             break;                     \

  378.         case 64:                       \

  379.             DISPATCH_WITH_N_OC(N, 64); \

  380.             break;                     \

  381.     }

  382. #define DISPATCH()              \

  383.     switch (N) {                \

  384.         case 1:                 \

  385.             DISPATCH_WITH_N(1); \

  386.             break;              \

  387.         case 2:                 \

  388.             DISPATCH_WITH_N(2); \

  389.             break;              \

  390.     }

  391. 

  392.     DISPATCH();

  393. 

  394. #undef DISPATCH

  395. #undef DISPATCH_WITH_N

  396. #undef DISPATCH_WITH_N_OC

  397.     local_xcorr_generic(kparam);

  398. }

  399. 

  400. #undef FUNC_NAME

  401. 

  402. #define FUNC_NAME CONCAT_STR(local_conv_, MEGDNN_SIMD_NAME)

  403. 

  404. void FUNC_NAME(const LocalKParam& kparam) {

  405.     auto N = kparam.n, OC = kparam.oc;

  406. #define DISPATCH_WITH_N_OC(N, OC)      \

  407.     do {                               \

  408.         local_conv_tpl<N, OC>(kparam); \

  409.         return;                        \

  410.     } while (0)

  411. 

  412. #define DISPATCH_WITH_N(N)             \

  413.     switch (OC) {                      \

  414.         case 16:                       \

  415.             DISPATCH_WITH_N_OC(N, 16); \

  416.             break;                     \

  417.         case 32:                       \

  418.             DISPATCH_WITH_N_OC(N, 32); \

  419.             break;                     \

  420.         case 48:                       \

  421.             DISPATCH_WITH_N_OC(N, 48); \

  422.             break;                     \

  423.         case 64:                       \

  424.             DISPATCH_WITH_N_OC(N, 64); \

  425.             break;                     \

  426.     }

  427. #define DISPATCH()              \

  428.     switch (N) {                \

  429.         case 1:                 \

  430.             DISPATCH_WITH_N(1); \

  431.             break;              \

  432.         case 2:                 \

  433.             DISPATCH_WITH_N(2); \

  434.             break;              \

  435.     }

  436. 

  437.     DISPATCH();

  438. 

  439. #undef DISPATCH

  440. #undef DISPATCH_WITH_N

  441. #undef DISPATCH_WITH_N_OC

  442.     local_conv_generic(kparam);

  443. }

  444. 

  445. #undef FUNC_NAME

  446. 

  447. }  // namespace megdnn

  448. 

  449. #include "src/common/macro_helper_epilogue.h"