wangwei
/
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: b8821edb3d
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'b8821edb3d'
${ noResults }
MegEngine/dnn/src/aarch64/elemwise/sigmoid.h

291 lines
14 kB
Raw Blame History 

  1. #pragma once

  2. #include "src/arm_common/elemwise_helper/elemwise_op.h"

  3. namespace megdnn {

  4. namespace elemwise {

  5. #if MEGDNN_AARCH64

  6. template <>

  7. struct OpCallerUnary<arm_common::SigmoidOp<float, float>, VEC> {

  8.     static void run(const float* src, float* dst, DType, DType, size_t nr_elems) {

  9.         size_t x6_iter = nr_elems / (4 * 6);

  10.         size_t offset = x6_iter * 4 * 6;

  11. 

  12.         float32x4_t lower_range;

  13.         float32x4_t upper_range;

  14.         float32x4_t alpha_9;

  15.         float32x4_t alpha_7;

  16.         float32x4_t alpha_5;

  17.         float32x4_t alpha_3;

  18.         float32x4_t alpha_1;

  19.         float32x4_t beta_10;

  20.         float32x4_t beta_8;

  21.         float32x4_t beta_6;

  22.         float32x4_t beta_4;

  23.         float32x4_t beta_2;

  24.         float32x4_t beta_0;

  25.         float32x4_t one_half;

  26. 

  27.         const float* const_ptr = &(arm_common::sigmoid_constants.lower_range);

  28.         if (x6_iter > 0) {

  29.             /**

  30.              * q0 - q5   : squared

  31.              * q6 - q11  : p

  32.              * q12- q17  : val(temp), q

  33.              * q18- q31  : const

  34.              */

  35.             asm volatile(

  36.                     "ld1r {%[lower_range].4s}, [%[const_ptr]], #4\n"

  37.                     "ld1r {%[upper_range].4s}, [%[const_ptr]], #4\n"

  38.                     "ld1r {%[alpha_9].4s},     [%[const_ptr]], #4\n"

  39.                     "ld1r {%[alpha_7].4s},     [%[const_ptr]], #4\n"

  40.                     "ld1r {%[alpha_5].4s},     [%[const_ptr]], #4\n"

  41.                     "ld1r {%[alpha_3].4s},     [%[const_ptr]], #4\n"

  42.                     "ld1r {%[alpha_1].4s},     [%[const_ptr]], #4\n"

  43.                     "ld1r {%[beta_10].4s},     [%[const_ptr]], #4\n"

  44.                     "ld1r {%[beta_8].4s},      [%[const_ptr]], #4\n"

  45.                     "ld1r {%[beta_6].4s},      [%[const_ptr]], #4\n"

  46.                     "ld1r {%[beta_4].4s},      [%[const_ptr]], #4\n"

  47.                     "ld1r {%[beta_2].4s},      [%[const_ptr]], #4\n"

  48.                     "ld1r {%[beta_0].4s},      [%[const_ptr]], #4\n"

  49.                     "ld1r {%[one_half].4s},    [%[const_ptr]], #4\n"

  50. 

  51.                     "1:\n"

  52.                     "ldr  q12, [%[a_ptr]]     \n"

  53.                     "ldr  q13, [%[a_ptr], #16]\n"

  54.                     "ldr  q14, [%[a_ptr], #32]\n"

  55.                     "ldr  q15, [%[a_ptr], #48]\n"

  56.                     "ldr  q16, [%[a_ptr], #64]\n"

  57.                     "ldr  q17, [%[a_ptr], #80]\n"

  58.                     // auto val = vmaxq_f32(vdupq_n_f32(sigmoid_constants.lower_range),

  59.                     // src);

  60.                     "fmax v12.4s, v12.4s, %[lower_range].4s\n"

  61.                     "fmax v13.4s, v13.4s, %[lower_range].4s\n"

  62.                     "fmax v14.4s, v14.4s, %[lower_range].4s\n"

  63.                     "fmax v15.4s, v15.4s, %[lower_range].4s\n"

  64.                     "fmax v16.4s, v16.4s, %[lower_range].4s\n"

  65.                     "fmax v17.4s, v17.4s, %[lower_range].4s\n"

  66.                     "add %[a_ptr], %[a_ptr], #96\n"

  67. 

  68.                     //  val = vminq_f32(vdupq_n_f32(sigmoid_constants.upper_range), val);

  69.                     "fmin v12.4s, v12.4s, %[upper_range].4s\n"

  70.                     "fmin v13.4s, v13.4s, %[upper_range].4s\n"

  71.                     "fmin v14.4s, v14.4s, %[upper_range].4s\n"

  72.                     "fmin v15.4s, v15.4s, %[upper_range].4s\n"

  73.                     "fmin v16.4s, v16.4s, %[upper_range].4s\n"

  74.                     "fmin v17.4s, v17.4s, %[upper_range].4s\n"

  75. 

  76.                     //! auto squared = vmulq_f32(val, val);

  77.                     "fmul v0.4s, v12.4s, v12.4s\n"

  78.                     "fmul v1.4s, v13.4s, v13.4s\n"

  79.                     "fmul v2.4s, v14.4s, v14.4s\n"

  80.                     "fmul v3.4s, v15.4s, v15.4s\n"

  81.                     "fmul v4.4s, v16.4s, v16.4s\n"

  82.                     "fmul v5.4s, v17.4s, v17.4s\n"

  83.                     //    auto p = fma_ps_f32(

  84.                     //             vdupq_n_f32(sigmoid_constants.alpha_7), squared,

  85.                     //             vdupq_n_f32(sigmoid_constants.alpha_9));

  86.                     "fmul v6.4s,  v0.4s,  %[alpha_9].4s\n"

  87.                     "fmul v7.4s,  v1.4s,  %[alpha_9].4s\n"

  88.                     "fmul v8.4s,  v2.4s,  %[alpha_9].4s\n"

  89.                     "fmul v9.4s,  v3.4s,  %[alpha_9].4s\n"

  90.                     "fmul v10.4s, v4.4s,  %[alpha_9].4s\n"

  91.                     "fmul v11.4s, v5.4s,  %[alpha_9].4s\n"

  92.                     "fadd v6.4s,  v6.4s,  %[alpha_7].4s\n"

  93.                     "fadd v7.4s,  v7.4s,  %[alpha_7].4s\n"

  94.                     "fadd v8.4s,  v8.4s,  %[alpha_7].4s\n"

  95.                     "fadd v9.4s,  v9.4s,  %[alpha_7].4s\n"

  96.                     "fadd v10.4s, v10.4s, %[alpha_7].4s\n"

  97.                     "fadd v11.4s, v11.4s, %[alpha_7].4s\n"

  98. 

  99.                     // p = fma_ps_f32(vdupq_n_f32(sigmoid_constants.alpha_5), p, squared);

  100.                     "fmul v6.4s,   v6.4s,  v0.4s\n"

  101.                     "fmul v7.4s,   v7.4s,  v1.4s\n"

  102.                     "fmul v8.4s,   v8.4s,  v2.4s\n"

  103.                     "fmul v9.4s,   v9.4s,  v3.4s\n"

  104.                     "fmul v10.4s,  v10.4s, v4.4s\n"

  105.                     "fmul v11.4s,  v11.4s, v5.4s\n"

  106.                     "fadd v6.4s,  v6.4s,  %[alpha_5].4s\n"

  107.                     "fadd v7.4s,  v7.4s,  %[alpha_5].4s\n"

  108.                     "fadd v8.4s,  v8.4s,  %[alpha_5].4s\n"

  109.                     "fadd v9.4s,  v9.4s,  %[alpha_5].4s\n"

  110.                     "fadd v10.4s, v10.4s, %[alpha_5].4s\n"

  111.                     "fadd v11.4s, v11.4s, %[alpha_5].4s\n"

  112. 

  113.                     // p = fma_ps_f32(vdupq_n_f32(sigmoid_constants.alpha_3), p, squared);

  114.                     "fmul v6.4s,   v6.4s,  v0.4s\n"

  115.                     "fmul v7.4s,   v7.4s,  v1.4s\n"

  116.                     "fmul v8.4s,   v8.4s,  v2.4s\n"

  117.                     "fmul v9.4s,   v9.4s,  v3.4s\n"

  118.                     "fmul v10.4s,  v10.4s, v4.4s\n"

  119.                     "fmul v11.4s,  v11.4s, v5.4s\n"

  120.                     "fadd v6.4s,  v6.4s,  %[alpha_3].4s\n"

  121.                     "fadd v7.4s,  v7.4s,  %[alpha_3].4s\n"

  122.                     "fadd v8.4s,  v8.4s,  %[alpha_3].4s\n"

  123.                     "fadd v9.4s,  v9.4s,  %[alpha_3].4s\n"

  124.                     "fadd v10.4s, v10.4s, %[alpha_3].4s\n"

  125.                     "fadd v11.4s, v11.4s, %[alpha_3].4s\n"

  126. 

  127.                     // p = fma_ps_f32(vdupq_n_f32(sigmoid_constants.alpha_1), p, squared);

  128.                     "fmul v6.4s,   v6.4s,  v0.4s\n"

  129.                     "fmul v7.4s,   v7.4s,  v1.4s\n"

  130.                     "fmul v8.4s,   v8.4s,  v2.4s\n"

  131.                     "fmul v9.4s,   v9.4s,  v3.4s\n"

  132.                     "fmul v10.4s,  v10.4s, v4.4s\n"

  133.                     "fmul v11.4s,  v11.4s, v5.4s\n"

  134.                     "fadd v6.4s,  v6.4s,  %[alpha_1].4s\n"

  135.                     "fadd v7.4s,  v7.4s,  %[alpha_1].4s\n"

  136.                     "fadd v8.4s,  v8.4s,  %[alpha_1].4s\n"

  137.                     "fadd v9.4s,  v9.4s,  %[alpha_1].4s\n"

  138.                     "fadd v10.4s, v10.4s, %[alpha_1].4s\n"

  139.                     "fadd v11.4s, v11.4s, %[alpha_1].4s\n"

  140. 

  141.                     //     p = vmulq_f32(p, val);

  142.                     "fmul v6.4s,   v6.4s,  v12.4s\n"

  143.                     "fmul v7.4s,   v7.4s,  v13.4s\n"

  144.                     "fmul v8.4s,   v8.4s,  v14.4s\n"

  145.                     "fmul v9.4s,   v9.4s,  v15.4s\n"

  146.                     "fmul v10.4s,  v10.4s, v16.4s\n"

  147.                     "fmul v11.4s,  v11.4s, v17.4s\n"

  148. 

  149.                     //     auto q = fma_ps_f32(

  150.                     //             vdupq_n_f32(sigmoid_constants.beta_8), squared,

  151.                     //             vdupq_n_f32(sigmoid_constants.beta_10));

  152.                     "fmul v12.4s, v0.4s,  %[beta_10].4s\n"

  153.                     "fmul v13.4s, v1.4s,  %[beta_10].4s\n"

  154.                     "fmul v14.4s, v2.4s,  %[beta_10].4s\n"

  155.                     "fmul v15.4s, v3.4s,  %[beta_10].4s\n"

  156.                     "fmul v16.4s, v4.4s,  %[beta_10].4s\n"

  157.                     "fmul v17.4s, v5.4s,  %[beta_10].4s\n"

  158.                     "fadd v12.4s, v12.4s, %[beta_8].4s\n"

  159.                     "fadd v13.4s, v13.4s, %[beta_8].4s\n"

  160.                     "fadd v14.4s, v14.4s, %[beta_8].4s\n"

  161.                     "fadd v15.4s, v15.4s, %[beta_8].4s\n"

  162.                     "fadd v16.4s, v16.4s, %[beta_8].4s\n"

  163.                     "fadd v17.4s, v17.4s, %[beta_8].4s\n"

  164. 

  165.                     //     q = fma_ps_f32(vdupq_n_f32(sigmoid_constants.beta_6), q,

  166.                     //     squared);

  167.                     "fmul v12.4s, v12.4s,  v0.4s\n"

  168.                     "fmul v13.4s, v13.4s,  v1.4s\n"

  169.                     "fmul v14.4s, v14.4s,  v2.4s\n"

  170.                     "fmul v15.4s, v15.4s,  v3.4s\n"

  171.                     "fmul v16.4s, v16.4s,  v4.4s\n"

  172.                     "fmul v17.4s, v17.4s,  v5.4s\n"

  173.                     "fadd v12.4s, v12.4s, %[beta_6].4s\n"

  174.                     "fadd v13.4s, v13.4s, %[beta_6].4s\n"

  175.                     "fadd v14.4s, v14.4s, %[beta_6].4s\n"

  176.                     "fadd v15.4s, v15.4s, %[beta_6].4s\n"

  177.                     "fadd v16.4s, v16.4s, %[beta_6].4s\n"

  178.                     "fadd v17.4s, v17.4s, %[beta_6].4s\n"

  179. 

  180.                     //     q = fma_ps_f32(vdupq_n_f32(sigmoid_constants.beta_4), q,

  181.                     //     squared);

  182.                     "fmul v12.4s, v12.4s,  v0.4s\n"

  183.                     "fmul v13.4s, v13.4s,  v1.4s\n"

  184.                     "fmul v14.4s, v14.4s,  v2.4s\n"

  185.                     "fmul v15.4s, v15.4s,  v3.4s\n"

  186.                     "fmul v16.4s, v16.4s,  v4.4s\n"

  187.                     "fmul v17.4s, v17.4s,  v5.4s\n"

  188.                     "fadd v12.4s, v12.4s, %[beta_4].4s\n"

  189.                     "fadd v13.4s, v13.4s, %[beta_4].4s\n"

  190.                     "fadd v14.4s, v14.4s, %[beta_4].4s\n"

  191.                     "fadd v15.4s, v15.4s, %[beta_4].4s\n"

  192.                     "fadd v16.4s, v16.4s, %[beta_4].4s\n"

  193.                     "fadd v17.4s, v17.4s, %[beta_4].4s\n"

  194. 

  195.                     //     q = fma_ps_f32(vdupq_n_f32(sigmoid_constants.beta_2), q,

  196.                     //     squared);

  197.                     "fmul v12.4s, v12.4s,  v0.4s\n"

  198.                     "fmul v13.4s, v13.4s,  v1.4s\n"

  199.                     "fmul v14.4s, v14.4s,  v2.4s\n"

  200.                     "fmul v15.4s, v15.4s,  v3.4s\n"

  201.                     "fmul v16.4s, v16.4s,  v4.4s\n"

  202.                     "fmul v17.4s, v17.4s,  v5.4s\n"

  203.                     "fadd v12.4s, v12.4s, %[beta_2].4s\n"

  204.                     "fadd v13.4s, v13.4s, %[beta_2].4s\n"

  205.                     "fadd v14.4s, v14.4s, %[beta_2].4s\n"

  206.                     "fadd v15.4s, v15.4s, %[beta_2].4s\n"

  207.                     "fadd v16.4s, v16.4s, %[beta_2].4s\n"

  208.                     "fadd v17.4s, v17.4s, %[beta_2].4s\n"

  209. 

  210.                     // q = fma_ps_f32(vdupq_n_f32(sigmoid_constants.beta_0), q, squared);

  211.                     "fmul v12.4s, v12.4s,  v0.4s\n"

  212.                     "fmul v13.4s, v13.4s,  v1.4s\n"

  213.                     "fmul v14.4s, v14.4s,  v2.4s\n"

  214.                     "fmul v15.4s, v15.4s,  v3.4s\n"

  215.                     "fmul v16.4s, v16.4s,  v4.4s\n"

  216.                     "fmul v17.4s, v17.4s,  v5.4s\n"

  217.                     "fadd v12.4s, v12.4s, %[beta_0].4s\n"

  218.                     "fadd v13.4s, v13.4s, %[beta_0].4s\n"

  219.                     "fadd v14.4s, v14.4s, %[beta_0].4s\n"

  220.                     "fadd v15.4s, v15.4s, %[beta_0].4s\n"

  221.                     "fadd v16.4s, v16.4s, %[beta_0].4s\n"

  222.                     "fadd v17.4s, v17.4s, %[beta_0].4s\n"

  223. 

  224.                     // vaddq_f32(div_ps_f32(p, q),

  225.                     // vdupq_n_f32(sigmoid_constants.one_half));

  226.                     "fdiv v12.4s, v6.4s,  v12.4s\n"

  227.                     "fdiv v13.4s, v7.4s,  v13.4s\n"

  228.                     "fdiv v14.4s, v8.4s,  v14.4s\n"

  229.                     "fdiv v15.4s, v9.4s,  v15.4s\n"

  230.                     "fdiv v16.4s, v10.4s, v16.4s\n"

  231.                     "fdiv v17.4s, v11.4s, v17.4s\n"

  232.                     "subs %w[x6_iter], %w[x6_iter], #1\n"

  233.                     "fadd v12.4s, v12.4s, %[one_half].4s\n"

  234.                     "fadd v13.4s, v13.4s, %[one_half].4s\n"

  235.                     "fadd v14.4s, v14.4s, %[one_half].4s\n"

  236.                     "fadd v15.4s, v15.4s, %[one_half].4s\n"

  237.                     "fadd v16.4s, v16.4s, %[one_half].4s\n"

  238.                     "fadd v17.4s, v17.4s, %[one_half].4s\n"

  239. 

  240.                     // save it

  241.                     "str  q12, [%[d_ptr]]     \n"

  242.                     "str  q13, [%[d_ptr], #16]\n"

  243.                     "str  q14, [%[d_ptr], #32]\n"

  244.                     "str  q15, [%[d_ptr], #48]\n"

  245.                     "str  q16, [%[d_ptr], #64]\n"

  246.                     "str  q17, [%[d_ptr], #80]\n"

  247.                     "add %[d_ptr], %[d_ptr], #96\n"

  248. 

  249.                     "bne 1b\n"

  250. 

  251.                     "2:\n"

  252.                     : [a_ptr] "+r"(src), [d_ptr] "+r"(dst), [const_ptr] "+r"(const_ptr),

  253.                       [x6_iter] "+r"(x6_iter), [lower_range] "=w"(lower_range),

  254.                       [alpha_9] "=w"(alpha_9), [upper_range] "=w"(upper_range),

  255.                       [alpha_7] "=w"(alpha_7), [alpha_5] "=w"(alpha_5),

  256.                       [alpha_3] "=w"(alpha_3), [alpha_1] "=w"(alpha_1),

  257.                       [beta_10] "=w"(beta_10), [beta_8] "=w"(beta_8),

  258.                       [beta_6] "=w"(beta_6), [beta_4] "=w"(beta_4),

  259.                       [beta_2] "=w"(beta_2), [beta_0] "=w"(beta_0),

  260.                       [one_half] "=w"(one_half)

  261.                     :

  262.                     : "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", "v10",

  263.                       "v11", "v12", "v13", "v14", "v15", "v16", "v17", "x1", "x2", "x8",

  264.                       "x9", "cc", "memory");

  265.         }

  266. 

  267.         using Op = arm_common::SigmoidOp<float, float>;

  268.         Op op;

  269.         ParamElemVisitorV2<typename Op::src_ctype> vis2;

  270.         size_t i = offset;

  271.         for (; i + Op::SIMD_WIDTH * 2 <= nr_elems; i += Op::SIMD_WIDTH * 2) {

  272.             op(vis2(src, src + Op::SIMD_WIDTH), dst);

  273.             src += Op::SIMD_WIDTH * 2;

  274.             dst += Op::SIMD_WIDTH * 2;

  275.         }

  276.         for (; i + Op::SIMD_WIDTH * 1 <= nr_elems; i += Op::SIMD_WIDTH * 1) {

  277.             op(vld1q_f32(src), dst);

  278.             src += Op::SIMD_WIDTH;

  279.             dst += Op::SIMD_WIDTH;

  280.         }

  281.         for (; i < nr_elems; i++) {

  282.             op(*src, dst);

  283.             src++;

  284.             dst++;

  285.         }

  286.     }

  287. };

  288. #endif

  289. 

  290. }  // namespace elemwise

  291. }  // namespace megdnn