[Autotvm] Support override

python/tvm/autotvm/task/topi_integration.py

@@ -284,7 +284,7 @@ def get(allow_duplicate=False):
         return TaskExtractEnv.current
 
 
-def register_topi_compute(topi_compute, target_keys, template_keys, func=None):
+def register_topi_compute(topi_compute, target_keys, template_keys, func=None, override=False):
     """Register a tunable template for a topi compute function.
 
     After the registration, this topi compute will become a configuration dispatcher. It uses
@@ -333,7 +333,7 @@ def config_dispatcher(*args, **kwargs):
 
             config_dispatcher = _REGISTERED_DISPATCHER[target_key][topi_compute]
 
-            @config_dispatcher.register(template_keys)
+            @config_dispatcher.register(template_keys, override=override)
             def template_call(cfg, *args, **kwargs):
                 """call the topi func and attach workload to compute node"""
                 assert not kwargs, "Do not support kwargs in template function call"
@@ -372,7 +372,7 @@ def template_call(cfg, *args, **kwargs):
     return _decorator
 
 
-def register_topi_schedule(topi_schedule, target_keys, template_keys, func=None):
+def register_topi_schedule(topi_schedule, target_keys, template_keys, func=None, override=False):
     """Register a tunable template for a topi schedule function.
 
     After the registration. This topi schedule will become a configuration dispatcher. It dispatches
@@ -438,7 +438,7 @@ def traverse(tensors):
 
             config_dispatcher = _REGISTERED_DISPATCHER[target_key][topi_schedule]
 
-            @config_dispatcher.register(template_keys)
+            @config_dispatcher.register(template_keys, override=override)
             def template_call(cfg, outs, *args, **kwargs):
                 """call the schedule func"""
                 if f == topi_schedule.fdefault:

tests/python/unittest/test_autotvm_dispatch_context.py

@@ -61,6 +61,17 @@ def simple_template(a, b):
     simple_template(2, 3)
 
 
+def test_override():
+    from topi import generic, nn
+    @autotvm.register_topi_compute(nn.dense, 'cpu', ['direct'], override=True)
+    def dense(cfg, data, weight, bias=None, out_dtype=None):
+        pass
+
+    @autotvm.register_topi_schedule(generic.schedule_dense, 'cpu', ['direct'], override=True)
+    def schedule_dense(cfg, outs):
+        pass
+
 if __name__ == "__main__":
     test_dispatch()
     test_fallback()
+    test_override()