intrin_rule_llvm.cc 1.77 KB
Newer Older
1 2 3 4 5 6
/*!
 *  Copyright (c) 2017 by Contributors
 * \file intrin_rule_llvm.cc
 */
#ifdef TVM_LLVM_VERSION

7
#include "./intrin_rule_llvm.h"
8 9 10 11 12

namespace tvm {
namespace codegen {
namespace llvm {

13 14 15
TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
.set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 0>);

16
TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
17
.set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
18

19
TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.fma")
20
.set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd, 1>);
21

22
TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.log")
23
.set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
24

ziheng committed
25
TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.sqrt")
26
.set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt, 1>);
ziheng committed
27

28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48
TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.tanh")
.set_body([](const TVMArgs& targs, TVMRetValue* rv) {
  Expr e = targs[0];
  const ir::Call* call = e.as<ir::Call>();
  CHECK(call != nullptr);
  const Expr& x = call->args[0];
  Expr one = make_const(x.type(), 1);
  Expr two = make_const(x.type(), 2);
  Expr neg_two = make_const(x.type(), -2);

  Expr exp_neg2x = ir::Call::make(
      x.type(), "exp", {neg_two * x}, ir::Call::PureIntrinsic);
  Expr exp_pos2x = ir::Call::make(
      x.type(), "exp", {two * x}, ir::Call::PureIntrinsic);

  Expr tanh_pos = (one - exp_neg2x) / (one + exp_neg2x);
  Expr tanh_neg = (exp_pos2x - one) / (exp_pos2x + one);
  *rv = ir::Select::make(
      x >= make_zero(x.type()), tanh_pos, tanh_neg);
});

49
TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.pow")
50
.set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::pow, 1>);
51

52 53 54
TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.popcount")
.set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::ctpop, 1>);

55 56 57 58 59
}  // namespace llvm
}  // namespace codegen
}  // namespace tvm

#endif  // LLVM_VERSION