[AutoTVM] Added @functools.wraps to function decorators

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"""