junrushao1994 commented on a change in pull request #7809:
URL: https://github.com/apache/tvm/pull/7809#discussion_r614581355
##########
File path: python/tvm/ir/op.py
##########
@@ -115,3 +127,50 @@ def _register(v):
return v
return _register(value) if value is not None else _register
+
+
+def register_op_intrin_lowering(
+ op_name,
+ f=None,
+ target="default",
+ level=10,
+ override=False,
+):
+ """Register Op lowering function
+
+ Parameters
+ ----------
+ op_name : str or function
+ The op name
+
+ f : function, optional
+ The function to be registered.
+
+ target : str
+ The target string for given intrinsic lowering function
+
+ level : int
+ The priority level
+
+ override: boolean optional
+ Whether override existing entry.
+
+ Returns
+ -------
+ fregister : function
+ Register op lowering function if f is not specified.
+ """
+ if not isinstance(op_name, str):
+ raise ValueError("expect string op name")
+
+ def _register(myf):
+ """internal intrinsic lowering registration function"""
+ assert isinstance(target, str)
+ if not isinstance(myf, PackedFuncBase):
+ myf = convert_to_tvm_func(myf)
+ _ffi_api.RegisterOpLowerIntrinsic(op_name, myf.handle, target, level,
override)
+ return myf
Review comment:
I assume the logic could be a bit simpler? Perhaps you might want to
refer the some previous implementations, like "def register_compute" :-)
##########
File path: src/ir/op.cc
##########
@@ -122,6 +124,19 @@ TVM_REGISTER_GLOBAL("ir.RegisterOpAttr")
}
});
+TVM_REGISTER_GLOBAL("ir.RegisterOpLowerIntrinsic")
Review comment:
I assume this should be put on `src/tir/op.cc`?
##########
File path: src/ir/op.cc
##########
@@ -122,6 +124,19 @@ TVM_REGISTER_GLOBAL("ir.RegisterOpAttr")
}
});
+TVM_REGISTER_GLOBAL("ir.RegisterOpLowerIntrinsic")
+ .set_body_typed([](String name, TVMFunctionHandle f, String target =
"default", int plevel = 10,
+ int can_override = 0) {
Review comment:
Forget about TVMFunctionHandle, which is only used in the C API but
nowhere else :-)
Let's embrace the PackedFunc magic! It is a nicer wrapper of a function from
either C++ or python (or potentially other languages).
BTW, don't need to set default arguments of a lambda because they are not
used any way.
```suggestion
.set_body_typed([](String name, PackedFunc f, String target, int plevel,
bool override_) {
```
##########
File path: python/tvm/ir/op.py
##########
@@ -115,3 +127,50 @@ def _register(v):
return v
return _register(value) if value is not None else _register
+
+
+def register_op_intrin_lowering(
+ op_name,
+ f=None,
+ target="default",
+ level=10,
+ override=False,
+):
Review comment:
Let's avoid providing too many defaults :-)
```suggestion
def register_op_intrin_lowering(
op_name,
f,
target,
level=10,
override=False,
):
```
--
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]
