OpenI
/
MegEngine
Not watched
1
0
Fork 0
Code
Releases 31 Wiki Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: 404ef808fa
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from '404ef808fa'
${ noResults }
MegEngine/src/jit/impl/mlir/ir/each_mode.cpp

481 lines
19 kB
Raw Blame History 

  1. /**

  2.  * \file src/jit/impl/mlir/ir/each_mode.cpp

  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. 

  13. #include "megbrain_build_config.h"

  14. #if MGB_JIT && MGB_JIT_MLIR

  15. 

  16. #include "./common.h"

  17. #include "./each_mode.h"

  18. #include "./numerical.h"

  19. #include "./types.h"

  20. 

  21. #include "megbrain/common.h"

  22. #include "megbrain/exception.h"

  23. #include "megbrain/jit/mlir/ir/dialect.h"

  24. 

  25. #include <mlir/Dialect/StandardOps/IR/Ops.h>

  26. 

  27. namespace mgb {

  28. namespace jit {

  29. 

  30. using Mode = megdnn::param::Elemwise::Mode;

  31. 

  32. template <Mode mode>

  33. mlir::Value lower_mode(mlir::OpBuilder& builder, mlir::Location loc,

  34.                        ValueRange operands);

  35. 

  36. /* ===================== trivial implementations ===================== */

  37. 

  38. #define cb(mode, fun)                                             \

  39.     template <>                                                   \

  40.     mlir::Value lower_mode<Mode::mode>(mlir::OpBuilder & builder, \

  41.                                        mlir::Location loc,        \

  42.                                        ValueRange operands) {     \

  43.         ValueBuilderHelper helper(builder, loc);                  \

  44.         return helper.fun(operands);                              \

  45.     }

  46. 

  47. //! unary

  48. cb(ABS, abs);

  49. cb(CEIL, ceil);

  50. cb(COS, cos);

  51. cb(EXP, exp);

  52. cb(FLOOR, floor);

  53. cb(LOG, log);

  54. cb(NEGATE, neg);

  55. cb(SIN, sin);

  56. cb(TANH, tanh);

  57. 

  58. //! binary

  59. cb(ADD, add);

  60. cb(MAX, max);

  61. cb(MIN, min);

  62. cb(MOD, mod);

  63. cb(MUL, mul);

  64. cb(SUB, sub);

  65. cb(TRUE_DIV, div);

  66. 

  67. #undef cb

  68. 

  69. /* ===================== unary op ===================== */

  70. 

  71. //! ACOS: pi / 2 - arctan2(x, sqrt(1 - x * x))

  72. template <>

  73. mlir::Value lower_mode<Mode::ACOS>(mlir::OpBuilder& builder, mlir::Location loc,

  74.                                    ValueRange operands) {

  75.     ValueBuilderHelper helper(builder, loc);

  76.     auto x = operands[0];

  77.     auto one_minus_x_2 = helper.sub(helper.const_f32(1.f), helper.mul(x, x));

  78.     auto asin = atan2_approx(helper, x, helper.sqrt(one_minus_x_2));

  79.     auto pi_over_2 = helper.const_f32(1.57079637f);

  80.     return helper.sub(pi_over_2, asin);

  81. }

  82. 

  83. //! ASIN: arctan2(x, sqrt(1 - x * x))

  84. template <>

  85. mlir::Value lower_mode<Mode::ASIN>(mlir::OpBuilder& builder, mlir::Location loc,

  86.                                    ValueRange operands) {

  87.     ValueBuilderHelper helper(builder, loc);

  88.     auto x = operands[0];

  89.     auto one_minus_x_2 = helper.sub(helper.const_f32(1.f), helper.mul(x, x));

  90.     return atan2_approx(helper, x, helper.sqrt(one_minus_x_2));

  91. }

  92. 

  93. //! ERFCINV: inverse of complementary gauss error function

  94. //! https://github.com/scipy/scipy/blob/master/scipy/special/cephes/erfinv.c

  95. template <>

  96. mlir::Value lower_mode<Mode::ERFCINV>(mlir::OpBuilder& builder,

  97.                                       mlir::Location loc, ValueRange operands) {

  98.     ValueBuilderHelper helper(builder, loc);

  99.     auto minus_sqrt2 = helper.const_f32(-1.4142135623f);

  100.     auto x = helper.mul(helper.const_f32(0.5f), operands[0]);

  101.     return helper.div(ndtri_approx(helper, x), minus_sqrt2);

  102. }

  103. 

  104. //! ERFC: complementary error function

  105. template <>

  106. mlir::Value lower_mode<Mode::ERFC>(mlir::OpBuilder& builder, mlir::Location loc,

  107.                                    ValueRange operands) {

  108.     ValueBuilderHelper helper(builder, loc);

  109.     return helper.sub(helper.const_f32(1.f), erf_approx(helper, operands[0]));

  110. }

  111. 

  112. //! ERFINV: inverse of gauss error function

  113. //! https://github.com/scipy/scipy/blob/master/scipy/special/cephes/erfinv.c

  114. template <>

  115. mlir::Value lower_mode<Mode::ERFINV>(mlir::OpBuilder& builder,

  116.                                      mlir::Location loc, ValueRange operands) {

  117.     ValueBuilderHelper helper(builder, loc);

  118.     auto sqrt2 = helper.const_f32(1.4142135623f);

  119.     auto x = helper.mul(helper.const_f32(0.5f),

  120.                         helper.add(operands[0], helper.const_f32(1.f)));

  121.     return helper.div(ndtri_approx(helper, x), sqrt2);

  122. }

  123. 

  124. //! ERF: gauss error function

  125. template <>

  126. mlir::Value lower_mode<Mode::ERF>(mlir::OpBuilder& builder, mlir::Location loc,

  127.                                   ValueRange operands) {

  128.     ValueBuilderHelper helper(builder, loc);

  129.     return erf_approx(helper, operands[0]);

  130. }

  131. 

  132. //! EXPM1: exp(x) - 1

  133. template <>

  134. mlir::Value lower_mode<Mode::EXPM1>(mlir::OpBuilder& builder,

  135.                                     mlir::Location loc, ValueRange operands) {

  136.     ValueBuilderHelper helper(builder, loc);

  137.     return helper.sub(helper.exp(operands[0]), helper.const_f32(1.f));

  138. }

  139. 

  140. //! FAST_TANH: x * (27.f + x * x) / (27.f + 9.f * x * x);

  141. template <>

  142. mlir::Value lower_mode<Mode::FAST_TANH>(mlir::OpBuilder& builder,

  143.                                         mlir::Location loc,

  144.                                         ValueRange operands) {

  145.     ValueBuilderHelper helper(builder, loc);

  146.     auto square = helper.mul(operands[0], operands[0]);

  147.     return helper.div(

  148.             helper.mul(operands[0], helper.add(helper.const_f32(27.f), square)),

  149.             helper.add(helper.const_f32(27.f),

  150.                        helper.mul(helper.const_f32(9.f), square)));

  151. }

  152. 

  153. //! H_SWISH: x * clip(x + 3, 0, 6) / 6

  154. template <>

  155. mlir::Value lower_mode<Mode::H_SWISH>(mlir::OpBuilder& builder,

  156.                                       mlir::Location loc, ValueRange operands) {

  157.     ValueBuilderHelper helper(builder, loc);

  158. 

  159.     auto const_3 = helper.const_f32(3.f);

  160.     auto const_0 = helper.const_f32(0.f);

  161.     auto const_6 = helper.const_f32(6.f);

  162.     auto tmp = helper.add(operands[0], const_3);

  163.     return helper.div(helper.mul(operands[0],

  164.                                  helper.min(helper.max(tmp, const_0), const_6)),

  165.                       const_6);

  166. }

  167. 

  168. //! LOG1P: log(1 + p)

  169. template <>

  170. mlir::Value lower_mode<Mode::LOG1P>(mlir::OpBuilder& builder,

  171.                                     mlir::Location loc, ValueRange operands) {

  172.     ValueBuilderHelper helper(builder, loc);

  173.     return helper.log(helper.add(operands[0], helper.const_f32(1.f)));

  174. }

  175. 

  176. //! RELU: max(x, 0)

  177. template <>

  178. mlir::Value lower_mode<Mode::RELU>(mlir::OpBuilder& builder, mlir::Location loc,

  179.                                    ValueRange operands) {

  180.     ValueBuilderHelper helper(builder, loc);

  181.     return helper.max(operands[0], helper.const_f32(0.f));

  182. }

  183. 

  184. //! ROUND

  185. template <>

  186. mlir::Value lower_mode<Mode::ROUND>(mlir::OpBuilder& builder,

  187.                                     mlir::Location loc, ValueRange operands) {

  188.     ValueBuilderHelper helper(builder, loc);

  189.     return helper.select(

  190.             helper.gt(operands[0], helper.const_f32(0.f)),

  191.             helper.floor(helper.add(operands[0], helper.const_f32(0.5f))),

  192.             helper.ceil(helper.sub(operands[0], helper.const_f32(0.5f))));

  193. }

  194. 

  195. //! SIGMOID: 1.f / (expf(-y) + 1.f))

  196. template <>

  197. mlir::Value lower_mode<Mode::SIGMOID>(mlir::OpBuilder& builder,

  198.                                       mlir::Location loc, ValueRange operands) {

  199.     ValueBuilderHelper helper(builder, loc);

  200.     return helper.div(helper.const_f32(1.f),

  201.                       helper.add(helper.exp(helper.neg(operands[0])),

  202.                                  helper.const_f32(1.f)));

  203. }

  204. 

  205. /* ===================== binary op ===================== */

  206. 

  207. //! ABS_GRAD: x > 0 ? y : -y

  208. template <>

  209. mlir::Value lower_mode<Mode::ABS_GRAD>(mlir::OpBuilder& builder,

  210.                                        mlir::Location loc,

  211.                                        ValueRange operands) {

  212.     ValueBuilderHelper helper(builder, loc);

  213.     return helper.select(helper.gt(operands[0], helper.const_f32(0.f)),

  214.                          operands[1], helper.neg(operands[1]));

  215. }

  216. 

  217. //! ATAN2

  218. template <>

  219. mlir::Value lower_mode<Mode::ATAN2>(mlir::OpBuilder& builder,

  220.                                     mlir::Location loc, ValueRange operands) {

  221.     ValueBuilderHelper helper(builder, loc);

  222.     return atan2_approx(helper, operands[0], operands[1]);

  223. }

  224. 

  225. //! EQ: x == y ? 1 : 0

  226. template <>

  227. mlir::Value lower_mode<Mode::EQ>(mlir::OpBuilder& builder, mlir::Location loc,

  228.                                  ValueRange operands) {

  229.     ValueBuilderHelper helper(builder, loc);

  230.     return helper.select(helper.eq(operands[0], operands[1]),

  231.                          helper.const_f32(1.f), helper.const_f32(0.f));

  232. }

  233. 

  234. //! FAST_TANH_GRAD: ((-48.f * x * x) / (3.f + x * x) + 27.f + x * x) / (3.f + x

  235. //! * x) * y

  236. template <>

  237. mlir::Value lower_mode<Mode::FAST_TANH_GRAD>(mlir::OpBuilder& builder,

  238.                                              mlir::Location loc,

  239.                                              ValueRange operands) {

  240.     ValueBuilderHelper helper(builder, loc);

  241.     auto x_pow2 = helper.mul(operands[0], operands[0]);

  242.     auto deno = helper.add(helper.const_f32(3.f), x_pow2);

  243.     return helper.mul(

  244.             helper.div(

  245.                     helper.add(

  246.                             helper.add(helper.div(helper.mul(helper.const_f32(

  247.                                                                      -48.f),

  248.                                                              x_pow2),

  249.                                                   deno),

  250.                                        helper.const_f32(27.f)),

  251.                             x_pow2),

  252.                     helper.mul(deno, helper.const_f32(9.f))),

  253.             operands[1]);

  254. }

  255. 

  256. //! FLOOR_DIV: floor(x/y)

  257. template <>

  258. mlir::Value lower_mode<Mode::FLOOR_DIV>(mlir::OpBuilder& builder,

  259.                                         mlir::Location loc,

  260.                                         ValueRange operands) {

  261.     ValueBuilderHelper helper(builder, loc);

  262.     return helper.floor(helper.div(operands[0], operands[1]));

  263. }

  264. 

  265. //! FUSE_ADD_H_SWISH: (x+y) * min(max(x + y + 3, 0), 6) * (1/6)

  266. template <>

  267. mlir::Value lower_mode<Mode::FUSE_ADD_H_SWISH>(mlir::OpBuilder& builder,

  268.                                                mlir::Location loc,

  269.                                                ValueRange operands) {

  270.     ValueBuilderHelper helper(builder, loc);

  271.     auto sum = helper.add(operands[0], operands[1]);

  272. 

  273.     auto const_3 = helper.const_f32(3.f);

  274.     auto const_0 = helper.const_f32(0.f);

  275.     auto const_6 = helper.const_f32(6.f);

  276.     auto tmp = helper.add(sum, const_3);

  277.     return helper.div(

  278.             helper.mul(sum, helper.min(helper.max(tmp, const_0), const_6)),

  279.             const_6);

  280. }

  281. 

  282. //! FUSE_ADD_RELU: (x + y) <= ctype(0) ? ctype(0) : (x + y)

  283. template <>

  284. mlir::Value lower_mode<Mode::FUSE_ADD_RELU>(mlir::OpBuilder& builder,

  285.                                             mlir::Location loc,

  286.                                             ValueRange operands) {

  287.     ValueBuilderHelper helper(builder, loc);

  288.     auto sum = helper.add(operands[0], operands[1]);

  289.     return helper.max(sum, helper.const_f32(0.f));

  290. }

  291. 

  292. //! FUSE_ADD_SIGMOID: 1.f / (expf(-(x+y)) + 1.f))

  293. template <>

  294. mlir::Value lower_mode<Mode::FUSE_ADD_SIGMOID>(mlir::OpBuilder& builder,

  295.                                                mlir::Location loc,

  296.                                                ValueRange operands) {

  297.     ValueBuilderHelper helper(builder, loc);

  298.     return helper.div(helper.const_f32(1.f),

  299.                       helper.add(helper.exp(helper.neg(

  300.                                          helper.add(operands[0], operands[1]))),

  301.                                  helper.const_f32(1.f)));

  302. }

  303. 

  304. //! FUSE_ADD_TANH: tanh(x + y)

  305. template <>

  306. mlir::Value lower_mode<Mode::FUSE_ADD_TANH>(mlir::OpBuilder& builder,

  307.                                             mlir::Location loc,

  308.                                             ValueRange operands) {

  309.     ValueBuilderHelper helper(builder, loc);

  310.     return helper.tanh(helper.add(operands[0], operands[1]));

  311. }

  312. 

  313. //! H_SWISH_GRAD: x < -3.f ? 0.f : (x > 3.f ? y : (2.f * x + 3.f) / 6.f * y)

  314. template <>

  315. mlir::Value lower_mode<Mode::H_SWISH_GRAD>(mlir::OpBuilder& builder,

  316.                                            mlir::Location loc,

  317.                                            ValueRange operands) {

  318.     ValueBuilderHelper helper(builder, loc);

  319.     return helper.select(

  320.             helper.lt(operands[0], helper.const_f32(-3.f)),

  321.             helper.const_f32(0.f),

  322.             helper.select(

  323.                     helper.gt(operands[0], helper.const_f32(3.f)), operands[1],

  324.                     helper.mul(

  325.                             helper.div(

  326.                                     helper.add(helper.mul(helper.const_f32(2.f),

  327.                                                           operands[0]),

  328.                                                helper.const_f32(3.f)),

  329.                                     helper.const_f32(6.f)),

  330.                             operands[1])));

  331. }

  332. 

  333. //! LEQ: x <= y ? 1 : 0

  334. template <>

  335. mlir::Value lower_mode<Mode::LEQ>(mlir::OpBuilder& builder, mlir::Location loc,

  336.                                   ValueRange operands) {

  337.     ValueBuilderHelper helper(builder, loc);

  338.     return helper.select(helper.le(operands[0], operands[1]),

  339.                          helper.const_f32(1.f), helper.const_f32(0.f));

  340. }

  341. 

  342. //! LOG_SUM_EXP: log(exp(x) + exp(y))

  343. template <>

  344. mlir::Value lower_mode<Mode::LOG_SUM_EXP>(mlir::OpBuilder& builder,

  345.                                           mlir::Location loc,

  346.                                           ValueRange operands) {

  347.     ValueBuilderHelper helper(builder, loc);

  348.     return helper.log(

  349.             helper.add(helper.exp(operands[0]), helper.exp(operands[1])));

  350. }

  351. 

  352. //! LT: x < y ? 1 : 0

  353. template <>

  354. mlir::Value lower_mode<Mode::LT>(mlir::OpBuilder& builder, mlir::Location loc,

  355.                                  ValueRange operands) {

  356.     ValueBuilderHelper helper(builder, loc);

  357.     return helper.select(helper.lt(operands[0], operands[1]),

  358.                          helper.const_f32(1.f), helper.const_f32(0.f));

  359. }

  360. 

  361. //! POW: x^y = exp(y * log(x))

  362. template <>

  363. mlir::Value lower_mode<Mode::POW>(mlir::OpBuilder& builder, mlir::Location loc,

  364.                                   ValueRange operands) {

  365.     ValueBuilderHelper helper(builder, loc);

  366.     return helper.exp(helper.mul(operands[1], helper.log(operands[0])));

  367. }

  368. 

  369. //! SIGMOID_GRAD: x * (1 - x) * y

  370. template <>

  371. mlir::Value lower_mode<Mode::SIGMOID_GRAD>(mlir::OpBuilder& builder,

  372.                                            mlir::Location loc,

  373.                                            ValueRange operands) {

  374.     ValueBuilderHelper helper(builder, loc);

  375.     return helper.mul(helper.mul(operands[0], helper.sub(helper.const_f32(1.f),

  376.                                                          operands[0])),

  377.                       operands[1]);

  378. }

  379. 

  380. //! SWITCH_GT0: (x > 0) * y

  381. template <>

  382. mlir::Value lower_mode<Mode::SWITCH_GT0>(mlir::OpBuilder& builder,

  383.                                          mlir::Location loc,

  384.                                          ValueRange operands) {

  385.     ValueBuilderHelper helper(builder, loc);

  386.     return helper.select(helper.gt(operands[0], helper.const_f32(0.f)),

  387.                          operands[1], helper.const_f32(0.f));

  388. }

  389. 

  390. //! TANH_GRAD: (1 - x * x) * y

  391. template <>

  392. mlir::Value lower_mode<Mode::TANH_GRAD>(mlir::OpBuilder& builder,

  393.                                         mlir::Location loc,

  394.                                         ValueRange operands) {

  395.     ValueBuilderHelper helper(builder, loc);

  396.     return helper.mul(helper.sub(helper.const_f32(1.0f),

  397.                                  helper.mul(operands[0], operands[0])),

  398.                       operands[1]);

  399. }

  400. 

  401. /* ===================== ternary op ===================== */

  402. 

  403. //! COND_LEQ_MOV: x <= y ? z : ctype(0)

  404. template <>

  405. mlir::Value lower_mode<Mode::COND_LEQ_MOV>(mlir::OpBuilder& builder,

  406.                                            mlir::Location loc,

  407.                                            ValueRange operands) {

  408.     ValueBuilderHelper helper(builder, loc);

  409.     return helper.select(helper.le(operands[0], operands[1]), operands[2],

  410.                          helper.const_f32(0.f));

  411. }

  412. 

  413. //! FUSE_MUL_ADD3: x * y + z

  414. template <>

  415. mlir::Value lower_mode<Mode::FUSE_MUL_ADD3>(mlir::OpBuilder& builder,

  416.                                             mlir::Location loc,

  417.                                             ValueRange operands) {

  418.     ValueBuilderHelper helper(builder, loc);

  419.     return helper.add(helper.mul(operands[0], operands[1]), operands[2]);

  420. }

  421. 

  422. /* ===================== elemwise ===================== */

  423. 

  424. mlir::Value lower_elemwise_to_std(mlir::Operation* op, mlir::OpBuilder& builder,

  425.                                   mlir::Location loc, ValueRange operands) {

  426.     auto mode = llvm::dyn_cast<dialect::Elemwise>(op).mode();

  427.     switch (mode) {

  428. #define cb(_, _mode)  \

  429.     case Mode::_mode: \

  430.         return lower_mode<Mode::_mode>(builder, loc, operands);

  431.         MLIR_MGB_FOREACH_ELEMWISE_MODE_UNARY(cb);

  432.         MLIR_MGB_FOREACH_ELEMWISE_MODE_BINARY(cb);

  433.         MLIR_MGB_FOREACH_ELEMWISE_MODE_TERNARY(cb);

  434.         default:

  435.             return nullptr;

  436.     }

  437. #undef cb

  438. }

  439. 

  440. /* ===================== typecvt ===================== */

  441. 

  442. mlir::Value lower_typecvt_to_std(mlir::Operation* op, mlir::OpBuilder& builder,

  443.                                  mlir::Location loc, mlir::Value input) {

  444.     auto&& typecvt = llvm::dyn_cast<dialect::TypeCvt>(op);

  445.     megdnn::DType idtype = typecvt.idtype();

  446.     megdnn::DType odtype = typecvt.odtype();

  447. 

  448.     mlir::Type itype = input.getType();

  449.     mlir::Type otype = megdnn_dtype_to_mlir_type(odtype, builder.getContext());

  450. 

  451.     if (mlir::FPExtOp::areCastCompatible(itype, otype)) {

  452.         return builder.create<mlir::FPExtOp>(loc, otype, input);

  453.     } else if (mlir::FPTruncOp::areCastCompatible(itype, otype)) {

  454.         return builder.create<mlir::FPTruncOp>(loc, otype, input);

  455.     } else if (mlir::FPToSIOp::areCastCompatible(itype, otype) and

  456.                is_signed_int_dtype(odtype)) {

  457.         return builder.create<mlir::FPToSIOp>(loc, otype, input);

  458.     } else if (mlir::FPToUIOp::areCastCompatible(itype, otype) and

  459.                is_unsigned_int_dtype(odtype)) {

  460.         return builder.create<mlir::FPToUIOp>(loc, otype, input);

  461.     } else if (mlir::SIToFPOp::areCastCompatible(itype, otype) and

  462.                is_signed_int_dtype(idtype)) {

  463.         return builder.create<mlir::SIToFPOp>(loc, otype, input);

  464.     } else if (mlir::UIToFPOp::areCastCompatible(itype, otype) and

  465.                is_unsigned_int_dtype(idtype)) {

  466.         return builder.create<mlir::UIToFPOp>(loc, otype, input);

  467.     } else {

  468.         mgb_throw(InternalError, "cannot convert from %s to %s", idtype.name(),

  469.                   odtype.name());

  470.     }

  471. 

  472.     return nullptr;

  473. }

  474. 

  475. }  // namespace jit

  476. }  // namespace mgb

  477. 

  478. #endif  // MGB_JIT && MGB_JIT_MLIR

  479. 

  480. // vim: syntax=cpp.doxygen