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: 8b94f49328
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from '8b94f49328'
${ noResults }
MegEngine/dnn/src/cuda/elemwise_helper.cpp

309 lines
9.6 kB
Raw Blame History 

  1. /**

  2.  * \file dnn/src/cuda/elemwise_helper.cpp

  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

  10.  * implied.

  11.  */

  12. 

  13. #include "src/cuda/elemwise_helper.cuh"

  14. #include "src/cuda/query_blocksize.cuh"

  15. #include "src/cuda/utils.h"

  16. 

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

  18. 

  19. #include <limits>

  20. #include <mutex>

  21. #include <unordered_map>

  22. 

  23. #define _cb_check_ndim(n) megdnn::TensorShape::MAX_NDIM == n ||

  24. static_assert(MEGDNN_FOREACH_TENSOR_NDIM(_cb_check_ndim) false,

  25.               "bad foreach ndim");

  26. #undef _cb_check_ndim

  27. 

  28. namespace megdnn {

  29. namespace cuda {

  30. 

  31. // ParamElemVisitor::init impls

  32. namespace elemwise_intl {

  33. 

  34. #pragma GCC diagnostic push

  35. #pragma GCC diagnostic ignored "-Warray-bounds"

  36. template <int ndim, typename ctype>

  37. void ParamVisitorBase<ndim, ctype, BCAST_OTHER>::host_init(

  38.         const TensorND& rv, int /*grid_size*/, int /*block_size*/,

  39.         int /*packed_size*/) {

  40.     megdnn_assert(rv.layout.ndim && rv.layout.ndim <= ndim);

  41.     m_ptr = rv.ptr<ctype>();

  42.     for (size_t i = 0; i < rv.layout.ndim; ++i) {

  43.         m_stride[i] = rv.layout.stride[i];

  44.         if (i + 1 < rv.layout.ndim)

  45.             m_shape_highdim[i] = rv.layout.shape[i + 1];

  46.     }

  47.     for (int i = rv.layout.ndim - 1; i < ndim - 1; ++i) {

  48.         m_shape_highdim[i] = 1;

  49.     }

  50.     for (int i = rv.layout.ndim; i < ndim; ++i) {

  51.         m_stride[i] = 0;

  52.     }

  53. }

  54. #pragma GCC diagnostic pop

  55. 

  56. template <typename ctype>

  57. void ParamVisitorBase<3, ctype, BCAST_101>::host_init(const TensorND& rv,

  58.                                                       int grid_size,

  59.                                                       int block_size,

  60.                                                       int packed_size) {

  61.     uint32_t shape2, shape1;

  62.     int stride1;

  63.     if (rv.layout.ndim == 3) {

  64.         megdnn_assert(!rv.layout.stride[0] && !rv.layout.stride[2]);

  65.         shape1 = rv.layout[1];

  66.         shape2 = rv.layout[2];

  67.         stride1 = rv.layout.stride[1];

  68.     } else {

  69.         megdnn_assert(rv.layout.ndim == 2 && !rv.layout.stride[1]);

  70.         shape1 = rv.layout[0];

  71.         shape2 = rv.layout[1];

  72.         stride1 = rv.layout.stride[0];

  73.     }

  74.     m_ptr = rv.ptr<ctype>();

  75.     m_stride1 = stride1;

  76.     m_shape12.host_init(packed_size * grid_size * block_size, shape2, shape1);

  77. }

  78. 

  79. template <typename ctype>

  80. void ParamVisitorBase<2, ctype, BCAST_10>::host_init(const TensorND& rv,

  81.                                                      int grid_size,

  82.                                                      int block_size,

  83.                                                      int packed_size) {

  84.     megdnn_assert(rv.layout.ndim == NDIM && !rv.layout.stride[0]);

  85.     m_ptr = rv.ptr<ctype>();

  86.     m_stride1 = rv.layout.stride[1];

  87.     m_shape1.host_init(packed_size * grid_size * block_size,

  88.                        rv.layout.shape[1]);

  89. }

  90. 

  91. template <typename ctype>

  92. void ParamVisitorBase<2, ctype, BCAST_01>::host_init(const TensorND& rv,

  93.                                                      int grid_size,

  94.                                                      int block_size,

  95.                                                      int packed_size) {

  96.     megdnn_assert(rv.layout.ndim == NDIM && !rv.layout.stride[1]);

  97.     m_ptr = rv.ptr<ctype>();

  98.     m_stride0 = rv.layout.stride[0];

  99.     m_shape1.host_init(packed_size * grid_size * block_size,

  100.                        rv.layout.shape[1]);

  101. }

  102. 

  103. template <typename ctype>

  104. void ParamVisitorBase<1, ctype, BCAST_FULL>::host_init(const TensorND& rv,

  105.                                                        int /*grid_size*/,

  106.                                                        int /*block_size*/,

  107.                                                        int /*packed_size*/) {

  108.     megdnn_assert(rv.layout.ndim == NDIM && !rv.layout.stride[0]);

  109.     m_ptr = rv.ptr<ctype>();

  110. }

  111. 

  112. template <typename ctype>

  113. void ParamVectVisitor<4, ctype, BCAST_1010>::host_init(const TensorND& rv,

  114.                                                        int grid_size,

  115.                                                        int block_size) {

  116.     megdnn_assert(rv.layout.ndim == NDIM && !rv.layout.stride[0] &&

  117.                   !rv.layout.stride[2]);

  118.     m_ptr = rv.ptr<ctype>();

  119.     m_stride1 = rv.layout.stride[1];

  120.     m_stride3 = rv.layout.stride[3];

  121.     uint32_t shape1 = rv.layout.shape[1];

  122.     uint32_t shape2 = rv.layout.shape[2];

  123.     uint32_t shape3 = rv.layout.shape[3];

  124.     m_shape123.host_init(packed_size * grid_size * block_size, shape2 * shape3,

  125.                          shape1);

  126.     m_shape3.host_init(packed_size * grid_size * block_size, shape3);

  127. }

  128. 

  129. #define INST(ndim, ctype, brd) template class ParamVisitorBase<ndim, ctype, brd>

  130. #define INST_FOR_CTYPE                  \

  131.     MEGDNN_FOREACH_TENSOR_NDIM(ndim_cb) \

  132.     INST(3, ct, BCAST_101);             \

  133.     INST(2, ct, BCAST_10);              \

  134.     INST(2, ct, BCAST_01);              \

  135.     INST(1, ct, BCAST_FULL);

  136. 

  137. #define ndim_cb(_ndim) INST(_ndim, ct, BCAST_OTHER);

  138. 

  139. #define ct dt_byte

  140. INST_FOR_CTYPE

  141. #undef ct

  142. #define ct dt_int32

  143. INST_FOR_CTYPE

  144. #undef ct

  145. #define ct dt_float32

  146. INST_FOR_CTYPE

  147. #undef ct

  148. #define ct dt_float16

  149. INST_FOR_CTYPE

  150. #undef ct

  151. #define ct dt_bfloat16

  152. INST_FOR_CTYPE

  153. #undef ct

  154. #define ct dt_int8

  155. INST_FOR_CTYPE

  156. #undef ct

  157. #define ct dt_uint8

  158. INST_FOR_CTYPE

  159. #undef ct

  160. #define ct dt_int16

  161. INST_FOR_CTYPE

  162. #undef ct

  163. #define ct dt_quint8

  164. INST_FOR_CTYPE

  165. #undef ct

  166. #define ct dt_qint8

  167. INST_FOR_CTYPE

  168. #undef ct

  169. #define ct dt_qint32

  170. INST_FOR_CTYPE

  171. #undef ct

  172. #define ct dt_bool

  173. INST_FOR_CTYPE

  174. #undef ct

  175. 

  176. #undef INST_FOR_CTYPE

  177. #undef INST

  178. 

  179. #define INST(ndim, ctype, brd) template class ParamElemVisitor<ndim, ctype, brd>

  180. #define INST_FOR_CTYPE                  \

  181.     MEGDNN_FOREACH_TENSOR_NDIM(ndim_cb) \

  182.     INST(3, ct, BCAST_101);             \

  183.     INST(2, ct, BCAST_10);              \

  184.     INST(2, ct, BCAST_01);              \

  185.     INST(1, ct, BCAST_FULL);

  186. 

  187. #define ndim_cb(_ndim) INST(_ndim, ct, BCAST_OTHER);

  188. 

  189. #define ct dt_byte

  190. INST_FOR_CTYPE

  191. #undef ct

  192. #define ct dt_int32

  193. INST_FOR_CTYPE

  194. #undef ct

  195. #define ct dt_float32

  196. INST_FOR_CTYPE

  197. #undef ct

  198. #define ct dt_float16

  199. INST_FOR_CTYPE

  200. #undef ct

  201. #define ct dt_bfloat16

  202. INST_FOR_CTYPE

  203. #undef ct

  204. #define ct dt_int8

  205. INST_FOR_CTYPE

  206. #undef ct

  207. #define ct dt_uint8

  208. INST_FOR_CTYPE

  209. #undef ct

  210. #define ct dt_int16

  211. INST_FOR_CTYPE

  212. #undef ct

  213. #define ct dt_quint8

  214. INST_FOR_CTYPE

  215. #undef ct

  216. #define ct dt_qint8

  217. INST_FOR_CTYPE

  218. #undef ct

  219. #define ct dt_qint32

  220. INST_FOR_CTYPE

  221. #undef ct

  222. #define ct dt_bool

  223. INST_FOR_CTYPE

  224. #undef ct

  225. 

  226. #undef ndim_cb

  227. 

  228. #undef INST_FOR_CTYPE

  229. #undef INST

  230. 

  231. #define INST(dt_ibyte) template class ParamVectVisitor<4, dt_ibyte, BCAST_1010>

  232. INST(dt_int8);

  233. INST(dt_uint8);

  234. INST(dt_bool);

  235. INST(dt_qint8);

  236. INST(dt_quint8);

  237. #undef dt_ibyte

  238. 

  239. template <int ndim>

  240. void ParamElemVisitor4bitBase<ndim, BCAST_OTHER>::host_init(

  241.         const TensorND& rv, int /*grid_size*/, int /*block_size*/) {

  242.     m_ptr = reinterpret_cast<Storage*>(rv.raw_ptr);

  243.     auto min_stride = rv.layout.stride[0];

  244.     for (size_t i = 0; i < rv.layout.ndim; ++i) {

  245.         m_stride[i] = rv.layout.stride[i];

  246.         m_shape[i] = rv.layout.shape[i];

  247.         if (i + 1 < rv.layout.ndim) {

  248.             m_shape_highdim[i] = rv.layout.shape[i + 1];

  249.             if (rv.layout.stride[i + 1] == 1)

  250.                 m_align_shape_highdim[i] =

  251.                         (uint32_t)round_up((int)rv.layout.shape[i + 1], 2);

  252.             else

  253.                 m_align_shape_highdim[i] = rv.layout.shape[i + 1];

  254.         }

  255.         if (min_stride > rv.layout.stride[i]) {

  256.             min_stride = rv.layout.stride[i];

  257.         }

  258.     }

  259.     megdnn_assert(min_stride == 1 || min_stride == 2);

  260.     m_is_min_stride_2 = (min_stride == 2);

  261.     for (size_t i = rv.layout.ndim - 1; i < ndim - 1; ++i) {

  262.         m_shape_highdim[i] = 1;

  263.         m_align_shape_highdim[i] = 1;

  264.     }

  265.     for (size_t i = rv.layout.ndim; i < ndim; ++i) {

  266.         m_stride[i] = 0;

  267.         m_shape[i] = 1;

  268.     }

  269.     m_is_physical_contiguous = rv.layout.is_physical_contiguous();

  270. }

  271. 

  272. #define ndim_cb(_ndim) \

  273.     template class ParamElemVisitor4bitBase<_ndim, BCAST_OTHER>;

  274. MEGDNN_FOREACH_TENSOR_NDIM(ndim_cb)

  275. #undef ndim_cb

  276. 

  277. }  // namespace elemwise_intl

  278. 

  279. void elemwise_intl::get_launch_spec(const void* kern, size_t size,

  280.                                     int* grid_size, int* block_size) {

  281.         safe_size_in_kern(size);

  282.         auto config = query_launch_config_for_kernel(kern);

  283.         *block_size = config.block_size;

  284.         int a = size / (config.block_size * 2),

  285.             b = (size - 1) / (config.block_size * 3) + 1;

  286.         if (current_device_prop().major <= 3) {

  287.             // for Kepler, less blocks (more work per thread) is faster

  288.             *grid_size = b;

  289.         } else {

  290.             *grid_size = std::max(a, b);

  291.         }

  292.         if (!*grid_size) {

  293.             *block_size = std::min<int>(std::max<int>(size / 64, 1) * 32, 1024);

  294.             *grid_size = std::max<int>(size / *block_size, 1);

  295.         }

  296.         // because we unroll 3 times in the kernel

  297.         megdnn_assert(static_cast<size_t>(*block_size) * *grid_size * 3 >=

  298.                       size);

  299.     }

  300. 

  301.     void elemwise_intl::on_bad_ndim(int ndim) {

  302.         megdnn_throw(ssprintf("invalid ndim: %d", ndim));

  303.         MEGDNN_MARK_USED_VAR(ndim);

  304.     }

  305. }  // namespace cuda

  306. }  // namespace megdnn

  307. 

  308. // vim: ft=cpp syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}