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.

nn.h 64 kB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568
  1. /**
  2. * \file dnn/include/megdnn/oprs/nn.h
  3. * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
  4. *
  5. * Copyright (c) 2014-2020 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. #pragma once
  13. #include "megdnn/internal/opr_header_prologue.h"
  14. namespace megdnn {
  15. class SeparableConvBase : public OperatorBase {
  16. DEF_OPR_IMPL_CTOR(SeparableConvBase, OperatorBase);
  17. DEF_OPR_PARAM(SeparableConv);
  18. public:
  19. using Mode = Param::Mode;
  20. protected:
  21. void deduce_layout_fwd(const TensorLayout& src,
  22. const TensorLayout& filter_x,
  23. const TensorLayout& filter_y, TensorLayout& dst);
  24. void check_layout_fwd(const TensorLayout& src, const TensorLayout& filter_x,
  25. const TensorLayout& filter_y,
  26. const TensorLayout& dst);
  27. };
  28. class SeparableConvForward : public SeparableConvBase {
  29. DEF_OPR_IMPL(SeparableConvForward, SeparableConvBase, 3, 1);
  30. public:
  31. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in filter_x,
  32. _megdnn_tensor_in filter_y, _megdnn_tensor_out dst,
  33. _megdnn_workspace workspace) = 0;
  34. void deduce_layout(const TensorLayout& src, const TensorLayout& filter_x,
  35. const TensorLayout& filter_y, TensorLayout& dst);
  36. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  37. const TensorLayout& filter_x,
  38. const TensorLayout& filter_y,
  39. const TensorLayout& dst) = 0;
  40. protected:
  41. void check_exec(const TensorLayout& src, const TensorLayout& filter_x,
  42. const TensorLayout& filter_y, const TensorLayout& dst,
  43. size_t workspace_in_bytes);
  44. };
  45. using SeparableConv = SeparableConvForward;
  46. namespace detail {
  47. struct PreprocessedFilter {
  48. //! user data; its lifetime should be bound to MegDNN Convolution
  49. //! operator
  50. void* algorithm_id;
  51. TensorNDArray tensors;
  52. };
  53. } // namespace intl
  54. /**
  55. * \brief base class for convolution operation
  56. *
  57. * This operator is supposed to perform convolution on arbitrary input
  58. * dimensions. The input/output format is N, C, dims..., and kernel format can
  59. * take two forms:
  60. * 1. OC, IC, dims..., for conventional dense convolution
  61. * 2. GROUP, OC_PER_GRP, IC_PER_GRP, dims... for sparse group convolution
  62. *
  63. * Currently, only 2D images are supported.
  64. */
  65. template <typename Parameter>
  66. class ConvolutionBase : public OperatorBase {
  67. DEF_OPR_IMPL_CTOR(ConvolutionBase, OperatorBase);
  68. using Param = Parameter;
  69. public:
  70. Param& param() { return m_param; }
  71. const Param& param() const { return m_param; }
  72. protected:
  73. Param m_param;
  74. public:
  75. static constexpr size_t MAX_SPATIAL_DIM = 2;
  76. using Mode = typename Param::Mode;
  77. struct CanonizedFilterMeta {
  78. DType dtype;
  79. typename Param::Format format;
  80. uint32_t
  81. //! whether filter should be flipped (i.e. is CONVOLUTION)
  82. should_flip,
  83. group, //!< number of groups
  84. icpg, //!< input channels per group
  85. ocpg, //!< output channels per group
  86. spatial_ndim, stride[MAX_SPATIAL_DIM], padding[MAX_SPATIAL_DIM],
  87. //! spatial dim
  88. spatial[MAX_SPATIAL_DIM], dilation[MAX_SPATIAL_DIM],
  89. //! spatial dim with dilation applied
  90. dilated_spatial[MAX_SPATIAL_DIM];
  91. //! T should be a ConvolutionBase<Z>::CanonizedFilterMeta
  92. template <typename T>
  93. void copy_from(const T& b) {
  94. dtype = b.dtype;
  95. format = b.format;
  96. should_flip = b.should_flip;
  97. group = b.group;
  98. icpg = b.icpg;
  99. ocpg = b.ocpg;
  100. spatial_ndim = b.spatial_ndim;
  101. memcpy(stride, b.stride, sizeof(stride));
  102. memcpy(padding, b.padding, sizeof(padding));
  103. memcpy(spatial, b.spatial, sizeof(spatial));
  104. memcpy(dilation, b.dilation, sizeof(dilation));
  105. memcpy(dilated_spatial, b.dilated_spatial, sizeof(dilated_spatial));
  106. }
  107. bool operator==(const CanonizedFilterMeta& b) const {
  108. bool flag = true;
  109. flag = flag && (format == b.format);
  110. flag = flag && (dtype == b.dtype);
  111. flag = flag && (should_flip == b.should_flip);
  112. flag = flag && (group == b.group);
  113. flag = flag && (icpg == b.icpg);
  114. flag = flag && (ocpg == b.ocpg);
  115. flag = flag && (spatial_ndim == b.spatial_ndim);
  116. if (flag) {
  117. for (uint32_t i = 0; i < spatial_ndim; ++i) {
  118. flag = flag && (stride[i] == b.stride[i]);
  119. flag = flag && (padding[i] == b.padding[i]);
  120. flag = flag && (spatial[i] == b.spatial[i]);
  121. flag = flag && (dilation[i] == b.dilation[i]);
  122. flag = flag && (dilated_spatial[i] == b.dilated_spatial[i]);
  123. }
  124. }
  125. return flag;
  126. }
  127. };
  128. using PreprocessedFilter = detail::PreprocessedFilter;
  129. protected:
  130. // Check or deduce output DType
  131. void check_or_deduce_dtype_fwd(DType src, DType filter, DType& dst) const;
  132. CanonizedFilterMeta deduce_layout_fwd(const TensorLayout& src,
  133. const TensorLayout& filter,
  134. TensorLayout& dst) const;
  135. CanonizedFilterMeta check_layout_fwd(const TensorLayout& src,
  136. const TensorLayout& filter,
  137. const TensorLayout& dst) const;
  138. CanonizedFilterMeta make_canonized_filter_meta(
  139. size_t src_ndim, const TensorLayout& filter) const;
  140. };
  141. class MaskPropagate : public OperatorBase {
  142. DEF_OPR_IMPL(MaskPropagate, OperatorBase, 1, 1);
  143. DEF_OPR_PARAM(MaskPropagate);
  144. public:
  145. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst,
  146. _megdnn_workspace workspace) = 0;
  147. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  148. const TensorLayout& dst) = 0;
  149. void deduce_layout(const TensorLayout& src, TensorLayout& dst);
  150. };
  151. /**
  152. * \brief ConvolutionForward Operator with 0/1 Mask matrix
  153. */
  154. class MaskConvForward : public ConvolutionBase<param::Convolution> {
  155. DEF_OPR_IMPL(MaskConvForward, ConvolutionBase, 3, 1);
  156. public:
  157. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in filter,
  158. _megdnn_tensor_in mask, _megdnn_tensor_out dst,
  159. _megdnn_workspace worksapce) = 0;
  160. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  161. const TensorLayout& filter,
  162. const TensorLayout& mask,
  163. const TensorLayout& dst) = 0;
  164. void deduce_dtype(DType src, DType filter, DType mask, DType& dst);
  165. void deduce_layout(const TensorLayout& src, const TensorLayout& filter,
  166. const TensorLayout& mask, TensorLayout& dst);
  167. protected:
  168. CanonizedFilterMeta check_exec(const TensorLayout& src,
  169. const TensorLayout& filter,
  170. const TensorLayout& mask,
  171. const TensorLayout& dst,
  172. size_t workspace_in_bytes);
  173. };
  174. using MaskConvolution = MaskConvForward;
  175. /**
  176. * \brief ConvolutionForward operator.
  177. */
  178. class ConvolutionForward : public ConvolutionBase<param::Convolution>,
  179. public detail::MultiAlgoOpr<ConvolutionForward, 3> {
  180. DEF_OPR_IMPL(ConvolutionForward, ConvolutionBase, 2, 1);
  181. public:
  182. /**
  183. * \param[in] src (n, ic, ih, iw)
  184. * \param[in] filter (oc, ic, fh, fw)
  185. * \param[out] dst (n, oc, oh, ow)
  186. */
  187. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in filter,
  188. _megdnn_tensor_out dst,
  189. const PreprocessedFilter* preprocessed_filter,
  190. _megdnn_workspace workspace) = 0;
  191. virtual void exec_preprocess(const TensorLayout& src_layout,
  192. _megdnn_tensor_in filter,
  193. const TensorLayout& dst_layout,
  194. PreprocessedFilter* preprocessed_filter,
  195. _megdnn_workspace workspace) = 0;
  196. void deduce_dtype(DType src, DType filter, DType& dst);
  197. void deduce_layout(const TensorLayout& src, const TensorLayout& filter,
  198. TensorLayout& dst);
  199. virtual size_t get_workspace_in_bytes(
  200. const TensorLayout& src, const TensorLayout& filter,
  201. const TensorLayout& dst,
  202. const PreprocessedFilter* preprocessed_filter) = 0;
  203. virtual SmallVector<TensorLayout> deduce_preprocessed_filter_layout(
  204. const TensorLayout& src, const TensorLayout& filter,
  205. const TensorLayout& dst) = 0;
  206. virtual size_t get_preprocess_workspace_in_bytes(
  207. const TensorLayout& src, const TensorLayout& filter,
  208. const TensorLayout& dst) = 0;
  209. protected:
  210. CanonizedFilterMeta check_exec(
  211. const TensorLayout& src, const TensorLayout& filter,
  212. const TensorLayout& dst, size_t workspace_in_bytes,
  213. const PreprocessedFilter* preprocessed_filter);
  214. };
  215. using Convolution = ConvolutionForward;
  216. /**
  217. * \brief ConvolutionBackwardData operator.
  218. *
  219. * Calculating the gradient wrt. convolution input data.
  220. */
  221. class ConvolutionBackwardData
  222. : public ConvolutionBase<param::Convolution>,
  223. public detail::MultiAlgoOpr<ConvolutionBackwardData, 3> {
  224. DEF_OPR_IMPL(ConvolutionBackwardData, ConvolutionBase, 2, 1);
  225. public:
  226. /**
  227. * \param[in] filter (oc, ic, fh, fw)
  228. * \param[in] diff (n, oc, oh, ow)
  229. * \param[out] grad (n, ic, ih, iw)
  230. */
  231. virtual void exec(_megdnn_tensor_in filter, _megdnn_tensor_in diff,
  232. _megdnn_tensor_out grad, _megdnn_workspace workspace) = 0;
  233. virtual size_t get_workspace_in_bytes(const TensorLayout& filter,
  234. const TensorLayout& diff,
  235. const TensorLayout& grad) = 0;
  236. void deduce_dtype(DType filter, DType diff, DType& grad);
  237. void deduce_layout(const TensorLayout& filter, const TensorLayout& diff,
  238. TensorLayout& grad);
  239. protected:
  240. CanonizedFilterMeta check_exec(const TensorLayout& filter,
  241. const TensorLayout& diff,
  242. const TensorLayout& grad,
  243. size_t workspace_in_bytes);
  244. };
  245. /**
  246. * \brief ConvolutionBackwardFilter operator.
  247. *
  248. * Calculating the gradient wrt. convolution filter.
  249. */
  250. class ConvolutionBackwardFilter
  251. : public ConvolutionBase<param::Convolution>,
  252. public detail::MultiAlgoOpr<ConvolutionBackwardFilter, 3> {
  253. DEF_OPR_IMPL(ConvolutionBackwardFilter, ConvolutionBase, 2, 1);
  254. public:
  255. /**
  256. * \param[in] src (n, ic, ih, iw)
  257. * \param[in] diff (n, oc, oh, ow)
  258. * \param[out] grad (oc, ic, fh, fw)
  259. */
  260. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in diff,
  261. _megdnn_tensor_out grad, _megdnn_workspace workspace) = 0;
  262. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  263. const TensorLayout& diff,
  264. const TensorLayout& grad) = 0;
  265. protected:
  266. CanonizedFilterMeta check_exec(const TensorLayout& src,
  267. const TensorLayout& diff,
  268. const TensorLayout& grad,
  269. size_t workspace_in_bytes);
  270. };
  271. /**
  272. * \brief ConvolutionBias operator
  273. */
  274. class ConvBiasForward : public ConvolutionBase<param::ConvBias>,
  275. public detail::MultiAlgoOpr<ConvBiasForward, 5> {
  276. DEF_OPR_IMPL(ConvBiasForward, ConvolutionBase, 4, 1);
  277. public:
  278. /**
  279. * \param[in] src (n, ic, ih, iw) or (n, ih, iw, ic)
  280. * \param[in] filter (oc, ic, fh, fw) or (oc, fh, fw, ic) or (oc/4, fh, fw,
  281. * 4 * ic)
  282. * \param[in] bias (1, oc, 1, 1)
  283. * \param[in] z same as dst
  284. * \param[out] dst (n, oc, oh, ow) or (n, oh, ow, oc)
  285. *
  286. * \note if the format is NCHW_WINOGRAD, the filter layout is (alphah,
  287. * alphaw, oc, ic)
  288. */
  289. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in filter,
  290. _megdnn_tensor_in bias, _megdnn_tensor_in z,
  291. _megdnn_tensor_out dst,
  292. const PreprocessedFilter* preprocessed_filter,
  293. _megdnn_workspace workspace) = 0;
  294. virtual void exec_preprocess(const TensorLayout& src_layout,
  295. _megdnn_tensor_in filter,
  296. const TensorLayout& bias_layout,
  297. const TensorLayout& z_layout,
  298. const TensorLayout& dst_layout,
  299. PreprocessedFilter* preprocessed_filter,
  300. _megdnn_workspace workspace) = 0;
  301. void deduce_dtype(DType src, DType filter, DType bias, DType z, DType& dst);
  302. void deduce_layout(const TensorLayout& src, const TensorLayout& filter,
  303. const TensorLayout& bias, const TensorLayout& z,
  304. TensorLayout& dst);
  305. virtual size_t get_workspace_in_bytes(
  306. const TensorLayout& src, const TensorLayout& filter,
  307. const TensorLayout& bias, const TensorLayout& z,
  308. const TensorLayout& dst,
  309. const PreprocessedFilter* preprocessed_filter) = 0;
  310. virtual size_t get_preprocess_workspace_in_bytes(
  311. const TensorLayout& src, const TensorLayout& filter,
  312. const TensorLayout& bias, const TensorLayout& z,
  313. const TensorLayout& dst) = 0;
  314. virtual SmallVector<TensorLayout> deduce_preprocessed_filter_layout(
  315. const TensorLayout& src, const TensorLayout& filter,
  316. const TensorLayout& bias, const TensorLayout& z,
  317. const TensorLayout& dst) = 0;
  318. static void deduce_winograd_origin_layout_and_param(
  319. const Param::Format format, const size_t output_block_size,
  320. const TensorLayout& src_layout,
  321. const TensorLayout& winograd_filter_layout,
  322. TensorLayout& origin_layout, Param& origin_param);
  323. enum class BiasMode : uint32_t {
  324. NO_BIAS = 0, //!< no bias
  325. BROADCAST_CHANNEL_BIAS, //!< broadcast channel bias, [1, c, 1, 1]
  326. BIAS //!< [N, C, H, W]
  327. };
  328. //! param for winograd algos.
  329. struct WinogradParam {
  330. uint32_t channel_block_size;
  331. uint32_t output_block_size;
  332. uint32_t tile_size;
  333. bool operator==(const WinogradParam& rhs) const {
  334. return channel_block_size == rhs.channel_block_size &&
  335. output_block_size == rhs.output_block_size &&
  336. tile_size == rhs.tile_size;
  337. }
  338. std::string to_string() const;
  339. };
  340. static constexpr WinogradParam INVALID_WINOGRAD_PARAM = {0, 0, 0};
  341. struct DirectParam {
  342. std::string to_string() const { return ""; }
  343. };
  344. struct MatmulParam {
  345. std::string to_string() const { return ""; }
  346. };
  347. struct DefaultParam {
  348. std::string to_string() const { return ""; }
  349. };
  350. //! get algo name, the format is ParamTrait<T>::category:base:p.to_string()
  351. //! \warning: base must not contain :.
  352. template <typename T>
  353. static std::string algo_name(
  354. const std::string& base, const T& p,
  355. param::ConvBias::Format format = param::ConvBias::Format::NCHW);
  356. /*!
  357. * \brief parse algo_name and get WinogradParam from algo name.
  358. *
  359. * \param algo name string
  360. * \return WinogradParam parsed from algo name, use pattern
  361. * winograd:base:m:tile_size.
  362. *
  363. * \warning: INVALID_WINOGRAD_PARAM returns if the algo_name is not matched.
  364. */
  365. static WinogradParam parse_winograd_name(const std::string& algo_name);
  366. /**
  367. * @brief find if there is nchw_nchwxx conv kernel optimized for argment,
  368. * nchw44 used for arm, nchw88 used for x86
  369. *
  370. * @param src_dtype conv feature map data type
  371. * @param filter_dtype conv filter or weight data type
  372. * @param dst_dtype output data type
  373. * @param fm filter meta param
  374. * @param bias_mode bias mode, no_bias or broadcast or bias
  375. * @param nonline_mode identity or relu or h_swish or sigmoid
  376. * @return true, found a kernel
  377. * @return false, can`t found any kernel
  378. */
  379. static bool is_nchw_nchwxx_optimized(
  380. const DTypeEnum src_dtype, const DTypeEnum filter_dtype,
  381. const DTypeEnum dst_dtype,
  382. const ConvolutionBase<param::Convolution>::CanonizedFilterMeta& fm,
  383. const ConvBiasForward::BiasMode bias_mode,
  384. const param::ConvBias::NonlineMode nonline_mode);
  385. protected:
  386. CanonizedFilterMeta check_exec(
  387. const TensorLayout& src, const TensorLayout& filter,
  388. const TensorLayout& bias, const TensorLayout& z,
  389. const TensorLayout& dst, size_t workspace_in_bytes,
  390. const PreprocessedFilter* preprocessed_filter);
  391. };
  392. using ConvBias = ConvBiasForward;
  393. /**
  394. * \brief base class for Conv - Nonline - Pooling
  395. */
  396. class ConvPoolingBase : public OperatorBase {
  397. DEF_OPR_IMPL_CTOR(ConvPoolingBase, OperatorBase);
  398. /**
  399. * \ Param::Method: Two methods to fetch the input data.
  400. * Default methods is WITH_TEXTURE_OBJ.
  401. * If you want to use WITH_SHARED_MEM mode,
  402. * please make sure that the size of
  403. * [ all of the fliter kernels + a channel
  404. * of input data + a channel of output data]
  405. * should be no large than 38KB.
  406. * And the pooling mode should not be "MAX".
  407. */
  408. DEF_OPR_PARAM(ConvPooling);
  409. protected:
  410. virtual void deduce_layout(const TensorLayout& src,
  411. const TensorLayout& filter,
  412. const TensorLayout& bias, TensorLayout& dst) = 0;
  413. virtual void check_layout(const TensorLayout& src,
  414. const TensorLayout& filter,
  415. const TensorLayout& bias, TensorLayout& dst,
  416. size_t workspace_limit_in_bytes) = 0;
  417. };
  418. class ConvPoolingForward : public ConvPoolingBase {
  419. DEF_OPR_IMPL(ConvPoolingForward, ConvPoolingBase, 2, 1);
  420. public:
  421. /**
  422. * \param[in] src input tensor
  423. * \param[out] dst output tensor
  424. */
  425. virtual void exec(const _megdnn_in TensorND src,
  426. const _megdnn_in TensorND filter,
  427. const _megdnn_in TensorND bias, _megdnn_out TensorND dst,
  428. _megdnn_out Workspace workspace) = 0;
  429. virtual void deduce_layout(const TensorLayout& src,
  430. const TensorLayout& filter,
  431. const TensorLayout& bias, TensorLayout& dst) = 0;
  432. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  433. const TensorLayout& filter,
  434. const TensorLayout& bias,
  435. const TensorLayout& dst) = 0;
  436. protected:
  437. virtual void check_layout(const TensorLayout& src,
  438. const TensorLayout& filter,
  439. const TensorLayout& bias, TensorLayout& dst,
  440. size_t workspace_limit_in_bytes) = 0;
  441. };
  442. using ConvPooling = ConvPoolingForward;
  443. class GroupLocalBase : public OperatorBase {
  444. DEF_OPR_IMPL_CTOR(GroupLocalBase, OperatorBase);
  445. DEF_OPR_PARAM(Convolution);
  446. public:
  447. using Mode = Param::Mode;
  448. protected:
  449. void deduce_layout_fwd(const TensorLayout& src, const TensorLayout& filter,
  450. TensorLayout& dst);
  451. void check_layout_fwd(const TensorLayout& src, const TensorLayout& filter,
  452. const TensorLayout& dst);
  453. };
  454. class GroupLocalForward : public GroupLocalBase {
  455. DEF_OPR_IMPL(GroupLocalForward, GroupLocalBase, 2, 1);
  456. public:
  457. /**
  458. * \param[in] src (N, IC, IH, IW)
  459. * \param[in] filter (G, OH, OW, IC/G, FH, FW, OC/G)
  460. * \param[out] dst (N, OC, OH, OW)
  461. **/
  462. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in filter,
  463. _megdnn_tensor_out dst, _megdnn_workspace workspace) = 0;
  464. void deduce_layout(const TensorLayout& src, const TensorLayout& filter,
  465. TensorLayout& dst) {
  466. deduce_layout_fwd(src, filter, dst);
  467. }
  468. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  469. const TensorLayout& filter,
  470. const TensorLayout& dst) = 0;
  471. protected:
  472. void check_exec(const TensorLayout& src, const TensorLayout& filter,
  473. const TensorLayout& dst, size_t workspace_in_bytes);
  474. };
  475. using GroupLocal = GroupLocalForward;
  476. class GroupLocalBackwardData : public GroupLocalBase {
  477. DEF_OPR_IMPL(GroupLocalBackwardData, GroupLocalBase, 2, 1);
  478. public:
  479. virtual void exec(_megdnn_tensor_in filter, _megdnn_tensor_in diff,
  480. _megdnn_tensor_out grad, _megdnn_workspace workspace) = 0;
  481. virtual size_t get_workspace_in_bytes(const TensorLayout& filter,
  482. const TensorLayout& diff,
  483. const TensorLayout& grad) = 0;
  484. protected:
  485. void check_exec(const TensorLayout& filter, const TensorLayout& diff,
  486. const TensorLayout& grad, size_t workspace_in_bytes);
  487. };
  488. class GroupLocalBackwardFilter : public GroupLocalBase {
  489. DEF_OPR_IMPL(GroupLocalBackwardFilter, GroupLocalBase, 2, 1);
  490. public:
  491. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in diff,
  492. _megdnn_tensor_out grad, _megdnn_workspace workspace) = 0;
  493. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  494. const TensorLayout& diff,
  495. const TensorLayout& grad) = 0;
  496. protected:
  497. void check_exec(const TensorLayout& filter, const TensorLayout& diff,
  498. const TensorLayout& grad, size_t workspace_in_bytes);
  499. };
  500. class Images2NeibsBase : public OperatorBase {
  501. DEF_OPR_IMPL_CTOR(Images2NeibsBase, OperatorBase);
  502. DEF_OPR_PARAM(Images2Neibs);
  503. protected:
  504. void deduce_layout_fwd(const TensorLayout& src, TensorLayout& dst);
  505. void check_layout_fwd(const TensorLayout& filter, const TensorLayout& dst);
  506. };
  507. class Images2NeibsForward : public Images2NeibsBase {
  508. DEF_OPR_IMPL(Images2NeibsForward, Images2NeibsBase, 1, 1);
  509. public:
  510. /**
  511. * \param[in] src (N, C, IH, IW)
  512. * \param[out] dst (N, C, OH, OW, window_h, window_w)
  513. *
  514. * \see
  515. * http://deeplearning.net/software/theano/library/tensor/nnet/neighbours.html
  516. *
  517. * \f$ dst_{n, c, oh, ow, wh, ww} = src_{n, c, ih+wh, iw+fw}\f$,
  518. * where \f$ ih=-pad_h+oh*stride_h, iw=-pad_w+ow*stride_w\f$.
  519. */
  520. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst,
  521. _megdnn_workspace workspace) = 0;
  522. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  523. const TensorLayout& dst) = 0;
  524. void deduce_layout(const TensorLayout& src, TensorLayout& dst);
  525. protected:
  526. void check_exec(const TensorLayout& src, const TensorLayout& dst,
  527. size_t workspace_in_bytes);
  528. };
  529. using Images2Neibs = Images2NeibsForward;
  530. class Images2NeibsBackward : public Images2NeibsBase {
  531. DEF_OPR_IMPL(Images2NeibsBackward, Images2NeibsBase, 1, 1);
  532. public:
  533. /**
  534. * \param[in] diff the backpropagated gradient wrt. dst
  535. * \param[out] grad the backpropagated gradient wrt. src
  536. */
  537. virtual void exec(_megdnn_tensor_in diff, _megdnn_tensor_out grad,
  538. _megdnn_workspace workspace) = 0;
  539. virtual size_t get_workspace_in_bytes(const TensorLayout& diff,
  540. const TensorLayout& grad) = 0;
  541. protected:
  542. void check_exec(const TensorLayout& diff, const TensorLayout& grad,
  543. size_t workspace_in_bytes);
  544. };
  545. /**
  546. * \brief base class for Pooling
  547. */
  548. class PoolingBase : public OperatorBase {
  549. DEF_OPR_IMPL_CTOR(PoolingBase, OperatorBase);
  550. DEF_OPR_PARAM(Pooling);
  551. public:
  552. using Mode = Param::Mode;
  553. protected:
  554. void deduce_layout_fwd(const TensorLayout& src, TensorLayout& dst);
  555. void check_layout_fwd(const TensorLayout& src, const TensorLayout& dst);
  556. };
  557. class PoolingForward : public PoolingBase {
  558. DEF_OPR_IMPL(PoolingForward, PoolingBase, 1, 1);
  559. public:
  560. /**
  561. * \param[in] src input tensor
  562. * \param[out] dst output tensor
  563. */
  564. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst,
  565. _megdnn_workspace workspace) = 0;
  566. void deduce_layout(const TensorLayout& src, TensorLayout& dst);
  567. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  568. const TensorLayout& dst) = 0;
  569. protected:
  570. void check_exec(const TensorLayout& src, const TensorLayout& dst,
  571. size_t workspace_in_bytes);
  572. };
  573. using Pooling = PoolingForward;
  574. class PoolingBackward : public PoolingBase {
  575. DEF_OPR_IMPL(PoolingBackward, PoolingBase, 3, 1);
  576. public:
  577. /**
  578. * \param[in] src the `src' parameter in PoolingForward::exec
  579. * \param[in] dst the `dst' parameter in PoolingForward::exec
  580. * \param[in] diff the backpropagated gradient wrt. dst
  581. * \param[out] grad the backpropagated gradient wrt. src
  582. */
  583. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in dst,
  584. _megdnn_tensor_in diff, _megdnn_tensor_out grad,
  585. _megdnn_workspace workspace) = 0;
  586. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  587. const TensorLayout& dst,
  588. const TensorLayout& diff,
  589. const TensorLayout& grad) = 0;
  590. protected:
  591. void check_exec(const TensorLayout& src, const TensorLayout& dst,
  592. const TensorLayout& diff, const TensorLayout& grad,
  593. size_t workspace_in_bytes);
  594. };
  595. /**
  596. * \brief base class for AdaptivePooling
  597. */
  598. class AdaptivePoolingBase : public OperatorBase {
  599. DEF_OPR_IMPL_CTOR(AdaptivePoolingBase, OperatorBase);
  600. DEF_OPR_PARAM(AdaptivePooling);
  601. protected:
  602. param::Pooling deduce_pooling_param(const TensorLayout& src,
  603. const TensorLayout& dst);
  604. };
  605. class AdaptivePoolingForward : public AdaptivePoolingBase {
  606. DEF_OPR_IMPL(AdaptivePoolingForward, AdaptivePoolingBase, 1, 1);
  607. public:
  608. /**
  609. * \param[in] src input tensor
  610. * \param[out] dst output tensor
  611. */
  612. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst,
  613. _megdnn_workspace workspace) = 0;
  614. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  615. const TensorLayout& dst) = 0;
  616. };
  617. using AdaptivePooling = AdaptivePoolingForward;
  618. class AdaptivePoolingBackward : public AdaptivePoolingBase {
  619. DEF_OPR_IMPL(AdaptivePoolingBackward, AdaptivePoolingBase, 3, 1);
  620. public:
  621. /**
  622. * \param[in] src the `src' parameter in AdaptivePoolingForward::exec
  623. * \param[in] dst the `dst' parameter in AdaptivePoolingForward::exec
  624. * \param[in] diff the backpropagated gradient wrt. dst
  625. * \param[out] grad the backpropagated gradient wrt. src
  626. */
  627. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in dst,
  628. _megdnn_tensor_in diff, _megdnn_tensor_out grad,
  629. _megdnn_workspace workspace) = 0;
  630. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  631. const TensorLayout& dst,
  632. const TensorLayout& diff,
  633. const TensorLayout& grad) = 0;
  634. };
  635. /**
  636. * \brief base class for Local
  637. */
  638. class LocalBase : public OperatorBase {
  639. DEF_OPR_IMPL_CTOR(LocalBase, OperatorBase);
  640. DEF_OPR_PARAM(Convolution);
  641. public:
  642. using Mode = Param::Mode;
  643. protected:
  644. void deduce_layout_fwd(const TensorLayout& src, const TensorLayout& filter,
  645. TensorLayout& dst);
  646. void check_layout_fwd(const TensorLayout& src, const TensorLayout& filter,
  647. const TensorLayout& dst);
  648. };
  649. class LocalForward : public LocalBase {
  650. DEF_OPR_IMPL(LocalForward, LocalBase, 2, 1);
  651. public:
  652. /**
  653. * \param[in] src (n, ic, ih, iw)
  654. * \param[in] filter (oh, ow, ic, fh, fw, oc)
  655. * \param[out] dst (n, oc, oh, ow)
  656. */
  657. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in filter,
  658. _megdnn_tensor_out dst, _megdnn_workspace workspace) = 0;
  659. /**
  660. * \brief Deducing output tensor layouts from input tensor layouts.
  661. *
  662. * Be aware that the first and second dimension of `filter' are ignored
  663. * when deducing `dst' layout.
  664. */
  665. void deduce_layout(const TensorLayout& src, const TensorLayout& filter,
  666. TensorLayout& dst);
  667. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  668. const TensorLayout& filter,
  669. const TensorLayout& dst) = 0;
  670. protected:
  671. void check_exec(const TensorLayout& src, const TensorLayout& filter,
  672. const TensorLayout& dst, size_t workspace_in_bytes);
  673. };
  674. using Local = LocalForward;
  675. class LocalBackwardData : public LocalBase {
  676. DEF_OPR_IMPL(LocalBackwardData, LocalBase, 2, 1);
  677. public:
  678. /**
  679. * \param[in] filter (oh, ow, ic, fh, fw, oc)
  680. * \param[in] diff (n, oc, oh, ow)
  681. * \param[out] grad (n, ic, ih, iw)
  682. */
  683. virtual void exec(_megdnn_tensor_in filter, _megdnn_tensor_in diff,
  684. _megdnn_tensor_out grad, _megdnn_workspace workspace) = 0;
  685. virtual size_t get_workspace_in_bytes(const TensorLayout& filter,
  686. const TensorLayout& diff,
  687. const TensorLayout& grad) = 0;
  688. protected:
  689. void check_exec(const TensorLayout& filter, const TensorLayout& diff,
  690. const TensorLayout& grad, size_t workspace_in_bytes);
  691. };
  692. class LocalBackwardFilter : public LocalBase {
  693. DEF_OPR_IMPL(LocalBackwardFilter, LocalBase, 2, 1);
  694. public:
  695. /**
  696. * \param[in] src (n, ic, ih, iw)
  697. * \param[in] diff (n, oc, oh, ow)
  698. * \param[out] grad (oh, ow, ic, fh, fw, oc)
  699. */
  700. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in diff,
  701. _megdnn_tensor_out grad, _megdnn_workspace workspace) = 0;
  702. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  703. const TensorLayout& diff,
  704. const TensorLayout& grad) = 0;
  705. protected:
  706. void check_exec(const TensorLayout& src, const TensorLayout& diff,
  707. const TensorLayout& grad, size_t workspace_in_bytes);
  708. };
  709. class BNBase : public OperatorBase {
  710. DEF_OPR_IMPL_CTOR(BNBase, OperatorBase);
  711. DEF_OPR_PARAM(BN);
  712. protected:
  713. void check_param();
  714. };
  715. class BNForward : public BNBase {
  716. DEF_OPR_IMPL(BNForward, BNBase, 6, 5);
  717. public:
  718. /**
  719. * \dst[i] = gemma
  720. * *(x[i]-estimatedMean[k])/sqrt(epsilon+estimatedVariance[k]) + beta \where
  721. * epsilon is a very small value to avoid a "divide by zero" error.
  722. * \param[in] src (n, c, h, w)
  723. * \param[out] dst (n, c, h, w)
  724. * \param[out] mean (see m_param.ParamDim) Global mean.
  725. * \param[out] variance (see m_param.ParamDim) Global variance.
  726. * \Param[out] batch_mean (see m_param.ParamDim)
  727. * Optionally cached intermediate mean from forward pass
  728. * \Param[out] batch_inv_variance (see m_param.ParamDim)
  729. * Optionally cached intermediate variance from forward pass
  730. * src and dst must have the same shape.
  731. * src and dst must be contiguous.
  732. */
  733. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in bn_scale,
  734. _megdnn_tensor_in bn_bias, _megdnn_tensor_inout mean,
  735. _megdnn_tensor_inout variance,
  736. _megdnn_tensor_out batch_mean,
  737. _megdnn_tensor_out batch_inv_variance,
  738. _megdnn_tensor_out dst, _megdnn_workspace workspace) = 0;
  739. void deduce_layout(const TensorLayout& src, TensorLayout& bn_scale,
  740. TensorLayout& bn_bias, TensorLayout& mean,
  741. TensorLayout& variance, TensorLayout& batch_mean,
  742. TensorLayout& batch_inv_variance, TensorLayout& dst);
  743. virtual size_t get_workspace_in_bytes(
  744. const TensorLayout& src, const TensorLayout& bn_scale,
  745. const TensorLayout& bn_bias, const TensorLayout& mean,
  746. const TensorLayout& variance, const TensorLayout& batch_mean,
  747. const TensorLayout& batch_inv_variance,
  748. const TensorLayout& dst) = 0;
  749. protected:
  750. void check_exec(const TensorLayout& src, const TensorLayout& bn_scale,
  751. const TensorLayout& bn_bias, const TensorLayout& mean,
  752. const TensorLayout& variance,
  753. const TensorLayout& batch_mean,
  754. const TensorLayout& batch_inv_variance,
  755. const TensorLayout& dst, size_t workspace_in_bytes);
  756. };
  757. using BN = BNForward;
  758. class BNBackward : public BNBase {
  759. DEF_OPR_IMPL(BNBackward, BNBase, 5, 3);
  760. public:
  761. /**
  762. * \param[in] input data of forwarding propagate.
  763. * \param[in] dy the backpropagated gradient of y.
  764. * \param[out] dx the backpropagated gradient of x.
  765. * \param[out] d_bn_scale, the backpropagated gradient of bn_scale.
  766. * \param[out] d_bn_bias, the backpropagated gradient of bn_bias.
  767. * Optionally cached intermediate results from forward pass
  768. * \param[in] saved_batch_mean mean of the input batch.
  769. Calculated in the forwardpropagation.
  770. * \param[in] saved_batch_variance of the input batch.
  771. Calculated in the forwardpropagation.
  772. */
  773. virtual void exec(_megdnn_tensor_in x, _megdnn_tensor_in dy,
  774. _megdnn_tensor_in saved_batch_mean,
  775. _megdnn_tensor_in saved_batch_variance,
  776. _megdnn_tensor_in bn_scale, _megdnn_tensor_out d_bn_scale,
  777. _megdnn_tensor_out d_bn_bias, _megdnn_tensor_out dx,
  778. _megdnn_workspace workspace) = 0;
  779. virtual size_t get_workspace_in_bytes(
  780. const TensorLayout& x, const TensorLayout& dy,
  781. const TensorLayout& saved_batch_mean,
  782. const TensorLayout& saved_batch_variance,
  783. const TensorLayout& bn_scale, const TensorLayout& d_bn_scale,
  784. const TensorLayout& d_bn_bias, const TensorLayout& dx) = 0;
  785. protected:
  786. void check_exec(const TensorLayout& x, const TensorLayout& dy,
  787. const TensorLayout& saved_batch_mean,
  788. const TensorLayout& saved_batch_variance,
  789. const TensorLayout& bn_scale,
  790. const TensorLayout& d_bn_scale,
  791. const TensorLayout& d_bn_bias, const TensorLayout& dx,
  792. size_t workspace_in_bytes);
  793. };
  794. class LRNBase : public OperatorBase {
  795. DEF_OPR_IMPL_CTOR(LRNBase, OperatorBase);
  796. DEF_OPR_PARAM(LRN);
  797. protected:
  798. void check_param();
  799. };
  800. class LRNForward : public LRNBase {
  801. DEF_OPR_IMPL(LRNForward, LRNBase, 1, 1);
  802. public:
  803. /**
  804. * \see ImageNet Classification with Deep Convolutional Neural Networks
  805. * \param[in] src (n, c, h, w)
  806. * \param[out] dst (n, c, h, w)
  807. *
  808. * src and dst must have the same shape.
  809. * src and dst must be contiguous.
  810. */
  811. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst,
  812. _megdnn_workspace workspace) = 0;
  813. void deduce_layout(const TensorLayout& src, TensorLayout& dst);
  814. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  815. const TensorLayout& dst) = 0;
  816. protected:
  817. void check_exec(const TensorLayout& src, const TensorLayout& dst,
  818. size_t workspace_in_bytes);
  819. };
  820. using LRN = LRNForward;
  821. class LRNBackward : public LRNBase {
  822. DEF_OPR_IMPL(LRNBackward, LRNBase, 3, 1);
  823. public:
  824. /**
  825. * \param[in] src the `src' parameter in LRNForward::exec
  826. * \param[in] dst the `dst' parameter in LRNForward::exec
  827. * \param[in] diff the backpropagated gradient wrt. dst
  828. * \param[out] grad the backpropagated gradient wrt. src
  829. *
  830. * All tensors should be contiguous and of the same shape.
  831. */
  832. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in dst,
  833. _megdnn_tensor_in diff, _megdnn_tensor_out grad,
  834. _megdnn_workspace workspace) = 0;
  835. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  836. const TensorLayout& dst,
  837. const TensorLayout& diff,
  838. const TensorLayout& grad) = 0;
  839. protected:
  840. void check_exec(const TensorLayout& src, const TensorLayout& dst,
  841. const TensorLayout& diff, const TensorLayout& grad,
  842. size_t workspace_in_bytes);
  843. };
  844. class ROIPoolingBase : public OperatorBase {
  845. DEF_OPR_IMPL_CTOR(ROIPoolingBase, OperatorBase);
  846. DEF_OPR_PARAM(ROIPooling);
  847. protected:
  848. void check_layout_fwd(const TensorLayout& src, const TensorLayout& rois,
  849. const TensorLayout& dst, const TensorLayout& index);
  850. };
  851. class ROIPoolingForward : public ROIPoolingBase {
  852. DEF_OPR_IMPL(ROIPoolingForward, ROIPoolingBase, 2, 2);
  853. public:
  854. /**
  855. * \param[in] src (n, c, ih, iw)
  856. * \param[in] rois (m, 5)
  857. * \param[out] dst (m, c, oh, ow)
  858. * \param[out] index (m, c, oh, ow) if mode is MAX, (0) if mode is AVERAGE
  859. *
  860. * The internal implementation is akin to
  861. * https://github.com/rbgirshick/caffe-fast-rcnn .d
  862. * Note that rois(, 0) denotes the input image index. We store it as
  863. * a float, but it should be an integer instead.
  864. *
  865. * index is a temporary tensor to facilitate its backward operator.
  866. * It is used to store argmax indicex in MAX mode, and it is not used
  867. * in AVERAGE mode.
  868. */
  869. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in rois,
  870. _megdnn_tensor_out dst, _megdnn_tensor_out index,
  871. _megdnn_workspace workspace) = 0;
  872. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  873. const TensorLayout& rois,
  874. const TensorLayout& dst,
  875. const TensorLayout& index) = 0;
  876. protected:
  877. void check_exec(const TensorLayout& src, const TensorLayout& rois,
  878. const TensorLayout& dst, const TensorLayout& index,
  879. size_t workspace_in_bytes);
  880. };
  881. using ROIPooling = ROIPoolingForward;
  882. class ROIPoolingBackward : public ROIPoolingBase {
  883. DEF_OPR_IMPL(ROIPoolingBackward, ROIPoolingBase, 4, 1);
  884. public:
  885. /**
  886. * \param[in] diff the backpropagated gradient wrt. dst
  887. * \param[in] src the `src' parameter in ROIPoolingForward::exec
  888. * \param[in] rois the `rois' parameter in ROIPoolingForward::exec
  889. * \param[in] index the `index' parameter in ROIPoolingForward::exec
  890. * \param[out] grad the backpropagated gradient wrt. src
  891. */
  892. virtual void exec(_megdnn_tensor_in diff, _megdnn_tensor_in src,
  893. _megdnn_tensor_in rois, _megdnn_tensor_in index,
  894. _megdnn_tensor_out grad, _megdnn_workspace workspace) = 0;
  895. virtual size_t get_workspace_in_bytes(const TensorLayout& diff,
  896. const TensorLayout& src,
  897. const TensorLayout& rois,
  898. const TensorLayout& index,
  899. const TensorLayout& grad) = 0;
  900. protected:
  901. void check_exec(const TensorLayout& diff, const TensorLayout& src,
  902. const TensorLayout& rois, const TensorLayout& index,
  903. const TensorLayout& grad, size_t workspace_in_bytes);
  904. };
  905. class Convolution3DBase : public OperatorBase {
  906. DEF_OPR_IMPL_CTOR(Convolution3DBase, OperatorBase);
  907. DEF_OPR_PARAM(Convolution3D);
  908. public:
  909. static constexpr size_t MAX_SPATIAL_DIM = 3;
  910. using Mode = Param::Mode;
  911. struct CanonizedFilterMeta {
  912. DTypeEnum dtype_enum;
  913. Param::Format format;
  914. uint32_t
  915. //! whether filter should be flipped (i.e. is CONVOLUTION)
  916. should_flip,
  917. group, //!< number of groups
  918. icpg, //!< input channels per group
  919. ocpg, //!< output channels per group
  920. spatial_ndim, stride[MAX_SPATIAL_DIM], padding[MAX_SPATIAL_DIM],
  921. //! spatial dim
  922. spatial[MAX_SPATIAL_DIM], dilation[MAX_SPATIAL_DIM],
  923. //! spatial dim with dilation applied
  924. dilated_spatial[MAX_SPATIAL_DIM];
  925. } MEGDNN_PACKED;
  926. protected:
  927. CanonizedFilterMeta deduce_layout_fwd(const TensorLayout& src,
  928. const TensorLayout& filter,
  929. TensorLayout& dst) const;
  930. CanonizedFilterMeta check_layout_fwd(const TensorLayout& src,
  931. const TensorLayout& filter,
  932. const TensorLayout& dst) const;
  933. CanonizedFilterMeta make_canonized_filter_meta(
  934. size_t src_ndim, const TensorLayout& filter) const;
  935. };
  936. class Convolution3DForward
  937. : public Convolution3DBase,
  938. public detail::MultiAlgoOpr<Convolution3DForward, 3> {
  939. DEF_OPR_IMPL(Convolution3DForward, Convolution3DBase, 2, 1);
  940. public:
  941. /**
  942. * \param[in] src (n, ic, id, ih, iw)
  943. * \param[in] filter (oc, ic, fd, fh, fw)
  944. * \param[out] dst (n, oc, od, oh, ow)
  945. */
  946. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in filter,
  947. _megdnn_tensor_out dst, _megdnn_workspace workspace) = 0;
  948. void deduce_layout(const TensorLayout& src, const TensorLayout& filter,
  949. TensorLayout& dst);
  950. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  951. const TensorLayout& filter,
  952. const TensorLayout& dst) = 0;
  953. protected:
  954. CanonizedFilterMeta check_exec(const TensorLayout& src,
  955. const TensorLayout& filter,
  956. const TensorLayout& dst,
  957. size_t workspace_in_bytes);
  958. };
  959. using Convolution3D = Convolution3DForward;
  960. class Convolution3DBackwardData
  961. : public Convolution3DBase,
  962. public detail::MultiAlgoOpr<Convolution3DBackwardData, 3> {
  963. DEF_OPR_IMPL(Convolution3DBackwardData, Convolution3DBase, 2, 1);
  964. public:
  965. /**
  966. * \param[in] filter (oc, ic, fd, fh, fw)
  967. * \param[in] diff (n, oc, od, oh, ow)
  968. * \param[out] grad (n, ic, id, ih, iw)
  969. */
  970. virtual void exec(_megdnn_tensor_in filter, _megdnn_tensor_in diff,
  971. _megdnn_tensor_out grad, _megdnn_workspace workspace) = 0;
  972. virtual size_t get_workspace_in_bytes(const TensorLayout& filter,
  973. const TensorLayout& diff,
  974. const TensorLayout& grad) = 0;
  975. void deduce_layout(const TensorLayout& filter, const TensorLayout& diff,
  976. TensorLayout& grad);
  977. protected:
  978. CanonizedFilterMeta check_exec(const TensorLayout& filter,
  979. const TensorLayout& diff,
  980. const TensorLayout& grad,
  981. size_t workspace_in_bytes);
  982. };
  983. class Convolution3DBackwardFilter
  984. : public Convolution3DBase,
  985. public detail::MultiAlgoOpr<Convolution3DBackwardFilter, 3> {
  986. DEF_OPR_IMPL(Convolution3DBackwardFilter, Convolution3DBase, 2, 1);
  987. public:
  988. /**
  989. * \param[in] src (n, ic, id, ih, iw)
  990. * \param[in] diff (n, oc, od, oh, ow)
  991. * \param[out] grad (oc, ic, fd, fh, fw)
  992. */
  993. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in diff,
  994. _megdnn_tensor_out grad, _megdnn_workspace workspace) = 0;
  995. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  996. const TensorLayout& diff,
  997. const TensorLayout& grad) = 0;
  998. protected:
  999. CanonizedFilterMeta check_exec(const TensorLayout& src,
  1000. const TensorLayout& diff,
  1001. const TensorLayout& grad,
  1002. size_t workspace_in_bytes);
  1003. };
  1004. class LocalShareBase : public OperatorBase {
  1005. DEF_OPR_IMPL_CTOR(LocalShareBase, OperatorBase);
  1006. DEF_OPR_PARAM(LocalShare);
  1007. protected:
  1008. void deduce_layout_fwd(const TensorLayout& src, const TensorLayout& filter,
  1009. TensorLayout& dst);
  1010. void check_layout_fwd(const TensorLayout& src, const TensorLayout& filter,
  1011. const TensorLayout& dst);
  1012. };
  1013. class LocalShareForward : public LocalShareBase,
  1014. public detail::MultiAlgoOpr<LocalShareForward, 3> {
  1015. DEF_OPR_IMPL(LocalShareForward, LocalShareBase, 2, 1);
  1016. public:
  1017. /**
  1018. * \param[in] src (N, IC, IH, IW)
  1019. * \param[in] filter (G, spatial_groups_h, spatial_groups_w, IC / G,
  1020. * FH, FW, OC / G)
  1021. * \param[out] dst (N, OC, OH, OW)
  1022. */
  1023. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in filter,
  1024. _megdnn_tensor_out dst, _megdnn_workspace workspace) = 0;
  1025. /**
  1026. * \brief deduce layout of the ouput tensor
  1027. */
  1028. void deduce_layout(const TensorLayout& src, const TensorLayout& filter,
  1029. TensorLayout& dst);
  1030. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  1031. const TensorLayout& filter,
  1032. const TensorLayout& dst) = 0;
  1033. protected:
  1034. void check_exec(const TensorLayout& src, const TensorLayout& filter,
  1035. const TensorLayout& dst, size_t workspace_in_bytes);
  1036. };
  1037. using LocalShare = LocalShareForward;
  1038. class LocalShareBackwardData
  1039. : public LocalShareBase,
  1040. public detail::MultiAlgoOpr<LocalShareBackwardData, 3> {
  1041. DEF_OPR_IMPL(LocalShareBackwardData, LocalShareBase, 2, 1);
  1042. public:
  1043. /**
  1044. * \param[in] filter (G, spatial_groups_h, spatial_groups_w, IC / G,
  1045. * FH, FW, OC / G)
  1046. * \param[in] diff (N, OC, OH, OW)
  1047. * \param[out] grad (N, IC, IH, IW)
  1048. */
  1049. virtual void exec(_megdnn_tensor_in filter, _megdnn_tensor_in diff,
  1050. _megdnn_tensor_out grad, _megdnn_workspace workspace) = 0;
  1051. virtual size_t get_workspace_in_bytes(const TensorLayout& filter,
  1052. const TensorLayout& diff,
  1053. const TensorLayout& grad) = 0;
  1054. void deduce_layout(const TensorLayout& filter, const TensorLayout& diff,
  1055. TensorLayout& grad);
  1056. protected:
  1057. void check_exec(const TensorLayout& filter, const TensorLayout& diff,
  1058. const TensorLayout& grad, size_t workspace_in_bytes);
  1059. };
  1060. class LocalShareBackwardFilter
  1061. : public LocalShareBase,
  1062. public detail::MultiAlgoOpr<LocalShareBackwardFilter, 3> {
  1063. DEF_OPR_IMPL(LocalShareBackwardFilter, LocalShareBase, 2, 1);
  1064. public:
  1065. /**
  1066. * \param[in] src (N, IC, IH, IW)
  1067. * \param[in] diff (N, OC, OH, OW)
  1068. * \param[out] grad (G, spatial_groups_h, spatial_groups_w, IC / G,
  1069. * FH, FW, OC / G)
  1070. */
  1071. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in diff,
  1072. _megdnn_tensor_out grad, _megdnn_workspace workspace) = 0;
  1073. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  1074. const TensorLayout& diff,
  1075. const TensorLayout& grad) = 0;
  1076. protected:
  1077. void check_exec(const TensorLayout& src, const TensorLayout& diff,
  1078. const TensorLayout& grad, size_t workspace_in_bytes);
  1079. };
  1080. class ROIAlignBase : public OperatorBase {
  1081. DEF_OPR_IMPL_CTOR(ROIAlignBase, OperatorBase);
  1082. DEF_OPR_PARAM(ROIAlign);
  1083. protected:
  1084. void deduce_layout_fwd(const TensorLayout& src, const TensorLayout& rois,
  1085. TensorLayout& dst, TensorLayout& index);
  1086. void check_layout_fwd(const TensorLayout& src, const TensorLayout& rois,
  1087. const TensorLayout& dst, const TensorLayout& index);
  1088. };
  1089. class ROIAlignForward : public ROIAlignBase {
  1090. DEF_OPR_IMPL(ROIAlignForward, ROIAlignBase, 2, 2);
  1091. public:
  1092. /**
  1093. * \param[in] src (n, c, ih, iw)
  1094. * \param[in] rois (m, 5)
  1095. * \param[out] dst (m, c, oh, ow)
  1096. * \param[out] index (m, c, oh, ow) if mode is MAX, (0) if mode is AVERAGE
  1097. *
  1098. * Note that rois(, 0) denotes the input image index. We store it as
  1099. * a float, but it should be an integer instead.
  1100. *
  1101. * index is a temporary tensor to facilitate its backward operator.
  1102. * It is used to store argmax indicex in MAX mode, and it is not used
  1103. * in AVERAGE mode.
  1104. */
  1105. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in rois,
  1106. _megdnn_tensor_out dst, _megdnn_tensor_out index,
  1107. _megdnn_workspace workspace) = 0;
  1108. void deduce_layout(const TensorLayout& src, const TensorLayout& rois,
  1109. TensorLayout& dst, TensorLayout& index);
  1110. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  1111. const TensorLayout& rois,
  1112. const TensorLayout& dst,
  1113. const TensorLayout& index) = 0;
  1114. protected:
  1115. void check_exec(const TensorLayout& src, const TensorLayout& rois,
  1116. const TensorLayout& dst, const TensorLayout& index,
  1117. size_t workspace_in_bytes);
  1118. };
  1119. using ROIAlign = ROIAlignForward;
  1120. class ROIAlignBackward : public ROIAlignBase {
  1121. DEF_OPR_IMPL(ROIAlignBackward, ROIAlignBase, 3, 1);
  1122. public:
  1123. /**
  1124. * \param[in] diff the backpropagated gradient wrt. dst
  1125. * \param[in] rois the `rois' parameter in ROIAlignForward::exec
  1126. * \param[in] index the `index' parameter in ROIAlignForward::exec
  1127. * \param[out] grad the backpropagated gradient wrt. src
  1128. */
  1129. virtual void exec(_megdnn_tensor_in diff, _megdnn_tensor_in rois,
  1130. _megdnn_tensor_in index, _megdnn_tensor_out grad,
  1131. _megdnn_workspace workspace) = 0;
  1132. virtual size_t get_workspace_in_bytes(const TensorLayout& diff,
  1133. const TensorLayout& rois,
  1134. const TensorLayout& index,
  1135. const TensorLayout& grad) = 0;
  1136. protected:
  1137. void check_exec(const TensorLayout& diff, const TensorLayout& rois,
  1138. const TensorLayout& index, const TensorLayout& grad,
  1139. size_t workspace_in_bytes);
  1140. };
  1141. class DeformableConvBase : public OperatorBase {
  1142. DEF_OPR_IMPL_CTOR(DeformableConvBase, OperatorBase);
  1143. DEF_OPR_PARAM(Convolution);
  1144. public:
  1145. static constexpr size_t MAX_SPATIAL_DIM = 2;
  1146. struct CanonizedFilterMeta : Convolution::CanonizedFilterMeta {
  1147. uint32_t deformable_group;
  1148. };
  1149. protected:
  1150. CanonizedFilterMeta make_canonized_filter_meta(
  1151. size_t src_ndim, const TensorLayout& filter,
  1152. const TensorLayout& offset) const;
  1153. void deduce_layout_fwd(const TensorLayout& im, const TensorLayout& filter,
  1154. const TensorLayout& mask, const TensorLayout& offset,
  1155. TensorLayout& dst);
  1156. void check_layout_fwd(const TensorLayout& src, const TensorLayout& filter,
  1157. const TensorLayout& mask, const TensorLayout& offset,
  1158. const TensorLayout& dst);
  1159. };
  1160. class DeformableConvForward
  1161. : public DeformableConvBase,
  1162. public detail::MultiAlgoOpr<DeformableConvForward, 5> {
  1163. DEF_OPR_IMPL(DeformableConvForward, DeformableConvBase, 4, 1);
  1164. public:
  1165. /**
  1166. * \param[in] im (n, ic, ih, iw)
  1167. * \param[in] filter (oc, ic, fh, fw)
  1168. * \param[in] offset (dg, 2, fh, fw, oh, ow)
  1169. * \param[in] mask (dg, fh, fw, oh, ow)
  1170. * \param[out] dst (n, oc, oh, ow)
  1171. */
  1172. virtual void exec(_megdnn_tensor_in im, _megdnn_tensor_in filter,
  1173. _megdnn_tensor_in offset, _megdnn_tensor_in mask,
  1174. _megdnn_tensor_out dst, _megdnn_workspace workspace) = 0;
  1175. void deduce_layout(const TensorLayout& im, const TensorLayout& filter,
  1176. const TensorLayout& offset, const TensorLayout& mask,
  1177. TensorLayout& dst);
  1178. virtual size_t get_workspace_in_bytes(const TensorLayout& im,
  1179. const TensorLayout& filter,
  1180. const TensorLayout& offset,
  1181. const TensorLayout& mask,
  1182. const TensorLayout& dst) = 0;
  1183. protected:
  1184. CanonizedFilterMeta check_exec(const TensorLayout& im,
  1185. const TensorLayout& filter,
  1186. const TensorLayout& offset,
  1187. const TensorLayout& mask,
  1188. const TensorLayout& dst,
  1189. size_t workspace_in_bytes);
  1190. };
  1191. using DeformableConv = DeformableConvForward;
  1192. /**
  1193. * \brief DeformableConvBackwardFilter operator.
  1194. *
  1195. * Calculating the gradient wrt. convolution filter.
  1196. */
  1197. class DeformableConvBackwardFilter
  1198. : public DeformableConvBase,
  1199. public detail::MultiAlgoOpr<DeformableConvBackwardFilter, 5> {
  1200. DEF_OPR_IMPL(DeformableConvBackwardFilter, DeformableConvBase, 4, 1);
  1201. public:
  1202. /**
  1203. * \param[in] im (oc, ic, fh, fw)
  1204. * \param[in] offset (dg, 2, fh, fw, oh, ow)
  1205. * \param[in] mask (dg, fh, fw, oh, ow)
  1206. * \param[in] out_grad (n, oc, oh, ow)
  1207. * \param[out] filter_grad (oc, ic, ih, iw)
  1208. */
  1209. virtual void exec(_megdnn_tensor_in im, _megdnn_tensor_in offset,
  1210. _megdnn_tensor_in mask, _megdnn_tensor_in out_grad,
  1211. _megdnn_tensor_out filter_grad,
  1212. _megdnn_workspace workspace) = 0;
  1213. virtual size_t get_workspace_in_bytes(const TensorLayout& im,
  1214. const TensorLayout& offset,
  1215. const TensorLayout& mask,
  1216. const TensorLayout& out_grad,
  1217. const TensorLayout& filter_grad) = 0;
  1218. void deduce_layout(const TensorLayout& im, const TensorLayout& offset,
  1219. const TensorLayout& mask, const TensorLayout& out_grad,
  1220. TensorLayout& filter_grad);
  1221. protected:
  1222. CanonizedFilterMeta check_exec(const TensorLayout& im,
  1223. const TensorLayout& offset,
  1224. const TensorLayout& mask,
  1225. const TensorLayout& out_grad,
  1226. const TensorLayout& filter_grad,
  1227. size_t workspace_in_bytes);
  1228. };
  1229. /**
  1230. * \brief DeformableConvBackwardData operator.
  1231. *
  1232. * Calculating the gradient wrt. convolution input data, offset and mask.
  1233. */
  1234. class DeformableConvBackwardData
  1235. : public DeformableConvBase,
  1236. public detail::MultiAlgoOpr<DeformableConvBackwardData, 8> {
  1237. DEF_OPR_IMPL(DeformableConvBackwardData, DeformableConvBase, 5, 3);
  1238. public:
  1239. /**
  1240. * \param[in] im (oc, ic, fh, fw)
  1241. * \param[in] filter (oc, ic, fh, fw)
  1242. * \param[in] offset (dg, 2, fh, fw, oh, ow)
  1243. * \param[in] mask (dg, fh, fw, oh, ow)
  1244. * \param[in] out_grad (n, oc, oh, ow)
  1245. * \param[out] im_grad (n, ic, ih, iw)
  1246. * \param[out] offset_grad (dg, 2, fh, fw, oh, ow)
  1247. * \param[out] mask_grad (dg, fh, fw, oh, ow)
  1248. */
  1249. virtual void exec(_megdnn_tensor_in im, _megdnn_tensor_in filter,
  1250. _megdnn_tensor_in offset, _megdnn_tensor_in mask,
  1251. _megdnn_tensor_in out_grad, _megdnn_tensor_out im_grad,
  1252. _megdnn_tensor_out offset_grad,
  1253. _megdnn_tensor_out mask_grad,
  1254. _megdnn_workspace workspace) = 0;
  1255. virtual size_t get_workspace_in_bytes(
  1256. const TensorLayout& im, const TensorLayout& filter,
  1257. const TensorLayout& offset, const TensorLayout& mask,
  1258. const TensorLayout& out_grad, const TensorLayout& im_grad,
  1259. const TensorLayout& offset_grad, const TensorLayout& mask_grad) = 0;
  1260. void deduce_layout(const TensorLayout& im, const TensorLayout& filter,
  1261. const TensorLayout& offset, const TensorLayout& mask,
  1262. const TensorLayout& out_grad, TensorLayout& im_grad,
  1263. TensorLayout& offset_grad, TensorLayout& mask_grad);
  1264. protected:
  1265. CanonizedFilterMeta check_exec(
  1266. const TensorLayout& im, const TensorLayout& filter,
  1267. const TensorLayout& offset, const TensorLayout& mask,
  1268. const TensorLayout& out_grad, const TensorLayout& im_grad,
  1269. const TensorLayout& offset_grad, const TensorLayout& mask_grad,
  1270. size_t workspace_in_bytes);
  1271. };
  1272. class DeformablePSROIPoolingBase : public OperatorBase {
  1273. DEF_OPR_IMPL_CTOR(DeformablePSROIPoolingBase, OperatorBase);
  1274. DEF_OPR_PARAM(DeformablePSROIPooling);
  1275. protected:
  1276. void deduce_layout_fwd(const TensorLayout& data, const TensorLayout& trans,
  1277. const TensorLayout& rois, TensorLayout& out_data,
  1278. TensorLayout& out_count);
  1279. void check_layout_fwd(const TensorLayout& data, const TensorLayout& trans,
  1280. const TensorLayout& rois,
  1281. const TensorLayout& out_data,
  1282. const TensorLayout& out_count,
  1283. size_t workspace_in_bytes);
  1284. };
  1285. class DeformablePSROIPoolingForward : public DeformablePSROIPoolingBase {
  1286. DEF_OPR_IMPL(DeformablePSROIPoolingForward, DeformablePSROIPoolingBase, 3,
  1287. 2);
  1288. public:
  1289. /**
  1290. * \param[in] data (oc, ic, ih, iw)
  1291. * \param[in] rois (xx, xx, xx, xx)
  1292. * \param[in] trans (oc, ic, fh, fw)
  1293. * \param[out] out_data ( n, ic, ih, iw)
  1294. * \param[out] out_count ( n, ic, ih, iw)
  1295. */
  1296. virtual size_t get_workspace_in_bytes(const TensorLayout& data,
  1297. const TensorLayout& rois,
  1298. const TensorLayout& trans,
  1299. const TensorLayout& out_data,
  1300. const TensorLayout& out_count) = 0;
  1301. virtual void exec(_megdnn_tensor_in data, _megdnn_tensor_in rois,
  1302. _megdnn_tensor_in trans, _megdnn_tensor_out out_data,
  1303. _megdnn_tensor_out out_count,
  1304. _megdnn_workspace workspace) = 0;
  1305. void deduce_layout(const TensorLayout& data, const TensorLayout& rois,
  1306. const TensorLayout& trans, TensorLayout& out_data,
  1307. TensorLayout& out_count);
  1308. void check_exec(const TensorLayout& data, const TensorLayout& rois,
  1309. const TensorLayout& trans, const TensorLayout& out_data,
  1310. const TensorLayout& out_count, size_t workspace_in_bytes);
  1311. };
  1312. using DeformablePSROIPooling = DeformablePSROIPoolingForward;
  1313. class DeformablePSROIPoolingBackward : public DeformablePSROIPoolingBase {
  1314. DEF_OPR_IMPL(DeformablePSROIPoolingBackward, DeformablePSROIPoolingBase, 5,
  1315. 2);
  1316. public:
  1317. /**
  1318. * \param[in] data (oc, ic, ih, iw)
  1319. * \param[in] rois (xx, xx, xx, xx)
  1320. * \param[in] trans (oc, ic, fh, fw)
  1321. * \param[in] out_diff (xx, xx, xx, xx)
  1322. * \param[in] out_count (xx, xx, xx, xx)
  1323. * \param[out] data_diff ( n, ic, ih, iw)
  1324. * \param[out] trans_diff ( n, ic, ih, iw)
  1325. */
  1326. virtual void exec(_megdnn_tensor_in data, _megdnn_tensor_in rois,
  1327. _megdnn_tensor_in trans, _megdnn_tensor_in out_diff,
  1328. _megdnn_tensor_in out_count, _megdnn_tensor_out data_diff,
  1329. _megdnn_tensor_out trans_diff,
  1330. _megdnn_workspace workspace) = 0;
  1331. virtual size_t get_workspace_in_bytes(const TensorLayout& data,
  1332. const TensorLayout& rois,
  1333. const TensorLayout& trans,
  1334. const TensorLayout& out_diff,
  1335. const TensorLayout& out_count,
  1336. const TensorLayout& data_diff,
  1337. const TensorLayout& trans_diff) = 0;
  1338. void check_exec(const TensorLayout& data, const TensorLayout& rois,
  1339. const TensorLayout& trans, const TensorLayout& out_diff,
  1340. const TensorLayout& out_count,
  1341. const TensorLayout& data_diff,
  1342. const TensorLayout& trans_diff, size_t workspace_in_bytes);
  1343. };
  1344. class BatchConvBiasForward
  1345. : public ConvolutionBase<param::BatchConvBias>,
  1346. public detail::MultiAlgoOpr<BatchConvBiasForward, 5> {
  1347. DEF_OPR_IMPL(BatchConvBiasForward, ConvolutionBase, 4, 1);
  1348. public:
  1349. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_in filter,
  1350. _megdnn_tensor_in bias, _megdnn_tensor_in z,
  1351. _megdnn_tensor_out dst, _megdnn_workspace workspace) = 0;
  1352. void deduce_dtype(DType src, DType filter, DType bias, DType z, DType& dst);
  1353. void deduce_layout(const TensorLayout& src, const TensorLayout& filter,
  1354. const TensorLayout& bias, const TensorLayout& z,
  1355. TensorLayout& dst);
  1356. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  1357. const TensorLayout& filter,
  1358. const TensorLayout& bias,
  1359. const TensorLayout& z,
  1360. const TensorLayout& dst) = 0;
  1361. protected:
  1362. CanonizedFilterMeta check_exec(const TensorLayout& src,
  1363. const TensorLayout& filter,
  1364. const TensorLayout& bias,
  1365. const TensorLayout& z,
  1366. const TensorLayout& dst,
  1367. size_t workspace_in_bytes);
  1368. };
  1369. using BatchConvBias = BatchConvBiasForward;
  1370. } // namespace megdnn
  1371. #include "megdnn/internal/opr_header_epilogue.h"
  1372. // vim: syntax=cpp.doxygen

MegEngine 安装包中集成了使用 GPU 运行代码所需的 CUDA 环境,不用区分 CPU 和 GPU 版。 如果想要运行 GPU 程序,请确保机器本身配有 GPU 硬件设备并安装好驱动。 如果你想体验在云端 GPU 算力平台进行深度学习开发的感觉,欢迎访问 MegStudio 平台