From 94a0525737181782485d30839eefa3144a34dd5e Mon Sep 17 00:00:00 2001
From: Eddie Yan <eqy@cs.washington.edu>
Date: Tue, 27 Nov 2018 15:44:14 -0800
Subject: [PATCH 01/14] check in copy paste

---
 python/tvm/autotvm/task/__init__.py          |   1 +
 python/tvm/autotvm/task/relay_integration.py | 331 +++++++++++++++++++
 2 files changed, 332 insertions(+)
 create mode 100644 python/tvm/autotvm/task/relay_integration.py

diff --git a/python/tvm/autotvm/task/__init__.py b/python/tvm/autotvm/task/__init__.py
index 04bcec92fd57..f6ea07c272d0 100644
--- a/python/tvm/autotvm/task/__init__.py
+++ b/python/tvm/autotvm/task/__init__.py
@@ -14,3 +14,4 @@
 
 from .topi_integration import register_topi_compute, register_topi_schedule
 from .nnvm_integration import extract_from_graph, extract_from_multiple_graph
+from .relay_integration import extract_from_program, extract_from_multiple_program
diff --git a/python/tvm/autotvm/task/relay_integration.py b/python/tvm/autotvm/task/relay_integration.py
new file mode 100644
index 000000000000..277f46faae42
--- /dev/null
+++ b/python/tvm/autotvm/task/relay_integration.py
@@ -0,0 +1,331 @@
+# pylint: disable=unused-variable,invalid-name
+"""
+Decorator and utilities for the integration with TOPI and Relay
+99.9% copy-paste of implementation by @MerryMercy
+
+"""
+import warnings
+import logging
+
+
+from ... import tensor, placeholder, create_schedule, target as _target
+
+from ..util import get_const_tuple
+from .task import create, register
+
+logger = logging.getLogger('autotvm')
+
+def serialize_args(args):
+    """serialize arguments of a topi function to a hashable tuple.
+
+    Parameters
+    ----------
+    args: list of hashable or Tensor
+    """
+    ret = []
+    for t in args:
+        if isinstance(t, tensor.Tensor):
+            ret.append(('TENSOR', get_const_tuple(t.shape), t.dtype))
+        else:
+            ret.append(t)
+    return tuple(ret)
+
+
+def deserialize_args(args):
+    """The inverse function of :code:`serialize_args`.
+
+    Parameters
+    ----------
+    args: list of hashable or Tensor
+    """
+    ret = []
+    for t in args:
+        if isinstance(t, tuple) and t[0] == 'TENSOR':
+            ret.append(placeholder(shape=t[1], dtype=t[2]))
+        else:
+            ret.append(t)
+    return ret
+
+
+# Task extractor for relay program 
+class TaskExtractEnv:
+    """Global environment for extracting tuning tasks from relay program"""
+    current = None
+
+    def __init__(self):
+        import topi
+        from tvm import relay
+
+        # NOTE: To add more ops, you only need to change the following lists
+        # relay op -> topi compute
+        self.op2topi = {
+            relay.op.nn.conv2d: [topi.nn.conv2d, topi.nn.depthwise_conv2d_nchw,
+                              topi.nn.group_conv2d_nchw],
+            relay.op.nn.conv2d_transpose: [topi.nn.conv2d_transpose_nchw],
+            relay.op.nn.dense: [topi.nn.dense],
+        }
+
+        # topi compute -> autotvm task name
+        self.topi_to_task = {
+            topi.nn.conv2d: "topi_nn_conv2d",
+            topi.nn.depthwise_conv2d_nchw: "topi_nn_depthwise_conv2d_nchw",
+            topi.nn.group_conv2d_nchw: "topi_nn_group_conv2d_nchw",
+            topi.nn.conv2d_transpose_nchw: "topi_nn_conv2d_transpose_nchw",
+            topi.nn.dense: "topi_nn_dense",
+        }
+
+        self.topi_to_schedule = {
+            topi.nn.conv2d: [topi.generic.schedule_conv2d_nchw,
+                             topi.generic.schedule_conv2d_nhwc],
+            topi.nn.depthwise_conv2d_nchw: [topi.generic.schedule_depthwise_conv2d_nchw,
+                                            topi.generic.schedule_depthwise_conv2d_nhwc],
+            topi.nn.group_conv2d_nchw: [topi.generic.schedule_group_conv2d_nchw],
+            topi.nn.conv2d_transpose_nchw: [topi.generic.schedule_conv2d_transpose_nchw],
+            topi.nn.dense: [topi.generic.schedule_dense],
+        }
+
+        self._register_tracing()
+        self._register_topi_task()
+        self.task_collection = []
+        self.wanted_topi_funcs = list(self.topi_to_task.keys())
+        print("init env??")
+
+    def _register_tracing(self):
+        """Register tracing function to track the topi function call"""
+        # register topi compute for "tracing" target
+        for topi_compute in self.topi_to_task:
+            def _local_scope(compute_func):
+                """start a scope to hold the local function in for loop"""
+
+                @compute_func.register("tracing", )
+                def _tracing_topi_compute(*args, **kwargs):
+                    assert not kwargs, "Do not support extracting tuning tasks when" \
+                                       "kwargs is used in TOPI function call." \
+                                       "Please modify it to use only positional args."
+
+                    if compute_func in self.wanted_topi_funcs:  # record this call
+                        key = (self.topi_to_task[compute_func], serialize_args(args))
+                        if key not in self.task_collection:
+                            self.task_collection.append(key)
+
+                    return compute_func.fdefault(*args)
+            _local_scope(topi_compute)
+
+        # register topi schedule for "tracing" target
+        for topi_compute in self.topi_to_task:
+            for topi_schedule in self.topi_to_schedule[topi_compute]:
+                def _local_scope_(schedule_func):
+                    """start a scope to hold the local function in for loop"""
+
+                    @schedule_func.register("tracing", )
+                    def _tracing_topi_compute(outs):
+                        outs = [outs] if isinstance(outs, tensor.Tensor) else outs
+                        return create_schedule([x.op for x in outs])
+                _local_scope_(topi_schedule)
+
+    def _register_topi_task(self):
+        """register tuning wrapper for topi function"""
+        import topi
+
+        # Tuning wrapper for topi functions
+        @register("topi_nn_conv2d")
+        def _topi_nn_conv2d(*args, **kwargs):
+            assert not kwargs, "Do not support kwargs in template function call"
+            args = deserialize_args(args)
+            A, W = args[:2]
+            layout = args[-2]
+            assert layout == 'NCHW', "only support NCHW currently"
+            C = topi.nn.conv2d(*args, **kwargs)
+            s = topi.generic.schedule_conv2d_nchw([C])
+            return s, [A, W, C]
+
+        @register("topi_nn_depthwise_conv2d_nchw")
+        def _topi_nn_depthwise_conv2d_nchw(*args, **kwargs):
+            assert not kwargs, "Do not support kwargs in template function call"
+            args = deserialize_args(args)
+            A, W = args[:2]
+            C = topi.nn.depthwise_conv2d_nchw(*args, **kwargs)
+            s = topi.generic.schedule_depthwise_conv2d_nchw([C])
+            return s, [A, W, C]
+
+        @register("topi_nn_group_conv2d_nchw")
+        def _topi_nn_group_conv2d_nchw(*args, **kwargs):
+            assert not kwargs, "Do not support kwargs in template function call"
+            args = deserialize_args(args)
+            A, W = args[:2]
+            C = topi.nn.group_conv2d_nchw(*args, **kwargs)
+            s = topi.generic.schedule_group_conv2d_nchw([C])
+            return s, [A, W, C]
+
+        @register("topi_nn_conv2d_transpose_nchw")
+        def _topi_nn_conv2d_transpose_nchw(*args, **kwargs):
+            assert not kwargs, "Do not support kwargs in template function call"
+            args = deserialize_args(args)
+            A, W = args[:2]
+            C = topi.nn.conv2d_transpose_nchw(*args, **kwargs)
+            s = topi.generic.schedule_conv2d_transpose_nchw([C])
+            return s, [A, W, C]
+
+        @register("topi_nn_dense")
+        def _topi_nn_dense(*args, **kwargs):
+            assert not kwargs, "Do not support kwargs in template function call"
+            args = deserialize_args(args)
+            data, weight, bias = args
+            C = topi.nn.dense(*args, **kwargs)
+            s = topi.generic.schedule_dense([C])
+            if bias is not None:
+                return s, [data, weight, bias, C]
+            return s, [data, weight, C]
+
+    def reset(self, wanted_topi_funcs):
+        """Reset task collections
+
+        Parameters
+        ----------
+        wanted_topi_funcs: List of function
+            The topi function to be extracted
+        """
+        self.task_collection = []
+        self.wanted_topi_funcs = wanted_topi_funcs
+
+    def get_tasks(self):
+        """Get collected tasks
+
+        Returns
+        -------
+        tasks: List of tuple(name, args)
+            A list of tasks extracted from the relay program
+        """
+        return self.task_collection
+
+    @staticmethod
+    def get():
+        """Get the single instance of TaskExtractEnv
+
+        Returns
+        -------
+        env: TaskExtractEnv
+            The single instance of TaskExtractEnv
+        """
+        if not TaskExtractEnv.current:
+            TaskExtractEnv.current = TaskExtractEnv()
+        return TaskExtractEnv.current
+
+
+def extract_from_program(func, params, ops, target, target_host=None):
+    """ Extract tuning tasks from a relay program.
+
+    This function collects tuning tasks by building the graph
+    with a "tracing" target and tracing all the calls to topi.
+
+    Parameters
+    ----------
+    func: relay.expr.Function
+        The func to tune
+    params: dict of str to numpy array
+        The associated parameters of the graph
+    ops: List of relay op
+        List of relay ops to be tuned
+    dtype: str or dict of str to str
+        The input types to the graph
+    target: tvm.target.Target
+        The compilation target
+    target_host: tvm.target.Target
+        The host compilation target
+
+    Returns
+    -------
+    task: Array of autotvm.task.Task
+        collected tasks
+    """
+    from tvm import relay
+
+    env = TaskExtractEnv.get()
+
+    topi_funcs = []
+    for op_name in ops:
+        if op_name in env.op2topi:
+            topi_funcs.extend(env.op2topi[op_name])
+        else:
+            warnings.warn("Symbol %s is not tunable, ignored" % op_name)
+
+    # run compiler to collect all TOPI calls during compilation
+    env.reset(topi_funcs)
+
+    # disable logger temporarily
+    old_state = logger.disabled
+    logger.disabled = True
+
+    # use a "tracing" target to do a fake compile for collecting topi calls
+    tracing_target = _target.create("llvm -device=tracing")
+    relay.build(func, target=tracing_target, target_host=target_host, params=params)
+
+    logger.disabled = old_state
+
+    # create tasks for target
+    tasks = []
+    for task_name, args in env.get_tasks():
+        tasks.append(create(task_name, args,
+                            target=target, target_host=target_host,
+                            template_key='direct'))
+
+    return tasks
+
+
+def extract_from_multiple_program(funcs, params, ops, target, target_host=None):
+    """ Extract tuning tasks from multiple relay programs.
+
+    This function is the multiple graph version of extract_from_graph
+
+    Parameters
+    ----------
+    funcs: List of relay.expr.Function
+        The list of functions to tune
+    params: List of dict of str to numpy array
+        The input shape to the graph
+    ops: List of relay op
+        List of relay ops to be tuned
+    target: tvm.target.Target
+        The compilation target
+    target_host: tvm.target.Target
+        The host compilation target
+
+    Returns
+    -------
+    task: Array of autotvm.task.Task
+        collected tasks
+    """
+    from tvm import relay
+
+    env = TaskExtractEnv.get()
+
+    topi_funcs = []
+    for sym_name in symbols:
+        if sym_name in env.symbol2topi:
+            topi_funcs.extend(env.symbol2topi[sym_name])
+        else:
+            warnings.warn("Symbol %s is not tunable, ignored" % sym_name)
+
+    # run compiler to collect all TOPI calls during compilation
+    env.reset(topi_funcs)
+
+    # disable logger temporarily
+    old_state = logger.disabled
+    logger.disabled = True
+
+    # use a "tracing" target to do a fake compile for collecting topi calls
+    tracing_target = _target.create("llvm -device=tracing")
+
+    for func, param in zip(funcs, params):
+        relay.build(func, target=tracing_target, target_host=target_host, params=param)
+
+    logger.disabled = old_state
+
+    # create tasks for target
+    tasks = []
+    for task_name, args in env.get_tasks():
+        tasks.append(create(task_name, args,
+                            target=target, target_host=target_host,
+                            template_key='direct'))
+
+    return tasks

