| # 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. |
| """ |
| TopHub: Tensor Operator Hub |
| To get the best performance, we typically need auto-tuning for the specific devices. |
| TVM releases pre-tuned parameters in TopHub for some common networks and hardware targets. |
| TVM will download these parameters for you when you call relay.build. |
| """ |
| # pylint: disable=invalid-name |
| |
| import logging |
| import os |
| import sys |
| |
| from .task import ApplyHistoryBest |
| from ..target import Target |
| from ..contrib.download import download |
| from .record import load_from_file |
| from .util import EmptyContext |
| |
| # environment variable to read TopHub location |
| AUTOTVM_TOPHUB_LOC_VAR = "TOPHUB_LOCATION" |
| |
| # default location of TopHub |
| AUTOTVM_TOPHUB_DEFAULT_LOC = "https://raw.githubusercontent.com/uwsampl/tvm-distro/master/tophub" |
| |
| # value of AUTOTVM_TOPHUB_LOC_VAR to specify to not read from TopHub |
| AUTOTVM_TOPHUB_NONE_LOC = "NONE" |
| |
| # root path to store TopHub files |
| AUTOTVM_TOPHUB_ROOT_PATH = os.path.join(os.path.expanduser("~"), ".tvm", "tophub") |
| |
| # the version of each package |
| PACKAGE_VERSION = { |
| "arm_cpu": "v0.07", |
| "llvm": "v0.04", |
| "cuda": "v0.09", |
| "rocm": "v0.05", |
| "opencl": "v0.04", |
| "mali": "v0.06", |
| "intel_graphics": "v0.02", |
| "vta": "v0.09", |
| "amd_apu": "v0.01", |
| } |
| |
| logger = logging.getLogger("autotvm") |
| |
| |
| def _alias(name): |
| """convert alias for some packages""" |
| table = { |
| "vtacpu": "vta", |
| "metal": "opencl", |
| "webgpu": "opencl", |
| "vulkan": "opencl", |
| "nvptx": "cuda", |
| "amd_apu": "amd_apu", |
| } |
| return table.get(name, name) |
| |
| |
| def _get_tophub_location(): |
| location = os.getenv(AUTOTVM_TOPHUB_LOC_VAR, None) |
| return AUTOTVM_TOPHUB_DEFAULT_LOC if location is None else location |
| |
| |
| def context(target, extra_files=None): |
| """Return the dispatch context with pre-tuned parameters. |
| This function will load the corresponding *.log files in AUTOTVM_TOPHUB_ROOT_PATH. |
| If cannot find them, it will download them from TopHub github repo. |
| Users can also add their own files in argument `extra_files`. |
| |
| Parameters |
| ---------- |
| target: Target or List of Target |
| The compilation target |
| extra_files: list of str, optional |
| Extra log files to load |
| """ |
| tophub_location = _get_tophub_location() |
| if tophub_location == AUTOTVM_TOPHUB_NONE_LOC: |
| return EmptyContext() |
| |
| best_context = ApplyHistoryBest([]) |
| |
| targets = target if isinstance(target, (list, tuple)) else [target] |
| |
| for tgt in targets: |
| if isinstance(tgt, str): |
| tgt = Target(tgt) |
| |
| possible_names = [] |
| device = tgt.attrs.get("device", "") |
| if device != "": |
| possible_names.append(_alias(device)) |
| possible_names.append(tgt.kind.name) |
| |
| all_packages = list(PACKAGE_VERSION.keys()) |
| for name in possible_names: |
| name = _alias(name) |
| if name in all_packages: |
| if not check_backend(tophub_location, name): |
| continue |
| |
| filename = "%s_%s.log" % (name, PACKAGE_VERSION[name]) |
| best_context.load(os.path.join(AUTOTVM_TOPHUB_ROOT_PATH, filename)) |
| break # only load one file to avoid some fallback template mismatch problem |
| |
| if extra_files: |
| for filename in extra_files: |
| best_context.load(filename) |
| |
| return best_context |
| |
| |
| def check_backend(tophub_location, backend): |
| """Check whether have pre-tuned parameters of the certain target. |
| If not, will download it. |
| |
| Parameters |
| ---------- |
| backend: str |
| The name of backend. |
| |
| Returns |
| ---------- |
| success: bool |
| Whether the check is successful. |
| """ |
| backend = _alias(backend) |
| assert backend in PACKAGE_VERSION, 'Cannot find backend "%s" in TopHub' % backend |
| |
| version = PACKAGE_VERSION[backend] |
| package_name = "%s_%s.log" % (backend, version) |
| if os.path.isfile(os.path.join(AUTOTVM_TOPHUB_ROOT_PATH, package_name)): |
| return True |
| |
| # pylint: disable=import-outside-toplevel |
| if sys.version_info >= (3,): |
| import urllib.request as urllib2 |
| else: |
| import urllib2 |
| try: |
| download_package(tophub_location, package_name) |
| return True |
| except urllib2.URLError as e: |
| logging.warning("Failed to download tophub package for %s: %s", backend, e) |
| return False |
| |
| |
| def download_package(tophub_location, package_name): |
| """Download pre-tuned parameters of operators for a backend |
| |
| Parameters |
| ---------- |
| tophub_location: str |
| The location to download TopHub parameters from |
| |
| package_name: str |
| The name of package |
| """ |
| rootpath = AUTOTVM_TOPHUB_ROOT_PATH |
| |
| if not os.path.isdir(rootpath): |
| # make directory |
| splits = os.path.split(rootpath) |
| for j in range(1, len(splits) + 1): |
| path = os.path.join(*splits[:j]) |
| if not os.path.isdir(path): |
| os.mkdir(path) |
| |
| download_url = "{0}/{1}".format(tophub_location, package_name) |
| logger.info("Download pre-tuned parameters package from %s", download_url) |
| download(download_url, os.path.join(rootpath, package_name), True, verbose=0) |
| |
| |
| # global cache for load_reference_log |
| REFERENCE_LOG_CACHE = {} |
| |
| |
| def load_reference_log(backend, model, workload_name): |
| """Load reference log from TopHub to support fallback in template. |
| Template will use these reference logs to choose fallback config. |
| |
| Parameters |
| ---------- |
| backend: str |
| The backend name |
| model: str |
| The name of the device model |
| workload_name: str |
| The name of the workload. (The first item in the workload tuple) |
| """ |
| |
| backend = _alias(backend) |
| version = PACKAGE_VERSION[backend] |
| package_name = "%s_%s.log" % (backend, version) |
| filename = os.path.join(AUTOTVM_TOPHUB_ROOT_PATH, package_name) |
| |
| global REFERENCE_LOG_CACHE |
| key = (backend, model, workload_name) |
| |
| if key not in REFERENCE_LOG_CACHE: |
| tmp = [] |
| # If TOPHUB_LOCATION is not AUTOTVM_TOPHUB_NONE_LOC, |
| # Download the config file from tophub if not exists. |
| if not os.path.exists(filename): |
| tophub_location = _get_tophub_location() |
| if tophub_location != AUTOTVM_TOPHUB_NONE_LOC: |
| download_package(tophub_location, package_name) |
| if os.path.isfile(filename): # in case download failed |
| find = False |
| inp = None |
| counts = {} |
| for inp, res in load_from_file(filename): |
| counts[inp.target.model] = counts.get(inp.target.model, 0) + 1 |
| if model == inp.target.model: |
| find = True |
| break |
| # if device model is not find, use the device model with the most tuned workloads |
| if not find and counts: |
| model = max(counts.items(), key=lambda k: k[1])[0] |
| |
| for inp, res in load_from_file(filename): |
| if model == inp.target.model and inp.task.workload[0] == workload_name: |
| tmp.append((inp, res)) |
| REFERENCE_LOG_CACHE[key] = tmp |
| |
| return REFERENCE_LOG_CACHE[key] |
