[TIR][UX] allow override when register TensorIntrin (#12439)

* allow override when register TensorIntrin * lint
apache · Aug 15, 2022 · 55f1d7e · 55f1d7e
1 parent 25c4a73
commit 55f1d7e
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 6 deletions.
diff --git a/include/tvm/tir/function.h b/include/tvm/tir/function.h
@@ -225,10 +225,11 @@ class TensorIntrin : public ObjectRef {
    * up with its name.
    * \param name The name of the TensorIntrin to register
    * \param intrin The TensorIntrin to register.
+   * \param override Whether override existing intrinsic.
    * \throws This method throws an exception if the TensorIntrin with the specified name already
    *         exists.
    */
-  TVM_DLL static void Register(String name, TensorIntrin intrin);
+  TVM_DLL static void Register(String name, TensorIntrin intrin, bool override = false);
 
   /*!
    * \brief Look up TensorIntrin by name. Raises an exception if not found.

diff --git a/python/tvm/tir/function.py b/python/tvm/tir/function.py
@@ -226,7 +226,7 @@ def __init__(self, desc, impl):
         self.__init_handle_by_constructor__(_ffi_api.TensorIntrin, desc, impl)
 
     @staticmethod
-    def register(name: str, desc: PrimFunc, impl: PrimFunc):
+    def register(name: str, desc: PrimFunc, impl: PrimFunc, override: bool = False):
         """Register a tensor intrinsic with its name.
 
         Parameters
@@ -237,8 +237,12 @@ def register(name: str, desc: PrimFunc, impl: PrimFunc):
             The function to describe the computation.
         impl : PrimFunc
             The function of the implementation for the execution.
+        override: bool
+            Whether override existing intrinsic.
         """
-        return _ffi_api.TensorIntrinRegister(name, TensorIntrin(desc, impl))  # type: ignore
+        return _ffi_api.TensorIntrinRegister(
+            name, TensorIntrin(desc, impl), override
+        )  # type: ignore
 
     @staticmethod
     def get(name: str):

diff --git a/src/tir/ir/function.cc b/src/tir/ir/function.cc
@@ -88,10 +88,12 @@ TensorIntrin::TensorIntrin(PrimFunc desc, PrimFunc impl) {
   data_ = std::move(n);
 }
 
-void TensorIntrin::Register(String name, TensorIntrin intrin) {
+void TensorIntrin::Register(String name, TensorIntrin intrin, bool override) {
   TensorIntrinManager* manager = TensorIntrinManager::Global();
-  CHECK_EQ(manager->reg.count(name), 0)
-      << "ValueError: TensorIntrin '" << name << "' has already been registered";
+  if (!override) {
+    CHECK_EQ(manager->reg.count(name), 0)
+        << "ValueError: TensorIntrin '" << name << "' has already been registered";
+  }
   manager->reg.Set(name, intrin);
 }