From 53b877ffa0e4d7e29f757848725d27971f9bc766 Mon Sep 17 00:00:00 2001
From: Eddie Yan <eqy@cs.washington.edu>
Date: Tue, 27 Nov 2018 15:48:23 -0800
Subject: [PATCH 02/14] cleanup

---
 python/tvm/autotvm/task/relay_integration.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/python/tvm/autotvm/task/relay_integration.py b/python/tvm/autotvm/task/relay_integration.py
index 277f46faae42..706a525602c0 100644
--- a/python/tvm/autotvm/task/relay_integration.py
+++ b/python/tvm/autotvm/task/relay_integration.py
@@ -247,7 +247,7 @@ def extract_from_program(func, params, ops, target, target_host=None):
         if op_name in env.op2topi:
             topi_funcs.extend(env.op2topi[op_name])
         else:
-            warnings.warn("Symbol %s is not tunable, ignored" % op_name)
+            warnings.warn("Op %s is not tunable, ignored" % op_name)
 
     # run compiler to collect all TOPI calls during compilation
     env.reset(topi_funcs)
@@ -300,11 +300,11 @@ def extract_from_multiple_program(funcs, params, ops, target, target_host=None):
     env = TaskExtractEnv.get()
 
     topi_funcs = []
-    for sym_name in symbols:
-        if sym_name in env.symbol2topi:
-            topi_funcs.extend(env.symbol2topi[sym_name])
+    for op_name in ops:
+        if op_name in env.op2topi:
+            topi_funcs.extend(env.op2topi[op_name])
         else:
-            warnings.warn("Symbol %s is not tunable, ignored" % sym_name)
+            warnings.warn("Op %s is not tunable, ignored" % op_name)
 
     # run compiler to collect all TOPI calls during compilation
     env.reset(topi_funcs)

From ee79db790665dcc77f08d1921bd9ee99797d454b Mon Sep 17 00:00:00 2001
From: Eddie Yan <eqy@cs.washington.edu>
Date: Tue, 27 Nov 2018 15:50:09 -0800
Subject: [PATCH 03/14] fix lint

---
 python/tvm/autotvm/task/relay_integration.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/python/tvm/autotvm/task/relay_integration.py b/python/tvm/autotvm/task/relay_integration.py
index 706a525602c0..3b3cdea580f4 100644
--- a/python/tvm/autotvm/task/relay_integration.py
+++ b/python/tvm/autotvm/task/relay_integration.py
@@ -47,7 +47,7 @@ def deserialize_args(args):
     return ret
 
 
-# Task extractor for relay program 
+# Task extractor for relay program
 class TaskExtractEnv:
     """Global environment for extracting tuning tasks from relay program"""
     current = None
@@ -60,7 +60,7 @@ def __init__(self):
         # relay op -> topi compute
         self.op2topi = {
             relay.op.nn.conv2d: [topi.nn.conv2d, topi.nn.depthwise_conv2d_nchw,
-                              topi.nn.group_conv2d_nchw],
+                                 topi.nn.group_conv2d_nchw],
             relay.op.nn.conv2d_transpose: [topi.nn.conv2d_transpose_nchw],
             relay.op.nn.dense: [topi.nn.dense],
         }

From 28e9341e03c6fb75bbf0add1200860e00c219b97 Mon Sep 17 00:00:00 2001
From: Eddie Yan <eqy@cs.washington.edu>
Date: Tue, 27 Nov 2018 17:23:05 -0800
Subject: [PATCH 04/14] cleanup

---
 python/tvm/autotvm/task/relay_integration.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/python/tvm/autotvm/task/relay_integration.py b/python/tvm/autotvm/task/relay_integration.py
index 3b3cdea580f4..b38ed1ef076a 100644
--- a/python/tvm/autotvm/task/relay_integration.py
+++ b/python/tvm/autotvm/task/relay_integration.py
@@ -88,7 +88,6 @@ def __init__(self):
         self._register_topi_task()
         self.task_collection = []
         self.wanted_topi_funcs = list(self.topi_to_task.keys())
-        print("init env??")
 
     def _register_tracing(self):
         """Register tracing function to track the topi function call"""

From 0ba31420a5e0eafe1c22348e1092ba410fc6c7f7 Mon Sep 17 00:00:00 2001
From: Eddie Yan <eqy@cs.washington.edu>
Date: Wed, 28 Nov 2018 13:06:22 -0800
Subject: [PATCH 05/14] port nnvm tests

---
 python/tvm/autotvm/task/relay_integration.py  |  1 +
 .../relay/test_autotvm_task_extraction.py     | 56 +++++++++++++++++++
 2 files changed, 57 insertions(+)
 create mode 100644 tests/python/relay/test_autotvm_task_extraction.py

