[Target][Lowering] Update Op Intrinsic Lowering Mechanism And Intrinsic Lowering Pass

include/tvm/tir/op_attr_types.h

@@ -28,11 +28,12 @@
 #ifndef TVM_TIR_OP_ATTR_TYPES_H_
 #define TVM_TIR_OP_ATTR_TYPES_H_
 
+#include <tvm/ir/expr.h>
 #include <tvm/runtime/container.h>
+#include <tvm/runtime/packed_func.h>
 
 namespace tvm {
 namespace tir {
-
 /*!
  * \brief Global symbol of the op after lowering.
  */
@@ -43,6 +44,16 @@ using TGlobalSymbol = String;
  */
 using TVectorizable = bool;
 
+/*!
+ * \brief The intrinsic lowering function for given op.
+ */
+using FLowerIntrinsic = runtime::TypedPackedFunc<PrimExpr(PrimExpr)>;
+
+/*!
+ * \brief The legalization function for given tir op.
+ */
+using FLegalize = runtime::TypedPackedFunc<PrimExpr(PrimExpr)>;
+
 /*!
  * \brief The effect type of the call.
  */

python/tvm/ir/__init__.py

@@ -23,7 +23,7 @@
 from .tensor_type import TensorType
 from .type_relation import TypeCall, TypeRelation
 from .expr import BaseExpr, PrimExpr, RelayExpr, GlobalVar, Range
-from .op import Op, register_op_attr
+from .op import Op, register_op_attr, register_intrin_lowering
 from .function import CallingConv, BaseFunc
 from .adt import Constructor, TypeData
 from .module import IRModule

python/tvm/ir/op.py

@@ -115,3 +115,40 @@ def _register(v):
         return v
 
     return _register(value) if value is not None else _register
+
+
+def register_intrin_lowering(
+    op_name,
+    target,
+    *,
+    f=None,
+    level=10,
+):
+    """Register Op lowering function
+
+    Parameters
+    ----------
+    op_name : str
+        The op name
+
+    target : str
+        The target string for given intrinsic lowering function
+
+    f : function, optional
+        The function to be registered.
+
+    level : int
+        The priority level
+
+    Returns
+    -------
+    fregister : function
+        Register op lowering function if f is not specified.
+    """
+
+    def _register(f):
+        """internal register function"""
+        _ffi_api.RegisterOpLowerIntrinsic(op_name, f, target, level)
+        return f
+
+    return _register(f) if f is not None else _register

python/tvm/target/__init__.py

@@ -61,4 +61,3 @@
 from .generic_func import generic_func, get_native_generic_func, override_native_generic_func
 from . import datatype
 from . import codegen
-from .intrin import register_intrin_rule

python/tvm/target/intrin.py

@@ -15,54 +15,10 @@
 # specific language governing permissions and limitations
 # under the License.
 """Target dependent intrinsic registration."""
-import tvm._ffi
+from tvm.ir import register_intrin_lowering
 from tvm.tir import call_pure_extern
 
 
-# Intrinsic rule related code
-def register_intrin_rule(target, intrin, f=None, override=False):
-    """Register an intrinsic function generation rule.
-
-    Intrinsic generation rules are callback functions for
-    code generator to get device specific calls.
-    This function simply translates to.
-
-    :code:`register_func("tvm.intrin.rule.%s.%s" % (target, intrin), f, override)`
-
-    TVM may already pre-register intrinsic rules in the backend.
-    However, user can use this function to change the intrinsic translation
-    behavior or add new intrinsic rules during runtime.
-
-    Parameters
-    ----------
-    target : str
-        The name of codegen target.
-
-    intrin : str
-        The name of the intrinsic.
-
-    f : function, optional
-        The function to be registered.
-
-    override: boolean optional
-        Whether override existing entry.
-
-    Returns
-    -------
-    fregister : function
-        Register function if f is not specified.
-
-    Examples
-    --------
-    The following code registers exp expansion rule for opencl.
-
-    .. code-block:: python
-
-        register_intrin_rule("opencl", "exp", my_exp_rule, override=True)
-    """
-    return tvm._ffi.register_func("tvm.intrin.rule.%s.%s" % (target, intrin), f, override)
-
-
 def _rule_float_suffix(op):
     """Intrinsic rule: Add float suffix if it is float32.
 
@@ -81,7 +37,7 @@ def _rule_float_suffix(op):
 
     See Also
     --------
-    register_intrin_rule : The registration function for intrin rule.
+    register_intrin_lowering : The registration function for intrinsic lowering rule.
     """
     name = op.op.name
     assert name.startswith("tir.")
@@ -112,14 +68,14 @@ def _rule_float_direct(op):
 
     See Also
     --------
-    register_intrin_rule : The registration function for intrin rule.
+    register_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)
+register_intrin_lowering("tir.exp", target="opencl", f=_rule_float_direct, level=99)
 # default pattern for exp
-register_intrin_rule("default", "exp", _rule_float_suffix, override=True)
+register_intrin_lowering("tir.exp", target="default", f=_rule_float_suffix, level=99)

python/tvm/topi/arm_cpu/tensor_intrin.py

@@ -19,6 +19,7 @@
 
 import tvm
 from tvm import te
+from tvm.ir import register_intrin_lowering
 
 
 def gemm_4x4_int8_int8_int32(M, N, K, unroll, in_type):
@@ -1054,6 +1055,6 @@ def _q_multiply_shift_arm(op):
     return tvm.tir.Select(s < 0, out_1, out_2)
 
 
-tvm.target.intrin.register_intrin_rule(
-    "llvm.aarch64", "q_multiply_shift", _q_multiply_shift_arm, override=True
+register_intrin_lowering(
+    "tir.q_multiply_shift", target="llvm.aarch64", f=_q_multiply_shift_arm, level=99
 )

python/tvm/topi/cuda/nms.py

@@ -21,6 +21,7 @@
 from tvm import te
 from tvm.contrib import nvcc
 from tvm.contrib.thrust import can_use_thrust, can_use_rocthrust
+from tvm.ir import register_intrin_lowering
 from tvm.tir import if_then_else
 from .sort import argsort, argsort_thrust
 from .scan import exclusive_scan
@@ -51,11 +52,9 @@ def opencl_atomic_add_rule(op):
     raise RuntimeError("only support int32")
 
 
-tvm.target.intrin.register_intrin_rule("cuda", "atomic_add", cuda_atomic_add_rule, override=True)
+register_intrin_lowering("tir.atomic_add", target="cuda", f=cuda_atomic_add_rule, level=99)
 
-tvm.target.intrin.register_intrin_rule(
-    "opencl", "atomic_add", opencl_atomic_add_rule, override=True
-)
+register_intrin_lowering("tir.atomic_add", target="opencl", f=opencl_atomic_add_rule, level=99)
 
 
 def atomic_add(x, y):

src/ir/op.cc

@@ -26,6 +26,7 @@
 #include <tvm/runtime/container.h>
 #include <tvm/runtime/module.h>
 #include <tvm/runtime/packed_func.h>
+#include <tvm/tir/op_attr_types.h>
 
 #include <memory>
 
@@ -36,6 +37,7 @@ namespace tvm {
 using runtime::PackedFunc;
 using runtime::TVMArgs;
 using runtime::TVMRetValue;
+using tir::FLowerIntrinsic;
 
 using OpRegistry = AttrRegistry<OpRegEntry, Op>;
 
@@ -122,6 +124,12 @@ TVM_REGISTER_GLOBAL("ir.RegisterOpAttr")
       }
     });
 
+TVM_REGISTER_GLOBAL("ir.RegisterOpLowerIntrinsic")
+    .set_body_typed([](String name, PackedFunc f, String target, int plevel) {
+      tvm::OpRegEntry::RegisterOrGet(name).set_attr<FLowerIntrinsic>(target + ".FLowerIntrinsic", f,
+                                                                     plevel);
+    });
+
 // helper to get internal dev function in objectref.
 struct Op2ObjectPtr : public ObjectRef {
   static ObjectPtr<Object> Get(const Op& op) { return GetDataPtr<Object>(op); }