From 52d4cb8d639285e0793e9c344f20ffd1e13b40b9 Mon Sep 17 00:00:00 2001 From: Leandro Nunes Date: Fri, 19 Feb 2021 23:37:54 +0000 Subject: [PATCH] [TVMC] Add composite target passes for compilation and tuning (#7304) * Extend --target syntax to cover multiple targets for compilation and tuning * Add a new composite_target module to implement custom codegen passes into TVMC * Provide implementation to integrate TVMC, to target Arm Ethos-N NPU and Compute Library for the Arm Architecture (ACL) Change-Id: Iaee53fe22f0c14eb4e4c8ec47e72bade0c5e32cc --- python/tvm/driver/tvmc/autotuner.py | 9 +- python/tvm/driver/tvmc/common.py | 188 +++++++++++++++++- python/tvm/driver/tvmc/compiler.py | 23 ++- python/tvm/driver/tvmc/composite_target.py | 68 +++++++ python/tvm/relay/op/contrib/ethosn.py | 35 ++++ tests/python/driver/tvmc/test_common.py | 91 +++++++++ tests/python/driver/tvmc/test_compiler.py | 47 ++++- .../driver/tvmc/test_composite_target.py | 62 ++++++ 8 files changed, 506 insertions(+), 17 deletions(-) create mode 100644 python/tvm/driver/tvmc/composite_target.py create mode 100644 tests/python/driver/tvmc/test_composite_target.py diff --git a/python/tvm/driver/tvmc/autotuner.py b/python/tvm/driver/tvmc/autotuner.py index fe5bebcabcbc..187b7c5d2a31 100644 --- a/python/tvm/driver/tvmc/autotuner.py +++ b/python/tvm/driver/tvmc/autotuner.py @@ -29,7 +29,7 @@ from tvm.autotvm.tuner import RandomTuner from tvm.autotvm.tuner import XGBTuner -from . import common, frontends +from . import common, composite_target, frontends from .common import TVMCException from .main import register_parser @@ -241,9 +241,14 @@ def drive_tune(args): "need to provide an RPC tracker key (--rpc-key) for remote tuning" ) - target = common.target_from_cli(args.target) + target, extra_targets = common.target_from_cli(args.target) mod, params = frontends.load_model(args.FILE, args.model_format, shape_dict=args.input_shapes) + for codegen_from_cli in extra_targets: + codegen = composite_target.get_codegen_by_target(codegen_from_cli["name"]) + partition_function = codegen["pass_pipeline"] + mod = partition_function(mod, params) + # min_repeat_ms should be: # a. the value provided by the user, if any, or # b. 0ms in case target is "cpu"; otherwise 1000ms diff --git a/python/tvm/driver/tvmc/common.py b/python/tvm/driver/tvmc/common.py index 1845915bcbd1..71bf42ae1e5c 100644 --- a/python/tvm/driver/tvmc/common.py +++ b/python/tvm/driver/tvmc/common.py @@ -18,6 +18,7 @@ Common utility functions shared by TVMC modules. """ import re +import json import logging import os.path import argparse @@ -78,6 +79,168 @@ def convert_graph_layout(mod, desired_layout): ) +def validate_targets(parse_targets): + """ + Apply a series of validations in the targets provided via CLI. + """ + tvm_target_kinds = tvm.target.Target.list_kinds() + targets = [t["name"] for t in parse_targets] + + if len(targets) > len(set(targets)): + raise TVMCException("Duplicate target definitions are not allowed") + + if targets[-1] not in tvm_target_kinds: + tvm_target_names = ", ".join(tvm_target_kinds) + raise TVMCException( + f"The last target needs to be a TVM target. Choices: {tvm_target_names}" + ) + + tvm_targets = [t for t in targets if t in tvm_target_kinds] + if len(tvm_targets) > 1: + verbose_tvm_targets = ", ".join(tvm_targets) + raise TVMCException( + f"Only one of the following targets can be used at a time. " + "Found: {verbose_tvm_targets}." + ) + + +def tokenize_target(target): + """ + Extract a list of tokens from a target specification text. + + It covers some corner-cases that are not covered by the built-in + module 'shlex', such as the use of "+" as a punctuation character. + + + Example + ------- + + For the input `foo -op1=v1 -op2="v ,2", bar -op3=v-4` we + should obtain: + + ["foo", "-op1=v1", "-op2="v ,2"", ",", "bar", "-op3=v-4"] + + Parameters + ---------- + target : str + Target options sent via CLI arguments + + Returns + ------- + list of str + a list of parsed tokens extracted from the target string + """ + + target_pattern = ( + r"(\-{0,2}[\w\-]+\=?" + r"(?:[\w\+\-]+(?:,[\w\+\-])*|[\'][\w\+\-,\s]+[\']|[\"][\w\+\-,\s]+[\"])*|,)" + ) + + return re.findall(target_pattern, target) + + +def parse_target(target): + """ + Parse a plain string of targets provided via a command-line + argument. + + To send more than one codegen, a comma-separated list + is expected. Options start with -=. + + We use python standard library 'shlex' to parse the argument in + a POSIX compatible way, so that if options are defined as + strings with spaces or commas, for example, this is considered + and parsed accordingly. + + + Example + ------- + + For the input `--target="foo -op1=v1 -op2="v ,2", bar -op3=v-4"` we + should obtain: + + [ + { + name: "foo", + opts: {"op1":"v1", "op2":"v ,2"}, + raw: 'foo -op1=v1 -op2="v ,2"' + }, + { + name: "bar", + opts: {"op3":"v-4"}, + raw: 'bar -op3=v-4' + } + ] + + Parameters + ---------- + target : str + Target options sent via CLI arguments + + Returns + ------- + codegens : list of dict + This list preserves the order in which codegens were + provided via command line. Each Dict contains three keys: + 'name', containing the name of the codegen; 'opts' containing + a key-value for all options passed via CLI; 'raw', + containing the plain string for this codegen + """ + codegens = [] + + parsed_tokens = tokenize_target(target) + + split_codegens = [] + current_codegen = [] + split_codegens.append(current_codegen) + for token in parsed_tokens: + # every time there is a comma separating + # two codegen definitions, prepare for + # a new codegen + if token == ",": + current_codegen = [] + split_codegens.append(current_codegen) + else: + # collect a new token for the current + # codegen being parsed + current_codegen.append(token) + + # at this point we have a list of lists, + # each item on the first list is a codegen definition + # in the comma-separated values + for codegen_def in split_codegens: + # the first is expected to be the name + name = codegen_def[0] + raw_target = " ".join(codegen_def) + all_opts = codegen_def[1:] if len(codegen_def) > 1 else [] + opts = {} + for opt in all_opts: + try: + # deal with -- prefixed flags + if opt.startswith("--"): + opt_name = opt[2:] + opt_value = True + else: + opt = opt[1:] if opt.startswith("-") else opt + opt_name, opt_value = opt.split("=", maxsplit=1) + except ValueError: + raise ValueError(f"Error when parsing '{opt}'") + + opts[opt_name] = opt_value + + codegens.append({"name": name, "opts": opts, "raw": raw_target}) + + return codegens + + +def is_inline_json(target): + try: + json.loads(target) + return True + except json.decoder.JSONDecodeError: + return False + + def target_from_cli(target): """ Create a tvm.target.Target instance from a @@ -93,18 +256,33 @@ def target_from_cli(target): ------- tvm.target.Target an instance of target device information + extra_targets : list of dict + This list preserves the order in which extra targets were + provided via command line. Each Dict contains three keys: + 'name', containing the name of the codegen; 'opts' containing + a key-value for all options passed via CLI; 'raw', + containing the plain string for this codegen """ + extra_targets = [] if os.path.exists(target): with open(target) as target_file: - logger.info("using target input from file: %s", target) + logger.debug("target input is a path: %s", target) target = "".join(target_file.readlines()) + elif is_inline_json(target): + logger.debug("target input is inline JSON: %s", target) + else: + logger.debug("target input is plain text: %s", target) + try: + parsed_targets = parse_target(target) + except ValueError as ex: + raise TVMCException(f"Error parsing target string '{target}'.\nThe error was: {ex}") - # TODO(@leandron) We don't have an API to collect a list of supported - # targets yet - logger.debug("creating target from input: %s", target) + validate_targets(parsed_targets) + target = parsed_targets[-1]["raw"] + extra_targets = parsed_targets[:-1] if len(parsed_targets) > 1 else [] - return tvm.target.Target(target) + return tvm.target.Target(target), extra_targets def tracker_host_port_from_cli(rpc_tracker_str): diff --git a/python/tvm/driver/tvmc/compiler.py b/python/tvm/driver/tvmc/compiler.py index 282ae6a76b56..fc1805ee0ab4 100644 --- a/python/tvm/driver/tvmc/compiler.py +++ b/python/tvm/driver/tvmc/compiler.py @@ -28,7 +28,7 @@ from tvm.contrib import cc from tvm.contrib import utils -from . import common, frontends +from . import common, composite_target, frontends from .main import register_parser @@ -72,7 +72,7 @@ def add_compile_parser(subparsers): ) parser.add_argument( "--target", - help="compilation target as plain string, inline JSON or path to a JSON file", + help="compilation targets as comma separated string, inline JSON or path to a JSON file.", required=True, ) parser.add_argument( @@ -185,13 +185,21 @@ def compile_model( """ dump_code = [x.strip() for x in dump_code.split(",")] if dump_code else None mod, params = frontends.load_model(path, model_format, shape_dict) + config = {} if alter_layout: mod = common.convert_graph_layout(mod, alter_layout) - tvm_target = common.target_from_cli(target) + tvm_target, extra_targets = common.target_from_cli(target) target_host = tvm_target if not target_host else target_host + for codegen_from_cli in extra_targets: + codegen = composite_target.get_codegen_by_target(codegen_from_cli["name"]) + partition_function = codegen["pass_pipeline"] + mod = partition_function(mod, params) + if codegen["config_key"] is not None: + config[codegen["config_key"]] = codegen_from_cli["opts"] + if tuning_records and os.path.exists(tuning_records): logger.debug("tuning records file provided: %s", tuning_records) @@ -203,22 +211,21 @@ def compile_model( if use_autoscheduler: with auto_scheduler.ApplyHistoryBest(tuning_records): - with tvm.transform.PassContext( - opt_level=3, config={"relay.backend.use_auto_scheduler": True} - ): + config["relay.backend.use_auto_scheduler"] = True + with tvm.transform.PassContext(opt_level=3, config=config): logger.debug("building relay graph with autoscheduler") graph_module = relay.build( mod, target=target, params=params, target_host=target_host ) else: with autotvm.apply_history_best(tuning_records): - with tvm.transform.PassContext(opt_level=3): + with tvm.transform.PassContext(opt_level=3, config=config): logger.debug("building relay graph with tuning records") graph_module = relay.build( mod, tvm_target, params=params, target_host=target_host ) else: - with tvm.transform.PassContext(opt_level=3): + with tvm.transform.PassContext(opt_level=3, config=config): logger.debug("building relay graph (no tuning records provided)") graph_module = relay.build(mod, tvm_target, params=params, target_host=target_host) diff --git a/python/tvm/driver/tvmc/composite_target.py b/python/tvm/driver/tvmc/composite_target.py new file mode 100644 index 000000000000..7c08994e0e75 --- /dev/null +++ b/python/tvm/driver/tvmc/composite_target.py @@ -0,0 +1,68 @@ +# 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. +""" +Provides support to composite target on TVMC. +""" +import logging + +from tvm.relay.op.contrib.arm_compute_lib import partition_for_arm_compute_lib +from tvm.relay.op.contrib.ethosn import partition_for_ethosn + +from .common import TVMCException + + +# pylint: disable=invalid-name +logger = logging.getLogger("TVMC") + +# Global dictionary to map targets with the configuration key +# to be used in the PassContext (if any), and a function +# responsible for partitioning to that target. +REGISTERED_CODEGEN = { + "acl": { + "config_key": None, + "pass_pipeline": partition_for_arm_compute_lib, + }, + "ethos-n77": { + "config_key": "relay.ext.ethos-n.options", + "pass_pipeline": partition_for_ethosn, + }, +} + + +def get_codegen_names(): + """Return a list of all registered codegens. + + Returns + ------- + list of str + all registered targets + """ + return list(REGISTERED_CODEGEN.keys()) + + +def get_codegen_by_target(name): + """Return a codegen entry by name. + + Returns + ------- + dict + requested target information + """ + try: + return REGISTERED_CODEGEN[name] + except KeyError: + raise TVMCException("Composite target %s is not defined in TVMC." % name) diff --git a/python/tvm/relay/op/contrib/ethosn.py b/python/tvm/relay/op/contrib/ethosn.py index 3a05011242e7..478a1ec46f26 100644 --- a/python/tvm/relay/op/contrib/ethosn.py +++ b/python/tvm/relay/op/contrib/ethosn.py @@ -17,7 +17,11 @@ # pylint: disable=invalid-name, unused-argument """Arm(R) Ethos(TM) -N NPU supported operators.""" from enum import Enum + import tvm.ir +from tvm.relay import transform +from tvm.relay.build_module import bind_params_by_name + from ...dataflow_pattern import wildcard, is_op, is_constant from ... import qnn as _qnn from .register import register_pattern_table @@ -42,6 +46,37 @@ def ethosn_available(): return Available.SW_AND_HW if hw else Available.SW_ONLY +def partition_for_ethosn(mod, params=None): + """Partition the graph greedily offloading supported + operators to Arm Ethos-N NPU. + + Parameters + ---------- + mod : Module + The module to run passes on. + params : Optional[Dict[str, NDArray]] + Constant input parameters. + + Returns + ------- + ret : annotated and partitioned module. + """ + if params: + mod["main"] = bind_params_by_name(mod["main"], params) + + seq = tvm.transform.Sequential( + [ + transform.InferType(), + transform.MergeComposite(pattern_table()), + transform.AnnotateTarget("ethos-n"), + transform.MergeCompilerRegions(), + transform.PartitionGraph(), + ] + ) + + return seq(mod) + + @register_pattern_table("ethos-n") def pattern_table(): """Get the Ethos-N compiler pattern table.""" diff --git a/tests/python/driver/tvmc/test_common.py b/tests/python/driver/tvmc/test_common.py index f30949b54497..253f32d3f0aa 100644 --- a/tests/python/driver/tvmc/test_common.py +++ b/tests/python/driver/tvmc/test_common.py @@ -24,6 +24,8 @@ from tvm import relay from tvm.driver import tvmc +from tvm.driver.tvmc.common import TVMCException + def test_compile_tflite_module_nhwc_to_nchw(tflite_mobilenet_v1_1_quant): # some CI environments wont offer TFLite, so skip in case it is not present @@ -182,3 +184,92 @@ def test_shape_parser(): shape_string = "input:5,10 input2:10,10" with pytest.raises(argparse.ArgumentTypeError): tvmc.common.parse_shape_string(shape_string) + + +def test_target_from_cli__error_duplicate(): + with pytest.raises(TVMCException): + _ = tvmc.common.target_from_cli("llvm, llvm") + + +def test_target_from_cli__error_target_not_found(): + with pytest.raises(TVMCException): + _ = tvmc.common.target_from_cli("invalidtarget") + + +def test_target_from_cli__error_no_tvm_target(): + with pytest.raises(TVMCException): + _ = tvmc.common.target_from_cli("ethos-n77") + + +def test_tokenize_target_with_opts(): + tokens = tvmc.common.tokenize_target("foo -opt1=value1 --flag, bar -opt2=value2") + expected_tokens = ["foo", "-opt1=value1", "--flag", ",", "bar", "-opt2=value2"] + + assert len(tokens) == len(expected_tokens) + assert tokens == expected_tokens + + +def test_tokenize_target_with_plus_sign(): + tokens = tvmc.common.tokenize_target("foo -opt1=+value1 --flag, bar -opt2=test,+v") + expected_tokens = ["foo", "-opt1=+value1", "--flag", ",", "bar", "-opt2=test,+v"] + + assert len(tokens) == len(expected_tokens) + assert tokens == expected_tokens + + +def test_tokenize_target_with_commas(): + tokens = tvmc.common.tokenize_target("foo -opt1=v,a,l,u,e,1 --flag") + expected_tokens = ["foo", "-opt1=v,a,l,u,e,1", "--flag"] + + assert len(tokens) == len(expected_tokens) + assert tokens == expected_tokens + + +def test_tokenize_target_with_commas_and_single_quotes(): + tokens = tvmc.common.tokenize_target("foo -opt1='v, a, l, u, e', bar") + expected_tokens = ["foo", "-opt1='v, a, l, u, e'", ",", "bar"] + + assert len(tokens) == len(expected_tokens) + assert tokens == expected_tokens + + +def test_tokenize_target_with_commas_and_double_quotes(): + tokens = tvmc.common.tokenize_target('foo -opt1="v, a, l, u, e", bar') + expected_tokens = ["foo", '-opt1="v, a, l, u, e"', ",", "bar"] + + assert len(tokens) == len(expected_tokens) + assert tokens == expected_tokens + + +def test_tokenize_target_with_dashes(): + tokens = tvmc.common.tokenize_target("foo-bar1 -opt-1=t-e-s-t, baz") + expected_tokens = ["foo-bar1", "-opt-1=t-e-s-t", ",", "baz"] + + assert len(tokens) == len(expected_tokens) + assert tokens == expected_tokens + + +def test_parse_single_target_with_opts(): + targets = tvmc.common.parse_target("llvm -device=arm_cpu --system-lib") + + assert len(targets) == 1 + assert "device" in targets[0]["opts"] + assert "system-lib" in targets[0]["opts"] + + +def test_parse_multiple_target(): + targets = tvmc.common.parse_target("acl, llvm -device=arm_cpu --system-lib") + + assert len(targets) == 2 + assert "acl" == targets[0]["name"] + assert "llvm" == targets[1]["name"] + + +def test_parse_multiple_target_with_opts(): + targets = tvmc.common.parse_target("ethos-n77 -myopt=value, llvm -device=arm_cpu --system-lib") + + assert len(targets) == 2 + assert "ethos-n77" == targets[0]["name"] + assert "myopt" in targets[0]["opts"] + assert "value" == targets[0]["opts"]["myopt"] + assert "llvm" == targets[1]["name"] diff --git a/tests/python/driver/tvmc/test_compiler.py b/tests/python/driver/tvmc/test_compiler.py index 4cb342c2e967..ae859298facd 100644 --- a/tests/python/driver/tvmc/test_compiler.py +++ b/tests/python/driver/tvmc/test_compiler.py @@ -19,10 +19,13 @@ import shutil from os import path +from unittest import mock import pytest import tvm +from tvm.relay.op.contrib.ethosn import ethosn_available + from tvm.driver import tvmc @@ -73,7 +76,7 @@ def test_cross_compile_aarch64_tflite_module(tflite_mobilenet_v1_1_quant): graph, lib, params, dumps = tvmc.compiler.compile_model( tflite_mobilenet_v1_1_quant, - target="llvm -device=arm_cpu -mtriple=aarch64-linux-gnu -mattr=+neon", + target="llvm -device=arm_cpu -mtriple=aarch64-linux-gnu -mattr='+neon'", dump_code="asm", ) @@ -110,7 +113,7 @@ def test_cross_compile_aarch64_keras_module(keras_resnet50): graph, lib, params, dumps = tvmc.compiler.compile_model( keras_resnet50, - target="llvm -device=arm_cpu -mtriple=aarch64-linux-gnu -mattr=+neon", + target="llvm -device=arm_cpu -mtriple=aarch64-linux-gnu -mattr='+neon'", dump_code="asm", ) @@ -185,3 +188,43 @@ def test_compile_opencl(tflite_mobilenet_v1_0_25_128): assert type(lib) is tvm.runtime.module.Module assert type(params) is dict assert type(dumps) is dict + + +@pytest.mark.skipif( + not ethosn_available(), + reason="--target=ethos-n77 is not available. TVM built with 'USE_ETHOSN OFF'", +) +def test_compile_tflite_module_with_external_codegen(tflite_mobilenet_v1_1_quant): + pytest.importorskip("tflite") + + graph, lib, params, dumps = tvmc.compiler.compile_model( + tflite_mobilenet_v1_1_quant, target="ethos-n77, llvm", dump_code="relay" + ) + + # check for output types + assert type(graph) is str + assert type(lib) is tvm.runtime.module.Module + assert type(params) is dict + assert type(dumps) is dict + + +@mock.patch("tvm.relay.build") +@mock.patch("tvm.driver.tvmc.composite_target.get_codegen_by_target") +@mock.patch("tvm.driver.tvmc.frontends.load_model") +@mock.patch("tvm.transform.PassContext") +def test_compile_check_configs_composite_target(mock_pc, mock_fe, mock_ct, mock_relay): + mock_codegen = {} + mock_codegen["config_key"] = "relay.ext.mock.options" + mock_codegen["pass_pipeline"] = lambda *args: None + + mock_fe.return_value = (None, None) + mock_ct.return_value = mock_codegen + mock_relay.return_value = mock.MagicMock() + + graph, lib, params, dumps = tvmc.compiler.compile_model( + "no_file_needed", target="mockcodegen -testopt=value, llvm" + ) + + mock_pc.assert_called_once_with( + opt_level=3, config={"relay.ext.mock.options": {"testopt": "value"}} + ) diff --git a/tests/python/driver/tvmc/test_composite_target.py b/tests/python/driver/tvmc/test_composite_target.py new file mode 100644 index 000000000000..eda7cd9224fd --- /dev/null +++ b/tests/python/driver/tvmc/test_composite_target.py @@ -0,0 +1,62 @@ +# 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. +import argparse +import os +import shutil + +from inspect import isfunction +from os import path + +import pytest + +import tvm + +from tvm.driver import tvmc + +from tvm.driver.tvmc.common import TVMCException + + +def test_get_codegen_names(): + names = tvmc.composite_target.get_codegen_names() + + assert "ethos-n77" in names + assert len(names) > 0 + + +def test_valid_codegen(): + codegen = tvmc.composite_target.get_codegen_by_target("acl") + + assert codegen is not None + assert codegen["pass_pipeline"] is not None + + +def test_invalid_codegen(): + with pytest.raises(TVMCException): + _ = tvmc.composite_target.get_codegen_by_target("invalid") + + +def test_all_codegens_contain_pass_pipeline(): + for name in tvmc.composite_target.get_codegen_names(): + codegen = tvmc.composite_target.get_codegen_by_target(name) + assert "pass_pipeline" in codegen, f"{name} does not contain a pass_pipeline" + assert isfunction(codegen["pass_pipeline"]) + + +def test_all_pass_pipelines_are_functions(): + for name in tvmc.composite_target.get_codegen_names(): + codegen = tvmc.composite_target.get_codegen_by_target(name) + assert isfunction(codegen["pass_pipeline"]), f"pass_pipeline for {name} is not a function"