diff --git a/python/tvm/autotvm/task/relay_integration.py b/python/tvm/autotvm/task/relay_integration.py
index b38ed1ef076a..8125163e70fd 100644
--- a/python/tvm/autotvm/task/relay_integration.py
+++ b/python/tvm/autotvm/task/relay_integration.py
@@ -257,6 +257,7 @@ def extract_from_program(func, params, ops, target, target_host=None):
 
     # use a "tracing" target to do a fake compile for collecting topi calls
     tracing_target = _target.create("llvm -device=tracing")
+    relay.backend.compile_engine.get().clear()
     relay.build(func, target=tracing_target, target_host=target_host, params=params)
 
     logger.disabled = old_state
diff --git a/tests/python/relay/test_autotvm_task_extraction.py b/tests/python/relay/test_autotvm_task_extraction.py
new file mode 100644
index 000000000000..8c93e4a56642
--- /dev/null
+++ b/tests/python/relay/test_autotvm_task_extraction.py
@@ -0,0 +1,56 @@
+"""Test task extraction for autotvm"""
+import tvm.relay.testing
+from tvm import relay
+from tvm import autotvm
+
+def get_network(name, batch_size):
+    """Get the symbol definition and random weight of a network"""
+    input_shape = (batch_size, 3, 224, 224)
+
+    if name == 'resnet-18':
+        net, params = relay.testing.resnet.get_workload(num_layers=18, batch_size=batch_size)
+    elif name == 'mobilenet':
+        net, params = relay.testing.mobilenet.get_workload(batch_size=batch_size)
+    elif name == 'dcgan':
+        net, params = relay.testing.dcgan.get_workload(batch_size=batch_size)
+        input_shape = (batch_size, 100)
+    else:
+        raise ValueError("Unsupported network: " + name)
+
+    return net, params, input_shape
+
+def test_task_extraction():
+    target = 'llvm'
+
+    net, params, input_shape = get_network('resnet-18', batch_size=1)
+    tasks = autotvm.task.extract_from_program(net, target=target,
+                                            params=params,
+                                            ops=(relay.op.nn.conv2d,))
+    assert len(tasks) == 12
+
+    net, params, input_shape = get_network('resnet-18', batch_size=1)
+    tasks = autotvm.task.extract_from_program(net, target=target,
+                                            params=params,
+                                            ops=(relay.op.nn.dense,))
+    assert len(tasks) == 1
+
+    net, params, input_shape = get_network('resnet-18', batch_size=1)
+    tasks = autotvm.task.extract_from_program(net, target=target,
+                                            params=params,
+                                            ops=(relay.op.nn.conv2d, relay.op.nn.dense))
+    assert len(tasks) == 13
+
+    net, params, input_shape = get_network('mobilenet', batch_size=1)
+    tasks = autotvm.task.extract_from_program(net, target=target,
+                                            params=params,
+                                            ops=(relay.op.nn.conv2d, relay.op.nn.dense))
+    assert len(tasks) == 20
+
+    net, params, input_shape = get_network('dcgan', batch_size=1)
+    tasks = autotvm.task.extract_from_program(net, target=target,
+                                            params=params,
+                                            ops=(relay.op.nn.conv2d_transpose,))
+    assert len(tasks) == 4
+
+if __name__ == '__main__':
+    test_task_extraction()

From 5c162af2c1f78a6424650eb739871b66dd142e62 Mon Sep 17 00:00:00 2001
From: Eddie Yan <eqy@cs.washington.edu>
Date: Wed, 28 Nov 2018 14:07:05 -0800
Subject: [PATCH 06/14] wrap build call in thread

---
 python/tvm/autotvm/task/relay_integration.py | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/python/tvm/autotvm/task/relay_integration.py b/python/tvm/autotvm/task/relay_integration.py
