You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2020/10/19 09:29:19 UTC
[GitHub] [incubator-tvm] merrymercy opened a new pull request #6710: [AutoScheduler] Relay integration : Task extraction
merrymercy opened a new pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710
This pr implements the basic relay integration for auto-scheduler.
## How to integrate auto-scheduler in relay
1. register auto-scheduler as an implementation in the op strategy. We have a universal schedule function for all topi compute functions.
2. Use tracing to extract all tasks similar to autotvm.
If auto-scheduler is enabled by `auto_scheduler.enable_relay_integration`, then auto-scheduler is always prefered to any other registered implementation.
## Limitation
Ideally, we should be able to mix autotvm, auto-scheduler, and static libraries, then we can select the best implementation according to a profling records.
However, this is not trivial to implement, because auto-scheduler works on a subgraph level while the current OpStrategy is designed for single operator level.
The limitation of OpStrategy also makes auto-scheduler unable to support multiple compute declarations. For example, auto-scheduler cannot select direct vs. winograd based on profiling records, while autotvm can easilly do this.
## Minor todo
- [ ] Add an API to disable the cache in relay compilation, so we can extract the weigths (i.e., the number of appearance) of each task for the task-scheduler.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] merrymercy commented on a change in pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
merrymercy commented on a change in pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#discussion_r511954470
##########
File path: python/tvm/auto_scheduler/relay_integration.py
##########
@@ -0,0 +1,179 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=unused-variable,invalid-name
+
+"""
+Integrate auto_scheduler into relay. It implements the following items:
+1. Extract search tasks from a relay program
+2. Provide auto-scheduling for all TOPI compute functions
+"""
+
+import os
+import threading
+
+import tvm
+from tvm import te, transform
+from tvm.te.tensor import PlaceholderOp, ComputeOp
+from .dispatcher import DispatchContext, FallbackConfig
+from .workload_registry import register_workload_tensors
+from .compute_dag import ComputeDAG
+
+
+def call_all_topi_funcs(mod, params, target):
+ """Call all TOPI compute + schedule to extract tasks in a relay program"""
+ # pylint: disable=import-outside-toplevel
+ from tvm import relay
+ from tvm.relay.backend import graph_runtime_codegen
+
+ with transform.PassContext(opt_level=3):
+ opt_mod, _ = relay.optimize(mod, target, params)
+ grc = graph_runtime_codegen.GraphRuntimeCodegen(None, target)
+ grc.codegen(opt_mod["main"])
+
+
+def extract_tasks(mod, params, target):
+ """Extract tuning tasks from a relay program.
+
+ Parameters
+ ----------
+ mod: tvm.IRModule or relay.function.Function
+ The module or function to tune
+ params: dict of str to numpy array
+ The associated parameters of the program
+ target: tvm.target.Target
+ The compilation target
+
+ Returns
+ -------
+ wkl_keys: List[str]
+ The hash key of extracted workloads
+ wkl_weights: List[int]
+ The weight (i.e. the number of appearance) of extracted workload
+ """
+ # pylint: disable=import-outside-toplevel
+ from tvm import relay
Review comment:
It is due to a legacy issue. Removed it.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] merrymercy commented on pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
merrymercy commented on pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#issuecomment-711911420
cc @icemelon9 @tqchen @minminsun @jcf94 @comaniac @FrozenGene
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] FrozenGene commented on a change in pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
FrozenGene commented on a change in pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#discussion_r514002681
##########
File path: python/tvm/auto_scheduler/utils.py
##########
@@ -30,7 +30,7 @@
try:
import psutil
except ImportError:
- raise ImportError("psutil not found, try `pip install psutil` to fix this")
Review comment:
I prefer keeping this ImportError. During internal development, we have met several times our environment doesn't import this package, we could tune but meet process error and find this root cause is hard.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] merrymercy commented on a change in pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
merrymercy commented on a change in pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#discussion_r511955715
##########
File path: python/tvm/auto_scheduler/relay_integration.py
##########
@@ -0,0 +1,179 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=unused-variable,invalid-name
+
+"""
+Integrate auto_scheduler into relay. It implements the following items:
+1. Extract search tasks from a relay program
+2. Provide auto-scheduling for all TOPI compute functions
+"""
+
+import os
+import threading
+
+import tvm
+from tvm import te, transform
+from tvm.te.tensor import PlaceholderOp, ComputeOp
+from .dispatcher import DispatchContext, FallbackConfig
+from .workload_registry import register_workload_tensors
+from .compute_dag import ComputeDAG
+
+
+def call_all_topi_funcs(mod, params, target):
+ """Call all TOPI compute + schedule to extract tasks in a relay program"""
+ # pylint: disable=import-outside-toplevel
+ from tvm import relay
+ from tvm.relay.backend import graph_runtime_codegen
+
+ with transform.PassContext(opt_level=3):
+ opt_mod, _ = relay.optimize(mod, target, params)
+ grc = graph_runtime_codegen.GraphRuntimeCodegen(None, target)
+ grc.codegen(opt_mod["main"])
+
+
+def extract_tasks(mod, params, target):
+ """Extract tuning tasks from a relay program.
+
+ Parameters
+ ----------
+ mod: tvm.IRModule or relay.function.Function
+ The module or function to tune
+ params: dict of str to numpy array
+ The associated parameters of the program
+ target: tvm.target.Target
+ The compilation target
+
+ Returns
+ -------
+ wkl_keys: List[str]
+ The hash key of extracted workloads
+ wkl_weights: List[int]
+ The weight (i.e. the number of appearance) of extracted workload
+ """
+ # pylint: disable=import-outside-toplevel
+ from tvm import relay
+
+ env = TracingEnvironment()
+ with env:
+ # Run the compiler to collect all TOPI calls during compilation.
+
+ # Wrap build call in a new thread to avoid the conflict
+ # between python's multiprocessing and tvm's thread pool
+ build_thread = threading.Thread(target=call_all_topi_funcs, args=(mod, params, target))
+ build_thread.start()
+ build_thread.join()
+ relay.backend.compile_engine.get().clear()
+
+ # create tasks for target
+ wkl_keys = []
+ wkl_weights = []
+ for wkl_key, wkl_weight in env.wkl_key_collection.items():
+ wkl_keys.append(wkl_key)
+ wkl_weights.append(wkl_weight)
+
+ return wkl_keys, wkl_weights
+
+
+class TracingEnvironment:
+ """Global environment for tracing all topi function calls"""
+
+ current = None
+
+ def __init__(self):
+ self.relay_disable_build_cache = "false"
+ self.layout_rewrite_success_ct = 0
+ self.wkl_key_collection = {}
+
+ def __enter__(self):
+ self.relay_disable_build_cache = os.environ.get("TVM_RELAY_DISABLE_BUILD_CACHE", "false")
+ os.environ["TVM_RELAY_DISABLE_BUILD_CACHE"] = "true"
+ TracingEnvironment.current = self
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ os.environ["TVM_RELAY_DISABLE_BUILD_CACHE"] = self.relay_disable_build_cache
+ TracingEnvironment.current = None
+
+ def add_workload_key(self, key):
+ """Add the workload key of an Ansor search task
+
+ Parameters
+ ----------
+ key: str
+ The workload key of a task
+ """
+ if key in self.wkl_key_collection:
+ self.wkl_key_collection[key] += 1
+ else:
+ self.wkl_key_collection[key] = 1
+
+
+def traverse_to_get_io_tensors(outs):
+ """Traverse from a list of output tensors to get a whole computational DAG"""
+ inputs = []
+
+ visited = set()
+
+ def traverse(t):
+ if t in visited:
+ return
+ if isinstance(t.op, PlaceholderOp):
+ inputs.append(t)
+ elif isinstance(t.op, ComputeOp):
+ for x in t.op.input_tensors:
+ traverse(x)
+ visited.add(t)
+
+ for t in outs:
+ traverse(t)
+
+ return inputs + list(outs)
+
+
+# The suffix of implementations that use the auto-scheduler in the OpStrategy.
+auto_schedule_impl_suffix = ".auto_scheduler"
+
+
+def auto_schedule_topi(outs):
Review comment:
We cannot call it `auto_schedule_te`. We have another function `auto_schedule` (https://github.com/apache/incubator-tvm/blob/e59c603515befb02035e237794aa0645dbfbaf09/python/tvm/auto_scheduler/auto_schedule.py#L161) for your use case.
But this function is designed to be used as a TOPI schedule function for Relay, because it does a lot of other things.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] merrymercy edited a comment on pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
merrymercy edited a comment on pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#issuecomment-716901922
@comaniac I only added conv2d_nhwc and dense in this PR to test the functionality.
I will send another PR with more ops and a tutorial on tuning networks for CUDA.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] merrymercy commented on a change in pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
merrymercy commented on a change in pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#discussion_r514835117
##########
File path: python/tvm/auto_scheduler/utils.py
##########
@@ -157,6 +157,9 @@ def __reduce__(self):
def kill_child_processes(parent_pid, sig=signal.SIGTERM):
"""kill all child processes recursively"""
+ if not psutil:
+ raise ImportError("psutil not found, try `pip install psutil` to fix this")
Review comment:
Yes, it is for testing. The goal is to minimize the dependency to run tvm.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] FrozenGene commented on a change in pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
FrozenGene commented on a change in pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#discussion_r511743107
##########
File path: tests/python/relay/test_auto_scheduler_task_extraction.py
##########
@@ -0,0 +1,60 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Test task extraction for auto-scheduler"""
+import tvm.relay.testing
+from tvm import auto_scheduler, relay
+
+
+def get_network(name, batch_size=1):
+ """Get the symbol definition and random weight of a network"""
+
+ # auto-scheduler prefer NHWC layout
+ layout = "NHWC"
Review comment:
How about move this to the parameter of `get_network`? We prefer NHWC but we could add testing of NCHW so that we could make sure we could work on NCHW layout too.
##########
File path: python/tvm/auto_scheduler/dispatcher.py
##########
@@ -0,0 +1,288 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+The global context that dispatches best configurations to workloads
+"""
+# pylint: disable=invalid-name
+
+from __future__ import absolute_import as _abs
Review comment:
I think we could remove it as we require Python 3.6 now
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] merrymercy commented on a change in pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
merrymercy commented on a change in pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#discussion_r511955715
##########
File path: python/tvm/auto_scheduler/relay_integration.py
##########
@@ -0,0 +1,179 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=unused-variable,invalid-name
+
+"""
+Integrate auto_scheduler into relay. It implements the following items:
+1. Extract search tasks from a relay program
+2. Provide auto-scheduling for all TOPI compute functions
+"""
+
+import os
+import threading
+
+import tvm
+from tvm import te, transform
+from tvm.te.tensor import PlaceholderOp, ComputeOp
+from .dispatcher import DispatchContext, FallbackConfig
+from .workload_registry import register_workload_tensors
+from .compute_dag import ComputeDAG
+
+
+def call_all_topi_funcs(mod, params, target):
+ """Call all TOPI compute + schedule to extract tasks in a relay program"""
+ # pylint: disable=import-outside-toplevel
+ from tvm import relay
+ from tvm.relay.backend import graph_runtime_codegen
+
+ with transform.PassContext(opt_level=3):
+ opt_mod, _ = relay.optimize(mod, target, params)
+ grc = graph_runtime_codegen.GraphRuntimeCodegen(None, target)
+ grc.codegen(opt_mod["main"])
+
+
+def extract_tasks(mod, params, target):
+ """Extract tuning tasks from a relay program.
+
+ Parameters
+ ----------
+ mod: tvm.IRModule or relay.function.Function
+ The module or function to tune
+ params: dict of str to numpy array
+ The associated parameters of the program
+ target: tvm.target.Target
+ The compilation target
+
+ Returns
+ -------
+ wkl_keys: List[str]
+ The hash key of extracted workloads
+ wkl_weights: List[int]
+ The weight (i.e. the number of appearance) of extracted workload
+ """
+ # pylint: disable=import-outside-toplevel
+ from tvm import relay
+
+ env = TracingEnvironment()
+ with env:
+ # Run the compiler to collect all TOPI calls during compilation.
+
+ # Wrap build call in a new thread to avoid the conflict
+ # between python's multiprocessing and tvm's thread pool
+ build_thread = threading.Thread(target=call_all_topi_funcs, args=(mod, params, target))
+ build_thread.start()
+ build_thread.join()
+ relay.backend.compile_engine.get().clear()
+
+ # create tasks for target
+ wkl_keys = []
+ wkl_weights = []
+ for wkl_key, wkl_weight in env.wkl_key_collection.items():
+ wkl_keys.append(wkl_key)
+ wkl_weights.append(wkl_weight)
+
+ return wkl_keys, wkl_weights
+
+
+class TracingEnvironment:
+ """Global environment for tracing all topi function calls"""
+
+ current = None
+
+ def __init__(self):
+ self.relay_disable_build_cache = "false"
+ self.layout_rewrite_success_ct = 0
+ self.wkl_key_collection = {}
+
+ def __enter__(self):
+ self.relay_disable_build_cache = os.environ.get("TVM_RELAY_DISABLE_BUILD_CACHE", "false")
+ os.environ["TVM_RELAY_DISABLE_BUILD_CACHE"] = "true"
+ TracingEnvironment.current = self
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ os.environ["TVM_RELAY_DISABLE_BUILD_CACHE"] = self.relay_disable_build_cache
+ TracingEnvironment.current = None
+
+ def add_workload_key(self, key):
+ """Add the workload key of an Ansor search task
+
+ Parameters
+ ----------
+ key: str
+ The workload key of a task
+ """
+ if key in self.wkl_key_collection:
+ self.wkl_key_collection[key] += 1
+ else:
+ self.wkl_key_collection[key] = 1
+
+
+def traverse_to_get_io_tensors(outs):
+ """Traverse from a list of output tensors to get a whole computational DAG"""
+ inputs = []
+
+ visited = set()
+
+ def traverse(t):
+ if t in visited:
+ return
+ if isinstance(t.op, PlaceholderOp):
+ inputs.append(t)
+ elif isinstance(t.op, ComputeOp):
+ for x in t.op.input_tensors:
+ traverse(x)
+ visited.add(t)
+
+ for t in outs:
+ traverse(t)
+
+ return inputs + list(outs)
+
+
+# The suffix of implementations that use the auto-scheduler in the OpStrategy.
+auto_schedule_impl_suffix = ".auto_scheduler"
+
+
+def auto_schedule_topi(outs):
Review comment:
We cannot call it `auto_schedule_te`. We have another function `auto_schedule` (https://github.com/apache/incubator-tvm/blob/e59c603515befb02035e237794aa0645dbfbaf09/python/tvm/auto_scheduler/auto_schedule.py#L161) for your use case.
But this function is designed to be used as a TOPI schedule function for Relay.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] merrymercy commented on a change in pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
merrymercy commented on a change in pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#discussion_r514835117
##########
File path: python/tvm/auto_scheduler/utils.py
##########
@@ -157,6 +157,9 @@ def __reduce__(self):
def kill_child_processes(parent_pid, sig=signal.SIGTERM):
"""kill all child processes recursively"""
+ if not psutil:
+ raise ImportError("psutil not found, try `pip install psutil` to fix this")
Review comment:
Yes, it is for testing. The goal is to minimize the dependency to run tvm.
If users do not use auto-scheduler, then they do not need to install this package.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] merrymercy commented on a change in pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
merrymercy commented on a change in pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#discussion_r511993769
##########
File path: python/tvm/auto_scheduler/dispatcher.py
##########
@@ -0,0 +1,288 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+The global context that dispatches best configurations to workloads
+"""
+# pylint: disable=invalid-name
+
+from __future__ import absolute_import as _abs
+
+import logging
+from pathlib import Path
+
+import numpy as np
+
+from tvm.tir.expr import FloatImm
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class DispatchContext(object):
+ """
+ Base class of dispatch context.
+ """
+
+ current = None
+
+ def __init__(self):
+ self._old_ctx = DispatchContext.current
+
+ def query(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ If cannot find the result inside this context, this function will query it
+ from the upper contexts.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : str
+ The current workload
+
+ Returns
+ -------
+ cfg : State
+ The schedule configuration for the workload
+ """
+ ret = self._query_inside(target, workload)
+ if ret is None:
+ ret = self._old_ctx.query(target, workload)
+ return ret
+
+ def update(self, target, workload, state):
+ """
+ Update the config for a workload
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+ state : State
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def _query_inside(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ This function only query config inside this context.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+
+ Returns
+ -------
+ cfg : State or str
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def __enter__(self):
+ self._old_ctx = DispatchContext.current
+ DispatchContext.current = self
+ return self
+
+ def __exit__(self, ptype, value, trace):
+ DispatchContext.current = self._old_ctx
+
+
+class ApplyHistoryBest(DispatchContext):
+ """
+ Apply the history best config
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+
+ def __init__(self, records, n_lines=None):
+ super(ApplyHistoryBest, self).__init__()
+
+ self.best_by_targetkey = {}
+ self.best_by_model = {}
+ self._best_user_defined = {}
+
+ if records:
+ self.load(records, n_lines)
+
+ def load(self, records, n_lines=None):
+ """Load records to this dispatch context
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+ # pylint: disable=import-outside-toplevel
+ from . import load_from_file
+
+ if isinstance(records, Path):
+ records = str(records)
+
+ if isinstance(records, str):
+ records = load_from_file(records)
+ if not records:
+ return
+
+ best_by_targetkey = self.best_by_targetkey
+ best_by_model = self.best_by_model
+
+ counter = 0
+ for inp, res in records:
+ if n_lines is not None and counter >= n_lines:
+ break
+ counter += 1
+ if res.error_no != 0:
+ continue
+
+ # use target keys in tvm target system as key to build best map
+ for k in inp.task.target.keys:
+ key = (k, inp.task.workload_key)
+ if key not in best_by_targetkey:
+ best_by_targetkey[key] = (inp, res)
+ else:
+ _, other_res = best_by_targetkey[key]
+ other_costs = [x.value for x in other_res.costs if isinstance(x, FloatImm)]
+ costs = [x.value for x in res.costs if isinstance(x, FloatImm)]
+ if np.mean(other_costs) > np.mean(costs):
+ best_by_targetkey[key] = (inp, res)
+
+ # use model as key to build best map
+ key = (inp.task.target.model, inp.task.workload_key)
+ if key not in best_by_model:
+ if inp.task.target.model != "unknown":
+ best_by_model[key] = (inp, res)
+ else:
+ _, other_res = best_by_model[key]
+ other_costs = [x.value for x in other_res.costs if isinstance(x, FloatImm)]
+ costs = [x.value for x in res.costs if isinstance(x, FloatImm)]
+ if np.mean(other_costs) > np.mean(costs):
+ best_by_model[key] = (inp, res)
+
+ logger.debug("Finish loading %d records", counter)
+
+ def _query_inside(self, target, workload):
+ if target is None:
+ raise RuntimeError(
+ "Need a target context to find the history best. "
+ "Hint: If your target is llvm, use `with tvm.target.create('llvm'):`"
+ " above the dispatcher call. So does other target. "
+ )
+
+ # first try matching by model
+ key = (target.model, workload)
+ if key in self._best_user_defined:
+ return self._best_user_defined[key]
+ if key in self.best_by_model:
+ return self.best_by_model[key][0].state
+
+ # then try matching by target key
+ for k in target.keys:
+ key = (k, workload)
+ if key in self._best_user_defined:
+ return self._best_user_defined[key]
+ if key in self.best_by_targetkey:
+ return self.best_by_targetkey[key][0].state
+
+ return None
+
+ def update(self, target, workload, state):
+ model = target.model
+ key = (model, workload)
+ self._best_user_defined[key] = state
+
+ for k in target.keys:
+ key = (k, workload)
+ self._best_user_defined[key] = state
+
+
+class FallbackConfig:
+ pass
+
+
+class FallbackContext(DispatchContext):
+ """
+ A fallback dispatch context.
+ This is used as the root context.
+ """
+
+ def __init__(self):
+ super(FallbackContext, self).__init__()
+ self.memory = {}
+ self.silent = False
+
+ # a set to prevent print duplicated message
+ self.messages = set()
+
+ def _query_inside(self, target, workload):
+ key = (str(target), workload)
+ if key in self.memory:
+ return self.memory[key]
+
+ if not self.silent:
+ msg = (
+ "Cannot find config for target=%s, workload=%s. A fallback configuration "
+ "is used, which may bring great performance regression." % (target, workload)
+ )
+ if msg not in self.messages:
+ self.messages.add(msg)
+ logger.warning(msg)
+ cfg = FallbackConfig()
Review comment:
Removed the class `FallbackConfig`. The function now returns `None` if nothing is found.
We cannot return `init_state` because the dispatcher does not save ComputeDAG.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] comaniac commented on a change in pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#discussion_r507927667
##########
File path: python/tvm/auto_scheduler/dispatcher.py
##########
@@ -0,0 +1,288 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+The global context that dispatches best configurations to workloads
+"""
+# pylint: disable=invalid-name
+
+from __future__ import absolute_import as _abs
+
+import logging
+from pathlib import Path
+
+import numpy as np
+
+from tvm.tir.expr import FloatImm
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class DispatchContext(object):
+ """
+ Base class of dispatch context.
+ """
+
+ current = None
+
+ def __init__(self):
+ self._old_ctx = DispatchContext.current
+
+ def query(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ If cannot find the result inside this context, this function will query it
+ from the upper contexts.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : str
+ The current workload
+
+ Returns
+ -------
+ cfg : State
+ The schedule configuration for the workload
+ """
+ ret = self._query_inside(target, workload)
+ if ret is None:
+ ret = self._old_ctx.query(target, workload)
+ return ret
+
+ def update(self, target, workload, state):
+ """
+ Update the config for a workload
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+ state : State
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def _query_inside(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ This function only query config inside this context.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+
+ Returns
+ -------
+ cfg : State or str
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def __enter__(self):
+ self._old_ctx = DispatchContext.current
+ DispatchContext.current = self
+ return self
+
+ def __exit__(self, ptype, value, trace):
+ DispatchContext.current = self._old_ctx
+
+
+class ApplyHistoryBest(DispatchContext):
+ """
+ Apply the history best config
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+
+ def __init__(self, records, n_lines=None):
+ super(ApplyHistoryBest, self).__init__()
+
+ self.best_by_targetkey = {}
+ self.best_by_model = {}
+ self._best_user_defined = {}
+
+ if records:
Review comment:
In what situation `records` could be `False`?
##########
File path: python/tvm/auto_scheduler/dispatcher.py
##########
@@ -0,0 +1,288 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+The global context that dispatches best configurations to workloads
+"""
+# pylint: disable=invalid-name
+
+from __future__ import absolute_import as _abs
Review comment:
nit: do we need `_abs`?
##########
File path: python/tvm/auto_scheduler/dispatcher.py
##########
@@ -0,0 +1,288 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+The global context that dispatches best configurations to workloads
+"""
+# pylint: disable=invalid-name
+
+from __future__ import absolute_import as _abs
+
+import logging
+from pathlib import Path
+
+import numpy as np
+
+from tvm.tir.expr import FloatImm
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class DispatchContext(object):
+ """
+ Base class of dispatch context.
+ """
+
+ current = None
+
+ def __init__(self):
+ self._old_ctx = DispatchContext.current
+
+ def query(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ If cannot find the result inside this context, this function will query it
+ from the upper contexts.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : str
+ The current workload
+
+ Returns
+ -------
+ cfg : State
+ The schedule configuration for the workload
+ """
+ ret = self._query_inside(target, workload)
+ if ret is None:
+ ret = self._old_ctx.query(target, workload)
+ return ret
+
+ def update(self, target, workload, state):
+ """
+ Update the config for a workload
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+ state : State
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def _query_inside(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ This function only query config inside this context.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+
+ Returns
+ -------
+ cfg : State or str
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def __enter__(self):
+ self._old_ctx = DispatchContext.current
+ DispatchContext.current = self
+ return self
+
+ def __exit__(self, ptype, value, trace):
+ DispatchContext.current = self._old_ctx
+
+
+class ApplyHistoryBest(DispatchContext):
+ """
+ Apply the history best config
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+
+ def __init__(self, records, n_lines=None):
+ super(ApplyHistoryBest, self).__init__()
+
+ self.best_by_targetkey = {}
+ self.best_by_model = {}
+ self._best_user_defined = {}
+
+ if records:
+ self.load(records, n_lines)
+
+ def load(self, records, n_lines=None):
+ """Load records to this dispatch context
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+ # pylint: disable=import-outside-toplevel
+ from . import load_from_file
Review comment:
ditto. Why lazy import is necessary here?
##########
File path: python/tvm/auto_scheduler/dispatcher.py
##########
@@ -0,0 +1,288 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+The global context that dispatches best configurations to workloads
+"""
+# pylint: disable=invalid-name
+
+from __future__ import absolute_import as _abs
+
+import logging
+from pathlib import Path
+
+import numpy as np
+
+from tvm.tir.expr import FloatImm
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class DispatchContext(object):
+ """
+ Base class of dispatch context.
+ """
+
+ current = None
+
+ def __init__(self):
+ self._old_ctx = DispatchContext.current
+
+ def query(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ If cannot find the result inside this context, this function will query it
+ from the upper contexts.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : str
+ The current workload
+
+ Returns
+ -------
+ cfg : State
+ The schedule configuration for the workload
+ """
+ ret = self._query_inside(target, workload)
+ if ret is None:
+ ret = self._old_ctx.query(target, workload)
+ return ret
+
+ def update(self, target, workload, state):
+ """
+ Update the config for a workload
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+ state : State
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def _query_inside(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ This function only query config inside this context.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+
+ Returns
+ -------
+ cfg : State or str
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def __enter__(self):
+ self._old_ctx = DispatchContext.current
+ DispatchContext.current = self
+ return self
+
+ def __exit__(self, ptype, value, trace):
+ DispatchContext.current = self._old_ctx
+
+
+class ApplyHistoryBest(DispatchContext):
+ """
+ Apply the history best config
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+
+ def __init__(self, records, n_lines=None):
+ super(ApplyHistoryBest, self).__init__()
+
+ self.best_by_targetkey = {}
+ self.best_by_model = {}
+ self._best_user_defined = {}
+
+ if records:
+ self.load(records, n_lines)
+
+ def load(self, records, n_lines=None):
+ """Load records to this dispatch context
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+ # pylint: disable=import-outside-toplevel
+ from . import load_from_file
+
+ if isinstance(records, Path):
+ records = str(records)
+
+ if isinstance(records, str):
+ records = load_from_file(records)
+ if not records:
+ return
+
+ best_by_targetkey = self.best_by_targetkey
+ best_by_model = self.best_by_model
+
+ counter = 0
+ for inp, res in records:
+ if n_lines is not None and counter >= n_lines:
+ break
+ counter += 1
+ if res.error_no != 0:
+ continue
+
+ # use target keys in tvm target system as key to build best map
+ for k in inp.task.target.keys:
+ key = (k, inp.task.workload_key)
+ if key not in best_by_targetkey:
+ best_by_targetkey[key] = (inp, res)
+ else:
+ _, other_res = best_by_targetkey[key]
+ other_costs = [x.value for x in other_res.costs if isinstance(x, FloatImm)]
+ costs = [x.value for x in res.costs if isinstance(x, FloatImm)]
+ if np.mean(other_costs) > np.mean(costs):
+ best_by_targetkey[key] = (inp, res)
+
+ # use model as key to build best map
+ key = (inp.task.target.model, inp.task.workload_key)
+ if key not in best_by_model:
+ if inp.task.target.model != "unknown":
+ best_by_model[key] = (inp, res)
+ else:
+ _, other_res = best_by_model[key]
+ other_costs = [x.value for x in other_res.costs if isinstance(x, FloatImm)]
+ costs = [x.value for x in res.costs if isinstance(x, FloatImm)]
+ if np.mean(other_costs) > np.mean(costs):
Review comment:
ditto.
##########
File path: tests/python/relay/test_auto_scheduler_task_extraction.py
##########
@@ -0,0 +1,60 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Test task extraction for auto-scheduler"""
+import tvm.relay.testing
+from tvm import auto_scheduler, relay
+
+
+def get_network(name, batch_size=1):
+ """Get the symbol definition and random weight of a network"""
+
+ # auto-scheduler prefer NHWC layout
+ layout = "NHWC"
+ image_shape = (224, 224, 3)
+
+ if name == "resnet-18":
+ mod, params = relay.testing.resnet.get_workload(
+ num_layers=18, batch_size=batch_size, layout=layout, image_shape=image_shape
+ )
+ elif name == "resnet3d-18":
+ mod, params = relay.testing.resnet_3d.get_workload(
+ num_layers=18, batch_size=batch_size, layout=layout, image_shape=image_shape
+ )
+ elif name == "mobilenet":
+ mod, params = relay.testing.mobilenet.get_workload(
+ batch_size=batch_size, layout=layout, image_shape=image_shape
+ )
+ elif name == "dcgan":
+ mod, params = relay.testing.dcgan.get_workload(batch_size=batch_size, layout=layout)
+ else:
+ raise ValueError("Unsupported network: " + name)
+
+ return mod, params
+
+
+def test_task_extraction():
+ auto_scheduler.enable_relay_integration()
+
+ mod, params = get_network("resnet-18")
+ target = tvm.target.Target("cuda")
Review comment:
Should we disable this test for non-CUDA environment?
##########
File path: python/tvm/auto_scheduler/dispatcher.py
##########
@@ -0,0 +1,288 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+The global context that dispatches best configurations to workloads
+"""
+# pylint: disable=invalid-name
+
+from __future__ import absolute_import as _abs
+
+import logging
+from pathlib import Path
+
+import numpy as np
+
+from tvm.tir.expr import FloatImm
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class DispatchContext(object):
+ """
+ Base class of dispatch context.
+ """
+
+ current = None
+
+ def __init__(self):
+ self._old_ctx = DispatchContext.current
+
+ def query(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ If cannot find the result inside this context, this function will query it
+ from the upper contexts.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : str
+ The current workload
+
+ Returns
+ -------
+ cfg : State
+ The schedule configuration for the workload
+ """
+ ret = self._query_inside(target, workload)
+ if ret is None:
+ ret = self._old_ctx.query(target, workload)
+ return ret
+
+ def update(self, target, workload, state):
+ """
+ Update the config for a workload
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+ state : State
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def _query_inside(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ This function only query config inside this context.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+
+ Returns
+ -------
+ cfg : State or str
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def __enter__(self):
+ self._old_ctx = DispatchContext.current
+ DispatchContext.current = self
+ return self
+
+ def __exit__(self, ptype, value, trace):
+ DispatchContext.current = self._old_ctx
+
+
+class ApplyHistoryBest(DispatchContext):
+ """
+ Apply the history best config
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+
+ def __init__(self, records, n_lines=None):
+ super(ApplyHistoryBest, self).__init__()
+
+ self.best_by_targetkey = {}
+ self.best_by_model = {}
+ self._best_user_defined = {}
+
+ if records:
+ self.load(records, n_lines)
+
+ def load(self, records, n_lines=None):
+ """Load records to this dispatch context
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+ # pylint: disable=import-outside-toplevel
+ from . import load_from_file
+
+ if isinstance(records, Path):
+ records = str(records)
+
+ if isinstance(records, str):
+ records = load_from_file(records)
+ if not records:
+ return
+
+ best_by_targetkey = self.best_by_targetkey
+ best_by_model = self.best_by_model
+
+ counter = 0
+ for inp, res in records:
+ if n_lines is not None and counter >= n_lines:
+ break
+ counter += 1
+ if res.error_no != 0:
+ continue
+
+ # use target keys in tvm target system as key to build best map
+ for k in inp.task.target.keys:
+ key = (k, inp.task.workload_key)
+ if key not in best_by_targetkey:
+ best_by_targetkey[key] = (inp, res)
+ else:
+ _, other_res = best_by_targetkey[key]
+ other_costs = [x.value for x in other_res.costs if isinstance(x, FloatImm)]
+ costs = [x.value for x in res.costs if isinstance(x, FloatImm)]
+ if np.mean(other_costs) > np.mean(costs):
+ best_by_targetkey[key] = (inp, res)
+
+ # use model as key to build best map
+ key = (inp.task.target.model, inp.task.workload_key)
+ if key not in best_by_model:
+ if inp.task.target.model != "unknown":
+ best_by_model[key] = (inp, res)
+ else:
+ _, other_res = best_by_model[key]
+ other_costs = [x.value for x in other_res.costs if isinstance(x, FloatImm)]
+ costs = [x.value for x in res.costs if isinstance(x, FloatImm)]
+ if np.mean(other_costs) > np.mean(costs):
+ best_by_model[key] = (inp, res)
+
+ logger.debug("Finish loading %d records", counter)
+
+ def _query_inside(self, target, workload):
+ if target is None:
+ raise RuntimeError(
+ "Need a target context to find the history best. "
+ "Hint: If your target is llvm, use `with tvm.target.create('llvm'):`"
+ " above the dispatcher call. So does other target. "
+ )
+
+ # first try matching by model
+ key = (target.model, workload)
+ if key in self._best_user_defined:
+ return self._best_user_defined[key]
+ if key in self.best_by_model:
+ return self.best_by_model[key][0].state
+
+ # then try matching by target key
+ for k in target.keys:
+ key = (k, workload)
+ if key in self._best_user_defined:
+ return self._best_user_defined[key]
+ if key in self.best_by_targetkey:
+ return self.best_by_targetkey[key][0].state
+
+ return None
+
+ def update(self, target, workload, state):
+ model = target.model
+ key = (model, workload)
+ self._best_user_defined[key] = state
+
+ for k in target.keys:
+ key = (k, workload)
+ self._best_user_defined[key] = state
+
+
+class FallbackConfig:
+ pass
+
+
+class FallbackContext(DispatchContext):
+ """
+ A fallback dispatch context.
+ This is used as the root context.
+ """
+
+ def __init__(self):
+ super(FallbackContext, self).__init__()
+ self.memory = {}
+ self.silent = False
+
+ # a set to prevent print duplicated message
+ self.messages = set()
+
+ def _query_inside(self, target, workload):
+ key = (str(target), workload)
+ if key in self.memory:
+ return self.memory[key]
+
+ if not self.silent:
+ msg = (
+ "Cannot find config for target=%s, workload=%s. A fallback configuration "
+ "is used, which may bring great performance regression." % (target, workload)
+ )
+ if msg not in self.messages:
+ self.messages.add(msg)
+ logger.warning(msg)
+ cfg = FallbackConfig()
Review comment:
Should we simply return an `init_state` or an undefined state so that we don't need `FallbackConfig` anymore? This return type is also inconsistent as the type description of `query`.
##########
File path: python/tvm/auto_scheduler/relay_integration.py
##########
@@ -0,0 +1,179 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=unused-variable,invalid-name
+
+"""
+Integrate auto_scheduler into relay. It implements the following items:
+1. Extract search tasks from a relay program
+2. Provide auto-scheduling for all TOPI compute functions
+"""
+
+import os
+import threading
+
+import tvm
+from tvm import te, transform
+from tvm.te.tensor import PlaceholderOp, ComputeOp
+from .dispatcher import DispatchContext, FallbackConfig
+from .workload_registry import register_workload_tensors
+from .compute_dag import ComputeDAG
+
+
+def call_all_topi_funcs(mod, params, target):
+ """Call all TOPI compute + schedule to extract tasks in a relay program"""
+ # pylint: disable=import-outside-toplevel
+ from tvm import relay
+ from tvm.relay.backend import graph_runtime_codegen
+
+ with transform.PassContext(opt_level=3):
+ opt_mod, _ = relay.optimize(mod, target, params)
+ grc = graph_runtime_codegen.GraphRuntimeCodegen(None, target)
+ grc.codegen(opt_mod["main"])
+
+
+def extract_tasks(mod, params, target):
+ """Extract tuning tasks from a relay program.
+
+ Parameters
+ ----------
+ mod: tvm.IRModule or relay.function.Function
+ The module or function to tune
+ params: dict of str to numpy array
+ The associated parameters of the program
+ target: tvm.target.Target
+ The compilation target
+
+ Returns
+ -------
+ wkl_keys: List[str]
+ The hash key of extracted workloads
+ wkl_weights: List[int]
+ The weight (i.e. the number of appearance) of extracted workload
+ """
+ # pylint: disable=import-outside-toplevel
+ from tvm import relay
Review comment:
Is any particular reason to do late import?
##########
File path: python/tvm/auto_scheduler/dispatcher.py
##########
@@ -0,0 +1,288 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+The global context that dispatches best configurations to workloads
+"""
+# pylint: disable=invalid-name
+
+from __future__ import absolute_import as _abs
+
+import logging
+from pathlib import Path
+
+import numpy as np
+
+from tvm.tir.expr import FloatImm
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class DispatchContext(object):
+ """
+ Base class of dispatch context.
+ """
+
+ current = None
+
+ def __init__(self):
+ self._old_ctx = DispatchContext.current
+
+ def query(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ If cannot find the result inside this context, this function will query it
+ from the upper contexts.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : str
+ The current workload
+
+ Returns
+ -------
+ cfg : State
+ The schedule configuration for the workload
+ """
+ ret = self._query_inside(target, workload)
+ if ret is None:
+ ret = self._old_ctx.query(target, workload)
+ return ret
+
+ def update(self, target, workload, state):
+ """
+ Update the config for a workload
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+ state : State
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def _query_inside(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ This function only query config inside this context.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+
+ Returns
+ -------
+ cfg : State or str
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def __enter__(self):
+ self._old_ctx = DispatchContext.current
+ DispatchContext.current = self
+ return self
+
+ def __exit__(self, ptype, value, trace):
+ DispatchContext.current = self._old_ctx
+
+
+class ApplyHistoryBest(DispatchContext):
+ """
+ Apply the history best config
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+
+ def __init__(self, records, n_lines=None):
+ super(ApplyHistoryBest, self).__init__()
+
+ self.best_by_targetkey = {}
+ self.best_by_model = {}
+ self._best_user_defined = {}
+
+ if records:
+ self.load(records, n_lines)
+
+ def load(self, records, n_lines=None):
+ """Load records to this dispatch context
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+ # pylint: disable=import-outside-toplevel
+ from . import load_from_file
+
+ if isinstance(records, Path):
+ records = str(records)
+
+ if isinstance(records, str):
+ records = load_from_file(records)
+ if not records:
+ return
+
+ best_by_targetkey = self.best_by_targetkey
+ best_by_model = self.best_by_model
+
+ counter = 0
+ for inp, res in records:
+ if n_lines is not None and counter >= n_lines:
+ break
+ counter += 1
+ if res.error_no != 0:
+ continue
+
+ # use target keys in tvm target system as key to build best map
+ for k in inp.task.target.keys:
+ key = (k, inp.task.workload_key)
+ if key not in best_by_targetkey:
+ best_by_targetkey[key] = (inp, res)
+ else:
+ _, other_res = best_by_targetkey[key]
+ other_costs = [x.value for x in other_res.costs if isinstance(x, FloatImm)]
+ costs = [x.value for x in res.costs if isinstance(x, FloatImm)]
+ if np.mean(other_costs) > np.mean(costs):
Review comment:
Should we use mean or median?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] merrymercy commented on pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
merrymercy commented on pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#issuecomment-716901922
@comaniac I only added conv2d_nhwc and dense in this PR for testing the functionality.
I will send another PR with more ops and a tutorial on tuning networks for CUDA.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] merrymercy commented on a change in pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
merrymercy commented on a change in pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#discussion_r511990647
##########
File path: python/tvm/auto_scheduler/dispatcher.py
##########
@@ -0,0 +1,288 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+The global context that dispatches best configurations to workloads
+"""
+# pylint: disable=invalid-name
+
+from __future__ import absolute_import as _abs
+
+import logging
+from pathlib import Path
+
+import numpy as np
+
+from tvm.tir.expr import FloatImm
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class DispatchContext(object):
+ """
+ Base class of dispatch context.
+ """
+
+ current = None
+
+ def __init__(self):
+ self._old_ctx = DispatchContext.current
+
+ def query(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ If cannot find the result inside this context, this function will query it
+ from the upper contexts.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : str
+ The current workload
+
+ Returns
+ -------
+ cfg : State
+ The schedule configuration for the workload
+ """
+ ret = self._query_inside(target, workload)
+ if ret is None:
+ ret = self._old_ctx.query(target, workload)
+ return ret
+
+ def update(self, target, workload, state):
+ """
+ Update the config for a workload
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+ state : State
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def _query_inside(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ This function only query config inside this context.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+
+ Returns
+ -------
+ cfg : State or str
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def __enter__(self):
+ self._old_ctx = DispatchContext.current
+ DispatchContext.current = self
+ return self
+
+ def __exit__(self, ptype, value, trace):
+ DispatchContext.current = self._old_ctx
+
+
+class ApplyHistoryBest(DispatchContext):
+ """
+ Apply the history best config
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+
+ def __init__(self, records, n_lines=None):
+ super(ApplyHistoryBest, self).__init__()
+
+ self.best_by_targetkey = {}
+ self.best_by_model = {}
+ self._best_user_defined = {}
+
+ if records:
+ self.load(records, n_lines)
+
+ def load(self, records, n_lines=None):
+ """Load records to this dispatch context
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+ # pylint: disable=import-outside-toplevel
+ from . import load_from_file
+
+ if isinstance(records, Path):
+ records = str(records)
+
+ if isinstance(records, str):
+ records = load_from_file(records)
+ if not records:
+ return
+
+ best_by_targetkey = self.best_by_targetkey
+ best_by_model = self.best_by_model
+
+ counter = 0
+ for inp, res in records:
+ if n_lines is not None and counter >= n_lines:
+ break
+ counter += 1
+ if res.error_no != 0:
+ continue
+
+ # use target keys in tvm target system as key to build best map
+ for k in inp.task.target.keys:
+ key = (k, inp.task.workload_key)
+ if key not in best_by_targetkey:
+ best_by_targetkey[key] = (inp, res)
+ else:
+ _, other_res = best_by_targetkey[key]
+ other_costs = [x.value for x in other_res.costs if isinstance(x, FloatImm)]
+ costs = [x.value for x in res.costs if isinstance(x, FloatImm)]
+ if np.mean(other_costs) > np.mean(costs):
Review comment:
All other parts of code use mean, so we use mean here to keep consistency.
We can send another PR to do replacement in all placeces.
##########
File path: python/tvm/auto_scheduler/dispatcher.py
##########
@@ -0,0 +1,288 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+The global context that dispatches best configurations to workloads
+"""
+# pylint: disable=invalid-name
+
+from __future__ import absolute_import as _abs
+
+import logging
+from pathlib import Path
+
+import numpy as np
+
+from tvm.tir.expr import FloatImm
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class DispatchContext(object):
+ """
+ Base class of dispatch context.
+ """
+
+ current = None
+
+ def __init__(self):
+ self._old_ctx = DispatchContext.current
+
+ def query(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ If cannot find the result inside this context, this function will query it
+ from the upper contexts.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : str
+ The current workload
+
+ Returns
+ -------
+ cfg : State
+ The schedule configuration for the workload
+ """
+ ret = self._query_inside(target, workload)
+ if ret is None:
+ ret = self._old_ctx.query(target, workload)
+ return ret
+
+ def update(self, target, workload, state):
+ """
+ Update the config for a workload
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+ state : State
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def _query_inside(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ This function only query config inside this context.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+
+ Returns
+ -------
+ cfg : State or str
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def __enter__(self):
+ self._old_ctx = DispatchContext.current
+ DispatchContext.current = self
+ return self
+
+ def __exit__(self, ptype, value, trace):
+ DispatchContext.current = self._old_ctx
+
+
+class ApplyHistoryBest(DispatchContext):
+ """
+ Apply the history best config
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+
+ def __init__(self, records, n_lines=None):
+ super(ApplyHistoryBest, self).__init__()
+
+ self.best_by_targetkey = {}
+ self.best_by_model = {}
+ self._best_user_defined = {}
+
+ if records:
+ self.load(records, n_lines)
+
+ def load(self, records, n_lines=None):
+ """Load records to this dispatch context
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+ # pylint: disable=import-outside-toplevel
+ from . import load_from_file
+
+ if isinstance(records, Path):
+ records = str(records)
+
+ if isinstance(records, str):
+ records = load_from_file(records)
+ if not records:
+ return
+
+ best_by_targetkey = self.best_by_targetkey
+ best_by_model = self.best_by_model
+
+ counter = 0
+ for inp, res in records:
+ if n_lines is not None and counter >= n_lines:
+ break
+ counter += 1
+ if res.error_no != 0:
+ continue
+
+ # use target keys in tvm target system as key to build best map
+ for k in inp.task.target.keys:
+ key = (k, inp.task.workload_key)
+ if key not in best_by_targetkey:
+ best_by_targetkey[key] = (inp, res)
+ else:
+ _, other_res = best_by_targetkey[key]
+ other_costs = [x.value for x in other_res.costs if isinstance(x, FloatImm)]
+ costs = [x.value for x in res.costs if isinstance(x, FloatImm)]
+ if np.mean(other_costs) > np.mean(costs):
Review comment:
All other parts of code use mean, so we use mean here to keep consistency.
We can send another PR to do the replacement in all places.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] masahi commented on a change in pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#discussion_r511300388
##########
File path: python/tvm/auto_scheduler/relay_integration.py
##########
@@ -0,0 +1,179 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=unused-variable,invalid-name
+
+"""
+Integrate auto_scheduler into relay. It implements the following items:
+1. Extract search tasks from a relay program
+2. Provide auto-scheduling for all TOPI compute functions
+"""
+
+import os
+import threading
+
+import tvm
+from tvm import te, transform
+from tvm.te.tensor import PlaceholderOp, ComputeOp
+from .dispatcher import DispatchContext, FallbackConfig
+from .workload_registry import register_workload_tensors
+from .compute_dag import ComputeDAG
+
+
+def call_all_topi_funcs(mod, params, target):
+ """Call all TOPI compute + schedule to extract tasks in a relay program"""
+ # pylint: disable=import-outside-toplevel
+ from tvm import relay
+ from tvm.relay.backend import graph_runtime_codegen
+
+ with transform.PassContext(opt_level=3):
+ opt_mod, _ = relay.optimize(mod, target, params)
+ grc = graph_runtime_codegen.GraphRuntimeCodegen(None, target)
+ grc.codegen(opt_mod["main"])
+
+
+def extract_tasks(mod, params, target):
+ """Extract tuning tasks from a relay program.
+
+ Parameters
+ ----------
+ mod: tvm.IRModule or relay.function.Function
+ The module or function to tune
+ params: dict of str to numpy array
+ The associated parameters of the program
+ target: tvm.target.Target
+ The compilation target
+
+ Returns
+ -------
+ wkl_keys: List[str]
+ The hash key of extracted workloads
+ wkl_weights: List[int]
+ The weight (i.e. the number of appearance) of extracted workload
+ """
+ # pylint: disable=import-outside-toplevel
+ from tvm import relay
+
+ env = TracingEnvironment()
+ with env:
+ # Run the compiler to collect all TOPI calls during compilation.
+
+ # Wrap build call in a new thread to avoid the conflict
+ # between python's multiprocessing and tvm's thread pool
+ build_thread = threading.Thread(target=call_all_topi_funcs, args=(mod, params, target))
+ build_thread.start()
+ build_thread.join()
+ relay.backend.compile_engine.get().clear()
+
+ # create tasks for target
+ wkl_keys = []
+ wkl_weights = []
+ for wkl_key, wkl_weight in env.wkl_key_collection.items():
+ wkl_keys.append(wkl_key)
+ wkl_weights.append(wkl_weight)
+
+ return wkl_keys, wkl_weights
+
+
+class TracingEnvironment:
+ """Global environment for tracing all topi function calls"""
+
+ current = None
+
+ def __init__(self):
+ self.relay_disable_build_cache = "false"
+ self.layout_rewrite_success_ct = 0
+ self.wkl_key_collection = {}
+
+ def __enter__(self):
+ self.relay_disable_build_cache = os.environ.get("TVM_RELAY_DISABLE_BUILD_CACHE", "false")
+ os.environ["TVM_RELAY_DISABLE_BUILD_CACHE"] = "true"
+ TracingEnvironment.current = self
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ os.environ["TVM_RELAY_DISABLE_BUILD_CACHE"] = self.relay_disable_build_cache
+ TracingEnvironment.current = None
+
+ def add_workload_key(self, key):
+ """Add the workload key of an Ansor search task
+
+ Parameters
+ ----------
+ key: str
+ The workload key of a task
+ """
+ if key in self.wkl_key_collection:
+ self.wkl_key_collection[key] += 1
+ else:
+ self.wkl_key_collection[key] = 1
+
+
+def traverse_to_get_io_tensors(outs):
+ """Traverse from a list of output tensors to get a whole computational DAG"""
+ inputs = []
+
+ visited = set()
+
+ def traverse(t):
+ if t in visited:
+ return
+ if isinstance(t.op, PlaceholderOp):
+ inputs.append(t)
+ elif isinstance(t.op, ComputeOp):
+ for x in t.op.input_tensors:
+ traverse(x)
+ visited.add(t)
+
+ for t in outs:
+ traverse(t)
+
+ return inputs + list(outs)
+
+
+# The suffix of implementations that use the auto-scheduler in the OpStrategy.
+auto_schedule_impl_suffix = ".auto_scheduler"
+
+
+def auto_schedule_topi(outs):
Review comment:
Is it better to call it `auto_schedule_te` (and replace all `topi` in this file with `te`)? I imagine there are cases where inputs to this function don't come from topi, such as TE from gradient ops generated by Relay.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] merrymercy commented on a change in pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
merrymercy commented on a change in pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#discussion_r511954470
##########
File path: python/tvm/auto_scheduler/relay_integration.py
##########
@@ -0,0 +1,179 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=unused-variable,invalid-name
+
+"""
+Integrate auto_scheduler into relay. It implements the following items:
+1. Extract search tasks from a relay program
+2. Provide auto-scheduling for all TOPI compute functions
+"""
+
+import os
+import threading
+
+import tvm
+from tvm import te, transform
+from tvm.te.tensor import PlaceholderOp, ComputeOp
+from .dispatcher import DispatchContext, FallbackConfig
+from .workload_registry import register_workload_tensors
+from .compute_dag import ComputeDAG
+
+
+def call_all_topi_funcs(mod, params, target):
+ """Call all TOPI compute + schedule to extract tasks in a relay program"""
+ # pylint: disable=import-outside-toplevel
+ from tvm import relay
+ from tvm.relay.backend import graph_runtime_codegen
+
+ with transform.PassContext(opt_level=3):
+ opt_mod, _ = relay.optimize(mod, target, params)
+ grc = graph_runtime_codegen.GraphRuntimeCodegen(None, target)
+ grc.codegen(opt_mod["main"])
+
+
+def extract_tasks(mod, params, target):
+ """Extract tuning tasks from a relay program.
+
+ Parameters
+ ----------
+ mod: tvm.IRModule or relay.function.Function
+ The module or function to tune
+ params: dict of str to numpy array
+ The associated parameters of the program
+ target: tvm.target.Target
+ The compilation target
+
+ Returns
+ -------
+ wkl_keys: List[str]
+ The hash key of extracted workloads
+ wkl_weights: List[int]
+ The weight (i.e. the number of appearance) of extracted workload
+ """
+ # pylint: disable=import-outside-toplevel
+ from tvm import relay
Review comment:
Otherwise, it will lead to mutual import.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] comaniac commented on a change in pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#discussion_r512258985
##########
File path: python/tvm/auto_scheduler/dispatcher.py
##########
@@ -0,0 +1,288 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+The global context that dispatches best configurations to workloads
+"""
+# pylint: disable=invalid-name
+
+from __future__ import absolute_import as _abs
+
+import logging
+from pathlib import Path
+
+import numpy as np
+
+from tvm.tir.expr import FloatImm
+
+logger = logging.getLogger("auto_scheduler")
+
+
+class DispatchContext(object):
+ """
+ Base class of dispatch context.
+ """
+
+ current = None
+
+ def __init__(self):
+ self._old_ctx = DispatchContext.current
+
+ def query(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ If cannot find the result inside this context, this function will query it
+ from the upper contexts.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : str
+ The current workload
+
+ Returns
+ -------
+ cfg : State
+ The schedule configuration for the workload
+ """
+ ret = self._query_inside(target, workload)
+ if ret is None:
+ ret = self._old_ctx.query(target, workload)
+ return ret
+
+ def update(self, target, workload, state):
+ """
+ Update the config for a workload
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+ state : State
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def _query_inside(self, target, workload):
+ """
+ Query the context to get the specific config for a workload.
+ This function only query config inside this context.
+
+ Parameters
+ ----------
+ target: Target
+ The current target
+ workload : Workload
+ The current workload.
+
+ Returns
+ -------
+ cfg : State or str
+ The schedule configuration for the workload
+ """
+ raise NotImplementedError()
+
+ def __enter__(self):
+ self._old_ctx = DispatchContext.current
+ DispatchContext.current = self
+ return self
+
+ def __exit__(self, ptype, value, trace):
+ DispatchContext.current = self._old_ctx
+
+
+class ApplyHistoryBest(DispatchContext):
+ """
+ Apply the history best config
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+
+ def __init__(self, records, n_lines=None):
+ super(ApplyHistoryBest, self).__init__()
+
+ self.best_by_targetkey = {}
+ self.best_by_model = {}
+ self._best_user_defined = {}
+
+ if records:
+ self.load(records, n_lines)
+
+ def load(self, records, n_lines=None):
+ """Load records to this dispatch context
+
+ Parameters
+ ----------
+ records : str or iterator of (MeasureInput, MeasureResult)
+ Collection of tuning records.
+ If is str, then it should be the filename of a records log file.
+ Each row of this file is an encoded record pair.
+ Otherwise, it is an iterator.
+ n_lines: int (optional)
+ if it is not None, only load the first `n_lines` lines of log
+ """
+ # pylint: disable=import-outside-toplevel
+ from . import load_from_file
+
+ if isinstance(records, Path):
+ records = str(records)
+
+ if isinstance(records, str):
+ records = load_from_file(records)
+ if not records:
+ return
+
+ best_by_targetkey = self.best_by_targetkey
+ best_by_model = self.best_by_model
+
+ counter = 0
+ for inp, res in records:
+ if n_lines is not None and counter >= n_lines:
+ break
+ counter += 1
+ if res.error_no != 0:
+ continue
+
+ # use target keys in tvm target system as key to build best map
+ for k in inp.task.target.keys:
+ key = (k, inp.task.workload_key)
+ if key not in best_by_targetkey:
+ best_by_targetkey[key] = (inp, res)
+ else:
+ _, other_res = best_by_targetkey[key]
+ other_costs = [x.value for x in other_res.costs if isinstance(x, FloatImm)]
+ costs = [x.value for x in res.costs if isinstance(x, FloatImm)]
+ if np.mean(other_costs) > np.mean(costs):
Review comment:
Make sense. I'll help create an issue in the Ansor repo as a reminder.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [incubator-tvm] FrozenGene commented on a change in pull request #6710: [AutoScheduler] Relay integration : Task extraction
Posted by GitBox <gi...@apache.org>.
FrozenGene commented on a change in pull request #6710:
URL: https://github.com/apache/incubator-tvm/pull/6710#discussion_r514004465
##########
File path: python/tvm/auto_scheduler/utils.py
##########
@@ -157,6 +157,9 @@ def __reduce__(self):
def kill_child_processes(parent_pid, sig=signal.SIGTERM):
"""kill all child processes recursively"""
+ if not psutil:
+ raise ImportError("psutil not found, try `pip install psutil` to fix this")
Review comment:
Ah...you move it here. But i am curious why we should need it here. For testing?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org