tqchen commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r619197339
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMIntrin<::llvm::Intrinsic::prefetch,
4>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2,
1>);
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs,
TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call =
e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr ln10 = make_const(x.dtype(),
2.302585093);
+ PrimExpr ret = exp(x * ln10);
+ *rv = ret;
+ }));
+
+TVM_REGISTER_OP("tir.fma").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd,
3>);
+
+TVM_REGISTER_OP("tir.log").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
+
+TVM_REGISTER_OP("tir.log2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log2,
1>);
+
+TVM_REGISTER_OP("tir.log10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::log10, 1>);
+
+TVM_REGISTER_OP("tir.sqrt")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt,
1>);
+
+TVM_REGISTER_OP("tir.floor")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
+
+TVM_REGISTER_OP("tir.ceil")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil,
1>);
+
+TVM_REGISTER_OP("tir.trunc")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::trunc, 1>);
+
+TVM_REGISTER_OP("tir.fabs")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::fabs,
1>);
+
+TVM_REGISTER_OP("tir.round")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::round, 1>);
+
+TVM_REGISTER_OP("tir.nearbyint")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
+
+TVM_REGISTER_OP("tir.tanh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs,
TVMRetValue* rv) {
Review comment:
[](PrimExpr expr) -> PrimExpr
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMIntrin<::llvm::Intrinsic::prefetch,
4>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2,
1>);
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs,
TVMRetValue* rv) {
Review comment:
Consider convert this to TypedPackedFunc as well, by using a
`[](PrimExpr expr) -> PrimExpr` signature
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMIntrin<::llvm::Intrinsic::prefetch,
4>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2,
1>);
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs,
TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call =
e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr ln10 = make_const(x.dtype(),
2.302585093);
+ PrimExpr ret = exp(x * ln10);
+ *rv = ret;
+ }));
+
+TVM_REGISTER_OP("tir.fma").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd,
3>);
+
+TVM_REGISTER_OP("tir.log").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
+
+TVM_REGISTER_OP("tir.log2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log2,
1>);
+
+TVM_REGISTER_OP("tir.log10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::log10, 1>);
+
+TVM_REGISTER_OP("tir.sqrt")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt,
1>);
+
+TVM_REGISTER_OP("tir.floor")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
+
+TVM_REGISTER_OP("tir.ceil")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil,
1>);
+
+TVM_REGISTER_OP("tir.trunc")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::trunc, 1>);
+
+TVM_REGISTER_OP("tir.fabs")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::fabs,
1>);
+
+TVM_REGISTER_OP("tir.round")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::round, 1>);
+
+TVM_REGISTER_OP("tir.nearbyint")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
+
+TVM_REGISTER_OP("tir.tanh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs,
TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call =
e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr one = make_const(x.dtype(), 1);
+ PrimExpr two = make_const(x.dtype(), 2);
+ PrimExpr neg_two = make_const(x.dtype(), -2);
+
+ PrimExpr exp_neg2x = exp(neg_two * x);
+ PrimExpr exp_pos2x = exp(two * x);
+
+ PrimExpr tanh_pos = (one - exp_neg2x) / (one
+ exp_neg2x);
+ PrimExpr tanh_neg = (exp_pos2x - one) /
(exp_pos2x + one);
+ *rv = tir::Select(x >= make_zero(x.dtype()),
tanh_pos, tanh_neg);
+ }));
+
+TVM_REGISTER_OP("tir.pow").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::pow, 2>);
+
+TVM_REGISTER_OP("tir.popcount")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::ctpop, 1>);
+
+TVM_REGISTER_OP("tir.tan").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", PackedFunc([](const TVMArgs& targs, TVMRetValue*
rv) {
Review comment:
NOTE: we should change these to FLegalize later once we start to
introduce legalize registries logics. And perhaps move to generic legalize
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMIntrin<::llvm::Intrinsic::prefetch,
4>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2,
1>);
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs,
TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call =
e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr ln10 = make_const(x.dtype(),
2.302585093);
+ PrimExpr ret = exp(x * ln10);
+ *rv = ret;
+ }));
+
+TVM_REGISTER_OP("tir.fma").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd,
3>);
+
+TVM_REGISTER_OP("tir.log").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
+
+TVM_REGISTER_OP("tir.log2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log2,
1>);
+
+TVM_REGISTER_OP("tir.log10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::log10, 1>);
+
+TVM_REGISTER_OP("tir.sqrt")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt,
1>);
+
+TVM_REGISTER_OP("tir.floor")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
+
+TVM_REGISTER_OP("tir.ceil")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil,
1>);
+
+TVM_REGISTER_OP("tir.trunc")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::trunc, 1>);
+
+TVM_REGISTER_OP("tir.fabs")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::fabs,
1>);
+
+TVM_REGISTER_OP("tir.round")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::round, 1>);
+
+TVM_REGISTER_OP("tir.nearbyint")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
+
+TVM_REGISTER_OP("tir.tanh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs,
TVMRetValue* rv) {
Review comment:
NOTE: we should change these to FLegalize later once we start to
introduce legalize registries logics. And perhaps move to generic legalize
##########
File path: src/tir/transforms/lower_intrin.cc
##########
@@ -42,28 +42,38 @@ class IntrinInjecter : public
tvm::arith::IRMutatorWithAnalyzer {
IntrinInjecter(arith::Analyzer* analyzer, std::string target, std::string
mtriple = "")
: IRMutatorWithAnalyzer(analyzer) {
- patterns_.push_back("tvm.intrin.rule." + target + ".");
+ patterns_.push_back(target + ".FLowerIntrinsic");
bool is_llvm_aarch64 = (mtriple.find("aarch64") != std::string::npos);
if (is_llvm_aarch64) {
- patterns_.push_back("tvm.intrin.rule." + target + "." + "aarch64.");
+ patterns_.push_back(target + ".aarch64.FLowerIntrinsic");
}
- patterns_.push_back("tvm.intrin.rule.default.");
- fma_ = runtime::Registry::Get(patterns_[0] + "fma");
+ patterns_.push_back("default.FLowerIntrinsic");
+ fma_ = runtime::Registry::Get("tvm.intrin.rule." + target + ".fma");
if (target == "stackvm") {
support_bitwise_op_ = false;
}
}
PrimExpr VisitExpr_(const CallNode* op) final {
if (auto* ptr_op = op->op.as<OpNode>()) {
- // Still use legacy string based rewriting
- // TODO(tvm-team): migrate the pattern application from global function
look up
- // to an OpAttrMap<PackedFunc>
- std::string name = ptr_op->name;
- PrimExpr r = ApplyPattern(name, GetRef<PrimExpr>(op));
- if (r.defined()) return r;
+ for (const std::string& pattern : patterns_)
+ if (Op::HasAttrMap(pattern)) {
+ auto f_lower_intrin_map = Op::GetAttrMap<FLowerIntrinsic>(pattern);
Review comment:
Let us restructure this. Because the list of patterns are known in
constructor time, it is better to cache the attr_map early. So we do not
involve string based lookup in this internal CallNode.
The string based lookup is done once during constructor time, then the rest
of lookups are all index based. cc @junrushao1994
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMIntrin<::llvm::Intrinsic::prefetch,
4>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2,
1>);
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs,
TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call =
e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr ln10 = make_const(x.dtype(),
2.302585093);
+ PrimExpr ret = exp(x * ln10);
+ *rv = ret;
+ }));
+
+TVM_REGISTER_OP("tir.fma").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd,
3>);
+
+TVM_REGISTER_OP("tir.log").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
+
+TVM_REGISTER_OP("tir.log2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log2,
1>);
+
+TVM_REGISTER_OP("tir.log10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::log10, 1>);
+
+TVM_REGISTER_OP("tir.sqrt")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt,
1>);
+
+TVM_REGISTER_OP("tir.floor")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
+
+TVM_REGISTER_OP("tir.ceil")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil,
1>);
+
+TVM_REGISTER_OP("tir.trunc")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::trunc, 1>);
+
+TVM_REGISTER_OP("tir.fabs")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::fabs,
1>);
+
+TVM_REGISTER_OP("tir.round")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::round, 1>);
+
+TVM_REGISTER_OP("tir.nearbyint")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
+
+TVM_REGISTER_OP("tir.tanh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs,
TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call =
e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr one = make_const(x.dtype(), 1);
+ PrimExpr two = make_const(x.dtype(), 2);
+ PrimExpr neg_two = make_const(x.dtype(), -2);
+
+ PrimExpr exp_neg2x = exp(neg_two * x);
+ PrimExpr exp_pos2x = exp(two * x);
+
+ PrimExpr tanh_pos = (one - exp_neg2x) / (one
+ exp_neg2x);
+ PrimExpr tanh_neg = (exp_pos2x - one) /
(exp_pos2x + one);
+ *rv = tir::Select(x >= make_zero(x.dtype()),
tanh_pos, tanh_neg);
+ }));
+
+TVM_REGISTER_OP("tir.pow").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::pow, 2>);
+
+TVM_REGISTER_OP("tir.popcount")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::ctpop, 1>);
+
+TVM_REGISTER_OP("tir.tan").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", PackedFunc([](const TVMArgs& targs, TVMRetValue*
rv) {
Review comment:
[](PrimExpr expr) -> PrimExpr
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMIntrin<::llvm::Intrinsic::prefetch,
4>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2,
1>);
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs,
TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call =
e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr ln10 = make_const(x.dtype(),
2.302585093);
+ PrimExpr ret = exp(x * ln10);
+ *rv = ret;
+ }));
+
+TVM_REGISTER_OP("tir.fma").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd,
3>);
+
+TVM_REGISTER_OP("tir.log").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
+
+TVM_REGISTER_OP("tir.log2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log2,
1>);
+
+TVM_REGISTER_OP("tir.log10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::log10, 1>);
+
+TVM_REGISTER_OP("tir.sqrt")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt,
1>);
+
+TVM_REGISTER_OP("tir.floor")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
+
+TVM_REGISTER_OP("tir.ceil")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil,
1>);
+
+TVM_REGISTER_OP("tir.trunc")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::trunc, 1>);
+
+TVM_REGISTER_OP("tir.fabs")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::fabs,
1>);
+
+TVM_REGISTER_OP("tir.round")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::round, 1>);
+
+TVM_REGISTER_OP("tir.nearbyint")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
+
+TVM_REGISTER_OP("tir.tanh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs,
TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call =
e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr one = make_const(x.dtype(), 1);
+ PrimExpr two = make_const(x.dtype(), 2);
+ PrimExpr neg_two = make_const(x.dtype(), -2);
+
+ PrimExpr exp_neg2x = exp(neg_two * x);
+ PrimExpr exp_pos2x = exp(two * x);
+
+ PrimExpr tanh_pos = (one - exp_neg2x) / (one
+ exp_neg2x);
+ PrimExpr tanh_neg = (exp_pos2x - one) /
(exp_pos2x + one);
+ *rv = tir::Select(x >= make_zero(x.dtype()),
tanh_pos, tanh_neg);
+ }));
+
+TVM_REGISTER_OP("tir.pow").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::pow, 2>);
+
+TVM_REGISTER_OP("tir.popcount")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::ctpop, 1>);
+
+TVM_REGISTER_OP("tir.tan").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", PackedFunc([](const TVMArgs& targs, TVMRetValue*
rv) {
PrimExpr e = targs[0];
const tir::CallNode* call = e.as<tir::CallNode>();
ICHECK(call != nullptr);
const PrimExpr& x = call->args[0];
- PrimExpr ln10 = make_const(x.dtype(), 2.302585093);
- PrimExpr ret = exp(x * ln10);
- *rv = ret;
- });
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.fma")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd, 3>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.log")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.log2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::log2, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.log10")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::log10, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.sqrt")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.floor")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.ceil")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.trunc")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::trunc, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.fabs")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::fabs, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.round")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::round, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.nearbyint")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.tanh")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
- PrimExpr e = targs[0];
- const tir::CallNode* call = e.as<tir::CallNode>();
- ICHECK(call != nullptr);
- const PrimExpr& x = call->args[0];
- PrimExpr one = make_const(x.dtype(), 1);
- PrimExpr two = make_const(x.dtype(), 2);
- PrimExpr neg_two = make_const(x.dtype(), -2);
-
- PrimExpr exp_neg2x = exp(neg_two * x);
- PrimExpr exp_pos2x = exp(two * x);
-
- PrimExpr tanh_pos = (one - exp_neg2x) / (one + exp_neg2x);
- PrimExpr tanh_neg = (exp_pos2x - one) / (exp_pos2x + one);
- *rv = tir::Select(x >= make_zero(x.dtype()), tanh_pos, tanh_neg);
- });
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.pow")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::pow, 2>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.popcount")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::ctpop, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.tan").set_body([](const TVMArgs&
targs, TVMRetValue* rv) {
- PrimExpr e = targs[0];
- const tir::CallNode* call = e.as<tir::CallNode>();
- ICHECK(call != nullptr);
- const PrimExpr& x = call->args[0];
- PrimExpr tan_x = sin(x) / cos(x);
- *rv = tan_x;
-});
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.cos")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::cos, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.cosh")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+ PrimExpr tan_x = sin(x) / cos(x);
+ *rv = tan_x;
+ }));
+
+TVM_REGISTER_OP("tir.cos").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::cos, 1>);
+
+TVM_REGISTER_OP("tir.cosh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs,
TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call =
e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr two = make_const(x.dtype(), 2);
+ PrimExpr neg_one = make_const(x.dtype(), -1);
+ PrimExpr exp_negx = exp(neg_one * x);
+ PrimExpr exp_posx = exp(x);
+ PrimExpr ret = (exp_posx + exp_negx) / two;
+ *rv = ret;
+ }));
+
+TVM_REGISTER_OP("tir.sin").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::sin, 1>);
+
+TVM_REGISTER_OP("tir.sinh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs,
TVMRetValue* rv) {
Review comment:
[](PrimExpr expr) -> PrimExpr
##########
File path: src/target/llvm/intrin_rule_llvm.cc
##########
@@ -25,155 +25,175 @@
#include "intrin_rule_llvm.h"
#include <tvm/tir/op.h>
+#include <tvm/tir/op_attr_types.h>
namespace tvm {
namespace codegen {
namespace llvm {
+using tir::FLowerIntrinsic;
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.prefetch")
- .set_body(DispatchLLVMIntrin<::llvm::Intrinsic::prefetch, 4>);
+TVM_REGISTER_OP("tir.prefetch")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMIntrin<::llvm::Intrinsic::prefetch,
4>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
+TVM_REGISTER_OP("tir.exp").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::exp, 1>);
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2, 1>);
+TVM_REGISTER_OP("tir.exp2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::exp2,
1>);
// TODO(tvm-team): migrate the legalization transformations as a separate
// set of rules in TIR that can be shared across backends.
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.exp10")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+TVM_REGISTER_OP("tir.exp10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs,
TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call =
e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr ln10 = make_const(x.dtype(),
2.302585093);
+ PrimExpr ret = exp(x * ln10);
+ *rv = ret;
+ }));
+
+TVM_REGISTER_OP("tir.fma").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd,
3>);
+
+TVM_REGISTER_OP("tir.log").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
+
+TVM_REGISTER_OP("tir.log2")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::log2,
1>);
+
+TVM_REGISTER_OP("tir.log10")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::log10, 1>);
+
+TVM_REGISTER_OP("tir.sqrt")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt,
1>);
+
+TVM_REGISTER_OP("tir.floor")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
+
+TVM_REGISTER_OP("tir.ceil")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil,
1>);
+
+TVM_REGISTER_OP("tir.trunc")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::trunc, 1>);
+
+TVM_REGISTER_OP("tir.fabs")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ DispatchLLVMPureIntrin<::llvm::Intrinsic::fabs,
1>);
+
+TVM_REGISTER_OP("tir.round")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::round, 1>);
+
+TVM_REGISTER_OP("tir.nearbyint")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
+
+TVM_REGISTER_OP("tir.tanh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs,
TVMRetValue* rv) {
+ using tir::make_const;
+ using tir::make_zero;
+ PrimExpr e = targs[0];
+ const tir::CallNode* call =
e.as<tir::CallNode>();
+ ICHECK(call != nullptr);
+ const PrimExpr& x = call->args[0];
+ PrimExpr one = make_const(x.dtype(), 1);
+ PrimExpr two = make_const(x.dtype(), 2);
+ PrimExpr neg_two = make_const(x.dtype(), -2);
+
+ PrimExpr exp_neg2x = exp(neg_two * x);
+ PrimExpr exp_pos2x = exp(two * x);
+
+ PrimExpr tanh_pos = (one - exp_neg2x) / (one
+ exp_neg2x);
+ PrimExpr tanh_neg = (exp_pos2x - one) /
(exp_pos2x + one);
+ *rv = tir::Select(x >= make_zero(x.dtype()),
tanh_pos, tanh_neg);
+ }));
+
+TVM_REGISTER_OP("tir.pow").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::pow, 2>);
+
+TVM_REGISTER_OP("tir.popcount")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+
DispatchLLVMPureIntrin<::llvm::Intrinsic::ctpop, 1>);
+
+TVM_REGISTER_OP("tir.tan").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", PackedFunc([](const TVMArgs& targs, TVMRetValue*
rv) {
PrimExpr e = targs[0];
const tir::CallNode* call = e.as<tir::CallNode>();
ICHECK(call != nullptr);
const PrimExpr& x = call->args[0];
- PrimExpr ln10 = make_const(x.dtype(), 2.302585093);
- PrimExpr ret = exp(x * ln10);
- *rv = ret;
- });
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.fma")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::fmuladd, 3>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.log")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::log, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.log2")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::log2, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.log10")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::log10, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.sqrt")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::sqrt, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.floor")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::floor, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.ceil")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::ceil, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.trunc")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::trunc, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.fabs")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::fabs, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.round")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::round, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.nearbyint")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::nearbyint, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.tanh")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
- PrimExpr e = targs[0];
- const tir::CallNode* call = e.as<tir::CallNode>();
- ICHECK(call != nullptr);
- const PrimExpr& x = call->args[0];
- PrimExpr one = make_const(x.dtype(), 1);
- PrimExpr two = make_const(x.dtype(), 2);
- PrimExpr neg_two = make_const(x.dtype(), -2);
-
- PrimExpr exp_neg2x = exp(neg_two * x);
- PrimExpr exp_pos2x = exp(two * x);
-
- PrimExpr tanh_pos = (one - exp_neg2x) / (one + exp_neg2x);
- PrimExpr tanh_neg = (exp_pos2x - one) / (exp_pos2x + one);
- *rv = tir::Select(x >= make_zero(x.dtype()), tanh_pos, tanh_neg);
- });
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.pow")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::pow, 2>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.popcount")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::ctpop, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.tan").set_body([](const TVMArgs&
targs, TVMRetValue* rv) {
- PrimExpr e = targs[0];
- const tir::CallNode* call = e.as<tir::CallNode>();
- ICHECK(call != nullptr);
- const PrimExpr& x = call->args[0];
- PrimExpr tan_x = sin(x) / cos(x);
- *rv = tan_x;
-});
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.cos")
- .set_body(DispatchLLVMPureIntrin<::llvm::Intrinsic::cos, 1>);
-
-TVM_REGISTER_GLOBAL("tvm.intrin.rule.llvm.cosh")
- .set_body([](const TVMArgs& targs, TVMRetValue* rv) {
- using tir::make_const;
- using tir::make_zero;
+ PrimExpr tan_x = sin(x) / cos(x);
+ *rv = tan_x;
+ }));
+
+TVM_REGISTER_OP("tir.cos").set_attr<FLowerIntrinsic>(
+ "llvm.FLowerIntrinsic", DispatchLLVMPureIntrin<::llvm::Intrinsic::cos, 1>);
+
+TVM_REGISTER_OP("tir.cosh")
+ .set_attr<FLowerIntrinsic>("llvm.FLowerIntrinsic",
+ PackedFunc([](const TVMArgs& targs,
TVMRetValue* rv) {
Review comment:
[](PrimExpr expr) -> PrimExpr
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