index 8125163e70fd..1450b65f96cb 100644
--- a/python/tvm/autotvm/task/relay_integration.py
+++ b/python/tvm/autotvm/task/relay_integration.py
@@ -4,6 +4,7 @@
 99.9% copy-paste of implementation by @MerryMercy
 
 """
+import threading
 import warnings
 import logging
 
@@ -258,8 +259,13 @@ def extract_from_program(func, params, ops, target, target_host=None):
     # use a "tracing" target to do a fake compile for collecting topi calls
     tracing_target = _target.create("llvm -device=tracing")
     relay.backend.compile_engine.get().clear()
-    relay.build(func, target=tracing_target, target_host=target_host, params=params)
-
+    # wrap build call in thread to avoid multiprocessing problems 
+    build_thread = threading.Thread(target=relay.build, args=(func,
+                                                              tracing_target,
+                                                              target_host,
+                                                              params))
+    build_thread.start()
+    build_thread.join()
     logger.disabled = old_state
 
     # create tasks for target
@@ -317,7 +323,13 @@ def extract_from_multiple_program(funcs, params, ops, target, target_host=None):
     tracing_target = _target.create("llvm -device=tracing")
 
     for func, param in zip(funcs, params):
-        relay.build(func, target=tracing_target, target_host=target_host, params=param)
+        # wrap build call in thread to avoid multiprocessing problems
+        build_thread = threading.Thread(target=relay.build, args=(func,
+                                                                  tracing_target,
+                                                                  target_host,
+                                                                  params))
+        build_thread.start()
+        build_thread.join()
 
     logger.disabled = old_state
 

From bba2b5145d3fc33afe06e8ae19254ca6a44677f0 Mon Sep 17 00:00:00 2001
From: Eddie Yan <eqy@cs.washington.edu>
Date: Wed, 28 Nov 2018 14:15:27 -0800
Subject: [PATCH 07/14] merge TaskExtractEnv

---
 python/tvm/autotvm/task/nnvm_integration.py  |  10 ++
 python/tvm/autotvm/task/relay_integration.py | 165 +------------------
 2 files changed, 11 insertions(+), 164 deletions(-)

diff --git a/python/tvm/autotvm/task/nnvm_integration.py b/python/tvm/autotvm/task/nnvm_integration.py
index 6a07194a594d..014275fd0969 100644
--- a/python/tvm/autotvm/task/nnvm_integration.py
+++ b/python/tvm/autotvm/task/nnvm_integration.py
@@ -54,6 +54,7 @@ class TaskExtractEnv:
     def __init__(self):
         import topi
         import nnvm
+        from tvm import relay
 
         # NOTE: To add more symbols, you only need to change the following lists
         # nnvm symbol -> topi compute
@@ -64,6 +65,15 @@ def __init__(self):
             nnvm.sym.dense: [topi.nn.dense],
         }
 
+        # NOTE: To add more ops, you only need to change the following lists
+        # relay op -> topi compute
+        self.op2topi = {
+            relay.op.nn.conv2d: [topi.nn.conv2d, topi.nn.depthwise_conv2d_nchw,
+                                 topi.nn.group_conv2d_nchw],
+            relay.op.nn.conv2d_transpose: [topi.nn.conv2d_transpose_nchw],
+            relay.op.nn.dense: [topi.nn.dense],
+        }
+
         # topi compute -> autotvm task name
         self.topi_to_task = {
             topi.nn.conv2d: "topi_nn_conv2d",
diff --git a/python/tvm/autotvm/task/relay_integration.py b/python/tvm/autotvm/task/relay_integration.py
index 1450b65f96cb..c50c83ee6f3b 100644
--- a/python/tvm/autotvm/task/relay_integration.py
+++ b/python/tvm/autotvm/task/relay_integration.py
@@ -13,6 +13,7 @@
 
 from ..util import get_const_tuple
 from .task import create, register
+from .nnvm_integration import TaskExtractEnv
 
 logger = logging.getLogger('autotvm')
 
@@ -48,170 +49,6 @@ def deserialize_args(args):
     return ret
 
 
-# Task extractor for relay program
-class TaskExtractEnv:
-    """Global environment for extracting tuning tasks from relay program"""
-    current = None
-
-    def __init__(self):
-        import topi
-        from tvm import relay
-
-        # NOTE: To add more ops, you only need to change the following lists
-        # relay op -> topi compute
-        self.op2topi = {
-            relay.op.nn.conv2d: [topi.nn.conv2d, topi.nn.depthwise_conv2d_nchw,
-                                 topi.nn.group_conv2d_nchw],
-            relay.op.nn.conv2d_transpose: [topi.nn.conv2d_transpose_nchw],
-            relay.op.nn.dense: [topi.nn.dense],
-        }
-
-        # topi compute -> autotvm task name
-        self.topi_to_task = {
-            topi.nn.conv2d: "topi_nn_conv2d",
-            topi.nn.depthwise_conv2d_nchw: "topi_nn_depthwise_conv2d_nchw",
-            topi.nn.group_conv2d_nchw: "topi_nn_group_conv2d_nchw",
-            topi.nn.conv2d_transpose_nchw: "topi_nn_conv2d_transpose_nchw",
-            topi.nn.dense: "topi_nn_dense",
-        }
-
-        self.topi_to_schedule = {
-            topi.nn.conv2d: [topi.generic.schedule_conv2d_nchw,
-                             topi.generic.schedule_conv2d_nhwc],
-            topi.nn.depthwise_conv2d_nchw: [topi.generic.schedule_depthwise_conv2d_nchw,
-                                            topi.generic.schedule_depthwise_conv2d_nhwc],
-            topi.nn.group_conv2d_nchw: [topi.generic.schedule_group_conv2d_nchw],
-            topi.nn.conv2d_transpose_nchw: [topi.generic.schedule_conv2d_transpose_nchw],
-            topi.nn.dense: [topi.generic.schedule_dense],
-        }
-
-        self._register_tracing()
-        self._register_topi_task()
-        self.task_collection = []
-        self.wanted_topi_funcs = list(self.topi_to_task.keys())
-
-    def _register_tracing(self):
-        """Register tracing function to track the topi function call"""
-        # register topi compute for "tracing" target
-        for topi_compute in self.topi_to_task:
-            def _local_scope(compute_func):
-                """start a scope to hold the local function in for loop"""
-
-                @compute_func.register("tracing", )
-                def _tracing_topi_compute(*args, **kwargs):
-                    assert not kwargs, "Do not support extracting tuning tasks when" \
-                                       "kwargs is used in TOPI function call." \
-                                       "Please modify it to use only positional args."
-
-                    if compute_func in self.wanted_topi_funcs:  # record this call
-                        key = (self.topi_to_task[compute_func], serialize_args(args))
-                        if key not in self.task_collection:
-                            self.task_collection.append(key)
-
-                    return compute_func.fdefault(*args)
-            _local_scope(topi_compute)
-
-        # register topi schedule for "tracing" target
-        for topi_compute in self.topi_to_task:
-            for topi_schedule in self.topi_to_schedule[topi_compute]:
-                def _local_scope_(schedule_func):
-                    """start a scope to hold the local function in for loop"""
-
-                    @schedule_func.register("tracing", )
-                    def _tracing_topi_compute(outs):
-                        outs = [outs] if isinstance(outs, tensor.Tensor) else outs
-                        return create_schedule([x.op for x in outs])
-                _local_scope_(topi_schedule)
-
-    def _register_topi_task(self):
-        """register tuning wrapper for topi function"""
-        import topi
-
-        # Tuning wrapper for topi functions
-        @register("topi_nn_conv2d")
-        def _topi_nn_conv2d(*args, **kwargs):
-            assert not kwargs, "Do not support kwargs in template function call"
-            args = deserialize_args(args)
-            A, W = args[:2]
-            layout = args[-2]
-            assert layout == 'NCHW', "only support NCHW currently"
-            C = topi.nn.conv2d(*args, **kwargs)
-            s = topi.generic.schedule_conv2d_nchw([C])
-            return s, [A, W, C]
-
-        @register("topi_nn_depthwise_conv2d_nchw")
-        def _topi_nn_depthwise_conv2d_nchw(*args, **kwargs):
-            assert not kwargs, "Do not support kwargs in template function call"
-            args = deserialize_args(args)
-            A, W = args[:2]
-            C = topi.nn.depthwise_conv2d_nchw(*args, **kwargs)
-            s = topi.generic.schedule_depthwise_conv2d_nchw([C])
-            return s, [A, W, C]
-
-        @register("topi_nn_group_conv2d_nchw")
-        def _topi_nn_group_conv2d_nchw(*args, **kwargs):
-            assert not kwargs, "Do not support kwargs in template function call"
-            args = deserialize_args(args)
-            A, W = args[:2]
-            C = topi.nn.group_conv2d_nchw(*args, **kwargs)
-            s = topi.generic.schedule_group_conv2d_nchw([C])
-            return s, [A, W, C]
-
-        @register("topi_nn_conv2d_transpose_nchw")
-        def _topi_nn_conv2d_transpose_nchw(*args, **kwargs):
-            assert not kwargs, "Do not support kwargs in template function call"
-            args = deserialize_args(args)
-            A, W = args[:2]
-            C = topi.nn.conv2d_transpose_nchw(*args, **kwargs)
-            s = topi.generic.schedule_conv2d_transpose_nchw([C])
-            return s, [A, W, C]
-
-        @register("topi_nn_dense")
-        def _topi_nn_dense(*args, **kwargs):
-            assert not kwargs, "Do not support kwargs in template function call"
-            args = deserialize_args(args)
-            data, weight, bias = args
-            C = topi.nn.dense(*args, **kwargs)
-            s = topi.generic.schedule_dense([C])
-            if bias is not None:
-                return s, [data, weight, bias, C]
-            return s, [data, weight, C]
-
-    def reset(self, wanted_topi_funcs):
-        """Reset task collections
-
-        Parameters
-        ----------
-        wanted_topi_funcs: List of function
-            The topi function to be extracted
-        """
-        self.task_collection = []
-        self.wanted_topi_funcs = wanted_topi_funcs
-
-    def get_tasks(self):
-        """Get collected tasks
-
-        Returns
-        -------
-        tasks: List of tuple(name, args)
-            A list of tasks extracted from the relay program
-        """
-        return self.task_collection
-
-    @staticmethod
-    def get():
-        """Get the single instance of TaskExtractEnv
-
-        Returns
-        -------
-        env: TaskExtractEnv
-            The single instance of TaskExtractEnv
-        """
-        if not TaskExtractEnv.current:
-            TaskExtractEnv.current = TaskExtractEnv()
-        return TaskExtractEnv.current
-
-
 def extract_from_program(func, params, ops, target, target_host=None):
     """ Extract tuning tasks from a relay program.
 

From a3774e36fee3a2244249a5bbfe28a8be70484226 Mon Sep 17 00:00:00 2001
From: Eddie Yan <eqy@cs.washington.edu>
Date: Wed, 28 Nov 2018 14:20:04 -0800
Subject: [PATCH 08/14] fix lint

---
 python/tvm/autotvm/task/relay_integration.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/python/tvm/autotvm/task/relay_integration.py b/python/tvm/autotvm/task/relay_integration.py
index c50c83ee6f3b..8c54c5cc698c 100644
--- a/python/tvm/autotvm/task/relay_integration.py
+++ b/python/tvm/autotvm/task/relay_integration.py
@@ -9,10 +9,10 @@
 import logging
 
 
-from ... import tensor, placeholder, create_schedule, target as _target
+from ... import tensor, placeholder, target as _target
 
 from ..util import get_const_tuple
-from .task import create, register
+from .task import create
 from .nnvm_integration import TaskExtractEnv
 
 logger = logging.getLogger('autotvm')
@@ -96,7 +96,7 @@ def extract_from_program(func, params, ops, target, target_host=None):
     # use a "tracing" target to do a fake compile for collecting topi calls
     tracing_target = _target.create("llvm -device=tracing")
     relay.backend.compile_engine.get().clear()
