From 8e4e6664a0199c1070d9b7dd505e0561bf0c6b46 Mon Sep 17 00:00:00 2001
From: Eric Lunderberg <elunderberg@octoml.ai>
Date: Thu, 10 Jun 2021 09:42:16 -0700
Subject: [PATCH] [AutoTVM] Added @functools.wraps to function decorators

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)>`.)
---
 python/tvm/autotvm/task/code_hash.py        | 3 +++
 python/tvm/autotvm/task/task.py             | 3 +++
 python/tvm/autotvm/task/topi_integration.py | 4 ++++
 3 files changed, 10 insertions(+)

diff --git a/python/tvm/autotvm/task/code_hash.py b/python/tvm/autotvm/task/code_hash.py
index 3331fc13c719..2bd053da7244 100644
--- a/python/tvm/autotvm/task/code_hash.py
+++ b/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)
diff --git a/python/tvm/autotvm/task/task.py b/python/tvm/autotvm/task/task.py
index 668832b8a86c..1f5827d7e9d0 100644
--- a/python/tvm/autotvm/task/task.py
+++ b/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)
diff --git a/python/tvm/autotvm/task/topi_integration.py b/python/tvm/autotvm/task/topi_integration.py
index 2558c7669ac9..32d8674640ed 100644
--- a/python/tvm/autotvm/task/topi_integration.py
+++ b/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"""