junrushao1994 commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r616961994
##########
File path: src/ir/op.cc
##########
@@ -122,6 +124,18 @@ TVM_REGISTER_GLOBAL("ir.RegisterOpAttr")
}
});
+TVM_REGISTER_GLOBAL("ir.RegisterOpLowerIntrinsic")
+ .set_body_typed([](String name, PackedFunc f, String target, int plevel,
int can_override) {
+ if (Op::HasAttrMap(target + ".FLowerIntrinsic") &&
+ OpRegistry::Global()->Get(name) != nullptr &&
+ Op::GetAttrMap<FLowerIntrinsic>(target +
".FLowerIntrinsic").count(Op::Get(name))) {
+ ICHECK(can_override) << "Op " << name << "'s intrinsic lowering
function " << target
Review comment:
Use `CHECK` here because it is a user-facing error message.
```suggestion
CHECK(can_override) << "Op " << name << "'s intrinsic lowering
function " << target
```
##########
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",
+
PackedFunc(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",
PackedFunc(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",
+
PackedFunc(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:
IIRC without this `PackedFunc` conversion the code still compiles
(because we do implicit conversion), so in this particular case (registering
with a lambda), we can remove this explicit conversion
##########
File path: src/ir/op.cc
##########
@@ -122,6 +124,18 @@ TVM_REGISTER_GLOBAL("ir.RegisterOpAttr")
}
});
+TVM_REGISTER_GLOBAL("ir.RegisterOpLowerIntrinsic")
+ .set_body_typed([](String name, PackedFunc f, String target, int plevel,
int can_override) {
+ if (Op::HasAttrMap(target + ".FLowerIntrinsic") &&
+ OpRegistry::Global()->Get(name) != nullptr &&
+ Op::GetAttrMap<FLowerIntrinsic>(target +
".FLowerIntrinsic").count(Op::Get(name))) {
+ ICHECK(can_override) << "Op " << name << "'s intrinsic lowering
function " << target
+ << ".FlowerIntrinsic is already registered";
+ }
+
tvm::OpRegEntry::RegisterOrGet(name).set_name().set_attr<FLowerIntrinsic>(
Review comment:
No need to set_name in this case
##########
File path: python/tvm/ir/op.py
##########
@@ -115,3 +127,43 @@ def _register(v):
return v
return _register(value) if value is not None else _register
+
+
+def register_op_intrin_lowering(
+ op_name,
+ target,
+ f=None,
+ level=10,
+ override=False,
+):
+ """Register Op lowering function
+
+ Parameters
+ ----------
+ op_name : str or function
+ The op name
Review comment:
Let's only allow the string case
##########
File path: python/tvm/target/intrin.py
##########
@@ -112,14 +68,18 @@ def _rule_float_direct(op):
See Also
--------
- register_intrin_rule : The registration function for intrin rule.
+ tvm.ir.op.register_op_intrin_lowering : The registration function for
intrinsic lowering rule.
"""
if str(op.dtype).startswith("float"):
return call_pure_extern(op.dtype, op.op.name[4:], *op.args)
return None
# opencl pattern for exp
-register_intrin_rule("opencl", "exp", _rule_float_direct, override=True)
+tvm.ir.op.register_op_intrin_lowering(
Review comment:
we should do like, `from tvm.ir.op import register_op_intrin_lowering`,
in the beginning of this file?
--
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:
[email protected]