-    # wrap build call in thread to avoid multiprocessing problems 
+    # wrap build call in thread to avoid multiprocessing problems
     build_thread = threading.Thread(target=relay.build, args=(func,
                                                               tracing_target,
                                                               target_host,

From 022960f39a44419f3f4162925a573ac0f7cd3c51 Mon Sep 17 00:00:00 2001
From: Eddie Yan <eqy@cs.washington.edu>
Date: Thu, 29 Nov 2018 17:46:28 -0800
Subject: [PATCH 09/14] try moving test to nnvm directory

---
 .../tests/python/compiler/test_autotvm_relay_task_extraction.py   | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename tests/python/relay/test_autotvm_task_extraction.py => nnvm/tests/python/compiler/test_autotvm_relay_task_extraction.py (100%)

diff --git a/tests/python/relay/test_autotvm_task_extraction.py b/nnvm/tests/python/compiler/test_autotvm_relay_task_extraction.py
similarity index 100%
rename from tests/python/relay/test_autotvm_task_extraction.py
rename to nnvm/tests/python/compiler/test_autotvm_relay_task_extraction.py

From 08590f8ec1d3934cbc38eaefc24aac01f2b6a1b1 Mon Sep 17 00:00:00 2001
From: Eddie Yan <eqy@cs.washington.edu>
Date: Fri, 30 Nov 2018 18:33:43 -0800
Subject: [PATCH 10/14] refactor, fix imports

---
 .../test_autotvm_relay_task_extraction.py     |  56 ----
 python/tvm/autotvm/task/nnvm_integration.py   | 241 +++---------------
 python/tvm/autotvm/task/relay_integration.py  |  36 ++-
 python/tvm/autotvm/task/topi_integration.py   | 192 +++++++++++++-
 4 files changed, 246 insertions(+), 279 deletions(-)
 delete mode 100644 nnvm/tests/python/compiler/test_autotvm_relay_task_extraction.py

diff --git a/nnvm/tests/python/compiler/test_autotvm_relay_task_extraction.py b/nnvm/tests/python/compiler/test_autotvm_relay_task_extraction.py
deleted file mode 100644
index 8c93e4a56642..000000000000
--- a/nnvm/tests/python/compiler/test_autotvm_relay_task_extraction.py
+++ /dev/null
@@ -1,56 +0,0 @@
-"""Test task extraction for autotvm"""
-import tvm.relay.testing
-from tvm import relay
-from tvm import autotvm
-
-def get_network(name, batch_size):
-    """Get the symbol definition and random weight of a network"""
-    input_shape = (batch_size, 3, 224, 224)
-
-    if name == 'resnet-18':
-        net, params = relay.testing.resnet.get_workload(num_layers=18, batch_size=batch_size)
-    elif name == 'mobilenet':
-        net, params = relay.testing.mobilenet.get_workload(batch_size=batch_size)
-    elif name == 'dcgan':
-        net, params = relay.testing.dcgan.get_workload(batch_size=batch_size)
-        input_shape = (batch_size, 100)
-    else:
-        raise ValueError("Unsupported network: " + name)
-
-    return net, params, input_shape
-
-def test_task_extraction():
-    target = 'llvm'
-
-    net, params, input_shape = get_network('resnet-18', batch_size=1)
-    tasks = autotvm.task.extract_from_program(net, target=target,
-                                            params=params,
-                                            ops=(relay.op.nn.conv2d,))
-    assert len(tasks) == 12
-
-    net, params, input_shape = get_network('resnet-18', batch_size=1)
-    tasks = autotvm.task.extract_from_program(net, target=target,
-                                            params=params,
-                                            ops=(relay.op.nn.dense,))
-    assert len(tasks) == 1
-
-    net, params, input_shape = get_network('resnet-18', batch_size=1)
-    tasks = autotvm.task.extract_from_program(net, target=target,
-                                            params=params,
-                                            ops=(relay.op.nn.conv2d, relay.op.nn.dense))
-    assert len(tasks) == 13
-
-    net, params, input_shape = get_network('mobilenet', batch_size=1)
-    tasks = autotvm.task.extract_from_program(net, target=target,
-                                            params=params,
-                                            ops=(relay.op.nn.conv2d, relay.op.nn.dense))
-    assert len(tasks) == 20
-
-    net, params, input_shape = get_network('dcgan', batch_size=1)
-    tasks = autotvm.task.extract_from_program(net, target=target,
-                                            params=params,
-                                            ops=(relay.op.nn.conv2d_transpose,))
-    assert len(tasks) == 4
-
-if __name__ == '__main__':
-    test_task_extraction()
diff --git a/python/tvm/autotvm/task/nnvm_integration.py b/python/tvm/autotvm/task/nnvm_integration.py
index 014275fd0969..770ac90520d8 100644
--- a/python/tvm/autotvm/task/nnvm_integration.py
+++ b/python/tvm/autotvm/task/nnvm_integration.py
@@ -7,218 +7,13 @@
 import logging
 
 
-from ... import tensor, placeholder, create_schedule, target as _target
+from ... import tensor, placeholder, target as _target
 
-from ..util import get_const_tuple
-from .task import create, register
+from .task import create
+from .topi_integration import TaskExtractEnv
 
 logger = logging.getLogger('autotvm')
 
-def serialize_args(args):
-    """serialize arguments of a topi function to a hashable tuple.
-
-    Parameters
-    ----------
-    args: list of hashable or Tensor
-    """
-    ret = []
-    for t in args:
-        if isinstance(t, tensor.Tensor):
-            ret.append(('TENSOR', get_const_tuple(t.shape), t.dtype))
-        else:
-            ret.append(t)
-    return tuple(ret)
-
-
-def deserialize_args(args):
-    """The inverse function of :code:`serialize_args`.
-
-    Parameters
-    ----------
-    args: list of hashable or Tensor
-    """
-    ret = []
-    for t in args:
-        if isinstance(t, tuple) and t[0] == 'TENSOR':
-            ret.append(placeholder(shape=t[1], dtype=t[2]))
-        else:
-            ret.append(t)
-    return ret
-
-
-# Task extractor for nnvm graph
-class TaskExtractEnv:
-    """Global environment for extracting tuning tasks from nnvm graph"""
-    current = None
-
-    def __init__(self):
-        import topi
-        import nnvm
-        from tvm import relay
-
-        # NOTE: To add more symbols, you only need to change the following lists
-        # nnvm symbol -> topi compute
-        self.symbol2topi = {
-            nnvm.sym.conv2d: [topi.nn.conv2d, topi.nn.depthwise_conv2d_nchw,
-                              topi.nn.group_conv2d_nchw],
-            nnvm.sym.conv2d_transpose: [topi.nn.conv2d_transpose_nchw],
-            nnvm.sym.dense: [topi.nn.dense],
-        }
-
-        # NOTE: To add more ops, you only need to change the following lists
-        # relay op -> topi compute
-        self.op2topi = {
-            relay.op.nn.conv2d: [topi.nn.conv2d, topi.nn.depthwise_conv2d_nchw,
-                                 topi.nn.group_conv2d_nchw],
-            relay.op.nn.conv2d_transpose: [topi.nn.conv2d_transpose_nchw],
-            relay.op.nn.dense: [topi.nn.dense],
-        }
-
-        # topi compute -> autotvm task name
-        self.topi_to_task = {
-            topi.nn.conv2d: "topi_nn_conv2d",
-            topi.nn.depthwise_conv2d_nchw: "topi_nn_depthwise_conv2d_nchw",
-            topi.nn.group_conv2d_nchw: "topi_nn_group_conv2d_nchw",
-            topi.nn.conv2d_transpose_nchw: "topi_nn_conv2d_transpose_nchw",
-            topi.nn.dense: "topi_nn_dense",
-        }
-
-        self.topi_to_schedule = {
-            topi.nn.conv2d: [topi.generic.schedule_conv2d_nchw,
-                             topi.generic.schedule_conv2d_nhwc],
-            topi.nn.depthwise_conv2d_nchw: [topi.generic.schedule_depthwise_conv2d_nchw,
-                                            topi.generic.schedule_depthwise_conv2d_nhwc],
-            topi.nn.group_conv2d_nchw: [topi.generic.schedule_group_conv2d_nchw],
-            topi.nn.conv2d_transpose_nchw: [topi.generic.schedule_conv2d_transpose_nchw],
-            topi.nn.dense: [topi.generic.schedule_dense],
-        }
-
-        self._register_tracing()
-        self._register_topi_task()
-        self.task_collection = []
-        self.wanted_topi_funcs = list(self.topi_to_task.keys())
-
-    def _register_tracing(self):
-        """Register tracing function to track the topi function call"""
-        # register topi compute for "tracing" target
-        for topi_compute in self.topi_to_task:
-            def _local_scope(compute_func):
-                """start a scope to hold the local function in for loop"""
-
-                @compute_func.register("tracing", )
-                def _tracing_topi_compute(*args, **kwargs):
-                    assert not kwargs, "Do not support extracting tuning tasks when" \
-                                       "kwargs is used in TOPI function call." \
-                                       "Please modify it to use only positional args."
-
-                    if compute_func in self.wanted_topi_funcs:  # record this call
-                        key = (self.topi_to_task[compute_func], serialize_args(args))
-                        if key not in self.task_collection:
-                            self.task_collection.append(key)
-
-                    return compute_func.fdefault(*args)
-            _local_scope(topi_compute)
-
-        # register topi schedule for "tracing" target
-        for topi_compute in self.topi_to_task:
-            for topi_schedule in self.topi_to_schedule[topi_compute]:
-                def _local_scope_(schedule_func):
-                    """start a scope to hold the local function in for loop"""
-
-                    @schedule_func.register("tracing", )
-                    def _tracing_topi_compute(outs):
-                        outs = [outs] if isinstance(outs, tensor.Tensor) else outs
-                        return create_schedule([x.op for x in outs])
-                _local_scope_(topi_schedule)
-
-    def _register_topi_task(self):
-        """register tuning wrapper for topi function"""
-        import topi
-
-        # Tuning wrapper for topi functions
-        @register("topi_nn_conv2d")
-        def _topi_nn_conv2d(*args, **kwargs):
-            assert not kwargs, "Do not support kwargs in template function call"
-            args = deserialize_args(args)
-            A, W = args[:2]
-            layout = args[-2]
-            assert layout == 'NCHW', "only support NCHW currently"
-            C = topi.nn.conv2d(*args, **kwargs)
-            s = topi.generic.schedule_conv2d_nchw([C])
-            return s, [A, W, C]
-
-        @register("topi_nn_depthwise_conv2d_nchw")
-        def _topi_nn_depthwise_conv2d_nchw(*args, **kwargs):
-            assert not kwargs, "Do not support kwargs in template function call"
-            args = deserialize_args(args)
-            A, W = args[:2]
-            C = topi.nn.depthwise_conv2d_nchw(*args, **kwargs)
-            s = topi.generic.schedule_depthwise_conv2d_nchw([C])
-            return s, [A, W, C]
-
-        @register("topi_nn_group_conv2d_nchw")
-        def _topi_nn_group_conv2d_nchw(*args, **kwargs):
-            assert not kwargs, "Do not support kwargs in template function call"
-            args = deserialize_args(args)
-            A, W = args[:2]
-            C = topi.nn.group_conv2d_nchw(*args, **kwargs)
-            s = topi.generic.schedule_group_conv2d_nchw([C])
-            return s, [A, W, C]
-
-        @register("topi_nn_conv2d_transpose_nchw")
-        def _topi_nn_conv2d_transpose_nchw(*args, **kwargs):
-            assert not kwargs, "Do not support kwargs in template function call"
-            args = deserialize_args(args)
-            A, W = args[:2]
-            C = topi.nn.conv2d_transpose_nchw(*args, **kwargs)
-            s = topi.generic.schedule_conv2d_transpose_nchw([C])
-            return s, [A, W, C]
-
-        @register("topi_nn_dense")
-        def _topi_nn_dense(*args, **kwargs):
-            assert not kwargs, "Do not support kwargs in template function call"
-            args = deserialize_args(args)
-            data, weight, bias = args
-            C = topi.nn.dense(*args, **kwargs)
-            s = topi.generic.schedule_dense([C])
-            if bias is not None:
-                return s, [data, weight, bias, C]
-            return s, [data, weight, C]
-
-    def reset(self, wanted_topi_funcs):
-        """Reset task collections
-
-        Parameters
-        ----------
-        wanted_topi_funcs: List of function
-            The topi function to be extracted
-        """
-        self.task_collection = []
-        self.wanted_topi_funcs = wanted_topi_funcs
-
-    def get_tasks(self):
-        """Get collected tasks
-
-        Returns
-        -------
-        tasks: List of tuple(name, args)
-            A list of tasks extracted from the nnvm graph
-        """
-        return self.task_collection
-
-    @staticmethod
-    def get():
-        """Get the single instance of TaskExtractEnv
-
-        Returns
-        -------
-        env: TaskExtractEnv
-            The single instance of TaskExtractEnv
-        """
-        if not TaskExtractEnv.current:
-            TaskExtractEnv.current = TaskExtractEnv()
-        return TaskExtractEnv.current
-
 
 def extract_from_graph(graph, shape, dtype, target, symbols, target_host=None):
     """ Extract tuning tasks from a nnvm graph.
@@ -247,13 +42,24 @@ def extract_from_graph(graph, shape, dtype, target, symbols, target_host=None):
         collected tasks
     """
     import nnvm.compiler
+    import nnvm
+    import topi
 
     env = TaskExtractEnv.get()
 
+    #NOTE: To add more symbols, you only need to change the following lists
+    #nnvm symbol -> topi compute
+    SYMBOL2TOPI = {
+        nnvm.sym.conv2d: [topi.nn.conv2d, topi.nn.depthwise_conv2d_nchw,
+                          topi.nn.group_conv2d_nchw],
+        nnvm.sym.conv2d_transpose: [topi.nn.conv2d_transpose_nchw],
+        nnvm.sym.dense: [topi.nn.dense],
+    }
+
     topi_funcs = []
     for sym_name in symbols:
-        if sym_name in env.symbol2topi:
-            topi_funcs.extend(env.symbol2topi[sym_name])
+        if sym_name in SYMBOL2TOPI:
+            topi_funcs.extend(SYMBOL2TOPI[sym_name])
         else:
             warnings.warn("Symbol %s is not tunable, ignored" % sym_name)
 
@@ -307,13 +113,24 @@ def extract_from_multiple_graph(graphs, shapes, dtypes, target, symbols, target_
         collected tasks
     """
     import nnvm.compiler
+    import nnvm
+    import topi
 
     env = TaskExtractEnv.get()
 
+    #NOTE: To add more symbols, you only need to change the following lists
+    #nnvm symbol -> topi compute
+    SYMBOL2TOPI = {
+        nnvm.sym.conv2d: [topi.nn.conv2d, topi.nn.depthwise_conv2d_nchw,
+                          topi.nn.group_conv2d_nchw],
+        nnvm.sym.conv2d_transpose: [topi.nn.conv2d_transpose_nchw],
+        nnvm.sym.dense: [topi.nn.dense],
+    }
+
     topi_funcs = []
     for sym_name in symbols:
-        if sym_name in env.symbol2topi:
-            topi_funcs.extend(env.symbol2topi[sym_name])
+        if sym_name in SYMBOL2TOPI:
+            topi_funcs.extend(SYMBOL2TOPI[sym_name])
         else:
             warnings.warn("Symbol %s is not tunable, ignored" % sym_name)
 
diff --git a/python/tvm/autotvm/task/relay_integration.py b/python/tvm/autotvm/task/relay_integration.py
index 8c54c5cc698c..82319ccf6aa8 100644
--- a/python/tvm/autotvm/task/relay_integration.py
+++ b/python/tvm/autotvm/task/relay_integration.py
@@ -11,12 +11,12 @@
 
 from ... import tensor, placeholder, target as _target
 
-from ..util import get_const_tuple
 from .task import create
-from .nnvm_integration import TaskExtractEnv
+from .topi_integration import TaskExtractEnv
 
 logger = logging.getLogger('autotvm')
 
+
 def serialize_args(args):
     """serialize arguments of a topi function to a hashable tuple.
 
@@ -75,14 +75,24 @@ def extract_from_program(func, params, ops, target, target_host=None):
     task: Array of autotvm.task.Task
         collected tasks
     """
+    env = TaskExtractEnv.get()
+    import tvm.relay.op
     from tvm import relay
+    import topi
 
-    env = TaskExtractEnv.get()
+    # NOTE: To add more ops, you only need to change the following lists
+    # relay op -> topi compute
+    OP2TOPI = {
+        tvm.relay.op.nn.conv2d: [topi.nn.conv2d, topi.nn.depthwise_conv2d_nchw,
+                             topi.nn.group_conv2d_nchw],
+        tvm.relay.op.nn.conv2d_transpose: [topi.nn.conv2d_transpose_nchw],
+        tvm.relay.op.nn.dense: [topi.nn.dense],
+    }
 
     topi_funcs = []
     for op_name in ops:
-        if op_name in env.op2topi:
-            topi_funcs.extend(env.op2topi[op_name])
+        if op_name in OP2TOPI:
+            topi_funcs.extend(OP2TOPI[op_name])
         else:
             warnings.warn("Op %s is not tunable, ignored" % op_name)
 
@@ -138,14 +148,24 @@ def extract_from_multiple_program(funcs, params, ops, target, target_host=None):
     task: Array of autotvm.task.Task
         collected tasks
     """
+    env = TaskExtractEnv.get()
+    import tvm.relay.op
     from tvm import relay
+    import topi
 
-    env = TaskExtractEnv.get()
+    # NOTE: To add more ops, you only need to change the following lists
+    # relay op -> topi compute
+    OP2TOPI = {
+        tvm.relay.op.nn.conv2d: [topi.nn.conv2d, topi.nn.depthwise_conv2d_nchw,
+                             topi.nn.group_conv2d_nchw],
+        tvm.relay.op.nn.conv2d_transpose: [topi.nn.conv2d_transpose_nchw],
+        tvm.relay.op.nn.dense: [topi.nn.dense],
+    }
 
     topi_funcs = []
     for op_name in ops:
-        if op_name in env.op2topi:
-            topi_funcs.extend(env.op2topi[op_name])
+        if op_name in OP2TOPI:
+            topi_funcs.extend(OP2TOPI[op_name])
         else:
             warnings.warn("Op %s is not tunable, ignored" % op_name)
 
diff --git a/python/tvm/autotvm/task/topi_integration.py b/python/tvm/autotvm/task/topi_integration.py
index f005ee0c9a54..412d7ae0e40b 100644
--- a/python/tvm/autotvm/task/topi_integration.py
+++ b/python/tvm/autotvm/task/topi_integration.py
@@ -11,16 +11,202 @@
 See tvm/topi/python/topi/arm_cpu/depthwise_conv2d.py for example usage.
 """
 
-from ... import _api_internal, tensor
-
-from .task import args_to_workload, dispatcher
+from ... import _api_internal, tensor, placeholder, create_schedule
 
+from .task import args_to_workload, dispatcher, register
+from ..util import get_const_tuple
 
 # A table that records all registered dispatcher for all targets
 _REGISTED_DISPATHCER = {
 }
 
 
+def serialize_args(args):
+    """serialize arguments of a topi function to a hashable tuple.
+
+    Parameters
+    ----------
+    args: list of hashable or Tensor
+    """
+    ret = []
+    for t in args:
+        if isinstance(t, tensor.Tensor):
+            ret.append(('TENSOR', get_const_tuple(t.shape), t.dtype))
+        else:
+            ret.append(t)
+    return tuple(ret)
+
+
+def deserialize_args(args):
+    """The inverse function of :code:`serialize_args`.
+
+    Parameters
+    ----------
+    args: list of hashable or Tensor
+    """
+    ret = []
+    for t in args:
+        if isinstance(t, tuple) and t[0] == 'TENSOR':
+            ret.append(placeholder(shape=t[1], dtype=t[2]))
+        else:
+            ret.append(t)
+    return ret
+
+
+# Task extractor for nnvm graph, relay program
+class TaskExtractEnv:
+    """Global environment for extracting tuning tasks from nnvm graph"""
+    current = None
+
+    def __init__(self):
+        import topi
+
+        # topi compute -> autotvm task name
+        self.topi_to_task = {
+            topi.nn.conv2d: "topi_nn_conv2d",
+            topi.nn.depthwise_conv2d_nchw: "topi_nn_depthwise_conv2d_nchw",
+            topi.nn.group_conv2d_nchw: "topi_nn_group_conv2d_nchw",
+            topi.nn.conv2d_transpose_nchw: "topi_nn_conv2d_transpose_nchw",
+            topi.nn.dense: "topi_nn_dense",
+        }
+
+        self.topi_to_schedule = {
+            topi.nn.conv2d: [topi.generic.schedule_conv2d_nchw,
+                             topi.generic.schedule_conv2d_nhwc],
+            topi.nn.depthwise_conv2d_nchw: [topi.generic.schedule_depthwise_conv2d_nchw,
+                                            topi.generic.schedule_depthwise_conv2d_nhwc],
+            topi.nn.group_conv2d_nchw: [topi.generic.schedule_group_conv2d_nchw],
+            topi.nn.conv2d_transpose_nchw: [topi.generic.schedule_conv2d_transpose_nchw],
+            topi.nn.dense: [topi.generic.schedule_dense],
+        }
+
+        self._register_tracing()
+        self._register_topi_task()
+        self.task_collection = []
+        self.wanted_topi_funcs = list(self.topi_to_task.keys())
+
+    def _register_tracing(self):
+        """Register tracing function to track the topi function call"""
+        # register topi compute for "tracing" target
+        for topi_compute in self.topi_to_task:
+            def _local_scope(compute_func):
+                """start a scope to hold the local function in for loop"""
+
+                @compute_func.register("tracing", )
+                def _tracing_topi_compute(*args, **kwargs):
+                    assert not kwargs, "Do not support extracting tuning tasks when" \
+                                       "kwargs is used in TOPI function call." \
+                                       "Please modify it to use only positional args."
+
+                    if compute_func in self.wanted_topi_funcs:  # record this call
+                        key = (self.topi_to_task[compute_func], serialize_args(args))
+                        if key not in self.task_collection:
+                            self.task_collection.append(key)
+
+                    return compute_func.fdefault(*args)
+            _local_scope(topi_compute)
+
+        # register topi schedule for "tracing" target
+        for topi_compute in self.topi_to_task:
+            for topi_schedule in self.topi_to_schedule[topi_compute]:
+                def _local_scope_(schedule_func):
+                    """start a scope to hold the local function in for loop"""
+
+                    @schedule_func.register("tracing", )
+                    def _tracing_topi_compute(outs):
+                        outs = [outs] if isinstance(outs, tensor.Tensor) else outs
+                        return create_schedule([x.op for x in outs])
+                _local_scope_(topi_schedule)
+
+    def _register_topi_task(self):
+        """register tuning wrapper for topi function"""
+        import topi
+
+        # Tuning wrapper for topi functions
+        @register("topi_nn_conv2d")
+        def _topi_nn_conv2d(*args, **kwargs):
+            assert not kwargs, "Do not support kwargs in template function call"
+            args = deserialize_args(args)
+            A, W = args[:2]
+            layout = args[-2]
+            assert layout == 'NCHW', "only support NCHW currently"
+            C = topi.nn.conv2d(*args, **kwargs)
+            s = topi.generic.schedule_conv2d_nchw([C])
+            return s, [A, W, C]
+
+        @register("topi_nn_depthwise_conv2d_nchw")
+        def _topi_nn_depthwise_conv2d_nchw(*args, **kwargs):
+            assert not kwargs, "Do not support kwargs in template function call"
+            args = deserialize_args(args)
+            A, W = args[:2]
+            C = topi.nn.depthwise_conv2d_nchw(*args, **kwargs)
+            s = topi.generic.schedule_depthwise_conv2d_nchw([C])
+            return s, [A, W, C]
+
+        @register("topi_nn_group_conv2d_nchw")
+        def _topi_nn_group_conv2d_nchw(*args, **kwargs):
+            assert not kwargs, "Do not support kwargs in template function call"
+            args = deserialize_args(args)
+            A, W = args[:2]
+            C = topi.nn.group_conv2d_nchw(*args, **kwargs)
+            s = topi.generic.schedule_group_conv2d_nchw([C])
+            return s, [A, W, C]
+
+        @register("topi_nn_conv2d_transpose_nchw")
+        def _topi_nn_conv2d_transpose_nchw(*args, **kwargs):
+            assert not kwargs, "Do not support kwargs in template function call"
+            args = deserialize_args(args)
+            A, W = args[:2]
+            C = topi.nn.conv2d_transpose_nchw(*args, **kwargs)
+            s = topi.generic.schedule_conv2d_transpose_nchw([C])
+            return s, [A, W, C]
+
+        @register("topi_nn_dense")
+        def _topi_nn_dense(*args, **kwargs):
+            assert not kwargs, "Do not support kwargs in template function call"
+            args = deserialize_args(args)
+            data, weight, bias = args
+            C = topi.nn.dense(*args, **kwargs)
+            s = topi.generic.schedule_dense([C])
+            if bias is not None:
+                return s, [data, weight, bias, C]
+            return s, [data, weight, C]
+
+    def reset(self, wanted_topi_funcs):
+        """Reset task collections
+
+        Parameters
+        ----------
+        wanted_topi_funcs: List of function
+            The topi function to be extracted
+        """
+        self.task_collection = []
+        self.wanted_topi_funcs = wanted_topi_funcs
+
+    def get_tasks(self):
+        """Get collected tasks
+
+        Returns
+        -------
+        tasks: List of tuple(name, args)
+            A list of tasks extracted from the nnvm graph
+        """
+        return self.task_collection
+
+    @staticmethod
+    def get():
+        """Get the single instance of TaskExtractEnv
+
+        Returns
+        -------
+        env: TaskExtractEnv
+            The single instance of TaskExtractEnv
+        """
+        if not TaskExtractEnv.current:
+            TaskExtractEnv.current = TaskExtractEnv()
+        return TaskExtractEnv.current
+
+
 def register_topi_compute(topi_compute, target_keys, template_keys, func=None):
     """Register a tunable template for a topi compute function.
 

From e902f586808e520f698c14877abf03482b95db0c Mon Sep 17 00:00:00 2001
From: Eddie Yan <eqy@cs.washington.edu>
Date: Sun, 2 Dec 2018 21:30:07 -0800
Subject: [PATCH 11/14] fix lint

---
 python/tvm/autotvm/task/nnvm_integration.py  | 2 +-
 python/tvm/autotvm/task/relay_integration.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/python/tvm/autotvm/task/nnvm_integration.py b/python/tvm/autotvm/task/nnvm_integration.py
index 770ac90520d8..cd7337586519 100644
--- a/python/tvm/autotvm/task/nnvm_integration.py
+++ b/python/tvm/autotvm/task/nnvm_integration.py
@@ -7,7 +7,7 @@
 import logging
 
 
-from ... import tensor, placeholder, target as _target
+from ... import target as _target
 
 from .task import create
 from .topi_integration import TaskExtractEnv
diff --git a/python/tvm/autotvm/task/relay_integration.py b/python/tvm/autotvm/task/relay_integration.py
index 82319ccf6aa8..7f4f9a919975 100644
--- a/python/tvm/autotvm/task/relay_integration.py
+++ b/python/tvm/autotvm/task/relay_integration.py
@@ -84,7 +84,7 @@ def extract_from_program(func, params, ops, target, target_host=None):
     # relay op -> topi compute
     OP2TOPI = {
         tvm.relay.op.nn.conv2d: [topi.nn.conv2d, topi.nn.depthwise_conv2d_nchw,
-                             topi.nn.group_conv2d_nchw],
+                                 topi.nn.group_conv2d_nchw],
         tvm.relay.op.nn.conv2d_transpose: [topi.nn.conv2d_transpose_nchw],
         tvm.relay.op.nn.dense: [topi.nn.dense],
     }
@@ -157,7 +157,7 @@ def extract_from_multiple_program(funcs, params, ops, target, target_host=None):
     # relay op -> topi compute
     OP2TOPI = {
         tvm.relay.op.nn.conv2d: [topi.nn.conv2d, topi.nn.depthwise_conv2d_nchw,
-                             topi.nn.group_conv2d_nchw],
+                                 topi.nn.group_conv2d_nchw],
         tvm.relay.op.nn.conv2d_transpose: [topi.nn.conv2d_transpose_nchw],
         tvm.relay.op.nn.dense: [topi.nn.dense],
     }

From ddc0a5ad4a975b9f80cee4f0c552d74b01f1c7db Mon Sep 17 00:00:00 2001
From: Eddie Yan <eqy@cs.washington.edu>
Date: Mon, 10 Dec 2018 09:37:46 -0800
Subject: [PATCH 12/14] trigger CI

---
 python/tvm/autotvm/task/relay_integration.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/python/tvm/autotvm/task/relay_integration.py b/python/tvm/autotvm/task/relay_integration.py
index 7f4f9a919975..21acf257f9ac 100644
--- a/python/tvm/autotvm/task/relay_integration.py
+++ b/python/tvm/autotvm/task/relay_integration.py
@@ -52,7 +52,7 @@ def deserialize_args(args):
 def extract_from_program(func, params, ops, target, target_host=None):
     """ Extract tuning tasks from a relay program.
 
-    This function collects tuning tasks by building the graph
+    This function collects tuning tasks by building the program
     with a "tracing" target and tracing all the calls to topi.
 
     Parameters
@@ -60,11 +60,11 @@ def extract_from_program(func, params, ops, target, target_host=None):
     func: relay.expr.Function
         The func to tune
     params: dict of str to numpy array
-        The associated parameters of the graph
+        The associated parameters of the program
     ops: List of relay op
         List of relay ops to be tuned
     dtype: str or dict of str to str
-        The input types to the graph
+        The input types to the program
     target: tvm.target.Target
         The compilation target
     target_host: tvm.target.Target
@@ -128,14 +128,14 @@ def extract_from_program(func, params, ops, target, target_host=None):
 def extract_from_multiple_program(funcs, params, ops, target, target_host=None):
     """ Extract tuning tasks from multiple relay programs.
 
-    This function is the multiple graph version of extract_from_graph
+    This function is the multiple program version of extract_from_program
 
     Parameters
     ----------
     funcs: List of relay.expr.Function
         The list of functions to tune
     params: List of dict of str to numpy array
-        The input shape to the graph
+        The associated parameters of the programs
     ops: List of relay op
         List of relay ops to be tuned
     target: tvm.target.Target

From 22488626182121bcaecceb41204030f2d58b4ca7 Mon Sep 17 00:00:00 2001
From: Eddie Yan <eqy@cs.washington.edu>
Date: Mon, 10 Dec 2018 10:00:38 -0800
Subject: [PATCH 13/14] fix import

---
 topi/python/topi/x86/conv2d.py           | 2 +-
 topi/python/topi/x86/depthwise_conv2d.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/topi/python/topi/x86/conv2d.py b/topi/python/topi/x86/conv2d.py
index 1a73736264bd..fe38b38d38e0 100644
--- a/topi/python/topi/x86/conv2d.py
+++ b/topi/python/topi/x86/conv2d.py
@@ -2,7 +2,7 @@
 """Conv2D schedule on x86"""
 import tvm
 from tvm import autotvm
-from tvm.autotvm.task.nnvm_integration import deserialize_args
+from tvm.autotvm.task.topi_integration import deserialize_args
 from tvm.autotvm.task import get_config
 from .. import generic, tag
 from .. import nn
diff --git a/topi/python/topi/x86/depthwise_conv2d.py b/topi/python/topi/x86/depthwise_conv2d.py
index 8f37a0316229..64858df91cdc 100644
--- a/topi/python/topi/x86/depthwise_conv2d.py
+++ b/topi/python/topi/x86/depthwise_conv2d.py
@@ -4,7 +4,7 @@
 from tvm import autotvm
 from tvm.autotvm.task import get_config
 from tvm.autotvm.task.space import SplitEntity
-from tvm.autotvm.task.nnvm_integration import deserialize_args
+from tvm.autotvm.task.topi_integration import deserialize_args
 from .. import generic, tag
 from ..nn.pad import pad
 from ..util import get_const_tuple

From 083cc12d5cc4380f399dd711039ce1e12eeb37cc Mon Sep 17 00:00:00 2001
From: Eddie Yan <eqy@cs.washington.edu>
Date: Mon, 10 Dec 2018 17:04:26 -0800
Subject: [PATCH 14/14] check in test

---
 .../relay/test_autotvm_task_extraction.py     | 56 +++++++++++++++++++
 1 file changed, 56 insertions(+)
 create mode 100644 tests/python/relay/test_autotvm_task_extraction.py

diff --git a/tests/python/relay/test_autotvm_task_extraction.py b/tests/python/relay/test_autotvm_task_extraction.py
new file mode 100644
index 000000000000..8c93e4a56642
--- /dev/null
+++ b/tests/python/relay/test_autotvm_task_extraction.py
@@ -0,0 +1,56 @@
+"""Test task extraction for autotvm"""
+import tvm.relay.testing
+from tvm import relay
+from tvm import autotvm
+
+def get_network(name, batch_size):
+    """Get the symbol definition and random weight of a network"""
+    input_shape = (batch_size, 3, 224, 224)
+
+    if name == 'resnet-18':
+        net, params = relay.testing.resnet.get_workload(num_layers=18, batch_size=batch_size)
+    elif name == 'mobilenet':
+        net, params = relay.testing.mobilenet.get_workload(batch_size=batch_size)
+    elif name == 'dcgan':
+        net, params = relay.testing.dcgan.get_workload(batch_size=batch_size)
+        input_shape = (batch_size, 100)
+    else:
+        raise ValueError("Unsupported network: " + name)
+
+    return net, params, input_shape
+
+def test_task_extraction():
+    target = 'llvm'
+
+    net, params, input_shape = get_network('resnet-18', batch_size=1)
+    tasks = autotvm.task.extract_from_program(net, target=target,
+                                            params=params,
+                                            ops=(relay.op.nn.conv2d,))
+    assert len(tasks) == 12
+
+    net, params, input_shape = get_network('resnet-18', batch_size=1)
+    tasks = autotvm.task.extract_from_program(net, target=target,
+                                            params=params,
+                                            ops=(relay.op.nn.dense,))
+    assert len(tasks) == 1
+
+    net, params, input_shape = get_network('resnet-18', batch_size=1)
+    tasks = autotvm.task.extract_from_program(net, target=target,
+                                            params=params,
+                                            ops=(relay.op.nn.conv2d, relay.op.nn.dense))
+    assert len(tasks) == 13
+
+    net, params, input_shape = get_network('mobilenet', batch_size=1)
+    tasks = autotvm.task.extract_from_program(net, target=target,
+                                            params=params,
+                                            ops=(relay.op.nn.conv2d, relay.op.nn.dense))
+    assert len(tasks) == 20
+
+    net, params, input_shape = get_network('dcgan', batch_size=1)
+    tasks = autotvm.task.extract_from_program(net, target=target,
+                                            params=params,
+                                            ops=(relay.op.nn.conv2d_transpose,))
+    assert len(tasks) == 4
+
+if __name__ == '__main__':
+    test_task_extraction()