nihui
/
ncnn
Not watched
1
0
Fork 0
Code
Releases 46 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
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: cfcb1cffa9
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'cfcb1cffa9'
${ noResults }
ncnn/src/layer/vulkan/shader/binaryop_broadcast.comp

554 lines
13 kB
Raw Blame History 

  1. // Tencent is pleased to support the open source community by making ncnn available.

  2. //

  3. // Copyright (C) 2019 THL A29 Limited, a Tencent company. All rights reserved.

  4. //

  5. // Licensed under the BSD 3-Clause License (the "License"); you may not use this file except

  6. // in compliance with the License. You may obtain a copy of the License at

  7. //

  8. // https://opensource.org/licenses/BSD-3-Clause

  9. //

  10. // Unless required by applicable law or agreed to in writing, software distributed

  11. // under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR

  12. // CONDITIONS OF ANY KIND, either express or implied. See the License for the

  13. // specific language governing permissions and limitations under the License.

  14. 

  15. #version 450

  16. 

  17. #if NCNN_fp16_storage

  18. #extension GL_EXT_shader_16bit_storage: require

  19. #endif

  20. #if NCNN_fp16_arithmetic

  21. #extension GL_EXT_shader_explicit_arithmetic_types_float16: require

  22. #endif

  23. 

  24. layout (constant_id = 0) const int op_type = 0;

  25. 

  26. #define shape_constant_id_offset 1

  27. layout (constant_id = shape_constant_id_offset + 0) const int adims = 0;

  28. layout (constant_id = shape_constant_id_offset + 1) const int aw = 0;

  29. layout (constant_id = shape_constant_id_offset + 2) const int ah = 0;

  30. layout (constant_id = shape_constant_id_offset + 3) const int ad = 0;

  31. layout (constant_id = shape_constant_id_offset + 4) const int ac = 0;

  32. layout (constant_id = shape_constant_id_offset + 5) const int acstep = 0;

  33. 

  34. layout (constant_id = shape_constant_id_offset + 6) const int bdims = 0;

  35. layout (constant_id = shape_constant_id_offset + 7) const int bw = 0;

  36. layout (constant_id = shape_constant_id_offset + 8) const int bh = 0;

  37. layout (constant_id = shape_constant_id_offset + 9) const int bd = 0;

  38. layout (constant_id = shape_constant_id_offset + 10) const int bc = 0;

  39. layout (constant_id = shape_constant_id_offset + 11) const int bcstep = 0;

  40. 

  41. layout (constant_id = shape_constant_id_offset + 12) const int outdims = 0;

  42. layout (constant_id = shape_constant_id_offset + 13) const int outw = 0;

  43. layout (constant_id = shape_constant_id_offset + 14) const int outh = 0;

  44. layout (constant_id = shape_constant_id_offset + 15) const int outd = 0;

  45. layout (constant_id = shape_constant_id_offset + 16) const int outc = 0;

  46. layout (constant_id = shape_constant_id_offset + 17) const int outcstep = 0;

  47. 

  48. #if NCNN_image_shader

  49. layout (binding = 0) uniform unfp sampler3D a_blob_3d;

  50. layout (binding = 1) uniform unfp sampler3D b_blob_3d;

  51. layout (binding = 2, imfmtc1) writeonly uniform unfp image3D top_blob_3d;

  52. #else

  53. layout (binding = 0) readonly buffer a_blob { sfp a_blob_data[]; };

  54. layout (binding = 1) readonly buffer b_blob { sfp b_blob_data[]; };

  55. layout (binding = 2) writeonly buffer top_blob { sfp top_blob_data[]; };

  56. #endif

  57. 

  58. layout (push_constant) uniform parameter

  59. {

  60.     int adims;

  61.     int aw;

  62.     int ah;

  63.     int ad;

  64.     int ac;

  65.     int acstep;

  66. 

  67.     int bdims;

  68.     int bw;

  69.     int bh;

  70.     int bd;

  71.     int bc;

  72.     int bcstep;

  73. 

  74.     int outdims;

  75.     int outw;

  76.     int outh;

  77.     int outd;

  78.     int outc;

  79.     int outcstep;

  80. } p;

  81. 

  82. void main()

  83. {

  84.     int gx = int(gl_GlobalInvocationID.x);

  85.     int gy = int(gl_GlobalInvocationID.y);

  86.     int gz = int(gl_GlobalInvocationID.z);

  87. 

  88.     if (gx >= psc(outw) || gy >= psc(outh) * psc(outd) || gz >= psc(outc))

  89.         return;

  90. 

  91. #if NCNN_image_shader

  92.     int ax = gx;

  93.     int ay = gy;

  94.     int az = gz;

  95.     int bx = gx;

  96.     int by = gy;

  97.     int bz = gz;

  98. 

  99.     if (psc(adims) == 4)

  100.     {

  101.         int yd = gy / psc(outh);

  102.         int yh = gy % psc(outh);

  103. 

  104.         if (psc(bdims) == 3)

  105.         {

  106.             // type 28

  107.             bx = yh;

  108.             by = yd;

  109.             bz = gz;

  110.         }

  111. 

  112.         if (psc(bdims) == 2)

  113.         {

  114.             // type 27

  115.             bx = yd;

  116.             by = gz;

  117.             bz = 0;

  118.         }

  119. 

  120.         if (psc(bdims) == 1)

  121.         {

  122.             if (psc(bw) == 1)

  123.             {

  124.                 // type 25

  125.                 bx = 0;

  126.                 by = 0;

  127.                 bz = 0;

  128.             }

  129.             else

  130.             {

  131.                 // type 26

  132.                 bx = gz;

  133.                 by = 0;

  134.                 bz = 0;

  135.             }

  136.         }

  137.     }

  138.     else if (psc(adims) == 3)

  139.     {

  140.         if (psc(bdims) == 4)

  141.         {

  142.             int yd = gy / psc(outh);

  143.             int yh = gy % psc(outh);

  144. 

  145.             // type 23

  146.             ax = yh;

  147.             ay = yd;

  148.             az = gz;

  149.         }

  150. 

  151.         if (psc(bdims) == 3)

  152.         {

  153.             if (psc(bw) == 1 && psc(bh) == 1)

  154.             {

  155.                 // special type 1

  156.                 bx = 0;

  157.                 by = 0;

  158.             }

  159. 

  160.             if (psc(bw) == psc(aw) && psc(bh) == psc(ah) && psc(bc) == 1)

  161.             {

  162.                 // special type 2

  163.                 bz = 0;

  164.             }

  165. 

  166.             if (psc(aw) == 1 && psc(ah) == 1)

  167.             {

  168.                 // special type 3

  169.                 ax = 0;

  170.                 ay = 0;

  171.             }

  172. 

  173.             if (psc(bw) == psc(aw) && psc(bh) == psc(ah) && psc(ac) == 1)

  174.             {

  175.                 // special type 4

  176.                 az = 0;

  177.             }

  178. 

  179.             if (psc(aw) != 1 && psc(bw) == 1 && psc(bh) == psc(ah) && psc(bc) == psc(ac))

  180.             {

  181.                 // special type 5

  182.                 bx = 0;

  183.             }

  184. 

  185.             if (psc(bw) == psc(aw) && psc(ah) != 1 && psc(bh) == 1 && psc(bc) == psc(ac))

  186.             {

  187.                 // special type 6

  188.                 by = 0;

  189.             }

  190. 

  191.             if (psc(bw) != 1 && psc(aw) == 1 && psc(bh) == psc(ah) && psc(bc) == psc(ac))

  192.             {

  193.                 // special type 7

  194.                 ax = 0;

  195.             }

  196. 

  197.             if (psc(bw) == psc(aw) && psc(bh) != 1 && psc(ah) == 1 && psc(bc) == psc(ac))

  198.             {

  199.                 // special type 8

  200.                 ay = 0;

  201.             }

  202.         }

  203. 

  204.         if (psc(bdims) == 2)

  205.         {

  206.             // type 18

  207.             bx = gy;

  208.             by = gz;

  209.             bz = 0;

  210.         }

  211. 

  212.         if (psc(bdims) == 1)

  213.         {

  214.             if (psc(bw) == 1)

  215.             {

  216.                 // type 16

  217.                 bx = 0;

  218.                 by = 0;

  219.                 bz = 0;

  220.             }

  221.             else

  222.             {

  223.                 // type 17

  224.                 bx = gz;

  225.                 by = 0;

  226.                 bz = 0;

  227.             }

  228.         }

  229.     }

  230.     else if (psc(adims) == 2)

  231.     {

  232.         if (psc(bdims) == 4)

  233.         {

  234.             int yd = gy / psc(outh);

  235.             int yh = gy % psc(outh);

  236. 

  237.             // type 22

  238.             ax = yd;

  239.             ay = gz;

  240.             az = 0;

  241.         }

  242. 

  243.         if (psc(bdims) == 3)

  244.         {

  245.             // type 14

  246.             ax = gy;

  247.             ay = gz;

  248.             az = 0;

  249.         }

  250. 

  251.         if (psc(bdims) == 1)

  252.         {

  253.             if (psc(bw) == 1)

  254.             {

  255.                 // type 11

  256.                 bx = 0;

  257.                 by = 0;

  258.                 bz = 0;

  259.             }

  260.             else

  261.             {

  262.                 // type 12

  263.                 bx = gy;

  264.                 by = 0;

  265.                 bz = 0;

  266.             }

  267.         }

  268.     }

  269.     else if (psc(adims) == 1)

  270.     {

  271.         if (psc(aw) == 1)

  272.         {

  273.             // type 2 3 4 20

  274.             ax = 0;

  275.             ay = 0;

  276.             az = 0;

  277.         }

  278.         else

  279.         {

  280.             if (psc(bdims) == 4)

  281.             {

  282.                 // type 21

  283.                 ax = gz;

  284.                 ay = 0;

  285.                 az = 0;

  286.             }

  287. 

  288.             if (psc(bdims) == 3)

  289.             {

  290.                 // type 9

  291.                 ax = gz;

  292.                 ay = 0;

  293.                 az = 0;

  294.             }

  295. 

  296.             if (psc(bdims) == 2)

  297.             {

  298.                 // type 8

  299.                 ax = gy;

  300.                 ay = 0;

  301.                 az = 0;

  302.             }

  303. 

  304.             if (psc(bdims) == 1)

  305.             {

  306.                 if (psc(bw) == 1)

  307.                 {

  308.                     // type 6

  309.                     bx = 0;

  310.                     by = 0;

  311.                     bz = 0;

  312.                 }

  313.             }

  314.         }

  315.     }

  316. 

  317.     afp v1 = image3d_ld1(a_blob_3d, ivec3(ax, ay, az));

  318.     afp v2 = image3d_ld1(b_blob_3d, ivec3(bx, by, bz));

  319. #else

  320.     const int gi = gz * psc(outcstep) + gy * psc(outw) + gx;

  321. 

  322.     int ai;

  323.     int bi;

  324. 

  325.     if (psc(adims) == 4)

  326.     {

  327.         int yd = gy / psc(outh);

  328.         int yh = gy % psc(outh);

  329. 

  330.         if (psc(bdims) == 3)

  331.         {

  332.             // type 28

  333.             ai = gi;

  334.             bi = gz * psc(bcstep) + yd * psc(bw) + yh;

  335.         }

  336. 

  337.         if (psc(bdims) == 2)

  338.         {

  339.             // type 27

  340.             ai = gi;

  341.             bi = gz * psc(bw) + yd;

  342.         }

  343. 

  344.         if (psc(bdims) == 1)

  345.         {

  346.             if (psc(bw) == 1)

  347.             {

  348.                 // type 25

  349.                 ai = gi;

  350.                 bi = 0;

  351.             }

  352.             else

  353.             {

  354.                 // type 26

  355.                 ai = gi;

  356.                 bi = gz;

  357.             }

  358.         }

  359.     }

  360.     else if (psc(adims) == 3)

  361.     {

  362.         if (psc(bdims) == 4)

  363.         {

  364.             int yd = gy / psc(outh);

  365.             int yh = gy % psc(outh);

  366. 

  367.             // type 23

  368.             ai = gz * psc(acstep) + yd * psc(aw) + yh;

  369.             bi = gi;

  370.         }

  371. 

  372.         if (psc(bdims) == 3)

  373.         {

  374.             if (psc(bw) == 1 && psc(bh) == 1)

  375.             {

  376.                 // special type 1

  377.                 ai = gi;

  378.                 bi = gz * psc(bcstep);

  379.             }

  380. 

  381.             if (psc(bw) == psc(aw) && psc(bh) == psc(ah) && psc(bc) == 1)

  382.             {

  383.                 // special type 2

  384.                 ai = gi;

  385.                 bi = gy * psc(bw) + gx;

  386.             }

  387. 

  388.             if (psc(aw) == 1 && psc(ah) == 1)

  389.             {

  390.                 // special type 3

  391.                 ai = gz * psc(acstep);

  392.                 bi = gi;

  393.             }

  394. 

  395.             if (psc(bw) == psc(aw) && psc(bh) == psc(ah) && psc(ac) == 1)

  396.             {

  397.                 // special type 4

  398.                 ai = gy * psc(aw) + gx;

  399.                 bi = gi;

  400.             }

  401. 

  402.             if (psc(aw) != 1 && psc(bw) == 1 && psc(bh) == psc(ah) && psc(bc) == psc(ac))

  403.             {

  404.                 // special type 5

  405.                 ai = gi;

  406.                 bi = gz * psc(bcstep) + gy;

  407.             }

  408. 

  409.             if (psc(bw) == psc(aw) && psc(ah) != 1 && psc(bh) == 1 && psc(bc) == psc(ac))

  410.             {

  411.                 // special type 6

  412.                 ai = gi;

  413.                 bi = gz * psc(bcstep) + gx;

  414.             }

  415. 

  416.             if (psc(bw) != 1 && psc(aw) == 1 && psc(bh) == psc(ah) && psc(bc) == psc(ac))

  417.             {

  418.                 // special type 7

  419.                 ai = gz * psc(acstep) + gy;

  420.                 bi = gi;

  421.             }

  422. 

  423.             if (psc(bw) == psc(aw) && psc(bh) != 1 && psc(ah) == 1 && psc(bc) == psc(ac))

  424.             {

  425.                 // special type 8

  426.                 ai = gz * psc(acstep) + gx;

  427.                 bi = gi;

  428.             }

  429.         }

  430. 

  431.         if (psc(bdims) == 2)

  432.         {

  433.             // type 18

  434.             ai = gi;

  435.             bi = gz * psc(bw) + gy;

  436.         }

  437. 

  438.         if (psc(bdims) == 1)

  439.         {

  440.             if (psc(bw) == 1)

  441.             {

  442.                 // type 16

  443.                 ai = gi;

  444.                 bi = 0;

  445.             }

  446.             else

  447.             {

  448.                 // type 17

  449.                 ai = gi;

  450.                 bi = gz;

  451.             }

  452.         }

  453.     }

  454.     else if (psc(adims) == 2)

  455.     {

  456.         if (psc(bdims) == 4)

  457.         {

  458.             int yd = gy / psc(outh);

  459.             int yh = gy % psc(outh);

  460. 

  461.             // type 22

  462.             ai = gz * psc(aw) + yd;

  463.             bi = gi;

  464.         }

  465. 

  466.         if (psc(bdims) == 3)

  467.         {

  468.             // type 14

  469.             ai = gz * psc(aw) + gy;

  470.             bi = gi;

  471.         }

  472. 

  473.         if (psc(bdims) == 1)

  474.         {

  475.             if (psc(bw) == 1)

  476.             {

  477.                 // type 11

  478.                 ai = gi;

  479.                 bi = 0;

  480.             }

  481.             else

  482.             {

  483.                 // type 12

  484.                 ai = gi;

  485.                 bi = gy;

  486.             }

  487.         }

  488.     }

  489.     else if (psc(adims) == 1)

  490.     {

  491.         if (psc(aw) == 1)

  492.         {

  493.             // type 2 3 4 20

  494.             ai = 0;

  495.             bi = gi;

  496.         }

  497.         else

  498.         {

  499.             if (psc(bdims) == 4)

  500.             {

  501.                 // type 21

  502.                 ai = gz;

  503.                 bi = gi;

  504.             }

  505. 

  506.             if (psc(bdims) == 3)

  507.             {

  508.                 // type 9

  509.                 ai = gz;

  510.                 bi = gi;

  511.             }

  512. 

  513.             if (psc(bdims) == 2)

  514.             {

  515.                 // type 8

  516.                 ai = gy;

  517.                 bi = gi;

  518.             }

  519. 

  520.             if (psc(bdims) == 1)

  521.             {

  522.                 if (psc(bw) == 1)

  523.                 {

  524.                     // type 6

  525.                     ai = gi;

  526.                     bi = 0;

  527.                 }

  528.             }

  529.         }

  530.     }

  531. 

  532.     afp v1 = buffer_ld1(a_blob_data, ai);

  533.     afp v2 = buffer_ld1(b_blob_data, bi);

  534. #endif

  535. 

  536.     afp res;

  537. 

  538.     if (op_type == 0) res = v1 + v2;

  539.     if (op_type == 1) res = v1 - v2;

  540.     if (op_type == 2) res = v1 * v2;

  541.     if (op_type == 3) res = v1 / v2;

  542.     if (op_type == 4) res = max(v1, v2);

  543.     if (op_type == 5) res = min(v1, v2);

  544.     if (op_type == 6) res = pow(v1, v2);

  545.     if (op_type == 7) res = v2 - v1;

  546.     if (op_type == 8) res = v2 / v1;

  547. 

  548. #if NCNN_image_shader

  549.     image3d_st1(top_blob_3d, ivec3(gx, gy, gz), res);

  550. #else

  551.     buffer_st1(top_blob_data, gi, res);

  552. #endif

  553. }