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.

general.h 52 kB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407
  1. /**
  2. * \file dnn/include/megdnn/oprs/general.h
  3. * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
  4. *
  5. * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
  6. *
  7. * Unless required by applicable law or agreed to in writing,
  8. * software distributed under the License is distributed on an
  9. * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  10. */
  11. #pragma once
  12. #include "megdnn/internal/opr_header_prologue.h"
  13. #include "megdnn/thin/small_vector.h"
  14. namespace megdnn {
  15. /*!
  16. * \brief standard element-wise operator
  17. *
  18. * Inputs must have same dtype, and their shapes must broadcastable into a final
  19. * shape. They can have arbitrary layouts, but non-contiguous and non-broadcast
  20. * layouts may harm performance seriously.
  21. *
  22. * Output dtype is the same as input dtype (note that even for compare oprs this
  23. * is true, e.g. float == float returns value of float). Output layout must be
  24. * contiguous.
  25. */
  26. class ElemwiseForward: public OperatorBase {
  27. DEF_OPR_PARAM(Elemwise);
  28. DEF_OPR_IMPL(ElemwiseForward, OperatorBase, -1, 1);
  29. public:
  30. using Mode = Param::Mode;
  31. //! information about a mode
  32. struct ModeTrait {
  33. uint32_t arity; //!< number of inputs needed
  34. bool commutable; //!< whether arity == 2 and inputs commutable
  35. bool allow_int; //!< whether int inputs allowed
  36. bool allow_float; //!< whether float inputs allowed
  37. bool allow_bool; //!< whether bool inputs allowed
  38. const char* name; //!< name of the mode
  39. ModeTrait():
  40. arity(0), commutable(0), allow_int(0), allow_float(0), allow_bool(0),
  41. name(NULL)
  42. {}
  43. //! get trait from a mode; this function is thread safe
  44. static const ModeTrait& from_mode(Mode mode);
  45. };
  46. //! get trait of current mode
  47. const ModeTrait& mode_trait() const {
  48. return ModeTrait::from_mode(m_param.mode);
  49. }
  50. /**
  51. * \param[in] src input tensor
  52. * \param[out] dst output tensor
  53. *
  54. * src and dst should have the same shape;
  55. * layouts should be contiguous;
  56. * the underlying data pointer can point to the same memory region for
  57. * src and dst.
  58. */
  59. virtual void exec(_megdnn_in const TensorNDArray &src,
  60. _megdnn_tensor_out dst) = 0;
  61. //! deduce output shape (do not check whether arity matches)
  62. static void deduce_shape(
  63. const TensorShapeArray &src,
  64. TensorShape &dst);
  65. static void deduce_format(const TensorFormatArray& src,
  66. TensorFormat& dst);
  67. //! deduce output layout
  68. void deduce_layout(const TensorLayoutArray &src,
  69. TensorLayout &dst);
  70. protected:
  71. //! throw exception if incorrect layout; broadcast input shape to
  72. //! output shape
  73. void check_layout_and_broadcast(
  74. const TensorLayoutPtrArray &src, const TensorLayout &dst);
  75. private:
  76. void check_dtype(DType dtype);
  77. };
  78. using Elemwise = ElemwiseForward;
  79. /*!
  80. * \brief compute ``x**a`` where ``a`` is a constant from the Param
  81. *
  82. * This opr is usually not directly accessible by the end user and it is created
  83. * by mgb optimizer, aiming to work around numerical stability issues with pow.
  84. * For example ``powf(x, 2.f)`` with ``x < 0`` in fast math mode may return NaN.
  85. *
  86. * Like elemwise, this opr supports arbitrary strides. But it should only be
  87. * used with monotone strides. Input and output should have the same
  88. * float-category dtype.
  89. */
  90. class PowC : public OperatorBase {
  91. DEF_OPR_PARAM(PowC);
  92. DEF_OPR_IMPL(PowC, OperatorBase, 1, 1);
  93. public:
  94. void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst);
  95. //! compatible API for mgb; workspace is not used
  96. void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst,
  97. _megdnn_workspace) {
  98. return exec(src, dst);
  99. }
  100. size_t get_workspace_in_bytes(const TensorLayout&, const TensorLayout&) {
  101. // the impls should require no workspace; this can be later changed to a
  102. // virtual function if this situation changes
  103. return 0;
  104. }
  105. void deduce_layout(const TensorLayout& src, TensorLayout& dst) {
  106. dst.dtype = src.dtype;
  107. dst.init_contiguous_stride(src);
  108. }
  109. protected:
  110. /*!
  111. * Perform the computing where layouts have been verified.
  112. *
  113. * \p src can have arbitrary layout, and \p dst is contiguous. They have the
  114. * same shape and dtype.
  115. *
  116. * The implementation should not access param(). It should check \p exp_f
  117. * and \p exp_i for the exponent value. Exactly one of them would be
  118. * non-null.
  119. *
  120. * Note: \p exp_f and \p exp_i must be dereferenced before dispatching any
  121. * kernel. They are allocated on the caller's stack.
  122. */
  123. virtual void do_exec(_megdnn_tensor_in src, _megdnn_tensor_out dst,
  124. const float* exp_f, const int* exp_i) = 0;
  125. };
  126. /*!
  127. * \brief modify a tensor inplace by adding another tensor to it
  128. *
  129. * dst and delta can have arbitrary layout but must have the same shape.
  130. */
  131. class AddUpdateForward: public OperatorBase {
  132. DEF_OPR_PARAM(AddUpdate);
  133. DEF_OPR_IMPL(AddUpdateForward, OperatorBase, -1, 1);
  134. public:
  135. virtual void exec(
  136. _megdnn_tensor_inout dst, _megdnn_tensor_in delta) = 0;
  137. protected:
  138. void check_exec(const TensorLayout &dst, const TensorLayout &delta);
  139. };
  140. using AddUpdate = AddUpdateForward;
  141. class ReduceForward: public OperatorBase {
  142. DEF_OPR_PARAM(Reduce);
  143. DEF_OPR_IMPL(ReduceForward, OperatorBase, 1, 1);
  144. public:
  145. using Mode = Param::Mode;
  146. using DataType = Param::DataType;
  147. /**
  148. * \param[in] src input tensor
  149. * \param[out] dst output tensor
  150. *
  151. * src and dst should be contiguous.
  152. * src and dst should be of the same shape for all dimensions except
  153. * param().axis.
  154. * the param().axis-th dimension shape for dst should be one.
  155. */
  156. virtual void exec(_megdnn_tensor_in src,
  157. _megdnn_tensor_out dst,
  158. _megdnn_workspace workspace) = 0;
  159. void deduce_layout(const TensorLayout &src, TensorLayout &dst);
  160. virtual size_t get_workspace_in_bytes(const TensorLayout &src,
  161. const TensorLayout &dst) = 0;
  162. protected:
  163. void check_exec(const TensorLayout &src, const TensorLayout &dst,
  164. size_t workspace_in_bytes);
  165. };
  166. using Reduce = ReduceForward;
  167. class CorrelationBase : public OperatorBase {
  168. DEF_OPR_IMPL_CTOR(CorrelationBase, OperatorBase);
  169. DEF_OPR_PARAM(Correlation);
  170. protected:
  171. void deduce_layout_fwd(const TensorLayout& data1, const TensorLayout& data2,
  172. TensorLayout& dst);
  173. void check_layout_fwd(const TensorLayout& data1, const TensorLayout& data2,
  174. const TensorLayout& dst);
  175. };
  176. class CorrelationForward : public CorrelationBase {
  177. DEF_OPR_IMPL(CorrelationForward, CorrelationBase, 2, 1);
  178. public:
  179. /**
  180. * \param[in] data1 (n, c, ih, iw)
  181. * \param[in] data2 (n, c, ih, iw)
  182. * \param[out] dst (n, q, oh, ow), q is the number of neighborhood
  183. * */
  184. virtual void exec(_megdnn_tensor_in data1, _megdnn_tensor_in data2,
  185. _megdnn_tensor_out dst, _megdnn_workspace workspace) = 0;
  186. void deduce_layout(const TensorLayout& data1, const TensorLayout& data2,
  187. TensorLayout& dst);
  188. virtual size_t get_workspace_in_bytes(const TensorLayout& data1,
  189. const TensorLayout& data2,
  190. const TensorLayout& dst) = 0;
  191. protected:
  192. void check_exec(const TensorLayout& data1, const TensorLayout& data2,
  193. const TensorLayout& dst, size_t workspace_in_bytes);
  194. };
  195. using Correlation = CorrelationForward;
  196. class CorrelationBackwardData1 : public CorrelationBase {
  197. DEF_OPR_IMPL(CorrelationBackwardData1, CorrelationBase, 3, 1);
  198. public:
  199. /**
  200. * \param[in] diff the backpropagated gradient wrt. dst
  201. * \param[in] data1 the `data1' parameter in CorrelationForward::exec
  202. * \param[in] data2 the `data2' parameter in CorrelationForward::exec
  203. * \param[out] grad1 the backpropagated gradient wrt. data1
  204. */
  205. virtual void exec(_megdnn_tensor_in diff, _megdnn_tensor_in data1, _megdnn_tensor_in data2,
  206. _megdnn_tensor_out grad1, _megdnn_workspace workspace) = 0;
  207. void deduce_layout(const TensorLayout& diff1, const TensorLayout& data1,
  208. const TensorLayout& data2, TensorLayout& dst);
  209. virtual size_t get_workspace_in_bytes(const TensorLayout& diff,
  210. const TensorLayout& data1,
  211. const TensorLayout& data2,
  212. const TensorLayout& grad1) = 0;
  213. protected:
  214. void check_exec(const TensorLayout& diff, const TensorLayout& data1, const TensorLayout& data2,
  215. const TensorLayout& grad1, size_t workspace_in_bytes);
  216. };
  217. class CorrelationBackwardData2 : public CorrelationBase {
  218. DEF_OPR_IMPL(CorrelationBackwardData2, CorrelationBase, 3, 1);
  219. public:
  220. /**
  221. * \param[in] diff the backpropagated gradient wrt. dst
  222. * \param[in] data1 the `data1' parameter in CorrelationForward::exec
  223. * \param[in] data2 the `data2' parameter in CorrelationForward::exec
  224. * \param[out] grad2 the backpropagated gradient wrt. data2
  225. */
  226. virtual void exec(_megdnn_tensor_in diff, _megdnn_tensor_in data1, _megdnn_tensor_in data2,
  227. _megdnn_tensor_out grad2, _megdnn_workspace workspace) = 0;
  228. void deduce_layout(const TensorLayout& diff1, const TensorLayout& data1,
  229. const TensorLayout& data2, TensorLayout& dst);
  230. virtual size_t get_workspace_in_bytes(const TensorLayout& diff,
  231. const TensorLayout& data1,
  232. const TensorLayout& data2,
  233. const TensorLayout& grad2) = 0;
  234. protected:
  235. void check_exec(const TensorLayout& diff, const TensorLayout& data1, const TensorLayout& data2,
  236. const TensorLayout& grad2, size_t workspace_in_bytes);
  237. };
  238. class CumsumForward: public OperatorBase {
  239. DEF_OPR_PARAM(Cumsum);
  240. DEF_OPR_IMPL(CumsumForward, OperatorBase, 1, 1);
  241. public:
  242. /**
  243. * \param[in] src input tensor
  244. * \param[out] dst output tensor
  245. *
  246. * src and dst should be contiguous.
  247. * src and dst should have the same shape.
  248. *
  249. * The exclusive flag specifies whether the current element it taken
  250. * into account when calculating results.
  251. *
  252. * The reverse flag specifies whether cumsum is forward (
  253. * from 0 to n) or backward (from n downto 0).
  254. *
  255. * Example:
  256. * exclusive && reverse:
  257. * dst_i = src_{i+1} + src_{i+2} + ... + src_{n-1}
  258. * exclusive && !reverse
  259. * dst_i = src_0 + src_1 + ... + src_{i-1}
  260. * !exclusive && reverse:
  261. * dst_i = src_i + src_{i+1} + ... + src_{n-1}
  262. * !exclusive && !reverse:
  263. * dst_i = src_0 + src_1 + ... + src_i
  264. */
  265. virtual void exec(_megdnn_tensor_in src,
  266. _megdnn_tensor_out dst,
  267. _megdnn_workspace workspace) = 0;
  268. void deduce_layout(const TensorLayout &src, TensorLayout &dst);
  269. virtual size_t get_workspace_in_bytes(const TensorLayout &src,
  270. const TensorLayout &dst) = 0;
  271. protected:
  272. void check_exec(const TensorLayout &src, const TensorLayout &dst,
  273. size_t workspace_in_bytes);
  274. };
  275. using Cumsum = CumsumForward;
  276. // mxx can be max or min
  277. class ArgmxxBase: public OperatorBase {
  278. DEF_OPR_IMPL_CTOR(ArgmxxBase, OperatorBase);
  279. DEF_OPR_PARAM(Axis);
  280. protected:
  281. void check_layout_fwd(const TensorLayout &src,
  282. const TensorLayout &dst);
  283. };
  284. class ArgmaxForward: public ArgmxxBase {
  285. DEF_OPR_IMPL(ArgmaxForward, ArgmxxBase, 1, 1);
  286. public:
  287. /**
  288. * \param[in] src input tensor
  289. * \param[out] dst output tensor containing the argmax indices
  290. *
  291. * src and dst should be contiguous.
  292. * src and dst should be of the same shape for all dimensions except
  293. * param().axis.
  294. * the param().axis-th dimension shape for dst should be one.
  295. */
  296. virtual void exec(_megdnn_tensor_in src,
  297. _megdnn_tensor_out dst,
  298. _megdnn_workspace workspace) = 0;
  299. void deduce_layout(const TensorLayout &src,
  300. TensorLayout &dst);
  301. virtual size_t get_workspace_in_bytes(const TensorLayout &src,
  302. const TensorLayout &dst) = 0;
  303. protected:
  304. void check_exec(const TensorLayout &src,
  305. const TensorLayout &dst,
  306. size_t workspace_in_bytes);
  307. };
  308. using Argmax = ArgmaxForward;
  309. class ArgminForward: public ArgmxxBase {
  310. DEF_OPR_IMPL(ArgminForward, ArgmxxBase, 1, 1);
  311. public:
  312. /**
  313. * \param[in] src input tensor
  314. * \param[out] dst output tensor containing the argmax indices
  315. *
  316. * src and dst should be contiguous.
  317. * src and dst should be of the same shape for all dimensions except
  318. * param().axis.
  319. * the param().axis-th dimension shape for dst should be one.
  320. */
  321. virtual void exec(_megdnn_tensor_in src,
  322. _megdnn_tensor_out dst,
  323. _megdnn_workspace workspace) = 0;
  324. void deduce_layout(const TensorLayout &src,
  325. TensorLayout &dst);
  326. virtual size_t get_workspace_in_bytes(const TensorLayout &src,
  327. const TensorLayout &dst) = 0;
  328. protected:
  329. void check_exec(const TensorLayout &src,
  330. const TensorLayout &dst,
  331. size_t workspace_in_bytes);
  332. };
  333. using Argmin = ArgminForward;
  334. /*!
  335. * \brief take values from input according to given condition
  336. *
  337. * Output two tensors:
  338. * 1. values copied from *data*, with same dtype as *data*
  339. * 2. selected indices with dtype int32; note that it is 1-dimensional and
  340. * based on the flatten input.
  341. *
  342. * Require data and mask to have the same shape and both be contiguous.
  343. */
  344. class CondTake : public OperatorBase {
  345. DEF_OPR_IMPL(CondTake, OperatorBase, 2, 2);
  346. DEF_OPR_PARAM(CondTake);
  347. public:
  348. using Output = std::array<TensorND, 2>;
  349. using OutputDType = std::array<DType, 2>;
  350. OutputDType infer_dtype(DType data, DType mask);
  351. virtual size_t get_workspace_in_bytes(const TensorLayout& data) = 0;
  352. virtual Output exec(_megdnn_tensor_in data, _megdnn_tensor_in mask,
  353. _megdnn_workspace workspace,
  354. DynOutMallocPolicyCall malloc_policy) = 0;
  355. protected:
  356. //! check input layouts and get flattened size
  357. size_t check_exec_get_size(const TensorLayout& data,
  358. const TensorLayout& mask,
  359. size_t workspace_in_bytes);
  360. };
  361. class TransposeForward: public OperatorBase {
  362. DEF_OPR_IMPL(TransposeForward, OperatorBase, 1, 1);
  363. DEF_OPR_PARAM(Empty);
  364. public:
  365. /**
  366. * \param[in] src (m, n) stride[0] >= n && stride[1] == 1
  367. * \param[out] dst (n, m) stride[0] >= m && stride[1] == 1
  368. */
  369. virtual void exec(_megdnn_tensor_in src,
  370. _megdnn_tensor_out dst,
  371. _megdnn_workspace workspace) = 0;
  372. void deduce_layout(const TensorLayout &src, TensorLayout &dst);
  373. virtual size_t get_workspace_in_bytes(const TensorLayout &src,
  374. const TensorLayout &dst) = 0;
  375. protected:
  376. void check_exec(const TensorLayout &src, const TensorLayout &dst,
  377. size_t workspace_in_bytes);
  378. };
  379. using Transpose = TransposeForward;
  380. /**
  381. * Change a tensor to another layout that has the same dtype and total number of
  382. * elements, and non-overlapping stride.
  383. *
  384. * ON CPU:
  385. * This operator is optimized for some cases(e.g. both dst and last dim of src
  386. * are contiguous)
  387. *
  388. * ON CUDA:
  389. * More contiguous the input/output layouts, higher performance. There is also
  390. * special optimization for broadcast case.
  391. */
  392. class RelayoutForward: public OperatorBase {
  393. DEF_OPR_IMPL(RelayoutForward, OperatorBase, 1, 1);
  394. DEF_OPR_PARAM(Empty);
  395. public:
  396. /*!
  397. * \brief execute relayout opr
  398. *
  399. * This operator should be placed on the same computing device of *dst*.
  400. *
  401. * \param src_handle handle of input tensor; for CUDA d2d copy, the
  402. * src handle can be on a different GPU for copy tensor with
  403. * non-contig dims <= 2
  404. */
  405. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst,
  406. Handle *src_handle = nullptr) = 0;
  407. protected:
  408. //! check layout and collapse contiguous
  409. void check_layout_and_canonize(
  410. TensorLayout &src, TensorLayout &dst);
  411. };
  412. using Relayout = RelayoutForward;
  413. /**
  414. * \brief Base class for Concat and Split operators
  415. */
  416. class ConcatSplitBase: public OperatorBase {
  417. public:
  418. using Param = param::Axis;
  419. ConcatSplitBase(Handle *handle);
  420. const Param &param() const { return m_param; }
  421. Param &param() { return m_param; }
  422. protected:
  423. void check_layout_common(const TensorLayoutArray &srcs,
  424. const TensorLayout &dst);
  425. Param m_param;
  426. /**
  427. * \brief a helper function
  428. *
  429. * A = shape[0] * shape[1] * ... * shape[axis-1]
  430. * B = {srcs[0].shape[axis], srcs[1].shape[axis], ...}
  431. * C = shape[axis+1] * shape[axis+2] * ... * shape[ndim-1]
  432. */
  433. void get_ABC(const TensorShapeArray &srcs,
  434. size_t &A,
  435. size_t *B,
  436. size_t &C);
  437. thin_function<TensorLayout(const TensorND &tensor)> m_get_layout;
  438. thin_function<TensorShape(const TensorLayout &layout)> m_get_shape;
  439. };
  440. class ConcatForward: public ConcatSplitBase {
  441. DEF_OPR_IMPL(ConcatForward, ConcatSplitBase, 1, 1);
  442. public:
  443. /**
  444. * \param[in] srcs a vector containing all inputs to be concatenated
  445. * \param[out] dst the output tensor.
  446. *
  447. * All tensors in srcs and dst should be contiguous.
  448. * All tensors should have the same shape for all axes except
  449. * param().axis.
  450. * For the param().axis-th axis, the axis shape for dst should be the
  451. * sum of corresponding axis shapes for all srcs.
  452. */
  453. virtual void exec(_megdnn_in const TensorNDArray &srcs,
  454. _megdnn_tensor_out dst,
  455. _megdnn_workspace workspace) = 0;
  456. void deduce_layout(const TensorLayoutArray &srcs,
  457. TensorLayout &dst);
  458. virtual size_t get_workspace_in_bytes(
  459. const TensorLayoutArray &srcs,
  460. const TensorLayout &dst) = 0;
  461. protected:
  462. void check_exec(const TensorLayoutArray &srcs,
  463. const TensorLayout &dst,
  464. size_t workspace_in_bytes);
  465. };
  466. using Concat = ConcatForward;
  467. class SplitForward: public ConcatSplitBase {
  468. DEF_OPR_IMPL(SplitForward, ConcatSplitBase, 1, 1);
  469. public:
  470. /**
  471. * \param[in] src input tensor
  472. * \param[out] dsts a vector containing all splitted result
  473. *
  474. * All tensors in src and dsts should be contiguous.
  475. * All tensors should have the same shape for all axes except
  476. * param().axis.
  477. * For the param().axis-th axis, the axis shape for src should be the
  478. * sum of corresponding axis shapes for all dsts.
  479. */
  480. virtual void exec(_megdnn_tensor_in src,
  481. const TensorNDArray &dsts,
  482. _megdnn_workspace workspace) = 0;
  483. virtual size_t get_workspace_in_bytes(const TensorLayout &src,
  484. const TensorLayoutArray &dsts) = 0;
  485. protected:
  486. void check_exec(const TensorLayout &src,
  487. const TensorLayoutArray &dsts,
  488. size_t workspace_in_bytes);
  489. };
  490. using Split = SplitForward;
  491. /**
  492. * \brief Base class for ParamPackConcat and ParamPackSplit Operators.
  493. *
  494. * ParamPack oprs act like Concat and Split, but they also are optimized for a
  495. * large number of inputs and can handle alignment requirements. Axis is also
  496. * not supported.
  497. *
  498. * The offsets can be generated by gen_offsets().
  499. */
  500. class ParamPackConcatSplitBase : public OperatorBase {
  501. protected:
  502. void check_exec(const TensorLayout& concated, const TensorLayout& offsets,
  503. const TensorLayout& parts);
  504. public:
  505. using Param = megdnn::param::Empty;
  506. ParamPackConcatSplitBase(Handle* handle) : OperatorBase(handle) {}
  507. //! generate offsets to be used with ParamPackConcat and ParamPackSplit
  508. static std::vector<dt_int32> gen_offsets(const TensorShapeArray& shapes,
  509. size_t alignment,
  510. size_t dtype_size);
  511. };
  512. /**
  513. * \brief ParamPackConcat, used for calculating gradient of ParamPackSplit
  514. * Combine multiple gradient tensors into a single large tensor, use copy
  515. * strategy due to AddUpdate or other dynamic situation.
  516. */
  517. class ParamPackConcat: public ParamPackConcatSplitBase {
  518. DEF_OPR_IMPL(ParamPackConcat, ParamPackConcatSplitBase, 2, 1);
  519. public:
  520. /*
  521. * \param[in] srcs: TensorND on cpu. srcs[i] corresponding to the
  522. * address of i-th Tensor.
  523. * \param[in] offsets: with size `2 * srcs.shape[0]`.
  524. * offsets[i * 2] and offsets[i * 2 + 1] means
  525. * the begin and the end of srcs[i]'s offsets in dst
  526. * \param[out] dst: output TensorND, live on cpu or gpu
  527. */
  528. virtual void exec(_megdnn_tensor_in srcs, _megdnn_tensor_in offsets,
  529. _megdnn_tensor_out dst, _megdnn_workspace workspace) = 0;
  530. virtual size_t get_workspace_in_bytes(const TensorShapeArray& srcs,
  531. const TensorShape& offsets,
  532. const TensorShape& dst) = 0;
  533. };
  534. /**
  535. * \brief base class for Tile and Repeat
  536. */
  537. class TileRepeatBase: public OperatorBase {
  538. public:
  539. TileRepeatBase(Handle *handle): OperatorBase(handle) {}
  540. struct Param {
  541. TensorShape times;
  542. };
  543. Param &param() { return m_param; }
  544. const Param &param() const { return m_param; }
  545. protected:
  546. void check_layout_fwd(const TensorLayout &src,
  547. const TensorLayout &dst);
  548. void deduce_layout_fwd(const TensorLayout &src,
  549. TensorLayout &dst);
  550. /**
  551. * Assuming src/dst/times are already simplified on entrance.
  552. */
  553. size_t get_workspace_in_bytes_fwd(const TensorShape &src,
  554. const TensorShape &dst,
  555. const TensorShape &times,
  556. DType dtype);
  557. Param m_param;
  558. };
  559. class TileBase: public TileRepeatBase {
  560. public:
  561. TileBase(Handle *handle): TileRepeatBase(handle) {}
  562. protected:
  563. void simplify_shape(const TensorShape &src,
  564. const TensorShape &dst,
  565. const TensorShape &times,
  566. TensorShape &src2,
  567. TensorShape &dst2,
  568. TensorShape &times2);
  569. /**
  570. * This is a helper function that would facilitate other backends'
  571. * implementation.
  572. */
  573. size_t get_workspace_in_bytes_fwd(const TensorLayout &src,
  574. const TensorLayout &dst);
  575. };
  576. class TileForward: public TileBase {
  577. DEF_OPR_IMPL(TileForward, TileBase, 1, 1);
  578. public:
  579. /**
  580. * \brief Tile src times to get dst.
  581. * \param[in] src input tensor
  582. * \param[out] dst output tensor
  583. * \param[out] workspace temporary workspace
  584. *
  585. * src and dst must be contiguous.
  586. * dst.shape should be {src.shape[0]*param().times[0],
  587. * src.shape[1]*param().times[1], ...}
  588. *
  589. * \see http://docs.scipy.org/doc/numpy/reference/generated/numpy.tile.html
  590. *
  591. * Difference between Tile and Repeat:
  592. * Tiling `abc' twice yields `abcabc', whereas repeating `abc' twice
  593. * yields `aabbcc'.
  594. */
  595. virtual void exec(_megdnn_tensor_in src,
  596. _megdnn_tensor_out dst,
  597. _megdnn_workspace workspace) = 0;
  598. void deduce_layout(const TensorLayout &src,
  599. TensorLayout &dst);
  600. virtual size_t get_workspace_in_bytes(const TensorLayout &src,
  601. const TensorLayout &dst) = 0;
  602. protected:
  603. void check_exec(const TensorLayout &src, const TensorLayout &dst,
  604. size_t workspace_in_bytes);
  605. };
  606. using Tile = TileForward;
  607. class TileBackward: public TileBase {
  608. DEF_OPR_IMPL(TileBackward, TileBase, 1, 1);
  609. public:
  610. /**
  611. * \param[in] diff the backpropagated gradient wrt. dst
  612. * \param[out] grad the backpropagated gradient wrt. src
  613. * \param[out] workspace temporary workspace
  614. */
  615. virtual void exec(_megdnn_tensor_in diff,
  616. _megdnn_tensor_out grad,
  617. _megdnn_workspace workspace) = 0;
  618. virtual size_t get_workspace_in_bytes(const TensorLayout &diff,
  619. const TensorLayout &grad) = 0;
  620. protected:
  621. void check_exec(const TensorLayout &diff, const TensorLayout &grad,
  622. size_t workspace_in_bytes);
  623. };
  624. class RepeatBase: public TileRepeatBase {
  625. public:
  626. RepeatBase(Handle *handle): TileRepeatBase(handle) {}
  627. protected:
  628. void simplify_shape(const TensorShape &src,
  629. const TensorShape &dst,
  630. const TensorShape &times,
  631. TensorShape &src2,
  632. TensorShape &dst2,
  633. TensorShape &times2);
  634. /**
  635. * This is a helper function that would facilitate other backends'
  636. * implementation.
  637. */
  638. size_t get_workspace_in_bytes_fwd(const TensorLayout &src,
  639. const TensorLayout &dst);
  640. };
  641. class RepeatForward: public RepeatBase {
  642. DEF_OPR_IMPL(RepeatForward, RepeatBase, 1, 1);
  643. public:
  644. /**
  645. * \brief Repeat src times to get dst.
  646. * \param[in] src input tensor
  647. * \param[out] dst output tensor
  648. * \param[out] workspace temporary workspace
  649. *
  650. * src and dst must be contiguous.
  651. * dst.shape should be {src.shape[0]*param().times[0],
  652. * src.shape[1]*param().times[1], ...}
  653. *
  654. * \see http://docs.scipy.org/doc/numpy/reference/generated/numpy.repeat.html
  655. * \see TileForward
  656. */
  657. virtual void exec(_megdnn_tensor_in src,
  658. _megdnn_tensor_out dst,
  659. _megdnn_workspace workspace) = 0;
  660. void deduce_layout(const TensorLayout &src,
  661. TensorLayout &dst);
  662. virtual size_t get_workspace_in_bytes(const TensorLayout &src,
  663. const TensorLayout &dst) = 0;
  664. protected:
  665. void check_exec(const TensorLayout &src,
  666. const TensorLayout &dst,
  667. size_t workspace_in_bytes);
  668. };
  669. using Repeat = RepeatForward;
  670. class RepeatBackward: public RepeatBase {
  671. DEF_OPR_IMPL(RepeatBackward, RepeatBase, 1, 1);
  672. public:
  673. /**
  674. * \param[in] diff the backpropagated gradient wrt. dst
  675. * \param[out] grad the backpropagated gradient wrt. src
  676. * \param[out] workspace temporary workspace
  677. */
  678. virtual void exec(_megdnn_tensor_in diff,
  679. _megdnn_tensor_out grad,
  680. _megdnn_workspace workspace) = 0;
  681. virtual size_t get_workspace_in_bytes(const TensorLayout &diff,
  682. const TensorLayout &grad) = 0;
  683. protected:
  684. void check_exec(const TensorLayout &diff,
  685. const TensorLayout &grad,
  686. size_t workspace_in_bytes);
  687. };
  688. class ArgsortForward: public OperatorBase {
  689. DEF_OPR_IMPL(ArgsortForward, OperatorBase, 1, 2);
  690. DEF_OPR_PARAM(Argsort);
  691. public:
  692. using Order = Param::Order;
  693. /**
  694. * \param[in] src (m, n)
  695. * \param[out] dst (m, n)
  696. * \param[out] indices (m, n)
  697. *
  698. * src, dst and indices should be contiguous.
  699. * Performing m independent sorting on m arrays of length n.
  700. * Sorting arrays and storing the resulting array in `dst',
  701. * and the corresponding indices in `indices'.
  702. *
  703. * Indices range from 0 to n-1.
  704. *
  705. * Note that indices is a TensorND of type int.
  706. */
  707. virtual void exec(_megdnn_tensor_in src,
  708. _megdnn_tensor_out dst,
  709. _megdnn_tensor_out indices,
  710. _megdnn_workspace workspace) = 0;
  711. void deduce_layout(const TensorLayout &src,
  712. TensorLayout &dst,
  713. TensorLayout &indices);
  714. virtual size_t get_workspace_in_bytes(const TensorLayout &src,
  715. const TensorLayout &dst,
  716. const TensorLayout &indices) = 0;
  717. protected:
  718. void check_exec(const TensorLayout &src,
  719. const TensorLayout &dst,
  720. const TensorLayout &indices,
  721. size_t workspace_in_bytes);
  722. };
  723. using Argsort = ArgsortForward;
  724. /*!
  725. * \brief backward opr for Argsort
  726. *
  727. * Note: the name is kept for backward compatibility. This opr is actually a
  728. * batched value setter. It is used for gradient computing of Argsort and TopK.
  729. */
  730. class ArgsortBackward : public OperatorBase {
  731. DEF_OPR_IMPL(ArgsortBackward, OperatorBase, 2, 1);
  732. DEF_OPR_PARAM(Empty);
  733. public:
  734. /**
  735. * \param[in] diff (m, k) the backpropagated gradient wrt. dst
  736. * \param[in] indices (m, k) the `indices' parameter in
  737. * ArgsortForward::exec
  738. * \param[out] grad (m, n) the backpropagated gradient wrt. src
  739. *
  740. * Constraint: n >= k. Untouched values would be initialized as zero.
  741. */
  742. virtual void exec(_megdnn_tensor_in diff, _megdnn_tensor_in indices,
  743. _megdnn_tensor_out grad, _megdnn_workspace workspace) = 0;
  744. virtual size_t get_workspace_in_bytes(const TensorLayout& diff,
  745. const TensorLayout& indices,
  746. const TensorLayout& grad) = 0;
  747. protected:
  748. void check_exec(const TensorLayout& diff, const TensorLayout& indices,
  749. const TensorLayout& grad, size_t workspace_in_bytes);
  750. };
  751. class TopK : public OperatorBase {
  752. DEF_OPR_IMPL(TopK, OperatorBase, 1, 2);
  753. DEF_OPR_PARAM(TopK);
  754. protected:
  755. //! impl exec; inputs have been validated
  756. virtual void do_exec(int k, _megdnn_tensor_in data,
  757. _megdnn_tensor_out values, int32_t* indices,
  758. _megdnn_workspace workspace) = 0;
  759. public:
  760. /*!
  761. * \param[in] k if positive, compute the smallest top-k values; otherwise
  762. * compute the largest top-k values
  763. * \param[in] data (m, n) input data, where top-k is computed on the
  764. * second axis. The second dimension must be contiguous, and the first
  765. * dimension can have arbitrary stride.
  766. * \param[out] values (m, ) or (m, k) output values; its shape depends
  767. * on mode
  768. * \param[out] indices () or (m, ) or (m, k) output values; its shape
  769. * depends on mode
  770. */
  771. void exec(int k, _megdnn_tensor_in data, _megdnn_tensor_out values,
  772. _megdnn_tensor_out indices, _megdnn_workspace workspace);
  773. virtual size_t get_workspace_in_bytes(int k, const TensorLayout& data,
  774. const TensorLayout& values,
  775. const TensorLayout& indices) = 0;
  776. void deduce_layout(int k, const TensorLayout& data, TensorLayout& values,
  777. TensorLayout& indices);
  778. };
  779. /*!
  780. * \brief convert dtype of *src* to match dtype of *dst*; *src* may have
  781. * arbitrary layout and *dst* must be contiguous.
  782. */
  783. class TypeCvtForward: public OperatorBase {
  784. DEF_OPR_PARAM(Empty);
  785. DEF_OPR_IMPL(TypeCvtForward, OperatorBase, 1, 1);
  786. public:
  787. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst) = 0;
  788. protected:
  789. void check_exec(const TensorLayout &src, const TensorLayout &dst);
  790. };
  791. using TypeCvt = TypeCvtForward;
  792. class IndexingRemapBase: public OperatorBase {
  793. public:
  794. using Param = param::IndexingRemap;
  795. IndexingRemapBase(Handle *handle): OperatorBase(handle) {}
  796. Param &param() { return m_param; }
  797. const Param &param() const { return m_param; }
  798. protected:
  799. Param m_param;
  800. void check_layout_fwd(const TensorLayout &src,
  801. const TensorLayout &map,
  802. const TensorLayout &dst);
  803. };
  804. class IndexingRemapForward: public IndexingRemapBase {
  805. DEF_OPR_IMPL(IndexingRemapForward, IndexingRemapBase, 2, 1);
  806. public:
  807. /**
  808. * \param[in] src input tensor
  809. * \param[in] map input map
  810. * \param[out] dst output tensor
  811. *
  812. * Suppose:
  813. * the shape of src is \f$(s_0, s_1, ..., s_{m-1}\f$;
  814. * the shape of dst is \f$(d_0, d_1, ..., d_{n-1})\f$;
  815. * then:
  816. * the shape of map must be \f$(d_0, d_1, ..., d_{n-1}, m)\f$.
  817. *
  818. * The last dimension of map indicates the src indices for the
  819. * corresponding dst entry.
  820. *
  821. * src and dst can be non-contiguous in a non-overlapping manner.
  822. */
  823. virtual void exec(_megdnn_tensor_in src,
  824. _megdnn_tensor_in map,
  825. _megdnn_tensor_out dst,
  826. _megdnn_workspace workspace) = 0;
  827. void deduce_layout(const TensorLayout &src,
  828. const TensorLayout &map,
  829. TensorLayout &dst);
  830. virtual size_t get_workspace_in_bytes(const TensorLayout &src,
  831. const TensorLayout &map,
  832. const TensorLayout &dst) = 0;
  833. protected:
  834. void check_exec(const TensorLayout &src,
  835. const TensorLayout &map,
  836. const TensorLayout &dst,
  837. size_t workspace_in_bytes);
  838. };
  839. using IndexingRemap = IndexingRemapForward;
  840. // The using directives preserve backward compatibility.
  841. using TensorRemapForward = IndexingRemap;
  842. using TensorRemap = TensorRemapForward;
  843. class IndexingRemapBackward: public IndexingRemapBase {
  844. DEF_OPR_IMPL(IndexingRemapBackward, IndexingRemapBase, 2, 1);
  845. public:
  846. /**
  847. * \param[in] diff the backpropagated gradient wrt. dst
  848. * \param[in] map the `map' parameter in IndexingRemapForward::exec
  849. * \param[out] grad the backpropagated gradient wrt. src
  850. */
  851. virtual void exec(_megdnn_tensor_in diff,
  852. _megdnn_tensor_in map,
  853. _megdnn_tensor_out grad,
  854. _megdnn_workspace workspace) = 0;
  855. virtual size_t get_workspace_in_bytes(const TensorLayout &diff,
  856. const TensorLayout &map,
  857. const TensorLayout &grad) = 0;
  858. protected:
  859. void check_exec(const TensorLayout &diff,
  860. const TensorLayout &map,
  861. const TensorLayout &grad,
  862. size_t workspace_in_bytes);
  863. };
  864. // The using directives preserve backward compatibility.
  865. using TensorRemapBackward = IndexingRemapBackward;
  866. class Linspace: public OperatorBase {
  867. DEF_OPR_IMPL(Linspace, OperatorBase, 0, 1);
  868. DEF_OPR_PARAM(LinspaceFull);
  869. public:
  870. /**
  871. * \param[out] dst must be 1d.
  872. *
  873. * \see http://docs.scipy.org/doc/numpy/reference/generated/numpy.linspace.html
  874. */
  875. virtual void exec(_megdnn_tensor_out dst,
  876. _megdnn_workspace workspace) = 0;
  877. virtual size_t get_workspace_in_bytes(const TensorLayout &dst) = 0;
  878. protected:
  879. void check_exec(const TensorLayout &dst, size_t workspace_in_bytes);
  880. };
  881. class Eye: public OperatorBase {
  882. DEF_OPR_IMPL(Eye, OperatorBase, 0, 1);
  883. DEF_OPR_PARAM(Eye);
  884. public:
  885. /**
  886. * \see http://docs.scipy.org/doc/numpy/reference/generated/numpy.eye.html
  887. */
  888. virtual void exec(_megdnn_tensor_out dst,
  889. _megdnn_workspace workspace) = 0;
  890. virtual size_t get_workspace_in_bytes(const TensorLayout &dst) = 0;
  891. protected:
  892. void check_exec(const TensorLayout &dst, size_t workspace_in_bytes);
  893. };
  894. class IndexingOneHotBase: public OperatorBase {
  895. DEF_OPR_IMPL_CTOR(IndexingOneHotBase, OperatorBase);
  896. DEF_OPR_PARAM(Axis);
  897. protected:
  898. void deduce_layout_fwd(const TensorLayout &src,
  899. const TensorLayout &index,
  900. TensorLayout &dst);
  901. void check_layout_fwd(const TensorLayout &src,
  902. const TensorLayout &index,
  903. const TensorLayout &dst);
  904. };
  905. /*!
  906. * \brief Indexing for one-hot encoding
  907. *
  908. * Given src, axis and index,
  909. * for all valid (n-1)-dimensional subscript tuples i iterating through index:
  910. * dst[i[0], ..., i[axis-1], 0, i[axis], ..., i[n-2]] =
  911. * inp[i[0], ..., i[axis-1], index[i], i[axis], ..., i[n-2]]
  912. *
  913. * \param[in] src n-dimensional input data
  914. * \param[in] index (n-1)-dimensional index, must be int
  915. * \param[out] dst n-dimensional output data
  916. */
  917. class IndexingOneHotForward: public IndexingOneHotBase {
  918. DEF_OPR_IMPL(IndexingOneHotForward, IndexingOneHotBase, 2, 1);
  919. public:
  920. void deduce_layout(const TensorLayout &src,
  921. const TensorLayout &index, TensorLayout &dst) {
  922. deduce_layout_fwd(src, index, dst);
  923. }
  924. virtual void exec(_megdnn_tensor_in src,
  925. _megdnn_tensor_in index,
  926. _megdnn_tensor_out dst,
  927. _megdnn_workspace workspace) = 0;
  928. virtual size_t get_workspace_in_bytes(const TensorLayout &src,
  929. const TensorLayout &index,
  930. const TensorLayout &dst) = 0;
  931. protected:
  932. void check_exec(const TensorLayout &src,
  933. const TensorLayout &index, const TensorLayout &dst,
  934. size_t workspace_in_bytes);
  935. };
  936. using IndexingOneHot = IndexingOneHotForward;
  937. /*!
  938. * \brief set-subtensor corresponding to IndexingOneHotForward
  939. *
  940. * \param[in,out] data n-dimensional input and output data, whose sub part
  941. * corresponding to *index* would be replaced by *sub*
  942. * \param[in] index (n-1)-dimensional index, must be int
  943. * \param[in] sub n-dimensional sub tensor to be filled in *data*
  944. */
  945. class IndexingSetOneHotForward: public IndexingOneHotBase {
  946. DEF_OPR_IMPL(IndexingSetOneHotForward, IndexingOneHotBase, -1, 1);
  947. public:
  948. virtual void exec(_megdnn_tensor_inout data, _megdnn_tensor_in index,
  949. _megdnn_tensor_in sub, _megdnn_workspace workspace) = 0;
  950. virtual size_t get_workspace_in_bytes(const TensorLayout &data,
  951. const TensorLayout &index,
  952. const TensorLayout &sub) = 0;
  953. protected:
  954. void check_exec(const TensorLayout &data,
  955. const TensorLayout &index, const TensorLayout &sub,
  956. size_t workspace_in_bytes);
  957. };
  958. using IndexingSetOneHot = IndexingSetOneHotForward;
  959. /*!
  960. * \brief base class for indexing on multiple axes using vector indices
  961. *
  962. * Note that the indexing axes are required to be sorted in ascending order
  963. */
  964. class IndexingMultiAxisVecBase: public OperatorBase {
  965. DEF_OPR_IMPL_CTOR(IndexingMultiAxisVecBase, OperatorBase);
  966. DEF_OPR_PARAM(Empty);
  967. public:
  968. struct AxisIndexer {
  969. size_t axis;
  970. TensorND vec;
  971. };
  972. struct AxisIndexerLayoutOnly {
  973. size_t axis;
  974. TensorLayout layout;
  975. };
  976. using IndexDesc = std::vector<AxisIndexer>;
  977. using IndexDescLayoutOnly = std::vector<AxisIndexerLayoutOnly>;
  978. /*!
  979. * \brief convert IndexDesc to IndexDescLayoutOnly
  980. */
  981. static IndexDescLayoutOnly extract_index_layout(const IndexDesc &index);
  982. /*!
  983. * \brief get the axes on src that are not used in index
  984. * \param[out] out output buffer; suggested size is
  985. * TensorLayout::MAX_NDIM
  986. * \return number of elements written to *out*
  987. */
  988. static size_t get_nonindex_axes(size_t src_ndim, const IndexDesc &index,
  989. size_t *out);
  990. /*!
  991. * \brief get contiguous-collapsed layout for indexing on value
  992. * \param idx_axis indexer axis on value (i.e. ExecInfo::idx_axis)
  993. * \return a tensor layout and an axis to iterate over *value* and also
  994. * access *data*; stride of layout on that axis would be zero, and
  995. * strides on other axes correspond to the strides in *data*
  996. */
  997. static std::pair<TensorLayout, size_t> get_value_iter_optimized_layout(
  998. const TensorLayout &data, const TensorLayout &value,
  999. const IndexDesc &index, size_t idx_axis);
  1000. //! helper info for kernel implementation
  1001. struct ExecInfo {
  1002. //! axis in value used by indexer
  1003. size_t idx_axis;
  1004. ptrdiff_t value_stride;
  1005. void* error_tracker;
  1006. megcore::AsyncErrorInfo* error_info;
  1007. };
  1008. protected:
  1009. /*!
  1010. * \return axis on dst used by indexer (i.e. ExecInfo::idx_axis)
  1011. */
  1012. static size_t deduce_layout_fwd(
  1013. const TensorLayout &data,
  1014. const IndexDescLayoutOnly &index,
  1015. TensorLayout &dst);
  1016. static ExecInfo check_exec_noworkspace(
  1017. const TensorLayout &data, const TensorLayout &value,
  1018. const IndexDesc &index, IndexDescLayoutOnly &index_layout);
  1019. };
  1020. /*!
  1021. * \brief compute indexing result, like numpy advanced indexing
  1022. *
  1023. * src can have arbitrary layout, but dst must be dim1-contig
  1024. */
  1025. class IndexingMultiAxisVec: public IndexingMultiAxisVecBase {
  1026. DEF_OPR_IMPL(IndexingMultiAxisVec, IndexingMultiAxisVecBase, 0, 1);
  1027. public:
  1028. virtual void exec(_megdnn_tensor_in src,
  1029. const IndexDesc &index,
  1030. _megdnn_tensor_out dst, _megdnn_workspace workspace) = 0;
  1031. /*!
  1032. * \brief get workspace size based on output shape and indexing axes
  1033. */
  1034. size_t get_workspace_in_bytes(
  1035. const TensorShape &dst,
  1036. const size_t *axes, size_t nr_axes);
  1037. static void deduce_layout(
  1038. const TensorLayout &data,
  1039. const IndexDescLayoutOnly &index,
  1040. TensorLayout &dst) {
  1041. deduce_layout_fwd(data, index, dst);
  1042. }
  1043. protected:
  1044. virtual size_t get_workspace_in_bytes(size_t dst_idx_size) = 0;
  1045. ExecInfo check_exec(
  1046. const TensorLayout &src,
  1047. const IndexDesc &index,
  1048. const TensorLayout &dst,
  1049. size_t workspace_in_bytes);
  1050. };
  1051. /*!
  1052. * \brief base class for modifying data by given index
  1053. *
  1054. * data can have arbitrary layout, but value must be dim1-contig
  1055. */
  1056. class IndexingModifyMultiAxisVecBase: public IndexingMultiAxisVecBase {
  1057. DEF_OPR_IMPL_CTOR(IndexingModifyMultiAxisVecBase, IndexingMultiAxisVecBase);
  1058. public:
  1059. virtual void exec(
  1060. _megdnn_tensor_inout data, _megdnn_tensor_in value,
  1061. const IndexDesc &index,
  1062. _megdnn_workspace workspace) = 0;
  1063. /*!
  1064. * \brief get workspace size based on shape of value input and indexing
  1065. * axes
  1066. */
  1067. size_t get_workspace_in_bytes(
  1068. const TensorShape &value,
  1069. const size_t *axes, size_t nr_axes);
  1070. protected:
  1071. ExecInfo check_exec(
  1072. const TensorLayout &data, const TensorLayout &value,
  1073. const IndexDesc &index,
  1074. size_t workspace_in_bytes);
  1075. virtual size_t get_workspace_in_bytes(size_t value_idx_size) = 0;
  1076. };
  1077. //! set value to indexed locations; index values must be non-overlapping
  1078. class IndexingSetMultiAxisVec: public IndexingModifyMultiAxisVecBase {
  1079. DEF_OPR_IMPL(IndexingSetMultiAxisVec,
  1080. IndexingModifyMultiAxisVecBase, 0, 0);
  1081. };
  1082. //! add value to indexed locations; index values must be non-overlapping
  1083. class IndexingIncrMultiAxisVec: public IndexingModifyMultiAxisVecBase {
  1084. DEF_OPR_IMPL(IndexingIncrMultiAxisVec,
  1085. IndexingModifyMultiAxisVecBase, 0, 0);
  1086. };
  1087. class MeshBase : public OperatorBase {
  1088. DEF_OPR_PARAM(Empty);
  1089. DEF_OPR_IMPL_CTOR(MeshBase, OperatorBase);
  1090. public:
  1091. using AxisIndexer = IndexingMultiAxisVecBase::AxisIndexer;
  1092. using IndexDesc = IndexingMultiAxisVecBase::IndexDesc;
  1093. using AxisIndexerLayoutOnly =
  1094. IndexingMultiAxisVecBase::AxisIndexerLayoutOnly;
  1095. using IndexDescLayoutOnly = IndexingMultiAxisVecBase::IndexDescLayoutOnly;
  1096. size_t get_workspace_in_bytes(const TensorShape&, const size_t*, size_t) {
  1097. return 0;
  1098. }
  1099. protected:
  1100. virtual void check_exec(const TensorLayout& origin,
  1101. const TensorLayout& indexed, const IndexDesc& desc);
  1102. };
  1103. class NormalMeshBase : public MeshBase {
  1104. DEF_OPR_IMPL(NormalMeshBase, MeshBase, 0, 0);
  1105. protected:
  1106. virtual void check_exec(const TensorLayout& origin,
  1107. const TensorLayout& indexed,
  1108. const IndexDesc& desc) override final;
  1109. };
  1110. class NormalMeshModifyBase : public NormalMeshBase {
  1111. DEF_OPR_IMPL_CTOR(NormalMeshModifyBase, NormalMeshBase);
  1112. public:
  1113. virtual void exec(_megdnn_tensor_inout data, _megdnn_tensor_in value,
  1114. const IndexDesc& desc, _megdnn_workspace workspace) = 0;
  1115. };
  1116. class BatchedMeshBase : public MeshBase {
  1117. DEF_OPR_IMPL_CTOR(BatchedMeshBase, MeshBase);
  1118. protected:
  1119. virtual void check_exec(const TensorLayout& origin,
  1120. const TensorLayout& indexed,
  1121. const IndexDesc& desc) override final;
  1122. };
  1123. class BatchedMeshModifyBase : public BatchedMeshBase {
  1124. DEF_OPR_IMPL_CTOR(BatchedMeshModifyBase, BatchedMeshBase);
  1125. public:
  1126. virtual void exec(_megdnn_tensor_inout data, _megdnn_tensor_in value,
  1127. const IndexDesc& desc, _megdnn_workspace workspace) = 0;
  1128. };
  1129. class MeshIndexing : public NormalMeshBase {
  1130. DEF_OPR_IMPL(MeshIndexing, NormalMeshBase, 0, 0);
  1131. public:
  1132. virtual void exec(_megdnn_tensor_in src, const IndexDesc& desc,
  1133. _megdnn_tensor_out dst, _megdnn_workspace workspace) = 0;
  1134. static void deduce_layout(const TensorLayout& inp,
  1135. const IndexDescLayoutOnly& layouts,
  1136. TensorLayout& out_layout);
  1137. };
  1138. class IncrMeshIndexing : public NormalMeshModifyBase {
  1139. DEF_OPR_IMPL(IncrMeshIndexing, NormalMeshModifyBase, 0, 0);
  1140. };
  1141. class SetMeshIndexing : public NormalMeshModifyBase {
  1142. DEF_OPR_IMPL(SetMeshIndexing, NormalMeshModifyBase, 0, 0);
  1143. };
  1144. class BatchedMeshIndexing : public BatchedMeshBase {
  1145. DEF_OPR_IMPL(BatchedMeshIndexing, BatchedMeshBase, 0, 0);
  1146. public:
  1147. virtual void exec(_megdnn_tensor_in src, const IndexDesc& desc,
  1148. _megdnn_tensor_out dst,
  1149. _megdnn_workspace workspace) = 0;
  1150. static void deduce_layout(const TensorLayout& inp,
  1151. const IndexDescLayoutOnly& layouts,
  1152. TensorLayout& out_layout);
  1153. };
  1154. class BatchedIncrMeshIndexing : public BatchedMeshModifyBase {
  1155. DEF_OPR_IMPL(BatchedIncrMeshIndexing, BatchedMeshModifyBase, 0, 0);
  1156. };
  1157. class BatchedSetMeshIndexing : public BatchedMeshModifyBase {
  1158. DEF_OPR_IMPL(BatchedSetMeshIndexing, BatchedMeshModifyBase, 0, 0);
  1159. };
  1160. class RelayoutFormat : public OperatorBase {
  1161. DEF_OPR_PARAM(RelayoutFormat);
  1162. DEF_OPR_IMPL(RelayoutFormat, OperatorBase, 1, 1);
  1163. public:
  1164. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst,
  1165. _megdnn_workspace workspace) = 0;
  1166. void deduce_layout(const TensorLayout& src, TensorLayout& dst);
  1167. void deduce_format(TensorFormat src, TensorFormat& dst);
  1168. virtual size_t get_workspace_in_bytes(const TensorLayout& src,
  1169. const TensorLayout& dst) = 0;
  1170. protected:
  1171. void deduce_layout_fwd(const TensorLayout& src, TensorLayout& dst);
  1172. void check_layout_fwd(const TensorLayout& src, const TensorLayout& dst);
  1173. void check_exec(const TensorLayout& src, const TensorLayout& dst,
  1174. size_t workspace_in_bytes);
  1175. void deduce_exec_layout(const TensorLayout& src, const TensorLayout& dst,
  1176. TensorLayout& exec_workspace,
  1177. TensorLayout& exec_src, TensorLayout& exec_dst);
  1178. };
  1179. /*!
  1180. * \brief check whether input contains inf or nan value.
  1181. */
  1182. class CheckNonFinite: public OperatorBase {
  1183. DEF_OPR_PARAM(Empty);
  1184. DEF_OPR_IMPL(CheckNonFinite, OperatorBase, 1, 1);
  1185. public:
  1186. virtual size_t get_workspace_in_bytes(const TensorLayout &src,
  1187. const TensorLayout &dst) = 0;
  1188. void deduce_layout(const TensorLayout &src, TensorLayout &dst);
  1189. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst,
  1190. _megdnn_workspace workspace) = 0;
  1191. protected:
  1192. void check_exec(const TensorLayout &src, const TensorLayout &dst,
  1193. size_t workspace_in_bytes);
  1194. };
  1195. /*!
  1196. * \brief fill the tensor with a scalar value
  1197. */
  1198. class Fill: public OperatorBase {
  1199. DEF_OPR_PARAM(Fill);
  1200. DEF_OPR_IMPL(Fill, OperatorBase, 0, 1);
  1201. public:
  1202. virtual void exec(_megdnn_tensor_out dst,
  1203. _megdnn_workspace workspace) = 0;
  1204. virtual size_t get_workspace_in_bytes(const TensorLayout& dst) = 0;
  1205. protected:
  1206. void check_exec(const TensorLayout& dst, size_t workspace_in_bytes);
  1207. };
  1208. /*!
  1209. * \brief standard padding operator
  1210. * Inputs must have the same dtype, and the output tensor shape must greater or equal than
  1211. * input tensor in every dimensions, the extra space will be fulled with m which default to
  1212. * be 0.
  1213. */
  1214. class PaddingBase: public OperatorBase {
  1215. DEF_OPR_PARAM(Padding);
  1216. DEF_OPR_IMPL(PaddingBase, OperatorBase, 1, 1);
  1217. public:
  1218. using Mode = Param::PaddingMode;
  1219. protected:
  1220. SmallVector<size_t> get_offsets();
  1221. void check_exec(const TensorLayout& src, const TensorLayout& dst);
  1222. };
  1223. class PaddingForward: public PaddingBase {
  1224. DEF_OPR_IMPL(PaddingForward, PaddingBase, 1, 1);
  1225. public:
  1226. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst) = 0;
  1227. void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst,
  1228. _megdnn_workspace) {
  1229. return exec(src, dst);
  1230. }
  1231. virtual size_t get_workspace_in_bytes(const TensorLayout& src, const TensorLayout& dst) = 0;
  1232. void deduce_layout(const TensorLayout &src, TensorLayout &dst);
  1233. protected:
  1234. void forward_check_exec(const TensorLayout& src, const TensorLayout& dst);
  1235. };
  1236. using Padding = PaddingForward;
  1237. class PaddingBackward: public PaddingBase {
  1238. DEF_OPR_IMPL(PaddingBackward, PaddingBase, 1, 1);
  1239. public:
  1240. virtual void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst) = 0;
  1241. void exec(_megdnn_tensor_in src, _megdnn_tensor_out dst,
  1242. _megdnn_workspace) {
  1243. return exec(src, dst);
  1244. }
  1245. virtual size_t get_workspace_in_bytes(const TensorLayout& src, const TensorLayout& dst) = 0;
  1246. protected:
  1247. void backward_check_exec(const TensorLayout& src, const TensorLayout& dst);
  1248. };
  1249. } // namespace megdnn
  1250. #include "megdnn/internal/opr_header_epilogue.h"
  1251. // vim: syntax=cpp.doxygen

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