[AutoTVM] Added @functools.wraps to function decorators (#8237)

This helps in debugging, as the function name, arguments, and
docstrings show the function name from the source code instead of the
wrapper function.(e.g.
`<function tvm.topi.cuda.dense.dense_small_batch(cfg, data, weight, bias=None, out_dtype=None)>`
instead of
`<function tvm.autotvm.task.topi_integration.register_topi_compute.<locals>._decorate.<locals>.wrapper(*args, **kwargs)>`.)

Co-authored-by: Eric Lunderberg <elunderberg@octoml.ai>

python/tvm/autotvm/task/code_hash.py

@@ -19,6 +19,7 @@
 
 code hashing is used to check the consistence of schedule code and the parameters loaded from log
 """
+import functools
 import inspect
 import zlib
 
@@ -35,6 +36,7 @@ def attach_code_hash(s):
     """
 
     def decorator(func):
+        @functools.wraps(func)
         def wrapper(*args, **kwargs):
             func(*args, **kwargs)
             raw_hash = zlib.crc32("".join(inspect.getsourcelines(func)[0]).encode())
@@ -56,6 +58,7 @@ def attach_code_hash_to_arg(arg_idx=1):
     """
 
     def decorator(func):
+        @functools.wraps(func)
         def wrapper(*args, **kwargs):
             func(*args, **kwargs)
             assert isinstance(args[arg_idx], schedule.Schedule)

python/tvm/autotvm/task/task.py

@@ -21,6 +21,8 @@
 func is a state-less function, or a string that
 registers the standard task.
 """
+import functools
+
 import numpy as np
 
 from tvm import runtime
@@ -411,6 +413,7 @@ def matmul(N, L, M, dtype):
     """
 
     def _decorate(f):
+        @functools.wraps(f)
         def wrapper(*args, **kwargs):
             assert not kwargs, "Do not support kwargs in template function call"
             workload = args_to_workload(args, task_name)

python/tvm/autotvm/task/topi_integration.py

@@ -26,6 +26,8 @@
 
 See tvm/topi/python/topi/arm_cpu/depthwise_conv2d.py for example usage.
 """
+import functools
+
 import tvm.te._ffi_api
 from tvm.target import Target
 from tvm.te import tensor
@@ -149,6 +151,7 @@ def register_topi_compute(task_name, func=None):
     """
 
     def _decorate(topi_compute):
+        @functools.wraps(topi_compute)
         @_register_task_compute(task_name)
         def wrapper(*args, **kwargs):
             """wrapper function for topi compute"""
@@ -224,6 +227,7 @@ def register_topi_schedule(task_name, func=None):
     """
 
     def _decorate(topi_schedule):
+        @functools.wraps(topi_schedule)
         @_register_task_schedule(task_name)
         def wrapper(outs, *args, **kwargs):
             """wrapper function for topi schedule"""