From 97aa80beed2bfc8ee9656a88b755dc34988dd322 Mon Sep 17 00:00:00 2001 From: mbs-octoml Date: Thu, 28 Apr 2022 15:57:05 -0700 Subject: [PATCH 01/11] [Relay] Support 'external codegen targets'. (Part of Collage, https://github.com/apache/tvm-rfcs/blob/main/rfcs/0062-collage.md) This change prepares the VM and Relay target handling machinery to support external codegen targets in addition to 'regular' targets. This allows us to configure the build with Collage as follows: ``` host_target = tvm.target.Target("llvm") targets = [tvm.target.Target("cuda", host_target), tvm.target.Target("cutlass", host_target), tvm.target.Target("cudnn", host_target)] with tvm.transform.PassContext(...): exe = tvm.relay.vm.compile(module, target=targets) ``` Four changes are required: 1. I introduce four new target kinds for the external codegens currently supported by Collage. Others can be added as they are vetted for use by Collage. These are given a device type matching the external codegen's assumption (ie just CUDA currently), and given a target kind attribute "is_external_codegen" of True. The latter is needed by Collage to signal the target kind name represents and external codegen 'compiler' name. See the RFC for specifics. 2. I introduce the binary relation Target::IsExternalCodegenFor so that external codegen targets can be related back to the 'underlying' targets they are implicitly using in their codegen. 3. I rework the VMCompiler and BuildModule interfaces to accept an Array of 'raw targets' instead of a Map. This more general representation is needed because we may now have multiple targets of the same device type active simultaneously. I add new static methods on the Python Target to convert to this form in a way that mimics check_and_update_host_consist. 4. I rework CompilationConfig to work from Array directly, to not depend on the host_target argument (since dealt with on the Python side), and to understand that if we have two targets for the same device type the non-external codegen target takes precedence. The change to CompilationConfig seems neutral with respect to the recent discussions on compilation configuration representation and tvmc. I made a few attempts to remove Target.check_and_update_host_const entirely in favor of using CompilationConfig as the definitive target handling choke point but backed out once they became too large. --- cmake/modules/CUDA.cmake | 4 + include/tvm/target/compilation_config.h | 113 ++++---- include/tvm/target/target.h | 30 +- include/tvm/target/target_kind.h | 20 ++ python/tvm/autotvm/tophub.py | 5 +- python/tvm/relay/backend/vm.py | 88 +----- python/tvm/relay/build_module.py | 18 +- python/tvm/target/target.py | 83 +++++- src/relay/backend/aot_executor_codegen.cc | 48 +--- src/relay/backend/build_module.cc | 39 ++- src/relay/backend/contrib/cublas/target.cc | 44 +++ src/relay/backend/contrib/cudnn/target.cc | 42 +++ src/relay/backend/contrib/cutlass/target.cc | 43 +++ src/relay/backend/contrib/tensorrt/target.cc | 42 +++ src/relay/backend/graph_executor_codegen.cc | 31 +-- src/relay/backend/interpreter.cc | 12 +- src/relay/backend/te_compiler.cc | 58 +--- src/relay/backend/te_compiler.h | 20 +- src/relay/backend/utils.cc | 10 +- src/relay/backend/utils.h | 5 +- src/relay/backend/vm/compiler.cc | 69 +++-- src/relay/backend/vm/compiler.h | 35 ++- src/target/compilation_config.cc | 275 ++++++++++--------- src/target/target.cc | 27 +- 24 files changed, 632 insertions(+), 529 deletions(-) create mode 100644 src/relay/backend/contrib/cublas/target.cc create mode 100644 src/relay/backend/contrib/cudnn/target.cc create mode 100644 src/relay/backend/contrib/cutlass/target.cc create mode 100644 src/relay/backend/contrib/tensorrt/target.cc diff --git a/cmake/modules/CUDA.cmake b/cmake/modules/CUDA.cmake index 10117aa12233..056ed18d442e 100644 --- a/cmake/modules/CUDA.cmake +++ b/cmake/modules/CUDA.cmake @@ -41,6 +41,8 @@ if(USE_CUDA) if(USE_CUDNN) message(STATUS "Build with cuDNN support") include_directories(SYSTEM ${CUDA_CUDNN_INCLUDE_DIRS}) + tvm_file_glob(GLOB CUDNN_RELAY_CONTRIB_SRC src/relay/backend/contrib/cudnn/*.cc) + list(APPEND COMPILER_SRCS ${CUDNN_RELAY_CONTRIB_SRC}) tvm_file_glob(GLOB CONTRIB_CUDNN_SRCS src/runtime/contrib/cudnn/*.cc) list(APPEND RUNTIME_SRCS ${CONTRIB_CUDNN_SRCS}) list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUDNN_LIBRARY}) @@ -48,6 +50,8 @@ if(USE_CUDA) if(USE_CUBLAS) message(STATUS "Build with cuBLAS support") + tvm_file_glob(GLOB CUBLAS_RELAY_CONTRIB_SRC src/relay/backend/contrib/cublas/*.cc) + list(APPEND COMPILER_SRCS ${CUBLAS_RELAY_CONTRIB_SRC}) tvm_file_glob(GLOB CONTRIB_CUBLAS_SRCS src/runtime/contrib/cublas/*.cc) list(APPEND RUNTIME_SRCS ${CONTRIB_CUBLAS_SRCS}) list(APPEND TVM_RUNTIME_LINKER_LIBS ${CUDA_CUBLAS_LIBRARY}) diff --git a/include/tvm/target/compilation_config.h b/include/tvm/target/compilation_config.h index 1c47a0f806a3..c2d265596180 100644 --- a/include/tvm/target/compilation_config.h +++ b/include/tvm/target/compilation_config.h @@ -20,7 +20,6 @@ /*! * \file tvm/target/compilation_config.h * \brief A helper class to collect all the targets in canonical form necessary for compilation. - * CAUTION: Preliminary, currently only used to support device planning, very likely to change. */ #ifndef TVM_TARGET_COMPILATION_CONFIG_H_ @@ -32,40 +31,21 @@ namespace tvm { /*! * \brief Gathers the \p Targets and distinguished \p VirtualDevices in canonical form needed to - * compile a Relay module. Centralizes any setup and validation logic needed to transition - * from configuration options conveyed implicitly (eg in \p PassContexts) or explicitly - * (eg a a list of \p Targets) to the configuration. + * compile a Relay module for execution over possibly heterogeneous devices. Centralizes the + * validation and canonicalization logic needed to transition from targets supplied by the Python + * APIs to a single internal representation. Also holds a cache of canonical \p VirtualDevices + * so that structural equal virtual devices have pointer equal canonical virtual devices. * - * CAUTION: This is subject to change as we rework compilation options in general. See - * https://github.com/apache/tvm-rfcs/blob/main/rfcs/0028-command-line-registry-composition.md. - * So far this class is only focussed on carrying just the configuration needed by PlanDevices, - * and removing target-munging code duplication and inconsistencies between the three major build - * flows for the VM (relay/backend/vm/compile.cc), Graph/AOT (relay/backend/build_module.cc) and - * Interpreter (relay/backend/interpreter.cc). Over time we expect more global compiler - * configuration (eg for executor and runtime config, for system memory pool configuration, etc) - * to migrate into this class, and instances thereof to be attached to \p IRModules using a - * well-known attribute. + * TODO(mbs): This is subject to change as we rework compilation options in general. This class + * is probably better called a 'CompositeTarget', and may be better made a sub-class of Target or + * some other common-target-root class. */ class CompilationConfigNode : public Object { public: - /*! - * \brief The legacy targets map, mapping device type to the corresponding \p Target to use - * when compiling primitive functions. Does not include an entry for the host target, however - * each \p Target in this map will have it's \p host field set to the \p host_target. - * - * Currently we require at most one \p Target per \p DLDeviceType, though we want to get rid of - * that limitation. - * - * CAUTION: Since keys are \p Integers they are compared by object equality not integer - * value. - * - * TODO(mbs): Remove once codegen updated for new target conventions. - */ - TargetMap legacy_target_map; - /*! * \brief The host target. Used for 'scalar' data and code (such as shapes and shape * functions) and residual Relay expressions and data (such as conditionals and ADTs). + * Each \p primitive_target below will have this exact target object as its 'host'. * * Note that it is possible for a \p Target used for primitive operations to be structurally * equal to the host \p Target (up to the \p host field.) However the \p Target objects will @@ -74,16 +54,26 @@ class CompilationConfigNode : public Object { Target host_target; /*! - * \brief Vector of all available \p Targets for compiling primitive operators. May contain - * a \p Target for the same device type as for the \p host_target, however the \p host_target - * should be used for all host computations and data. Each \p Target will have \p host_target - * as its host. + * \brief Vector of all available \p Targets for compiling primitive tensor operators (kernels). + * May contain a \p Target for the same device type as for the \p host_target, however the \p + * host_target should be used for all host computations and data. Each \p Target will have \p + * host_target as its 'host'. + * + * It is possible to have multiple primitive targets for the same device type. However given + * primitive targets left and right where: + * - left appears before right in the array + * - left->kind->device_type == right->kind->device_type + * then: + * - right.IsExternalCodegenFor(left) must be true + * In this way the FindPrimitiveTargetOrFail method will find the 'most general' target for + * the requested device type. */ Array primitive_targets; /*! * \brief \p VirtualDevice for primitive operators which are not otherwise constrained to a - * particular device. + * particular device. Used by the PlanDevices pass to determine a virtual device for every + * sub-expression. */ VirtualDevice default_primitive_virtual_device = VirtualDevice::FullyUnconstrained(); @@ -94,25 +84,33 @@ class CompilationConfigNode : public Object { * \brief If defined then compile and/or run in 'homogenous execution mode'. In this mode all * primitives are compiled for this target only. * - * This is to support legacy passes which have not been adapted to hetrogeneous execution and + * This is to support legacy passes which have not been adapted to heterogeneous execution and * rely on an implicit global \p Target to be in scope. * - * TODO(mbs): Remove once all passes are 'hetrogeneous aware'. + * TODO(mbs): Remove once all passes are 'heterogeneous aware'. */ Target optional_homogeneous_target; void VisitAttrs(AttrVisitor* v); + /*! + * \brief Return the unique \p Target to use for \p device_type. Fail if no such target exists. + * + * This will be the first primitive target with matching device type. + */ + Target FindPrimitiveTargetOrFail(DLDeviceType device_type) const; + /*! * \brief Returns a \p VirtualDevice agreeing with \p virtual_device on all its constrained * fields, however: - * - If the target is null then it is filled in from the known available primitive targets by - * matching on device type. Fails if no such target is known. + * - If the target is null then it is filled in using \p FindPrimitiveTargetOrFail to match + * the device type. * - The returned object is unique for the field values w.r.t. all other \p VirtualDevices - * returned by this method. + * returned by this method. * * We call the result the 'canonical' \p VirtualDevice. Two canonical \p VirtualDevices are - * structurally equal if and only if they are pointer equal. + * structurally equal if and only if they are pointer equal. In this way we can build maps + * from virtual devices using just pointer equality. */ VirtualDevice CanonicalVirtualDevice(const VirtualDevice& virtual_device) const; @@ -121,31 +119,20 @@ class CompilationConfigNode : public Object { private: /*! - * \brief Establishes the default \p VirtualDevice for primitives and the \p VirtualDevice for the - * host given: - * - the vector of available primitive \p Targets. - * - any host \p Target. + * \brief Sets the primitive targets, the host target, the default primitive virtual device, and + * the host virtual device given: + * - the vector of 'raw' targets (in any order) supplied by one of the TVM entry points. * - any "relay.fallback_device_type" attribute on \p pass_ctx. * - whether the LLVM backend is available. - * If necessary, creates new default \p Targets to match the required devices. - * - * NOTE: The implementation is a bit convoluted since it tries to maintain backwards - * compatibility with legacy methods for conveying \p Targets. - * - * CAUTION: Recreated the primitive_targets so that they all have the given/constructed - * host_target as their host (cf CheckAndUpdateHostConsistency). + * Will look for a suitable host target in the given primitive targets, but if none found may + * reuse a raw target or create a default CPU target. */ - void EstablishDefaultVirtualDevices(const transform::PassContext& pass_ctx); + void Init(const transform::PassContext& pass_ctx, const Array& raw_targets); /*! - * \brief Returns a freshly constructed \p Target to represent \p device_type. + * \brief Returns a freshly constructed CPU \p Target. */ - static Target MakeDefaultTarget(DLDeviceType device_type); - - /*! - * \brief Return the \p Target to use for \p device_type. Fail if no such target exists. - */ - Target FindPrimitiveTargetOrFail(DLDeviceType device_type) const; + static Target MakeDefaultCPUTarget(); /*! * \brief A cache of constructed virtual devices. @@ -163,13 +150,11 @@ class CompilationConfigNode : public Object { class CompilationConfig : public ObjectRef { public: /*! - * \brief Constructs the compilation config given the available \p Targets in the - * \p legacy_target_map_arg and an optional \p optional_host_target_arg. May use - * 'relay.fallback_device_type' and the availability of the LLVM compilation module - * to decide on appropriate default devices. + * \brief Constructs the compilation config given the settings in \p pass_ctx and supplied + * \p raw_targets. See \p CompilationConfigNode::Init for details. */ - TVM_DLL CompilationConfig(const transform::PassContext& pass_ctx, TargetMap legacy_target_map_arg, - Target optional_host_target_arg); + TVM_DLL CompilationConfig(const transform::PassContext& pass_ctx, + const Array& raw_targets); TVM_DEFINE_OBJECT_REF_METHODS(CompilationConfig, ObjectRef, CompilationConfigNode); }; diff --git a/include/tvm/target/target.h b/include/tvm/target/target.h index 21760bdc8dbf..29103bd20866 100644 --- a/include/tvm/target/target.h +++ b/include/tvm/target/target.h @@ -177,7 +177,26 @@ class Target : public ObjectRef { */ static Target WithHost(const Target& target, const Target& host); + /*! + * \brief Returns true if \p this target represents an external codegen which is compatible + * with \p that target. In particular: + * - \p this has a true ::tvm::attr::kIsExternalCodegen attribute + * - \p that does not have a true ::tvm::attr::kIsExternalCodegen attribute + * - \p this and \p that have the same kind->device_type + * + * After partitioning, the external codegen compilation path may use \p that to guide it's + * compilation to a \p runtime::Module. Given on \p this, an appropriate \p that can be + * found using \p CompilationConfig::FindPrimitiveTargetOrFail(this->kind->device_type). + * + * The \p CollagePartition pass uses this method to guide it's search over candidate partitions + * using external codegen. + */ + bool IsExternalCodegenFor(const Target& that) const; + private: + Target(TargetKind kind, Optional host, String tag, Array keys, + Map attrs); + // enable with syntax. friend class TargetInternal; friend class With; @@ -194,8 +213,6 @@ class Target : public ObjectRef { TVM_DLL void ExitWithScope(); }; -using TargetMap = Map; - /*! * \brief Check and update host field of the given legacy target and target host pair. * Note that this function is for legacy target api compatibility issue only, not @@ -205,15 +222,6 @@ using TargetMap = Map; */ void CheckAndUpdateHostConsistency(Target* target, Target* host); -/*! - * \brief Check and update host field of the given legacy heterogeneous targets and - * target host.Note that this function is for legacy target api compatibility issue only, - * not recommended for other use. - * \param target_map The pointer to a Map objects with values being Target objects - * \param host The Target typed object for target host to be updated - */ -void CheckAndUpdateHostConsistency(TargetMap* target_map, Target* host); - /*! * \brief Check and update host field of the given legacy heterogeneous targets and * target host.Note that this function is for legacy target api compatibility issue only, diff --git a/include/tvm/target/target_kind.h b/include/tvm/target/target_kind.h index e802a3088d2d..395d3aab6757 100644 --- a/include/tvm/target/target_kind.h +++ b/include/tvm/target/target_kind.h @@ -384,6 +384,26 @@ inline TargetKindRegEntry& TargetKindRegEntry::set_name() { #define TVM_TARGET_KIND_REGISTER_VAR_DEF \ static DMLC_ATTRIBUTE_UNUSED ::tvm::TargetKindRegEntry& __make_##TargetKind +namespace attr { +// +// Distinguished TargetKind attribute names. +// + +/*! + * \brief A \p TargetKind attribute of type \p Bool. If true, then the target kind name also + * corresponds to an external codegen 'compiler' name. That name may be used: + * - To retrieve partitioning rules using \p get_partition_table. + * - To attach to Relay Functions under the \p attr::kCompiler attribute to indicate + * the function is to be compiled by the external codegen path. + * + * The \p CollagePartition pass uses this attribute to guide it's search over candidate partitions + * using external codegen. + * + * See also \p Target::IsExternalCodegenFor + */ +constexpr const char* kIsExternalCodegen = "is_external_codegen"; +} // namespace attr + /*! * \def TVM_REGISTER_TARGET_KIND * \brief Register a new target kind, or set attribute of the corresponding target kind. diff --git a/python/tvm/autotvm/tophub.py b/python/tvm/autotvm/tophub.py index f438bc197afe..0a51bb12b2a4 100644 --- a/python/tvm/autotvm/tophub.py +++ b/python/tvm/autotvm/tophub.py @@ -26,6 +26,7 @@ from os import getenv import sys from pathlib import Path +from tvm.ir.container import Array from .task import ApplyHistoryBest from ..target import Target @@ -87,7 +88,7 @@ def context(target, extra_files=None): Parameters ---------- target: Target or List of Target - The compilation target + The compilation targets extra_files: list of str, optional Extra log files to load """ @@ -97,7 +98,7 @@ def context(target, extra_files=None): best_context = ApplyHistoryBest([]) - targets = target if isinstance(target, (list, tuple)) else [target] + targets = target if isinstance(target, (Array, list, tuple)) else [target] for tgt in targets: if isinstance(tgt, str): diff --git a/python/tvm/relay/backend/vm.py b/python/tvm/relay/backend/vm.py index 25744408d87b..72613421f11f 100644 --- a/python/tvm/relay/backend/vm.py +++ b/python/tvm/relay/backend/vm.py @@ -65,18 +65,10 @@ def compile(mod, target=None, target_host=None, params=None): exec : tvm.runtime.vm.Executable The VM executable that contains both library code and bytecode. """ - if target_host is not None: - warnings.warn( - "target_host parameter is going to be deprecated. " - "Please pass in tvm.target.Target(target, host=target_host) instead." - ) - target, target_host = Target.check_and_update_host_consist( - target, target_host, target_is_dict_key=False - ) compiler = VMCompiler() if params: compiler.set_params(params) - compiler.lower(mod, target) + compiler.lower(mod, target, target_host) compiler.codegen() return compiler.get_exec() @@ -139,20 +131,10 @@ def lower(self, mod, target=None, target_host=None): By default, llvm is used if it is enabled, otherwise a stackvm intepreter is used. """ - if target_host is not None: - warnings.warn( - "target_host parameter is going to be deprecated. " - "Please pass in tvm.target.Target(target, host=target_host) instead." - ) - target = self._update_target(target) - target_host = self._update_target_host(target, target_host) - target, target_host = Target.check_and_update_host_consist( - target, target_host, target_is_dict_key=False - ) - - tophub_context = self._tophub_context(target) + raw_targets = Target.canonicalize_target_and_host(target, target_host) + tophub_context = self._tophub_context(raw_targets) with tophub_context: - self._lower(mod, target, target_host) + self._lower(mod, raw_targets) def codegen(self): """Generate the kernel library.""" @@ -185,20 +167,10 @@ def optimize(self, mod, target=None, target_host=None, params=None): params : dict The parameters of the final module. """ - if target_host is not None: - warnings.warn( - "target_host parameter is going to be deprecated. " - "Please pass in tvm.target.Target(target, host=target_host) instead." - ) - target = self._update_target(target) - target_host = self._update_target_host(target, target_host) - target, target_host = Target.check_and_update_host_consist( - target, target_host, target_is_dict_key=False - ) - + raw_targets = Target.canonicalize_target_and_host(target, target_host) if params: self.set_params(params) - return self._optimize(mod, target, target_host), self.get_params() + return self._optimize(mod, raw_targets), self.get_params() def get_exec(self): """Get the VM executable. @@ -210,56 +182,12 @@ def get_exec(self): """ return vm_rt.Executable(self._get_exec()) - def _update_target(self, target): - """Update target.""" - target = target if target else tvm.target.Target.current() - if target is None: - raise ValueError("Target is not set in env or passed as argument.") - - if isinstance(target, str): - target = {target: target} - elif isinstance(target, tvm.target.Target): - target = {target.kind.name: target} - elif not isinstance(target, dict): - raise TypeError( - "target is expected to be str, tvm.target.Target, " - + "or dict of str to str/tvm.target.Target, but received " - + "{}".format(type(target)) - ) - - tgts = {} - for dev, tgt in target.items(): - dev_type = tvm.tir.IntImm("int32", tvm.nd.device(dev).device_type) - if isinstance(tgt, str): - tgt = tvm.target.Target(tgt) - - tgts[dev_type] = tgt - - return tgts - - def _update_target_host(self, target, target_host): - """Update target host.""" - target_host = None if target_host == "" else target_host - if not target_host: - for _, tgt in target.items(): - if tgt.host is not None: - return tgt.host - for device_type, tgt in target.items(): - if device_type.value == tvm.nd.cpu(0).device_type: - target_host = tgt - break - if not target_host: - target_host = "llvm" if tvm.runtime.enabled("llvm") else "stackvm" - if isinstance(target_host, str): - target_host = tvm.target.Target(target_host) - return target_host - - def _tophub_context(self, target): + def _tophub_context(self, raw_targets): """Get the autotvm context.""" # If current dispatch context is fallback context (the default root context), # then load pre-tuned parameters from TopHub if isinstance(autotvm.DispatchContext.current, autotvm.FallbackContext): - tophub_context = autotvm.tophub.context(list(target.values())) + tophub_context = autotvm.tophub.context(raw_targets) else: tophub_context = autotvm.utils.EmptyContext() return tophub_context diff --git a/python/tvm/relay/build_module.py b/python/tvm/relay/build_module.py index 876145c63fc0..10030118f6ad 100644 --- a/python/tvm/relay/build_module.py +++ b/python/tvm/relay/build_module.py @@ -173,15 +173,7 @@ def build( params : dict The parameters of the final graph. """ - if target_host is not None: - warnings.warn( - "target_host parameter is going to be deprecated. " - "Please pass in tvm.target.Target(target, host=target_host) instead." - ) - target = build_target_by_device_type_map(target) - target, target_host = Target.check_and_update_host_consist( - target, target_host, target_is_dict_key=False - ) + raw_targets = Target.canonicalize_target_and_host(target, target_host) # Setup the params. if params: @@ -199,7 +191,7 @@ def build( mod_name = mangle_module_name(mod_name) - self._build(mod, target, target_host, executor, runtime, workspace_memory_pools, mod_name) + self._build(mod, raw_targets, executor, runtime, workspace_memory_pools, mod_name) autotvm.GLOBAL_SCOPE.silent = old_autotvm_silent # Get artifacts @@ -209,7 +201,7 @@ def build( return executor_config, mod, params - def optimize(self, mod, target=None, params=None): + def optimize(self, mod, target=None, target_host=None, params=None): """ Parameters ---------- @@ -233,12 +225,12 @@ def optimize(self, mod, target=None, params=None): params : dict The parameters of the final graph. """ - target = build_target_by_device_type_map(target) + raw_targets = Target.canonicalize_target_and_host(target, target_host) # Setup the params. if params: self._set_params(params) - mod = self._optimize(mod, target) + mod = self._optimize(mod, raw_targets) # Get artifacts params = self.get_params() diff --git a/python/tvm/target/target.py b/python/tvm/target/target.py index f75db92c39b0..6ffc898b8763 100644 --- a/python/tvm/target/target.py +++ b/python/tvm/target/target.py @@ -24,7 +24,7 @@ from tvm._ffi import register_func as _register_func from tvm.runtime import Object, convert from tvm.runtime.container import String -from tvm.ir.container import Map +from tvm.ir.container import Map, Array from . import _ffi_api @@ -218,6 +218,63 @@ def list_kinds(): """Returns the list of available target names.""" return list(_ffi_api.ListTargetKinds()) + @staticmethod + def canonicalize_target(target): + """Given a single target-like object, returns the TVM Target object representing it. Can convert from: + - None (to None). + - An existing TVM Target object. + - A string. + - A Python dictionary binding the target 'kind' and other attributes.""" + if target is None: + return None + elif isinstance(target, Target): + return target + else: + return Target(target) + + @staticmethod + def canonicalize_multi_targets(multi_targets): + """Given a single or collection of target-like objects, returns a TVM Array of Target objects representing + then. Can convert from: + - None (to None). + - A single target-like object in a form recognized by canonicalize_target. + - A Python list or TVM Array of target-like objects in a form recognized by canonicalize_target. + - A Python dict or TVM Map from TVM IntImm objects representing device types to + a target-like object in a form recognized by canonicalize_target.""" + if multi_targets is None: + return None + elif isinstance(multi_targets, (dict, Map)) and "kind" not in multi_targets: + # Convert legacy heterogeneous map representation to ordinary list of targets. + return Target.canonicalize_multi_targets([t for _, t in multi_targets.items()]) + elif isinstance(multi_targets, (list, Array)): + # Multiple Target results. + return convert([Target.canonicalize_target(t) for t in multi_targets]) + else: + # Single Target result. + return convert([Target.canonicalize_target(multi_targets)]) + + @staticmethod + def canonicalize_target_and_host(target, target_host=None): + """Returns a TVM Array capturing target and target_host. The given target can be in any + form recognized by Target.canonicalize_target or Target.canonicalize_multi_targets. If given + target_host can be in any form recognized by Target.canonicalize_target. If target_host is given + it will be set as the 'host' in each result Target object (and a warning given). + """ + # Convert target to Array, but not yet accounting for any host. + raw_targets = Target.canonicalize_multi_targets(target) + assert raw_targets is not None + # Convert host to Target, if given. + target_host = Target.canonicalize_target(target_host) + if target_host is None: + return raw_targets + else: + warnings.warn( + "target_host parameter is going to be deprecated. " + "Please pass in tvm.target.Target(target, host=target_host) instead." + ) + # Make sure the (canonical) host is captured in all the (canonical) targets. + return convert([Target(t, target_host) for t in raw_targets]) + @staticmethod def check_and_update_host_consist(target, host=None, target_is_dict_key=True): """A helper function that merges a legacy "target, target_host" pair, then returns @@ -560,7 +617,7 @@ def get_arch_version(cpu_ver): # Check for valid codegen cpu valid_hex = ["v65", "v66", "v67", "v67t", "v68", "v69"] try: - cpu_ver = cpu_ver[cpu_ver.index("v") :].lower() + cpu_ver = cpu_ver[cpu_ver.index("v"):].lower() assert cpu_ver in valid_hex except: msg = "{} is not a valid Hexagon version\nvalid versions include {}" @@ -628,7 +685,7 @@ def validate_hvx_length(codegen_hvx, sim_options): # If --hvx_length was specified, check HVX length of sim # vs codegen i = sim_options.index("hvx_length") + len("hvx_length") + 1 - sim_hvx = sim_options[i : i + 3] + sim_hvx = sim_options[i: i + 3] if sim_hvx != str(codegen_hvx): msg = "sim hvx {} and codegen hvx {} mismatch!".format(sim_hvx, codegen_hvx) # Set the stacklevel to the tvm.target.hexagon() call. @@ -656,9 +713,9 @@ def validate_hvx_length(codegen_hvx, sim_options): # Regex match for allowed cpus valid_cpu_str_regex = ( - r"(?P
--.*\s)?(--m)?"
-                + r"(?Pv6[25678])(?P[a-z])?"
-                + r"(?P_[0-9]+)?(?P_rev[0-9])?\s?(?P--.*)?"
+                    r"(?P
--.*\s)?(--m)?"
+                    + r"(?Pv6[25678])(?P[a-z])?"
+                    + r"(?P_[0-9]+)?(?P_rev[0-9])?\s?(?P--.*)?"
             )
             m = re.match(valid_cpu_str_regex, sim_options.lower())
             if not m:
@@ -667,13 +724,13 @@ def validate_hvx_length(codegen_hvx, sim_options):
             # Parse options into correct order
             cpu_attr = {x: str(m.groupdict()[x] or "") for x in m.groupdict()}
             sim_options = (
-                cpu_attr["base_version"]
-                + cpu_attr["sub_version"]
-                + cpu_attr["l2_size"]
-                + cpu_attr["rev"]
-                + " "
-                + cpu_attr["pre"]
-                + cpu_attr["post"]
+                    cpu_attr["base_version"]
+                    + cpu_attr["sub_version"]
+                    + cpu_attr["l2_size"]
+                    + cpu_attr["rev"]
+                    + " "
+                    + cpu_attr["pre"]
+                    + cpu_attr["post"]
             )
 
         return sim_cpu + " " + validate_hvx_length(hvx, sim_options)
diff --git a/src/relay/backend/aot_executor_codegen.cc b/src/relay/backend/aot_executor_codegen.cc
index 22d4b1c032f4..a0c0aa940a0b 100644
--- a/src/relay/backend/aot_executor_codegen.cc
+++ b/src/relay/backend/aot_executor_codegen.cc
@@ -763,7 +763,7 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     String run_func_name = runtime::get_name_mangled(mod_name, runtime::symbol::tvm_module_main);
     dict_attrs.Set("global_symbol", run_func_name);
     dict_attrs.Set("runner_function", Bool(true));
-    dict_attrs.Set(tvm::attr::kTarget, target_host_);
+    dict_attrs.Set(tvm::attr::kTarget, config_->host_target);
 
     tir::Stmt device_activations = GenerateAllDeviceHook("Activate");
     tir::Stmt device_deactivations = GenerateAllDeviceHook("Deactivate");
@@ -909,7 +909,7 @@ class AOTExecutorCodegen : public MixedModeVisitor {
         CalculateWorkspaceBytes(tir_main_func, workspace_byte_alignment);
     backend::FunctionInfo main_func_info =
         lowered_mod->GetAttr("main_func_info").value();
-    main_func_info->workspace_sizes.Set(target_host_, main_workspace_size_bytes);
+    main_func_info->workspace_sizes.Set(config_->host_target, main_workspace_size_bytes);
     function_metadata_.Set(runtime::symbol::tvm_module_main, main_func_info);
     return lowered_mod;
   }
@@ -929,10 +929,8 @@ class AOTExecutorCodegen : public MixedModeVisitor {
   Map main_buffer_map_;
   /*! \brief maps input and output variables to TensorType which describe them */
   Map io_tensor_types_;
-  /*! \brief target device */
-  tec::TargetMap targets_;
-  /*! \brief target host */
-  Target target_host_;
+  /*! \brief All available targets. */
+  CompilationConfig config_;
   /*!
    * \brief The type of kernel call to be emitted.
    * See CallType for more documentation.
@@ -967,16 +965,10 @@ class AOTExecutorCodegen : public MixedModeVisitor {
   std::unordered_map io_var_names_;
 
  public:
-  AOTExecutorCodegen(runtime::Module* mod, const tec::TargetMap& targets, Target target_host)
-      : mod_(mod), targets_(targets), target_host_(target_host) {}
+  AOTExecutorCodegen(runtime::Module* mod, CompilationConfig config) : mod_(mod), config_(config) {}
 
   LoweredOutput Codegen(IRModule mod, relay::Function func, String mod_name) {
     VLOG_CONTEXT << "AOT";
-    for (const auto& kv : targets_) {
-      VLOG(1) << "target: " << kv.second->ToDebugString();
-    }
-    ICHECK(target_host_.defined()) << "require a target_host to be given for AOT codegen";
-    VLOG(1) << "target host: " << target_host_->ToDebugString();
 
     Runtime runtime_config = mod->GetAttr(tvm::attr::kRuntime).value();
     Executor executor_config = mod->GetAttr(tvm::attr::kExecutor).value();
@@ -1015,9 +1007,6 @@ class AOTExecutorCodegen : public MixedModeVisitor {
                     << ") is not one of the expected values";
     }
 
-    // TODO(mbs): Plumb from compiler config
-    VirtualDevice host_virtual_device = VirtualDevice::ForTarget(target_host_);
-
     IRModule lowered_mod = tec::LowerTEPass(
         mod_name,
         [this, workspace_byte_alignment](BaseFunc func) {
@@ -1033,7 +1022,7 @@ class AOTExecutorCodegen : public MixedModeVisitor {
           // lowering process directly.
           tec::UpdateFunctionMetadata(func, this->function_metadata_, workspace_byte_alignment);
         },
-        host_virtual_device)(mod);
+        config_->host_virtual_device)(mod);
 
     auto lowered_main = lowered_mod->Lookup("main");
     auto lowered_main_func = GetRef(lowered_main.as());
@@ -1046,7 +1035,7 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     // TODO(@electriclilies, @jroesch, @Mousius): remove UpdateMainWorkspaceSize
     StaticMemoryPlan memory_plan(storage_device_map_);
     backend::FunctionInfo func_info =
-        tec::UpdateMainWorkspaceSize(lowered_mod, targets_, memory_plan->expr_to_storage_info);
+        tec::UpdateMainWorkspaceSize(lowered_mod, config_, memory_plan->expr_to_storage_info);
     lowered_mod = WithAttr(lowered_mod, "main_func_info", func_info);
 
     for (auto input : lowered_main_func->params) {
@@ -1219,8 +1208,8 @@ class AOTExecutorCodegenModule : public runtime::ModuleNode {
         ICHECK_EQ(args.num_args, 2) << "The expected of arguments are: "
                                     << "runtime::Module mod and  Map targets";
         void* mod = args[0];
-        TargetMap targets = args[1];
-        init(mod, targets);
+        CompilationConfig config = args[1];
+        init(mod, std::move(config));
       });
     } else if (name == "codegen") {
       return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
@@ -1267,22 +1256,9 @@ class AOTExecutorCodegenModule : public runtime::ModuleNode {
   const char* type_key() const final { return "RelayGraphRuntimeCodegenModule"; }
 
  private:
-  void init(void* mod, TargetMap tmp) {
-    tec::TargetMap targets;
-    Target target_host;
-    for (const auto& it : tmp) {
-      auto dev_type = it.first.as();
-      // TODO(tvm-team): AoT only works with kDLCPU device type. We can remove kDLHexagon
-      // here once we refactored kDLHexagon to kDLCPU.
-      if (!target_host.defined() && ((it.second->kind->device_type == kDLCPU) ||
-                                     (it.second->kind->device_type == kDLHexagon))) {
-        target_host = it.second;
-      }
-      ICHECK(dev_type);
-      targets[static_cast(dev_type->value)] = it.second;
-    }
-    codegen_ = std::make_shared(reinterpret_cast(mod),
-                                                    targets, target_host);
+  void init(void* mod, CompilationConfig config) {
+    codegen_ =
+        std::make_shared(reinterpret_cast(mod), config);
   }
 
   Array list_params_name() {
diff --git a/src/relay/backend/build_module.cc b/src/relay/backend/build_module.cc
index 99f0517d1b7f..e742aaafbe07 100644
--- a/src/relay/backend/build_module.cc
+++ b/src/relay/backend/build_module.cc
@@ -61,7 +61,7 @@ struct BuildOutput {
 };
 
 struct ExecutorCodegen {
-  void Init(runtime::Module* m, TargetMap targets) { CallFunc("init", m, targets); }
+  void Init(runtime::Module* m, CompilationConfig config) { CallFunc("init", m, config); }
 
   void Codegen(IRModule mod, const Function& func, String mod_name) {
     CallFunc("codegen", mod, func, mod_name);
@@ -190,8 +190,8 @@ class RelayBuildModule : public runtime::ModuleNode {
           [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->GetModule(); });
     } else if (name == "build") {
       return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
-        ICHECK_EQ(args.num_args, 7);
-        this->Build(args[0], args[1], args[2], args[3], args[4], args[5], args[6]);
+        ICHECK_EQ(args.num_args, 6);
+        this->Build(args[0], args[1], args[2], args[3], args[4], args[5]);
       });
     } else if (name == "list_params") {
       return PackedFunc(
@@ -228,8 +228,8 @@ class RelayBuildModule : public runtime::ModuleNode {
       });
     } else if (name == "optimize") {
       return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
-        ICHECK_EQ(args.num_args, 2);
-        *rv = this->Optimize(args[0], args[1]);
+        ICHECK_EQ(args.num_args, 3);
+        *rv = this->Optimize(args[0], args[1], args[2]);
       });
     } else {
       LOG(FATAL) << "Unknown packed function: " << name;
@@ -296,20 +296,20 @@ class RelayBuildModule : public runtime::ModuleNode {
    * \brief Build relay IRModule for graph executor
    *
    * \param mod Relay IRModule
-   * \param targets Target devices
-   * \param target_host Host target device
+   * \param raw_targets List of available targets for kernels.
    * \param executor Executor to target
    * \param runtime Runtime to codegen for
    * \param mod_name Name of the module
    */
-  void Build(IRModule mod, const TargetMap& targets, const tvm::Target& target_host,
-             const Executor& executor, const Runtime& runtime,
-             const WorkspaceMemoryPools& workspace_memory_pools, const String mod_name) {
+  void Build(IRModule mod, const Array& raw_targets, const Executor& executor,
+             const Runtime& runtime, const WorkspaceMemoryPools& workspace_memory_pools,
+             const String& mod_name) {
     VLOG_CONTEXT << "Build";
     executor_ = executor;
     runtime_ = runtime;
     workspace_memory_pools_ = workspace_memory_pools;
-    config_ = CompilationConfig(PassContext::Current(), targets, target_host);
+    config_ = CompilationConfig(PassContext::Current(), raw_targets);
+    VLOG(1) << "Using compilation config:" << std::endl << config_;
     BuildRelay(std::move(mod), mod_name);
   }
 
@@ -318,16 +318,15 @@ class RelayBuildModule : public runtime::ModuleNode {
    * \brief Optimize a Relay IRModule.
    *
    * \param relay_module The input IRModule where optmization will be applied on.
-   * \param targets The device type to `Target` mapping.
+   * \param raw_targets List of available targets for kernels.
    *
    * \return relay::IRModule The updated Relay IR module after optimization.
    */
-  IRModule Optimize(IRModule relay_module, const TargetMap& targets) {
+  IRModule Optimize(IRModule relay_module, const Array& raw_targets,
+                    const Target& optional_host_target) {
     VLOG_CONTEXT << "Optimize";
-    // TODO(mbs): executor_ will be whatever was left over from last Build. Note that
-    // the empty executor string will CHECK fail, so how are folks using this API?
-    config_ = CompilationConfig(transform::PassContext::Current(), targets,
-                                /*optional_host_target=*/Target());
+    config_ = CompilationConfig(PassContext ::Current(), raw_targets);
+    VLOG(1) << "Using compilation config:" << std::endl << config_;
     return OptimizeImpl(std::move(relay_module));
   }
 
@@ -336,8 +335,8 @@ class RelayBuildModule : public runtime::ModuleNode {
 
     backend::BindParamsInModule(relay_module, params_);
 
-    Array pass_seqs = GetPassPrefix(
-        /*is_homogenous=*/config_->optional_homogeneous_target.defined(), /*is_vm=*/false);
+    Array pass_seqs =
+        GetPassPrefix(/*is_homogenous=*/config_->primitive_targets.size() == 1, /*is_vm=*/false);
     transform::PassContext pass_ctx = PassContext::Current();
 
     if (config_->optional_homogeneous_target.defined()) {
@@ -418,7 +417,7 @@ class RelayBuildModule : public runtime::ModuleNode {
 
     // Generate code for the updated function.
     executor_codegen_ = MakeExecutorCodegen(executor_->name);
-    executor_codegen_->Init(nullptr, config_->legacy_target_map);
+    executor_codegen_->Init(nullptr, config_);
     executor_codegen_->Codegen(func_module, func, mod_name);
     executor_codegen_->UpdateOutput(&ret_);
     ret_.params = executor_codegen_->GetParams();
diff --git a/src/relay/backend/contrib/cublas/target.cc b/src/relay/backend/contrib/cublas/target.cc
new file mode 100644
index 000000000000..45d3aaa314ae
--- /dev/null
+++ b/src/relay/backend/contrib/cublas/target.cc
@@ -0,0 +1,44 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file src/relay/backend/contrib/cudnn/target.cc
+ * \brief Registers the "cublas" external codegen TargetKind.
+ */
+
+#include 
+
+namespace tvm {
+namespace relay {
+namespace contrib {
+
+/*!
+ * \brief This external codegen target can use the CuBLAS library linked into the TVM runtime.
+ *  - Patterns and custom compiler: python/tvm/relay/op/contrib/cublas.py
+ *  - Custom schedules: python/tvm/contrib/cublas.py
+ *  - Runtime: src/runtime/contrib/cublas/cublas.cc
+ *
+ * CuBLAS can also be used via the "-libs=cublas" Target option.
+ */
+TVM_REGISTER_TARGET_KIND("cublas", kDLCUDA)
+    .set_attr(tvm::attr::kIsExternalCodegen, Bool(true));
+
+}  // namespace contrib
+}  // namespace relay
+}  // namespace tvm
diff --git a/src/relay/backend/contrib/cudnn/target.cc b/src/relay/backend/contrib/cudnn/target.cc
new file mode 100644
index 000000000000..1f1117391209
--- /dev/null
+++ b/src/relay/backend/contrib/cudnn/target.cc
@@ -0,0 +1,42 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file src/relay/backend/contrib/cudnn/target.cc
+ * \brief Registers the "cudnn" external codegen TargetKind.
+ */
+
+#include 
+
+namespace tvm {
+namespace relay {
+namespace contrib {
+
+/*!
+ * \brief This external codegen target can use the CuDNN library linked into the TVM runtime.
+ *  - Patterns and custom compiler: python/tvm/relay/op/contrib/cudnn.py
+ *  - Custom schedules: python/tvm/contrib/cudnn.py
+ *  - Runtime: src/runtime/contrib/cudnn/ *.cc
+ */
+TVM_REGISTER_TARGET_KIND("cudnn", kDLCUDA)
+    .set_attr(tvm::attr::kIsExternalCodegen, Bool(true));
+
+}  // namespace contrib
+}  // namespace relay
+}  // namespace tvm
diff --git a/src/relay/backend/contrib/cutlass/target.cc b/src/relay/backend/contrib/cutlass/target.cc
new file mode 100644
index 000000000000..3a7384fb19cc
--- /dev/null
+++ b/src/relay/backend/contrib/cutlass/target.cc
@@ -0,0 +1,43 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file src/relay/backend/contrib/cutlass/target.cc
+ * \brief Registers the "cutlass" external codegen TargetKind.
+ */
+
+#include 
+
+namespace tvm {
+namespace relay {
+namespace contrib {
+
+/*!
+ * \brief This external codegen target can use the CUTLASS template library included in
+ * TVM's 3rdparty/cutlass.
+ *  - Patterns: python/tvm/relay/op/contrib/cutlass.py
+ *  - Custom compiler: python/tvm/contrib/cutlass/build.py,
+ *                     src/relay/backend/contrib/cutlass/codegen.cc
+ */
+TVM_REGISTER_TARGET_KIND("cutlass", kDLCUDA)
+    .set_attr(tvm::attr::kIsExternalCodegen, Bool(true));
+
+}  // namespace contrib
+}  // namespace relay
+}  // namespace tvm
diff --git a/src/relay/backend/contrib/tensorrt/target.cc b/src/relay/backend/contrib/tensorrt/target.cc
new file mode 100644
index 000000000000..85d127ab7115
--- /dev/null
+++ b/src/relay/backend/contrib/tensorrt/target.cc
@@ -0,0 +1,42 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file src/relay/backend/contrib/tensorrt/target.cc
+ * \brief Registers the "tensorrt" external codegen TargetKind.
+ */
+
+#include 
+
+namespace tvm {
+namespace relay {
+namespace contrib {
+
+/*!
+ * \brief This external codegen target can offload compilation to the TensorRT compiler.
+ *  - Patterns: python/tvm/relay/op/contrib/tensorrt.py
+ *  - Custom compiler: src/relay/backend/contrib/tensorrt/codegen.cc
+ *  - Runtime: src/runtime/contrib/tensorrt/ *.cc
+ */
+TVM_REGISTER_TARGET_KIND("tensorrt", kDLCUDA)
+    .set_attr(tvm::attr::kIsExternalCodegen, Bool(true));
+
+}  // namespace contrib
+}  // namespace relay
+}  // namespace tvm
diff --git a/src/relay/backend/graph_executor_codegen.cc b/src/relay/backend/graph_executor_codegen.cc
index 30bc0beeebce..580458da2385 100644
--- a/src/relay/backend/graph_executor_codegen.cc
+++ b/src/relay/backend/graph_executor_codegen.cc
@@ -190,8 +190,8 @@ class GraphOpNode : public GraphNode {
  */
 class GraphExecutorCodegen : public backend::MemoizedExprTranslator> {
  public:
-  GraphExecutorCodegen(runtime::Module* mod, const TargetMap& targets)
-      : mod_(mod), targets_(targets) {}
+  GraphExecutorCodegen(runtime::Module* mod, CompilationConfig config)
+      : mod_(mod), config_(std::move(config)) {}
 
   StorageInfo GetStorageInfo(const Expr& e) {
     size_t count = memory_plan_->expr_to_storage_info.count(e);
@@ -204,9 +204,6 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslatorfunction_metadata_);
         },
-        config->host_virtual_device)(mod);
+        config_->host_virtual_device)(mod);
 
     Optional main_func_info =
         lowered_mod->GetAttr("main_func_info");
@@ -610,8 +597,8 @@ class GraphExecutorCodegen : public backend::MemoizedExprTranslator> var_map_;
-  /*! \brief target device */
-  TargetMap targets_;
+  /*! \brief Available targets */
+  CompilationConfig config_;
   /*!
    * \brief parameters (i.e. ConstantNodes found in the graph).
    * These are take as inputs to the GraphExecutor.
@@ -639,9 +626,9 @@ class GraphExecutorCodegenModule : public runtime::ModuleNode {
         ICHECK_EQ(args.num_args, 2) << "The expected of arguments are: "
                                     << "runtime::Module mod and Map targets";
         void* mod = args[0];
-        TargetMap target_map = args[1];
+        CompilationConfig config = args[1];
         codegen_ = std::make_shared(reinterpret_cast(mod),
-                                                          target_map);
+                                                          std::move(config));
       });
     } else if (name == "codegen") {
       return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
diff --git a/src/relay/backend/interpreter.cc b/src/relay/backend/interpreter.cc
index c4b1673e0731..673a547d2df0 100644
--- a/src/relay/backend/interpreter.cc
+++ b/src/relay/backend/interpreter.cc
@@ -1020,10 +1020,8 @@ TypedPackedFunc)> EvalFunction(IRModule mod, Expr expr, De
           << PrettyPrint(expr);
 
   ICHECK_EQ(device.device_type, target->kind->device_type);
-  TargetMap targets;
-  targets.Set(device.device_type, target);
-  CompilationConfig config(transform::PassContext::Current(), targets,
-                           /*optional_host_target_arg=*/{});
+  Array raw_targets = {target};
+  CompilationConfig config(transform::PassContext::Current(), raw_targets);
 
   //
   // Step 1: Prepare mod.
@@ -1111,10 +1109,8 @@ ObjectRef Eval(Expr expr, Map type_definitions,
                std::unordered_set import_set, Device device, Target target,
                Map attrs) {
   ICHECK_EQ(device.device_type, target->kind->device_type);
-  TargetMap targets;
-  targets.Set(device.device_type, target);
-  CompilationConfig config(transform::PassContext::Current(), targets,
-                           /*optional_host_target_arg=*/{});
+  Array raw_targets = {target};
+  CompilationConfig config(transform::PassContext::Current(), raw_targets);
 
   std::pair mod_and_global =
       IRModule::FromExprInContext(expr, /*global_funcs=*/{}, type_definitions, import_set);
diff --git a/src/relay/backend/te_compiler.cc b/src/relay/backend/te_compiler.cc
index 4209b0a8bbe7..70d74ea92377 100644
--- a/src/relay/backend/te_compiler.cc
+++ b/src/relay/backend/te_compiler.cc
@@ -773,7 +773,7 @@ class LowerTensorExprMutator : public DeviceAwareExprMutator {
     } else {
       // The target corresponding to the call_node expression's annotation.
       VirtualDevice virtual_device = GetVirtualDevice(GetRef(call_node));
-      ICHECK(!virtual_device->IsFullyUnconstrained());
+      ICHECK(!virtual_device->IsFullyUnconstrained()) << PrettyPrint(GetRef(call_node));
       target = virtual_device->target;
       ICHECK(target.defined());
     }
@@ -803,36 +803,6 @@ class LowerTensorExprMutator : public DeviceAwareExprMutator {
   const Op& debug_op_;
 };
 
-Target GetTargetFromInteger(DLDeviceType dev_type, tec::TargetMap targets) {
-  if (targets.size() == 1) {
-    // The homogeneous execution case, return the only target.
-    const auto& it = targets.begin();
-    return (*it).second;
-  } else {
-    // The heterogeneous execution case, return the target associated with the
-    // given device type.
-    // If "dev_type" equals to 0, the device name only can be got from
-    // "targets", and it may not be "llvm", so here just set it to "unknown".
-    std::string dev_name = "unknown";
-    if (dev_type != 0) {
-      dev_name = runtime::DeviceName(dev_type);
-    }
-
-    if (targets.count(dev_type) == 0) {
-      std::stringstream msg;
-      msg << "No target is specified for provided device name: `" << dev_name << "`\n\n"
-          << dev_name << " mapped to device type (" << dev_type
-          << ") which was not found in the target map.\n"
-          << "Availible targets: \n";
-      for (auto target : targets) {
-        msg << "  " << target.first << "-> " << target.second << "\n";
-      }
-      LOG(FATAL) << msg.str();
-    }
-    return targets[dev_type];
-  }
-}
-
 Pass LowerTensorExpr(const String& module_name, TECompiler compiler, ProcessFn process_fn,
                      VirtualDevice host_virtual_device) {
   runtime::TypedPackedFunc pass_func =
@@ -844,15 +814,12 @@ Pass LowerTensorExpr(const String& module_name, TECompiler compiler, ProcessFn p
   return CreateFunctionPass(pass_func, 0, "LowerTensorExpr", {});
 }
 
-backend::FunctionInfo UpdateMainWorkspaceSize(const IRModule& mod, tec::TargetMap targets,
+backend::FunctionInfo UpdateMainWorkspaceSize(const IRModule& mod, const CompilationConfig& config,
                                               Map storage_info_map) {
   Function func = Downcast(mod->Lookup("main"));
 
   VLOG_CONTEXT << "UpdateMainWorkspaceSize";
   VLOG(1) << "calculating FunctionInfo for main:" << std::endl << PrettyPrint(func);
-  for (const auto& kv : targets) {
-    VLOG(1) << "  target " << kv.first << " = " << kv.second->str();
-  }
 
   // This is a Map>
   // TODO(mbs): Collapsing VirtualDevices to just device type.
@@ -952,25 +919,24 @@ backend::FunctionInfo UpdateMainWorkspaceSize(const IRModule& mod, tec::TargetMa
   Map relay_primfuncs;
 
   // Initialize all target workspaces to zero
-  for (const auto& kv : targets) {
-    auto tgt = kv.second;
-    workspace_sizes.Set(tgt, 0);
+  for (const auto& target : config->primitive_targets) {
+    workspace_sizes.Set(target, 0);
   }
 
   for (const auto& dev_and_size : device_workspace) {
-    auto tgt = tec::GetTargetFromInteger(dev_and_size.first, targets);
-    workspace_sizes.Set(tgt, dev_and_size.second);
-    relay_primfuncs.Set(tgt, func);
+    Target target = config->FindPrimitiveTargetOrFail(dev_and_size.first);
+    workspace_sizes.Set(target, dev_and_size.second);
+    relay_primfuncs.Set(target, func);
   }
   for (const auto& dev_and_size : device_io) {
-    auto tgt = tec::GetTargetFromInteger(dev_and_size.first, targets);
-    io_sizes.Set(tgt, dev_and_size.second);
+    Target target = config->FindPrimitiveTargetOrFail(dev_and_size.first);
+    io_sizes.Set(target, dev_and_size.second);
   }
 
   for (const auto& dev_and_size : device_consts) {
-    auto tgt = tec::GetTargetFromInteger(dev_and_size.first, targets);
-    ICHECK_EQ(constant_sizes.count(tgt), 0);
-    constant_sizes.Set(tgt, dev_and_size.second);
+    Target target = config->FindPrimitiveTargetOrFail(dev_and_size.first);
+    ICHECK_EQ(constant_sizes.count(target), 0);
+    constant_sizes.Set(target, dev_and_size.second);
   }
 
   backend::FunctionInfo func_info(std::move(workspace_sizes), std::move(io_sizes),
diff --git a/src/relay/backend/te_compiler.h b/src/relay/backend/te_compiler.h
index b6f2218e2319..0b2288d6a156 100644
--- a/src/relay/backend/te_compiler.h
+++ b/src/relay/backend/te_compiler.h
@@ -58,11 +58,6 @@ namespace tvm {
 namespace relay {
 namespace tec {
 
-// TODO(@jroesch, @chrisS) these should be a tvm::Map for uniformity sake
-// we should a version of context which works in Map
-using TargetMap = std::unordered_map;
-using DeviceMap =
-    std::unordered_map;
 using ProcessFn = std::function;
 
 /*!
@@ -160,25 +155,14 @@ void UpdateFunctionMetadata(BaseFunc relay_func,
                             Map& function_metadata,  // NOLINT(*)
                             Integer workspace_byte_alignment = 16);
 
-/*!
- * \brief Obtain the Target from the device type.
- * If homogenous compilation, this will return the only target.
- * If heterogeneous compilation, this will select the associated target using the
- * targets_ Map.
- *
- * \param dev_type
- * \return Target
- */
-Target GetTargetFromInteger(DLDeviceType dev_type, tec::TargetMap targets);
-
 /*!
  * \brief Update the "main" control function's metadata
  *
  * \param mod The module
- * \param targets Map of targets
+ * \param config All the available targets.
  * \return function_infos Function info for each function in the module
  */
-backend::FunctionInfo UpdateMainWorkspaceSize(const IRModule& mod, tec::TargetMap targets,
+backend::FunctionInfo UpdateMainWorkspaceSize(const IRModule& mod, const CompilationConfig& config,
                                               Map storage_info_map);
 
 /*! \brief Returns all the global \p PrimFunc functions in \p mod, but separated into an \p IRModule
diff --git a/src/relay/backend/utils.cc b/src/relay/backend/utils.cc
index 4a6fe90289fb..133f9a9fc387 100644
--- a/src/relay/backend/utils.cc
+++ b/src/relay/backend/utils.cc
@@ -27,9 +27,7 @@
 
 #include 
 #include 
-
-#include "te_compiler.h"
-#include "tvm/runtime/ndarray.h"
+#include 
 
 namespace tvm {
 namespace relay {
@@ -205,7 +203,7 @@ ExecutorCodegenMetadata::ExecutorCodegenMetadata(
 
 TVM_REGISTER_NODE_TYPE(ExecutorCodegenMetadataNode);
 
-Array GetPassPrefix(bool is_homegeneous, bool is_vm) {
+Array GetPassPrefix(bool is_homogeneous, bool is_vm) {
   Array pass_seqs;
   // TODO(mbs): Would be nice to get spans on all diagnostics, but since they arg forgotton
   // by most passes there's little utility in including this now. Plus we'd need to only do
@@ -218,7 +216,7 @@ Array GetPassPrefix(bool is_homegeneous, bool is_vm) {
   pass_seqs.push_back(relay::qnn::transform::Legalize());
 
   // Legalize pass is restricted to homogeneous execution for now.
-  if (is_homegeneous) {
+  if (is_homogeneous) {
     pass_seqs.push_back(transform::Legalize());
   }
 
@@ -254,7 +252,7 @@ Array GetPassPrefix(bool is_homegeneous, bool is_vm) {
   pass_seqs.push_back(transform::CanonicalizeOps());
 
   // Alter layout transformation is currently only applied to homogeneous execution.
-  if (is_homegeneous) {
+  if (is_homogeneous) {
     if (!is_vm) {
       pass_seqs.push_back(transform::InferType());
     }
diff --git a/src/relay/backend/utils.h b/src/relay/backend/utils.h
index a31ff605cafa..360f366a162e 100644
--- a/src/relay/backend/utils.h
+++ b/src/relay/backend/utils.h
@@ -241,6 +241,7 @@ struct ConstantUpdater : public ExprVisitor {
 
   void VisitExpr_(const ConstantNode* cn) final {
     std::string name = symbol_ + "_const_" + std::to_string(const_idx_++);
+    VLOG(1) << "Binding " << name << " to constant of type " << PrettyPrint(cn->checked_type());
     (*params_)[name] = cn->data;
   }
 
@@ -515,11 +516,11 @@ inline bool IsMetaScheduleEnabled() {
  * difference. This function unifies the shared optimization pass prefix between vm and graph
  * runtime, and returns the pass prefix given the backend type.
  *
- * \param is_homogenous True if all primitives are to be executed on the same device and target.
+ * \param is_homogeneous True if all primitives are to be executed on the same device and target.
  * \param is_vm True if passes are to be used for the vm executor.
  * \return An array of passes.
  */
-Array GetPassPrefix(bool is_homogenous, bool is_vm);
+Array GetPassPrefix(bool is_homogeneous, bool is_vm);
 
 /*! \brief Target hash function */
 struct TargetStrHash {
diff --git a/src/relay/backend/vm/compiler.cc b/src/relay/backend/vm/compiler.cc
index b63409154350..e6aeb0bc4a0f 100644
--- a/src/relay/backend/vm/compiler.cc
+++ b/src/relay/backend/vm/compiler.cc
@@ -827,8 +827,8 @@ class VMFunctionCompiler : DeviceAwareExprFunctor {
 PackedFunc VMCompiler::GetFunction(const std::string& name, const ObjectPtr& sptr_to_self) {
   if (name == "lower") {
     return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
-      ICHECK_EQ(args.num_args, 3);
-      this->Lower(args[0], args[1], args[2]);
+      ICHECK_EQ(args.num_args, 2);
+      this->Lower(args[0], args[1]);
     });
   } else if (name == "codegen") {
     return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
@@ -836,8 +836,10 @@ PackedFunc VMCompiler::GetFunction(const std::string& name, const ObjectPtrCodegen();
     });
   } else if (name == "get_executable") {
-    return PackedFunc(
-        [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = runtime::Module(exec_); });
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+      ICHECK_EQ(args.num_args, 0);
+      *rv = this->GetExecutable();
+    });
   } else if (name == "set_params") {
     return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
       Map params = args[0];
@@ -855,8 +857,8 @@ PackedFunc VMCompiler::GetFunction(const std::string& name, const ObjectPtrOptimizeModule(args[0], args[1], args[2]);
+      ICHECK_EQ(args.num_args, 2);
+      *rv = this->OptimizeModule(args[0], args[1]);
     });
   } else {
     LOG(FATAL) << "Unknown packed function: " << name;
@@ -868,16 +870,42 @@ void VMCompiler::SetParam(const std::string& name, runtime::NDArray data_in) {
   params_[name] = data_in;
 }
 
-void VMCompiler::Lower(IRModule mod, TargetMap targets, tvm::Target target_host) {
+void VMCompiler::Lower(IRModule mod, const Array& raw_targets) {
   VLOG_CONTEXT << "VM Lower";
+  Setup(raw_targets);
+  LowerImpl(std::move(mod));
+}
+
+IRModule VMCompiler::OptimizeModule(IRModule mod, const Array& raw_targets) {
+  VLOG_CONTEXT << "VM Optimize";
+  Setup(raw_targets);
+  return OptimizeModuleImpl(std::move(mod));
+}
+
+runtime::Module VMCompiler::GetExecutable() const {
+  if (exec_ == nullptr) {
+    LOG(WARNING) << "No executable to return. Did you forget to call VMCompiler::Lower?";
+  }
+  if (exec_->imports().empty()) {
+    LOG(WARNING) << "Executable is empty. Did you forget to call VMCompiler::Codegen?";
+  }
+  return runtime::Module(exec_);
+}
+
+void VMCompiler::Setup(const Array& raw_targets) {
+  ICHECK(exec_ == nullptr) << "Can't reuse VMComplier object for multiple modules";
   exec_ = make_object();
-  config_ = CompilationConfig(PassContext::Current(), std::move(targets), std::move(target_host));
+  ICHECK(!config_.defined());
+  config_ = CompilationConfig(PassContext::Current(), raw_targets);
+  VLOG(1) << "Using compilation config:" << std::endl << config_;
 
   // The first device is always for the host.
   CHECK(context_.virtual_devices_.empty());
-  VLOG(2) << "virtual_device[0] = " << config_->host_virtual_device << " (host)";
+  VLOG(1) << "virtual_device[0] = " << config_->host_virtual_device << " (host)";
   context_.virtual_devices_.push_back(config_->host_virtual_device);
+}
 
+void VMCompiler::LowerImpl(IRModule mod) {
   // Run the optimizations necessary to target the VM.
   context_.module = OptimizeModuleImpl(std::move(mod));
 
@@ -1022,26 +1050,13 @@ transform::Sequential VMCompiler::FuseAndLowerOperators(const VirtualDevice& hos
   return transform::Sequential(std::move(pass_seqs));
 }
 
-IRModule VMCompiler::OptimizeModule(IRModule mod, const TargetMap& targets,
-                                    const Target& target_host) {
-  config_ = CompilationConfig(PassContext::Current(), targets, target_host);
-  // The first device always corresponds to the host.
-  CHECK(context_.virtual_devices_.empty());
-  context_.virtual_devices_.push_back(config_->host_virtual_device);
-  // TODO(mbs): exec_ is not allocated. What is the API here?
-  CHECK(exec_ == nullptr);
-  return OptimizeModuleImpl(std::move(mod));
-}
-
 IRModule VMCompiler::OptimizeModuleImpl(IRModule mod) {
-  VLOG_CONTEXT << "VM Optimize";
   backend::BindParamsInModule(mod, params_);
-
   Array pass_seqs = relay::backend::GetPassPrefix(
-      /*is_homogenous=*/config_->optional_homogeneous_target.defined(), /*is_vm=*/true);
+      /*is_homogeneous=*/config_->optional_homogeneous_target.defined(), /*is_vm=*/true);
 
   // Always plan devices so the remaining passes don't need to distinguish homogeneous vs
-  // hetrogeneous execution.
+  // heterogeneous execution.
   pass_seqs.push_back(transform::PlanDevices(config_));
 
   pass_seqs.push_back(transform::FuseOps());
@@ -1163,12 +1178,10 @@ void VMCompiler::Codegen() {
 
 runtime::Module CreateVMCompiler() {
   auto exec = make_object();
-  return runtime::Module(exec);
+  return runtime::Module(std::move(exec));
 }
 
-TVM_REGISTER_GLOBAL("relay._vm._VMCompiler").set_body([](TVMArgs args, TVMRetValue* rv) {
-  *rv = CreateVMCompiler();
-});
+TVM_REGISTER_GLOBAL("relay._vm._VMCompiler").set_body_typed(CreateVMCompiler);
 
 }  // namespace vm
 }  // namespace relay
diff --git a/src/relay/backend/vm/compiler.h b/src/relay/backend/vm/compiler.h
index 906e5148b593..3cbafd73eae9 100644
--- a/src/relay/backend/vm/compiler.h
+++ b/src/relay/backend/vm/compiler.h
@@ -106,34 +106,39 @@ class VMCompiler : public runtime::ModuleNode {
    *
    * ----------------------------------------------------------------------------------
    * | This is the main entry point for the VM compilation flow.                      |
-   * |  - Preceded by \p SetParam for the global params.                             |
+   * |  - Preceded by \p SetParam for the global params.                              |
    * |  - Followed by \p Codegen() to finalize the executable.                        |
-   * |  - Then the result runtime::Module can be constructed from the internal exec_. |
+   * |  - Then the result runtime::Module can be constructed by GetExecutable.        |
    * ----------------------------------------------------------------------------------
    *
    * \param mod Relay Module
-   * \param targets For heterogeneous compilation, it is a dictionary indicating device type
-   *                to target mapping. For homogeneous compilation, it is a singleton build target.
-   * \param target_host Host compilation target, if target is device.
+   * \param raw_targets List of available targets for running kernels. Any host target should
+   * be conveyed by the 'host' target field.
    */
-  void Lower(IRModule mod, TargetMap targets, Target target_host);
+  void Lower(IRModule mod, const Array& raw_targets);
 
-  /*! \brief Generate the machine code for lowered functions. */
-  void Codegen();
-
- protected:
   /*
-   * \brief Perform a series of optimizations on the input IR module.
+   * \brief Perform a series of optimizations on the input IR module. Can be used instead
+   * of Lower if wish to stop and observe optimized IRModule. Otherwise not needed on
+   * regular compilation flow.
    *
    * \param mod The input IRModule.
-   * \param targets For heterogeneous compilation, it is a dictionary indicating device type
-   *                to target mapping. For homogeneous compilation, it is a singleton build target.
-   * \param target_host Host compilation target.
+   * \param raw_targets List of available target for running kernels.
    *
    * \return The optimized IRModule.
    */
-  IRModule OptimizeModule(IRModule mod, const TargetMap& targets, const Target& target_host);
+  IRModule OptimizeModule(IRModule mod, const Array& raw_targets);
+
+  /*! \brief Generate the machine code for lowered functions. */
+  void Codegen();
+
+  /*! \brief Returns the runtime::Module containing the compiled VM code. */
+  runtime::Module GetExecutable() const;
+
+ protected:
 
+  void Setup(const Array& raw_targets);
+  void LowerImpl(IRModule mod);
   IRModule OptimizeModuleImpl(IRModule mod);
 
   transform::Sequential MemoryOpt(const VirtualDevice& host_virtual_device);
diff --git a/src/target/compilation_config.cc b/src/target/compilation_config.cc
index a56e0ad0777c..49dae60ac8b7 100644
--- a/src/target/compilation_config.cc
+++ b/src/target/compilation_config.cc
@@ -30,7 +30,6 @@ namespace tvm {
 TVM_REGISTER_NODE_TYPE(CompilationConfigNode);
 
 void CompilationConfigNode::VisitAttrs(AttrVisitor* v) {
-  v->Visit("legacy_target_map", &legacy_target_map);
   v->Visit("host_target", &host_target);
   v->Visit("primitive_targets", &primitive_targets);
   v->Visit("default_primitive_virtual_device", &default_primitive_virtual_device);
@@ -39,6 +38,27 @@ void CompilationConfigNode::VisitAttrs(AttrVisitor* v) {
   // NOTE: The virtual_device_cache_ is not accessible via FFI.
 }
 
+Target CompilationConfigNode::FindPrimitiveTargetOrFail(DLDeviceType device_type) const {
+  if (device_type < 0 && primitive_targets.size() == 1) {
+    // In the homogenous case don't be fussy with device types.
+    return primitive_targets.front();
+  }
+  ICHECK_GT(device_type, 0);
+  auto itr = std::find_if(
+      primitive_targets.begin(), primitive_targets.end(),
+      [device_type](const Target& target) { return target->kind->device_type == device_type; });
+  if (itr == primitive_targets.end()) {
+    std::stringstream msg;
+    msg << "No target is specified for device '" << runtime::DeviceName(device_type)
+        << "' mapped to device type " << device_type << ". The available targets are:" << std::endl;
+    for (const auto& target : primitive_targets) {
+      msg << "  " << target->kind->device_type << "-> " << target->ToDebugString() << std::endl;
+    }
+    LOG(FATAL) << msg.str();
+  }
+  return *itr;
+}
+
 VirtualDevice CompilationConfigNode::CanonicalVirtualDevice(
     const VirtualDevice& virtual_device) const {
   if (virtual_device->target.defined()) {
@@ -53,81 +73,107 @@ VirtualDevice CompilationConfigNode::CanonicalVirtualDevice(
                                                     target, virtual_device->memory_scope));
 }
 
-void CompilationConfigNode::EstablishDefaultVirtualDevices(const transform::PassContext& pass_ctx) {
+void CompilationConfigNode::Init(const transform::PassContext& pass_ctx,
+                                 const Array& raw_targets) {
+  VLOG_CONTEXT << "CompilationConfig";
+  ICHECK_GT(raw_targets.size(), 0U);
+
   //
-  // Gather the hints as to what our default device type for the 'host' should be, and
-  // create an appropriate target if we don't already have one.
+  // Decide on the host target.
   //
-  DLDeviceType host_device_type;
-  if (host_target.defined()) {
-    CHECK(!host_target->host.defined()) << "Host targets are not expected to have hosts";
-    host_device_type = static_cast(host_target->kind->device_type);
-    VLOG(1) << "Using the given host target " << host_target->ToDebugString() << " of device type "
-            << host_device_type << " for the host target";
-    for (const auto& primitive_target : primitive_targets) {
-      if (primitive_target->host.defined() &&
-          !StructuralEqual()(primitive_target->host, host_target)) {
-        VLOG(1) << "The primitive target " << primitive_target->ToDebugString()
-                << " already has a host which disagrees with the desired host target. It "
-                << "will be ignored.";
-      }
-    }
-  } else if (primitive_targets.size() == 1 && primitive_targets.front()->host.defined()) {
-    host_target = primitive_targets.front()->GetHost().value();
-    CHECK(!host_target->host.defined()) << "Host targets are not expected to have hosts";
-    host_device_type = static_cast(host_target->kind->device_type);
-    VLOG(1) << "Using the host of the unique primitive target, namely "
-            << host_target->ToDebugString() << " of device type " << host_device_type
-            << " for the host target";
-  } else if (primitive_targets.size() == 1 &&
-             primitive_targets.front()->kind->device_type == kDLCPU) {
-    // In the homogenous case without an explicit host target just use the given target so long as
-    // it's a CPU.
-    host_device_type = kDLCPU;
-    host_target = primitive_targets.front();
-    VLOG(1) << "Using the unique primitive target " << host_target->ToDebugString()
-            << " of device type " << host_device_type << " for the host target";
+
+  // Any CPU-like targets?
+  auto cpu_itr = std::find_if(raw_targets.begin(), raw_targets.end(), [](const Target& target) {
+    // TODO(tvm-team): AoT only works with kDLCPU device type. We can remove kDLHexagon
+    // here once we refactored kDLHexagon to kDLCPU.
+    return target->kind->device_type == kDLCPU || target->kind->device_type == kDLHexagon;
+  });
+
+  // Any targets with a host?
+  auto has_host_itr = std::find_if(raw_targets.begin(), raw_targets.end(),
+                                   [](const Target& target) { return target->host.defined(); });
+  if (has_host_itr != raw_targets.end()) {
+    // RULE A: If any raw target has a host, use the first such host for all the primitive
+    // targets.
+    host_target = Target((*has_host_itr)->GetHost().value(), /*host=*/Target());
+    VLOG(1) << "The target " << (*has_host_itr)->ToDebugString() << " supplies a host target "
+            << host_target->ToDebugString() << " of device type " << host_target->kind->device_type;
+  } else if (cpu_itr != raw_targets.end()) {
+    // RULE B: If any raw target is for a CPU-like device then also use that as the host.
+    host_target = Target(*cpu_itr, /*host=*/Target());
+    VLOG(1) << "Using target " << host_target->ToDebugString() << " of CPU-like device type "
+            << host_target->kind->device_type << " as the host target";
   } else {
-    // Fallback.
-    host_device_type = kDLCPU;
-    // Even if the list of available targets already includes one for kDLCPU we won't use it
-    // in the hetrogeneous case since its options may not be appropriate for host code
-    // (eg shape functions). Instead, create a fresh default Target.
-    host_target = MakeDefaultTarget(host_device_type);
-    VLOG(1) << "Using the default target " << host_target->ToDebugString() << " of device type "
-            << host_device_type << " for the host target";
+    // RULE C: Otherwise, create a default CPU host target.
+    host_target = MakeDefaultCPUTarget();
+    VLOG(1) << "Created a default target " << host_target->ToDebugString() << " of device type "
+            << host_target->kind->device_type << " for the host target";
   }
   ICHECK(host_target.defined());
   ICHECK(!host_target->host.defined());
 
-  if (host_device_type != kDLCPU) {
-    // I think we're on thin ice here until we've audited the code base for assumed kDLCPU.
-    VLOG(1) << "The host target is not a CPU.";
+  if (host_target->kind->device_type != kDLCPU) {
+    // I think we're on thin ice here until we've audited the code base for assumed CPU hosts.
+    VLOG(1) << "The host target is not a CPU. This is probably not going to work.";
   }
 
   //
   // Establish the host VirtualDevice.
   //
-  host_virtual_device =
-      virtual_device_cache_.Unique(VirtualDevice(host_device_type,
-                                                 /*virtual_device_id=*/0, host_target));
+  host_virtual_device = virtual_device_cache_.Unique(
+      VirtualDevice(static_cast(host_target->kind->device_type),
+                    /*virtual_device_id=*/0, host_target));
+  ICHECK(host_virtual_device.defined());
+  ICHECK(host_virtual_device->target.defined());
 
   //
-  // Now that we've settled on a host, we can set it as the host on all primitive targets.
+  // Now that we've settled on a host, we can set it as the host on all the raw targets.
   //
-  Array new_primitve_targets;
-  new_primitve_targets.reserve(primitive_targets.size());
-  for (const auto& primitive_target : primitive_targets) {
-    new_primitve_targets.push_back(Target(primitive_target, host_target));
+  primitive_targets.clear();
+  primitive_targets.reserve(raw_targets.size());
+  for (const auto& raw_target : raw_targets) {
+    if (raw_target->host.defined() &&
+        !StructuralEqual()(raw_target->host, host_target)) {
+      VLOG(1) << "The target " << raw_target->ToDebugString()
+              << " already has a host which disagrees with the desired host target. It "
+              << "will be overridden.";
+    }
+    primitive_targets.push_back(Target(raw_target, host_target));
   }
-  primitive_targets = new_primitve_targets;
+  ICHECK_GT(primitive_targets.size(), 0U);
 
   //
-  // Gather the hints as to what our default device type for primitives should be.
+  // Check the primitive_targets are ordered correctly re Target::IsExternalCodegenFor.
+  //
+  std::unordered_set primitive_target_device_types;
+  for (const auto& target : primitive_targets) {
+    primitive_target_device_types.emplace(static_cast(target->kind->device_type));
+  }
+  for (DLDeviceType device_type : primitive_target_device_types) {
+    Target first_primitive_target;
+    for (const auto& current_primitive_target : primitive_targets) {
+      if (current_primitive_target->kind->device_type != device_type) {
+        continue;
+      }
+      if (!first_primitive_target.defined()) {
+        first_primitive_target = current_primitive_target;
+        continue;
+      }
+      CHECK(current_primitive_target.IsExternalCodegenFor(first_primitive_target))
+          << "The first given target for device type " << device_type << " is "
+          << first_primitive_target->ToDebugString() << ", however a later target "
+          << current_primitive_target->ToDebugString()
+          << " for the same device type is not an external codegen target.";
+    }
+  }
+
+  //
+  // Decide on the default device type for primitives.
   //
   DLDeviceType default_primitive_device_type;
   Optional opt_fallback_dev = pass_ctx->GetConfig("relay.fallback_device_type");
   if (opt_fallback_dev) {
+    // RULE D: Respect the PassContext setting if given.
     const int64_t v = opt_fallback_dev.value()->value;
     CHECK_GT(v, 0)
         << "The 'relay.fallback_device_type' pass attribute is set to an invalid device type " << v;
@@ -135,16 +181,13 @@ void CompilationConfigNode::EstablishDefaultVirtualDevices(const transform::Pass
     VLOG(1) << "Using the 'relay.fallback_device_type' pass attribute "
             << default_primitive_device_type
             << " as the default device type for all primitive operations";
-  } else if (primitive_targets.size() == 1) {
-    // In the homogeneous case there's no free choice.
-    default_primitive_device_type =
-        static_cast(primitive_targets.front()->kind->device_type);
-    VLOG(1) << "Using the device type " << default_primitive_device_type
-            << " of the unique primitive target as the default device type for all primitive "
-            << "operations";
+  } else if (primitive_target_device_types.size() == 1) {
+    // RULE E: Since only one device in use it is the defacto default.
+    default_primitive_device_type = *primitive_target_device_types.begin();
+    VLOG(1) << "All primitive targets have the device type " << default_primitive_device_type
+            << " so that is also the default device type for all primitive operations.";
   } else {
-    // Fallback. Note that we'll require a primitive Target of kDLCPU device_type to be given
-    // and won't manufacture one out of thin air.
+    // RULE F: Fallback to CPU.
     default_primitive_device_type = kDLCPU;
     VLOG(1) << "Using " << default_primitive_device_type
             << " as the default device type for all primitive operations";
@@ -152,95 +195,57 @@ void CompilationConfigNode::EstablishDefaultVirtualDevices(const transform::Pass
 
   //
   // Establish the default primitive VirtualDevice, choosing a known Target to match the device
-  // type.
+  // type. We do not create a default target, it must already exist as a primitive target.
   //
   default_primitive_virtual_device = virtual_device_cache_.Unique(VirtualDevice(
       default_primitive_device_type,
       /*virtual_device_id=*/0, FindPrimitiveTargetOrFail(default_primitive_device_type)));
+
+  ICHECK(default_primitive_virtual_device.defined());
+  ICHECK(default_primitive_virtual_device->target.defined());
+
+  // Legacy: Some passes only support homogenous compilation and expect the target to be
+  // given by the global target context. Make this easy to detect.
+  optional_homogeneous_target =
+      primitive_targets.size() == 1 ? *primitive_targets.begin() : Target();
 }
 
-/* static */ Target CompilationConfigNode::MakeDefaultTarget(DLDeviceType device_type) {
-  std::string name = runtime::DeviceName(device_type);
-  if (name == "cpu") {
-    if (runtime::Registry::Get("codegen.LLVMModuleCreate")) {
-      // LLVM is available.
-      // TODO(mbs): More robust extension mechanism?
-      return Target("llvm");
-    } else {
-      // LLVM is not available.
-      // TODO(mbs): Already deprecated?
-      return Target("stackvm");
-    }
+/* static */ Target CompilationConfigNode::MakeDefaultCPUTarget() {
+  if (runtime::Registry::Get("codegen.LLVMModuleCreate")) {
+    // LLVM is available.
+    // TODO(mbs): More robust extension mechanism?
+    return Target("llvm");
   } else {
-    return Target(name);
+    // LLVM is not available.
+    // TODO(mbs): Already deprecated?
+    return Target("stackvm");
   }
 }
 
-Target CompilationConfigNode::FindPrimitiveTargetOrFail(DLDeviceType device_type) const {
-  auto itr = std::find_if(
-      primitive_targets.begin(), primitive_targets.end(),
-      [device_type](const Target& target) { return target->kind->device_type == device_type; });
-  CHECK(itr != primitive_targets.end()) << "No target for device type " << device_type << " in the "
-                                        << primitive_targets.size() << " given by the targets list";
-  return *itr;
-}
+TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
+    .set_dispatch([](const ObjectRef& ref, ReprPrinter* p) {
+      auto* node = ref.as();
+      p->stream << "Primitive targets:";
+      for (const auto& target : node->primitive_targets) {
+        p->stream << std::endl
+                  << "  " << target->kind->device_type << " |-> " << target->ToDebugString();
+      }
+      p->stream << std::endl
+                << "Default primitive virtual device: " << node->default_primitive_virtual_device;
+      p->stream << std::endl << "Host virtual device: " << node->host_virtual_device;
+    });
 
 CompilationConfig::CompilationConfig(const transform::PassContext& pass_ctx,
-                                     TargetMap legacy_target_map_arg,
-                                     Target optional_host_target_arg) {
-  VLOG_CONTEXT << "CompilationConfig";
-
+                                     const Array& raw_targets) {
   auto node = make_object();
-
-  for (const auto& pair : legacy_target_map_arg) {
-    VLOG(0) << "Available primitive target " << pair.first << " = " << pair.second->ToDebugString();
-  }
-  if (optional_host_target_arg.defined()) {
-    VLOG(0) << "Available host target " << optional_host_target_arg->ToDebugString();
-  }
-
-  // Capture the arguments in our preferred representation.
-  for (const auto& pair : legacy_target_map_arg) {
-    node->primitive_targets.push_back(pair.second);
-  }
-  node->host_target = optional_host_target_arg;
-
-  // Complete the targets vector and establish default scopes. After this primitive_targets will
-  // contain the definitive list of all required targets, target_host will be defined, and
-  // all primitive targets will have host target_host.
-  node->EstablishDefaultVirtualDevices(pass_ctx);
-
-  // LEGACY: Reconstruct the target map from all the primitive targets.
-  // Note that we require pointer equality between targets in legacy_target_map and
-  // primitive_targets.
-  for (const auto& primitive_target : node->primitive_targets) {
-    node->legacy_target_map.Set(Integer(primitive_target->kind->device_type), primitive_target);
-  }
-
-  ICHECK(node->default_primitive_virtual_device->target.defined());
-  ICHECK(node->host_virtual_device->target.defined());
-  ICHECK_GT(node->primitive_targets.size(), 0U);
-
-  // Legacy: Some passes only support homogenous compilation and expect the target to be
-  // given by the global target context. Make this easy to detect.
-  node->optional_homogeneous_target =
-      node->legacy_target_map.size() == 1 ? (*node->legacy_target_map.begin()).second : Target();
-
-  for (const auto& target : node->primitive_targets) {
-    VLOG(1) << "Target " << target->ToDebugString() << " of device type "
-            << target->kind->device_type << " is available for primitives";
-  }
-  VLOG(1) << "Using default primitive virtual device " << node->default_primitive_virtual_device;
-  VLOG(1) << "Using host virtual device " << node->host_virtual_device;
-
+  node->Init(pass_ctx, raw_targets);
   data_ = std::move(node);
 }
 
 TVM_REGISTER_GLOBAL("target.MakeCompilationConfig")
-    .set_body_typed([](const transform::PassContext& pass_ctx, TargetMap legacy_target_map,
-                       Target optional_host_target) -> CompilationConfig {
-      return CompilationConfig(pass_ctx, std::move(legacy_target_map),
-                               std::move(optional_host_target));
+    .set_body_typed([](const transform::PassContext& pass_ctx,
+                       const Array& raw_targets) -> CompilationConfig {
+      return CompilationConfig(pass_ctx, raw_targets);
     });
 
 }  // namespace tvm
diff --git a/src/target/target.cc b/src/target/target.cc
index a5c493a582ab..e26fed6c74f9 100644
--- a/src/target/target.cc
+++ b/src/target/target.cc
@@ -74,16 +74,6 @@ void CheckAndUpdateHostConsistency(Target* target, Target* host) {
   *host = (*target)->GetHost().value_or(Target());
 }
 
-void CheckAndUpdateHostConsistency(TargetMap* targets, Target* host) {
-  Map new_targets;
-  for (auto& it : *targets) {
-    auto target = it.second;
-    CheckAndUpdateHostConsistency(&target, host);
-    new_targets.Set(it.first, target);
-  }
-  *targets = new_targets;
-}
-
 void CheckAndUpdateHostConsistency(Map* targets, Target* host) {
   Map new_targets;
   for (auto& it : *targets) {
@@ -493,6 +483,23 @@ Target::Target(Target target, Target host) {
   data_ = std::move(n);
 }
 
+Target::Target(TargetKind kind, Optional host, String tag, Array keys,
+               Map attrs) {
+  auto data = runtime::make_object();
+  data->kind = std::move(kind);
+  data->host = std::move(host);
+  data->tag = std::move(tag);
+  data->keys = std::move(keys);
+  data->attrs = std::move(attrs);
+  data_ = std::move(data);
+}
+
+bool Target::IsExternalCodegenFor(const Target& that) const {
+  TargetKindAttrMap attr_map = TargetKind::GetAttrMap(::tvm::attr::kIsExternalCodegen);
+  return get()->kind->device_type == that->kind->device_type &&
+         attr_map.get(get()->kind, Bool(false)) && !attr_map.get(that->kind, Bool(false));
+}
+
 std::vector TargetNode::GetKeys() const {
   std::vector result;
   for (auto& expr : keys) {

From 8976c9f0af9250abf3ed3969cd1509cb5c558a63 Mon Sep 17 00:00:00 2001
From: mbs-octoml 
Date: Fri, 29 Apr 2022 07:51:02 -0700
Subject: [PATCH 02/11] - Working on unit tests

---
 src/target/compilation_config.cc              |  1 +
 .../relay/transforms/device_domains_test.cc   |  5 +--
 tests/cpp/target/compilation_config_test.cc   | 41 ++++++++++---------
 3 files changed, 23 insertions(+), 24 deletions(-)

diff --git a/src/target/compilation_config.cc b/src/target/compilation_config.cc
index 49dae60ac8b7..de0e8d02d0ed 100644
--- a/src/target/compilation_config.cc
+++ b/src/target/compilation_config.cc
@@ -92,6 +92,7 @@ void CompilationConfigNode::Init(const transform::PassContext& pass_ctx,
   // Any targets with a host?
   auto has_host_itr = std::find_if(raw_targets.begin(), raw_targets.end(),
                                    [](const Target& target) { return target->host.defined(); });
+
   if (has_host_itr != raw_targets.end()) {
     // RULE A: If any raw target has a host, use the first such host for all the primitive
     // targets.
diff --git a/tests/cpp/relay/transforms/device_domains_test.cc b/tests/cpp/relay/transforms/device_domains_test.cc
index dac109d23ea2..c5b2f26315b2 100644
--- a/tests/cpp/relay/transforms/device_domains_test.cc
+++ b/tests/cpp/relay/transforms/device_domains_test.cc
@@ -47,11 +47,8 @@ IRModule TestModule() {
 TEST(DeviceDomains, SmokeTest) {
   VirtualDevice cpu = VirtualDevice::ForDeviceType(kDLCPU);
   VirtualDevice cuda = VirtualDevice::ForDeviceType(kDLCUDA);
-  TargetMap target_map;
-  target_map.Set(Integer(static_cast(kDLCPU)), Target("llvm"));
-  target_map.Set(Integer(static_cast(kDLCUDA)), Target("cuda"));
   transform::PassContext ctxt = transform::PassContext::Create();
-  CompilationConfig config(ctxt, target_map, /*optional_host_target=*/{});
+  CompilationConfig config(ctxt, {Target("llvm"), Target("cuda")});
   DeviceDomains domains(config);
   IRModule mod = TestModule();
   Function f = Downcast(mod->Lookup("f"));
diff --git a/tests/cpp/target/compilation_config_test.cc b/tests/cpp/target/compilation_config_test.cc
index 2b1041b47d0b..8a31113441c4 100644
--- a/tests/cpp/target/compilation_config_test.cc
+++ b/tests/cpp/target/compilation_config_test.cc
@@ -34,46 +34,46 @@ Target TestExtDevTarget() { return Target("ext_dev"); }
 
 CompilationConfig TestCompilationConfig() {
   transform::PassContext pass_ctx = transform::PassContext::Create();
-  TargetMap legacy_target_map;
-  legacy_target_map.Set(Integer(static_cast(kDLCUDA)), TestCudaTarget());
-  legacy_target_map.Set(Integer(static_cast(kDLCPU)), TestCpuTarget());
-  return CompilationConfig(pass_ctx, legacy_target_map, TestDefaultCpuTarget());
+  Target host_target = TestDefaultCpuTarget();
+  Array raw_targets = {Target::WithHost(TestCudaTarget(), host_target),
+                               Target::WithHost(TestCpuTarget(), host_target)};
+  return CompilationConfig(pass_ctx, raw_targets);
 }
 
-TEST(CompilationConfig, Constructor_Homogeneous_FallbackCPUHost) {
+TEST(CompilationConfig, Constructor_Homogeneous_WithHost) {
   transform::PassContext pass_ctx = transform::PassContext::Create();
   Target host_target = TestDefaultCpuTarget();
-  Target cuda_target = TestCudaTarget();
-  TargetMap legacy_target_map;
-  legacy_target_map.Set(Integer(static_cast(kDLCUDA)), cuda_target);
-  CompilationConfig config(pass_ctx, legacy_target_map, /*optional_host_target_arg=*/{});
+  Target cuda_target = Target::WithHost(TestCudaTarget(), host_target);
+  CompilationConfig config(pass_ctx, {cuda_target});
 
   VirtualDevice expected_default_primitive_virtual_device(
       kDLCUDA, 0, Target::WithHost(cuda_target, host_target));
   VirtualDevice expected_host_virtual_device(kDLCPU, 0, host_target);
 
-  ASSERT_EQ(config->legacy_target_map.size(), 1);
-  EXPECT_TRUE(StructuralEqual()((*config->legacy_target_map.begin()).second,
-                                Target::WithHost(cuda_target, host_target)));
   EXPECT_TRUE(config->host_target.defined());
+  // RULE A: Picked the host.
   EXPECT_TRUE(StructuralEqual()(config->host_target, host_target));
+  EXPECT_TRUE(StructuralEqual()(config->host_virtual_device, expected_host_virtual_device));
+
   ASSERT_EQ(config->primitive_targets.size(), 1);
-  EXPECT_TRUE(
-      StructuralEqual()(config->primitive_targets[0], Target::WithHost(cuda_target, host_target)));
+  EXPECT_TRUE(StructuralEqual()(config->primitive_targets[0], cuda_target));
+
+  // RULE E: Pick device type of sole target.
   EXPECT_TRUE(StructuralEqual()(config->default_primitive_virtual_device,
                                 expected_default_primitive_virtual_device));
-  EXPECT_TRUE(StructuralEqual()(config->host_virtual_device, expected_host_virtual_device));
+
+  // Homogeneous case.
   ASSERT_TRUE(config->optional_homogeneous_target.defined());
-  EXPECT_TRUE(StructuralEqual()(config->optional_homogeneous_target,
-                                Target::WithHost(cuda_target, host_target)));
+  EXPECT_TRUE(StructuralEqual()(config->optional_homogeneous_target, cuda_target));
 }
 
-TEST(CompilationConfig, Constructor_Homegenoous_InnerHost) {
+#if 0
+TEST(CompilationConfig, Constructor_Homogenoous_InnerHost) {
   transform::PassContext pass_ctx = transform::PassContext::Create();
   Target host_target = TestCpuTarget();
   Target cuda_target = Target::WithHost(TestCudaTarget(), host_target);
-  TargetMap legacy_target_map;
-  legacy_target_map.Set(Integer(static_cast(kDLCUDA)), cuda_target);
+  Array raw_targets =
+      legacy_target_map.Set(Integer(static_cast(kDLCUDA)), cuda_target);
   CompilationConfig config(pass_ctx, legacy_target_map, /*optional_host_target_arg=*/{});
 
   EXPECT_TRUE(StructuralEqual()(config->host_target, host_target));
@@ -223,6 +223,7 @@ TEST(CompilationConfig, CanonicalVirtualDevice_NoMatchingTarget) {
   VirtualDevice no_such_target(kDLMetal);
   EXPECT_ANY_THROW(config->CanonicalVirtualDevice(no_such_target));
 }
+#endif
 
 }  // namespace
 }  // namespace tvm

From 08623d1bbe32139bfc64ec798b61298a81f93519 Mon Sep 17 00:00:00 2001
From: mbs-octoml 
Date: Fri, 29 Apr 2022 13:20:13 -0700
Subject: [PATCH 03/11] - Fix two Debug-only failures

---
 CMakeLists.txt                 |  6 ++++++
 tests/cpp/aot_metadata_test.cc | 16 ++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 666fefbe0cd2..90cc0f95185d 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -647,6 +647,12 @@ if(GTEST_FOUND)
   target_link_libraries(cpptest PRIVATE ${TVM_TEST_LIBRARY_NAME} GTest::GTest GTest::Main GTest::gmock pthread dl)
   set_target_properties(cpptest PROPERTIES EXCLUDE_FROM_ALL 1)
   set_target_properties(cpptest PROPERTIES EXCLUDE_FROM_DEFAULT_BUILD 1)
+  if(USE_RELAY_DEBUG)
+    target_compile_definitions(cpptest PRIVATE "USE_RELAY_DEBUG")
+    target_compile_definitions(cpptest PRIVATE "TVM_LOG_DEBUG")
+  else()
+    target_compile_definitions(cpptest PRIVATE "NDEBUG")
+  endif()
   # For some reason, compile definitions are not propagated correctly, so we manually add them here
   target_compile_definitions(cpptest PUBLIC $)
   gtest_discover_tests(cpptest)
diff --git a/tests/cpp/aot_metadata_test.cc b/tests/cpp/aot_metadata_test.cc
index b1dea64aaa9c..f8ce614b24cf 100644
--- a/tests/cpp/aot_metadata_test.cc
+++ b/tests/cpp/aot_metadata_test.cc
@@ -327,6 +327,10 @@ TEST(DiscoverArraysVisitor, DiscoverArrays) {
                                    DiscoveredNameEq("kTvmgenMetadata_pools")}));
 }
 
+// In Debug builds the _type_key is no longer inlined but also has no
+// link-time definition.
+#define WITH_TYPE_KEY 0
+
 template ::value, bool> =
               true>
@@ -338,18 +342,30 @@ class TVMObjectIsInstanceMatcher : public MatcherInterfaceIsInstance();
     if (!result) {
+#if WITH_TYPE_KEY
       (*os) << "is an instance of type " << T::ContainerType::_type_key;
+#else
+      (*os) << "is not of expected instance type";
+#endif
     }
 
     return result;
   }
 
   void DescribeTo(std::ostream* os) const override {
+#if WITH_TYPE_KEY
     (*os) << "is an instance of type " << T::ContainerType::_type_key;
+#else
+    (*os) << "is not of expected instance type";
+#endif
   }
 
   void DescribeNegationTo(std::ostream* os) const override {
+#if WITH_TYPE_KEY
     (*os) << "is not an instance of type " << T::ContainerType::_type_key;
+#else
+    (*os) << "is not of expected instance type";
+#endif
   }
 };
 

From f9d698857d0df2d0db2aa51f90ec9565758832c7 Mon Sep 17 00:00:00 2001
From: mbs-octoml 
Date: Mon, 2 May 2022 16:18:27 -0700
Subject: [PATCH 04/11] - Use Array in
 GraphExecutorCodegen/AOTExecutorCodegen ifaces instead   of CompilationConfig
 (don't want to bake it into any official APIs). - Started unit tests.

---
 include/tvm/target/compilation_config.h       |   9 +
 .../relay/backend/graph_executor_codegen.py   |  11 +-
 python/tvm/target/target.py                   |  32 +--
 src/relay/backend/aot_executor_codegen.cc     |  13 +-
 src/relay/backend/build_module.cc             |  13 +-
 src/relay/backend/graph_executor_codegen.cc   |  10 +-
 src/relay/backend/vm/compiler.h               |   9 +-
 src/target/compilation_config.cc              |   7 +-
 tests/cpp/relay_build_module_test.cc          |   5 +-
 tests/cpp/runtime_test.cc                     |   5 +-
 tests/cpp/target/compilation_config_test.cc   | 251 +++++++++++-------
 tests/python/driver/tvmc/test_target.py       |   2 -
 tests/python/relay/test_build_module.py       |   9 +-
 tests/python/unittest/test_target_target.py   |  43 ++-
 14 files changed, 265 insertions(+), 154 deletions(-)

diff --git a/include/tvm/target/compilation_config.h b/include/tvm/target/compilation_config.h
index c2d265596180..bb278c911358 100644
--- a/include/tvm/target/compilation_config.h
+++ b/include/tvm/target/compilation_config.h
@@ -36,6 +36,15 @@ namespace tvm {
  * APIs to a single internal representation. Also holds a cache of canonical \p VirtualDevices
  * so that structural equal virtual devices have pointer equal canonical virtual devices.
  *
+ * The construction of \p CompilationConfig is idempotent, in that given the same \p PassContext
+ * \p ctx and an arbitrary \p Array \p raw_targets:
+ *
+ * \code
+ *   CompilationConfig(ctxt, raw_targets)
+ *      is structurally equal to
+ *   CompilationConfig(ctxt, CompilationConfig(ctxt, raw_targets)->primitive_targets)
+ * \endcode
+ *
  * TODO(mbs): This is subject to change as we rework compilation options in general. This class
  * is probably better called a 'CompositeTarget', and may be better made a sub-class of Target or
  * some other common-target-root class.
diff --git a/python/tvm/relay/backend/graph_executor_codegen.py b/python/tvm/relay/backend/graph_executor_codegen.py
index e13d73c1a68b..9ce6a056a80d 100644
--- a/python/tvm/relay/backend/graph_executor_codegen.py
+++ b/python/tvm/relay/backend/graph_executor_codegen.py
@@ -54,15 +54,8 @@ def __init__(self, mod, target):
         self._setup(mod, target)
 
     def _setup(self, mod, target):
-        tgts = {}
-        if isinstance(target, dict):
-            for dev, tgt in target.items():
-                if not isinstance(tgt, (str, Target)):
-                    raise Exception("Unknown target type")
-                tgts[dev] = Target(tgt)
-        elif isinstance(target, (str, Target)):
-            tgts[_expr.IntImm("int32", 0)] = Target(target)
-        self._init(mod, tgts)
+        raw_targets = Target.canonicalize_target_and_host(target)
+        self._init(mod, raw_targets)
 
     def codegen(self, ir_module, func):
         """Compile a single function into a graph.
diff --git a/python/tvm/target/target.py b/python/tvm/target/target.py
index 6ffc898b8763..7b38bff94006 100644
--- a/python/tvm/target/target.py
+++ b/python/tvm/target/target.py
@@ -221,10 +221,10 @@ def list_kinds():
     @staticmethod
     def canonicalize_target(target):
         """Given a single target-like object, returns the TVM Target object representing it. Can convert from:
-         - None (to None).
-         - An existing TVM Target object.
-         - A string.
-         - A Python dictionary binding the target 'kind' and other attributes."""
+        - None (to None).
+        - An existing TVM Target object.
+        - A string.
+        - A Python dictionary binding the target 'kind' and other attributes."""
         if target is None:
             return None
         elif isinstance(target, Target):
@@ -617,7 +617,7 @@ def get_arch_version(cpu_ver):
     # Check for valid codegen cpu
     valid_hex = ["v65", "v66", "v67", "v67t", "v68", "v69"]
     try:
-        cpu_ver = cpu_ver[cpu_ver.index("v"):].lower()
+        cpu_ver = cpu_ver[cpu_ver.index("v") :].lower()
         assert cpu_ver in valid_hex
     except:
         msg = "{} is not a valid Hexagon version\nvalid versions include {}"
@@ -685,7 +685,7 @@ def validate_hvx_length(codegen_hvx, sim_options):
                 # If --hvx_length was specified, check HVX length of sim
                 # vs codegen
                 i = sim_options.index("hvx_length") + len("hvx_length") + 1
-                sim_hvx = sim_options[i: i + 3]
+                sim_hvx = sim_options[i : i + 3]
                 if sim_hvx != str(codegen_hvx):
                     msg = "sim hvx {} and codegen hvx {} mismatch!".format(sim_hvx, codegen_hvx)
                     # Set the stacklevel to the tvm.target.hexagon() call.
@@ -713,9 +713,9 @@ def validate_hvx_length(codegen_hvx, sim_options):
 
             # Regex match for allowed cpus
             valid_cpu_str_regex = (
-                    r"(?P
--.*\s)?(--m)?"
-                    + r"(?Pv6[25678])(?P[a-z])?"
-                    + r"(?P_[0-9]+)?(?P_rev[0-9])?\s?(?P--.*)?"
+                r"(?P
--.*\s)?(--m)?"
+                + r"(?Pv6[25678])(?P[a-z])?"
+                + r"(?P_[0-9]+)?(?P_rev[0-9])?\s?(?P--.*)?"
             )
             m = re.match(valid_cpu_str_regex, sim_options.lower())
             if not m:
@@ -724,13 +724,13 @@ def validate_hvx_length(codegen_hvx, sim_options):
             # Parse options into correct order
             cpu_attr = {x: str(m.groupdict()[x] or "") for x in m.groupdict()}
             sim_options = (
-                    cpu_attr["base_version"]
-                    + cpu_attr["sub_version"]
-                    + cpu_attr["l2_size"]
-                    + cpu_attr["rev"]
-                    + " "
-                    + cpu_attr["pre"]
-                    + cpu_attr["post"]
+                cpu_attr["base_version"]
+                + cpu_attr["sub_version"]
+                + cpu_attr["l2_size"]
+                + cpu_attr["rev"]
+                + " "
+                + cpu_attr["pre"]
+                + cpu_attr["post"]
             )
 
         return sim_cpu + " " + validate_hvx_length(hvx, sim_options)
diff --git a/src/relay/backend/aot_executor_codegen.cc b/src/relay/backend/aot_executor_codegen.cc
index a0c0aa940a0b..399d84594de9 100644
--- a/src/relay/backend/aot_executor_codegen.cc
+++ b/src/relay/backend/aot_executor_codegen.cc
@@ -965,7 +965,8 @@ class AOTExecutorCodegen : public MixedModeVisitor {
   std::unordered_map io_var_names_;
 
  public:
-  AOTExecutorCodegen(runtime::Module* mod, CompilationConfig config) : mod_(mod), config_(config) {}
+  AOTExecutorCodegen(runtime::Module* mod, const Array& targets)
+      : mod_(mod), config_(transform::PassContext::Current(), targets) {}
 
   LoweredOutput Codegen(IRModule mod, relay::Function func, String mod_name) {
     VLOG_CONTEXT << "AOT";
@@ -1206,10 +1207,10 @@ class AOTExecutorCodegenModule : public runtime::ModuleNode {
     if (name == "init") {
       return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
         ICHECK_EQ(args.num_args, 2) << "The expected of arguments are: "
-                                    << "runtime::Module mod and  Map targets";
+                                    << "runtime::Module mod and Array targets";
         void* mod = args[0];
-        CompilationConfig config = args[1];
-        init(mod, std::move(config));
+        Array targets = args[1];
+        init(mod, targets);
       });
     } else if (name == "codegen") {
       return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
@@ -1256,9 +1257,9 @@ class AOTExecutorCodegenModule : public runtime::ModuleNode {
   const char* type_key() const final { return "RelayGraphRuntimeCodegenModule"; }
 
  private:
-  void init(void* mod, CompilationConfig config) {
+  void init(void* mod, const Array& targets) {
     codegen_ =
-        std::make_shared(reinterpret_cast(mod), config);
+        std::make_shared(reinterpret_cast(mod), targets);
   }
 
   Array list_params_name() {
diff --git a/src/relay/backend/build_module.cc b/src/relay/backend/build_module.cc
index e742aaafbe07..9ddddeb389f3 100644
--- a/src/relay/backend/build_module.cc
+++ b/src/relay/backend/build_module.cc
@@ -61,7 +61,9 @@ struct BuildOutput {
 };
 
 struct ExecutorCodegen {
-  void Init(runtime::Module* m, CompilationConfig config) { CallFunc("init", m, config); }
+  void Init(runtime::Module* m, const Array& raw_targets) {
+    CallFunc("init", m, raw_targets);
+  }
 
   void Codegen(IRModule mod, const Function& func, String mod_name) {
     CallFunc("codegen", mod, func, mod_name);
@@ -228,8 +230,8 @@ class RelayBuildModule : public runtime::ModuleNode {
       });
     } else if (name == "optimize") {
       return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
-        ICHECK_EQ(args.num_args, 3);
-        *rv = this->Optimize(args[0], args[1], args[2]);
+        ICHECK_EQ(args.num_args, 2);
+        *rv = this->Optimize(args[0], args[1]);
       });
     } else {
       LOG(FATAL) << "Unknown packed function: " << name;
@@ -322,8 +324,7 @@ class RelayBuildModule : public runtime::ModuleNode {
    *
    * \return relay::IRModule The updated Relay IR module after optimization.
    */
-  IRModule Optimize(IRModule relay_module, const Array& raw_targets,
-                    const Target& optional_host_target) {
+  IRModule Optimize(IRModule relay_module, const Array& raw_targets) {
     VLOG_CONTEXT << "Optimize";
     config_ = CompilationConfig(PassContext ::Current(), raw_targets);
     VLOG(1) << "Using compilation config:" << std::endl << config_;
@@ -417,7 +418,7 @@ class RelayBuildModule : public runtime::ModuleNode {
 
     // Generate code for the updated function.
     executor_codegen_ = MakeExecutorCodegen(executor_->name);
-    executor_codegen_->Init(nullptr, config_);
+    executor_codegen_->Init(nullptr, config_->primitive_targets);
     executor_codegen_->Codegen(func_module, func, mod_name);
     executor_codegen_->UpdateOutput(&ret_);
     ret_.params = executor_codegen_->GetParams();
diff --git a/src/relay/backend/graph_executor_codegen.cc b/src/relay/backend/graph_executor_codegen.cc
index 580458da2385..2734439cddbd 100644
--- a/src/relay/backend/graph_executor_codegen.cc
+++ b/src/relay/backend/graph_executor_codegen.cc
@@ -190,8 +190,8 @@ class GraphOpNode : public GraphNode {
  */
 class GraphExecutorCodegen : public backend::MemoizedExprTranslator> {
  public:
-  GraphExecutorCodegen(runtime::Module* mod, CompilationConfig config)
-      : mod_(mod), config_(std::move(config)) {}
+  GraphExecutorCodegen(runtime::Module* mod, const Array& targets)
+      : mod_(mod), config_(transform::PassContext::Current(), targets) {}
 
   StorageInfo GetStorageInfo(const Expr& e) {
     size_t count = memory_plan_->expr_to_storage_info.count(e);
@@ -624,11 +624,11 @@ class GraphExecutorCodegenModule : public runtime::ModuleNode {
     if (name == "init") {
       return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
         ICHECK_EQ(args.num_args, 2) << "The expected of arguments are: "
-                                    << "runtime::Module mod and Map targets";
+                                    << "runtime::Module mod and Array targets";
         void* mod = args[0];
-        CompilationConfig config = args[1];
+        Array targets = args[1];
         codegen_ = std::make_shared(reinterpret_cast(mod),
-                                                          std::move(config));
+                                                          std::move(targets));
       });
     } else if (name == "codegen") {
       return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
diff --git a/src/relay/backend/vm/compiler.h b/src/relay/backend/vm/compiler.h
index 3cbafd73eae9..b1c977e52679 100644
--- a/src/relay/backend/vm/compiler.h
+++ b/src/relay/backend/vm/compiler.h
@@ -136,12 +136,19 @@ class VMCompiler : public runtime::ModuleNode {
   runtime::Module GetExecutable() const;
 
  protected:
-
+  /*! \brief Builds the executor and compilation config to match \p raw_targets. */
   void Setup(const Array& raw_targets);
+
+  /*! \brief Internal implementation of \p Lower. */
   void LowerImpl(IRModule mod);
+
+  /*! \brief Internal implementation of \p OptimizeModule. */
   IRModule OptimizeModuleImpl(IRModule mod);
 
+  /*! \brief Returns the passes which layout memory. */
   transform::Sequential MemoryOpt(const VirtualDevice& host_virtual_device);
+
+  /*! \brief Returns the passes which fuse then lower Relay primitive operators. */
   transform::Sequential FuseAndLowerOperators(const VirtualDevice& host_virtual_device);
 
   /*!
diff --git a/src/target/compilation_config.cc b/src/target/compilation_config.cc
index de0e8d02d0ed..d9d7d4b9675c 100644
--- a/src/target/compilation_config.cc
+++ b/src/target/compilation_config.cc
@@ -76,7 +76,7 @@ VirtualDevice CompilationConfigNode::CanonicalVirtualDevice(
 void CompilationConfigNode::Init(const transform::PassContext& pass_ctx,
                                  const Array& raw_targets) {
   VLOG_CONTEXT << "CompilationConfig";
-  ICHECK_GT(raw_targets.size(), 0U);
+  CHECK_GT(raw_targets.size(), 0U) << "Require at least one target";
 
   //
   // Decide on the host target.
@@ -133,8 +133,7 @@ void CompilationConfigNode::Init(const transform::PassContext& pass_ctx,
   primitive_targets.clear();
   primitive_targets.reserve(raw_targets.size());
   for (const auto& raw_target : raw_targets) {
-    if (raw_target->host.defined() &&
-        !StructuralEqual()(raw_target->host, host_target)) {
+    if (raw_target->host.defined() && !StructuralEqual()(raw_target->host, host_target)) {
       VLOG(1) << "The target " << raw_target->ToDebugString()
               << " already has a host which disagrees with the desired host target. It "
               << "will be overridden.";
@@ -183,7 +182,7 @@ void CompilationConfigNode::Init(const transform::PassContext& pass_ctx,
             << default_primitive_device_type
             << " as the default device type for all primitive operations";
   } else if (primitive_target_device_types.size() == 1) {
-    // RULE E: Since only one device in use it is the defacto default.
+    // RULE E: Since only one device in use there's no choice to make.
     default_primitive_device_type = *primitive_target_device_types.begin();
     VLOG(1) << "All primitive targets have the device type " << default_primitive_device_type
             << " so that is also the default device type for all primitive operations.";
diff --git a/tests/cpp/relay_build_module_test.cc b/tests/cpp/relay_build_module_test.cc
index 859f587f5a11..4814a1c7e7db 100644
--- a/tests/cpp/relay_build_module_test.cc
+++ b/tests/cpp/relay_build_module_test.cc
@@ -124,12 +124,11 @@ TEST(Relay, BuildModule) {
   auto build_f = build_mod.GetFunction("build", false);
   auto json_f = build_mod.GetFunction("get_graph_json", false);
   auto mod_f = build_mod.GetFunction("get_module", false);
-  Map targets;
   Target llvm_tgt = Target("llvm");
-  targets.Set(0, llvm_tgt);
+  Array targets = {llvm_tgt};
   auto relay_mod = tvm::IRModule::FromExpr(func);
   ICHECK(relay_mod.defined()) << "Module must be defined";
-  build_f(relay_mod, targets, llvm_tgt, Executor::Create("graph"), Runtime::Create("cpp"),
+  build_f(relay_mod, targets, Executor::Create("graph"), Runtime::Create("cpp"),
           WorkspaceMemoryPools(), "");
   std::string json = json_f();
   tvm::runtime::Module mod = mod_f();
diff --git a/tests/cpp/runtime_test.cc b/tests/cpp/runtime_test.cc
index 57686baf7b46..33f44f4f3e54 100644
--- a/tests/cpp/runtime_test.cc
+++ b/tests/cpp/runtime_test.cc
@@ -110,12 +110,11 @@ TEST(Runtime, ZeroCopy) {
   auto build_f = build_mod.GetFunction("build", false);
   auto json_f = build_mod.GetFunction("get_graph_json", false);
   auto mod_f = build_mod.GetFunction("get_module", false);
-  Map targets;
   Target llvm_tgt = Target("llvm");
-  targets.Set(0, llvm_tgt);
+  Array targets = {llvm_tgt};
   auto relay_mod = tvm::IRModule::FromExpr(func);
   ICHECK(relay_mod.defined()) << "Module must be defined";
-  build_f(relay_mod, targets, llvm_tgt, Executor::Create("graph"), Runtime::Create("cpp"),
+  build_f(relay_mod, targets, Executor::Create("graph"), Runtime::Create("cpp"),
           WorkspaceMemoryPools(), "");
   // create graph executor
   std::string json = json_f();
diff --git a/tests/cpp/target/compilation_config_test.cc b/tests/cpp/target/compilation_config_test.cc
index 8a31113441c4..4568d11d6232 100644
--- a/tests/cpp/target/compilation_config_test.cc
+++ b/tests/cpp/target/compilation_config_test.cc
@@ -32,160 +32,226 @@ Target TestDefaultCpuTarget() { return Target("llvm"); }
 
 Target TestExtDevTarget() { return Target("ext_dev"); }
 
+TVM_REGISTER_TARGET_KIND("test_ext_codegen_1", kDLCUDA)
+    .set_attr(tvm::attr::kIsExternalCodegen, Bool(true));
+
+TVM_REGISTER_TARGET_KIND("test_ext_codegen_2", kDLCUDA)
+    .set_attr(tvm::attr::kIsExternalCodegen, Bool(true));
+
+Target TestExtCodegenTarget1() { return Target("test_ext_codegen_1"); }
+Target TestExtCodegenTarget2() { return Target("test_ext_codegen_2"); }
+
 CompilationConfig TestCompilationConfig() {
   transform::PassContext pass_ctx = transform::PassContext::Create();
   Target host_target = TestDefaultCpuTarget();
-  Array raw_targets = {Target::WithHost(TestCudaTarget(), host_target),
-                               Target::WithHost(TestCpuTarget(), host_target)};
-  return CompilationConfig(pass_ctx, raw_targets);
+  Target cuda_target = Target::WithHost(TestCudaTarget(), host_target);
+  Target cpu_target = Target::WithHost(TestCpuTarget(), host_target);
+  return CompilationConfig(pass_ctx, {cuda_target, cpu_target});
 }
 
-TEST(CompilationConfig, Constructor_Homogeneous_WithHost) {
+TEST(CompilationConfig, Constructor_Heterogeneous_RuleA_RuleF_ReplaceHost) {
   transform::PassContext pass_ctx = transform::PassContext::Create();
+
   Target host_target = TestDefaultCpuTarget();
   Target cuda_target = Target::WithHost(TestCudaTarget(), host_target);
-  CompilationConfig config(pass_ctx, {cuda_target});
+  Target ignored_target = TestExtDevTarget();
+  Target raw_cpu_target = Target::WithHost(TestCpuTarget(), ignored_target);
+  CompilationConfig config(pass_ctx, {cuda_target, raw_cpu_target});
 
-  VirtualDevice expected_default_primitive_virtual_device(
-      kDLCUDA, 0, Target::WithHost(cuda_target, host_target));
+  Target cpu_target = Target::WithHost(TestCpuTarget(), host_target);
+  VirtualDevice expected_default_primitive_virtual_device(kDLCPU, 0, cpu_target);
   VirtualDevice expected_host_virtual_device(kDLCPU, 0, host_target);
 
+  // Host is chosen as per Rule A.
   EXPECT_TRUE(config->host_target.defined());
-  // RULE A: Picked the host.
   EXPECT_TRUE(StructuralEqual()(config->host_target, host_target));
   EXPECT_TRUE(StructuralEqual()(config->host_virtual_device, expected_host_virtual_device));
 
-  ASSERT_EQ(config->primitive_targets.size(), 1);
+  ASSERT_EQ(config->primitive_targets.size(), 2);
   EXPECT_TRUE(StructuralEqual()(config->primitive_targets[0], cuda_target));
+  // The host is taken from first raw target and overwritten in second.
+  EXPECT_TRUE(StructuralEqual()(config->primitive_targets[1], cpu_target));
 
-  // RULE E: Pick device type of sole target.
+  // Default primitive virtual device chosen as per Rule F
   EXPECT_TRUE(StructuralEqual()(config->default_primitive_virtual_device,
                                 expected_default_primitive_virtual_device));
 
-  // Homogeneous case.
-  ASSERT_TRUE(config->optional_homogeneous_target.defined());
-  EXPECT_TRUE(StructuralEqual()(config->optional_homogeneous_target, cuda_target));
+  // Heterogeneous case.
+  ASSERT_FALSE(config->optional_homogeneous_target.defined());
 }
 
-#if 0
-TEST(CompilationConfig, Constructor_Homogenoous_InnerHost) {
+TEST(CompilationConfig, Constructor_Homogeneous_RuleA_RuleE) {
   transform::PassContext pass_ctx = transform::PassContext::Create();
-  Target host_target = TestCpuTarget();
+
+  Target host_target = TestDefaultCpuTarget();
   Target cuda_target = Target::WithHost(TestCudaTarget(), host_target);
-  Array raw_targets =
-      legacy_target_map.Set(Integer(static_cast(kDLCUDA)), cuda_target);
-  CompilationConfig config(pass_ctx, legacy_target_map, /*optional_host_target_arg=*/{});
+  CompilationConfig config(pass_ctx, {cuda_target});
+
+  VirtualDevice expected_default_primitive_virtual_device(kDLCUDA, 0, cuda_target);
+  VirtualDevice expected_host_virtual_device(kDLCPU, 0, host_target);
 
+  // Host is chose as per Rule A.
+  EXPECT_TRUE(config->host_target.defined());
   EXPECT_TRUE(StructuralEqual()(config->host_target, host_target));
-}
+  EXPECT_TRUE(StructuralEqual()(config->host_virtual_device, expected_host_virtual_device));
 
-TEST(CompilationConfig, Constructor_Homogenous_CPUHost) {
-  transform::PassContext pass_ctx = transform::PassContext::Create();
-  Target host_target = TestCpuTarget();
-  Target cpu_target = TestCpuTarget();
-  TargetMap legacy_target_map;
-  legacy_target_map.Set(Integer(static_cast(kDLCPU)),
-                        Target::WithHost(cpu_target, host_target));
-  CompilationConfig config(pass_ctx, legacy_target_map, /*optional_host_target_arg=*/{});
+  ASSERT_EQ(config->primitive_targets.size(), 1);
+  EXPECT_TRUE(StructuralEqual()(config->primitive_targets[0], cuda_target));
+
+  // Default primitive virtual device chose as per rule E.
+  EXPECT_TRUE(StructuralEqual()(config->default_primitive_virtual_device,
+                                expected_default_primitive_virtual_device));
 
-  EXPECT_TRUE(StructuralEqual()(config->host_target, cpu_target));
+  // Homogeneous case.
   ASSERT_TRUE(config->optional_homogeneous_target.defined());
-  EXPECT_TRUE(StructuralEqual()(config->optional_homogeneous_target,
-                                Target::WithHost(cpu_target, cpu_target)));
+  EXPECT_TRUE(StructuralEqual()(config->optional_homogeneous_target, cuda_target));
 }
 
-TEST(CompilationConfig, Constructor_Hetrogeneous_FallbackCPUHost) {
+TEST(CompilationConfig, Constructor_Heterogeneous_RuleB_RuleD) {
   transform::PassContext pass_ctx = transform::PassContext::Create();
   pass_ctx->config.Set("relay.fallback_device_type", Integer(static_cast(kDLCUDA)));
-  Target host_target = TestDefaultCpuTarget();
-  Target cuda_target = TestCudaTarget();
-  Target cpu_target = TestCpuTarget();
-  TargetMap legacy_target_map;
-  legacy_target_map.Set(Integer(static_cast(kDLCPU)),
-                        Target::WithHost(cpu_target, host_target));
-  legacy_target_map.Set(Integer(static_cast(kDLCUDA)),
-                        Target::WithHost(cuda_target, host_target));
-  CompilationConfig config(pass_ctx, legacy_target_map, /*optional_host_target_arg=*/{});
-
-  VirtualDevice expected_default_primitive_virtual_device(
-      kDLCUDA, 0, Target::WithHost(cuda_target, host_target));
+
+  Target raw_cuda_target = TestCudaTarget();
+  Target raw_cpu_target = TestCpuTarget();
+  CompilationConfig config(pass_ctx, {raw_cuda_target, raw_cpu_target});
+
+  Target host_target = TestCpuTarget();
+  Target cuda_target = Target::WithHost(TestCudaTarget(), host_target);
+  Target cpu_target = Target::WithHost(TestCpuTarget(), host_target);
+
+  VirtualDevice expected_default_primitive_virtual_device(kDLCUDA, 0, cuda_target);
   VirtualDevice expected_host_virtual_device(kDLCPU, 0, host_target);
 
-  ASSERT_EQ(config->legacy_target_map.size(), 2);
-  for (const auto& pair : config->legacy_target_map) {
-    if (pair.first->value == kDLCPU) {
-      EXPECT_TRUE(StructuralEqual()(pair.second, Target::WithHost(cpu_target, host_target)));
-    } else if (pair.first->value == kDLCUDA) {
-      EXPECT_TRUE(StructuralEqual()(pair.second, Target::WithHost(cuda_target, host_target)));
-    }
-  }
+  // Host is chosen as per Rule B.
   EXPECT_TRUE(config->host_target.defined());
   EXPECT_TRUE(StructuralEqual()(config->host_target, host_target));
+  EXPECT_TRUE(StructuralEqual()(config->host_virtual_device, expected_host_virtual_device));
+
+  ASSERT_EQ(config->primitive_targets.size(), 2);
+  EXPECT_TRUE(StructuralEqual()(config->primitive_targets[0], cuda_target));
+  EXPECT_TRUE(StructuralEqual()(config->primitive_targets[1], cpu_target));
+
+  // Default primitive virtual device chosen as per Rule D
   EXPECT_TRUE(StructuralEqual()(config->default_primitive_virtual_device,
                                 expected_default_primitive_virtual_device));
-  EXPECT_TRUE(StructuralEqual()(config->host_virtual_device, expected_host_virtual_device));
-  EXPECT_FALSE(config->optional_homogeneous_target.defined());
+
+  // Heterogeneous case.
+  ASSERT_FALSE(config->optional_homogeneous_target.defined());
 }
 
-TEST(CompilationConfig, Constructor_Hetrogeneous_ExplicitHost) {
+TEST(CompilationConfig, Constructor_Homogeneous_RuleC_RuleE) {
   transform::PassContext pass_ctx = transform::PassContext::Create();
-  pass_ctx->config.Set("relay.fallback_device_type", Integer(static_cast(kDLCUDA)));
-  Target host_target = TestCpuTarget();
-  Target cuda_target = TestCudaTarget();
-  Target cpu_target = TestCpuTarget();
-  TargetMap legacy_target_map;
-  legacy_target_map.Set(Integer(static_cast(kDLCPU)),
-                        Target::WithHost(cpu_target, host_target));
-  legacy_target_map.Set(Integer(static_cast(kDLCUDA)),
-                        Target::WithHost(cuda_target, host_target));
-  CompilationConfig config(pass_ctx, legacy_target_map, host_target);
-
-  VirtualDevice expected_default_primitive_virtual_device(
-      kDLCUDA, 0, Target::WithHost(cuda_target, host_target));
+
+  Target raw_cuda_target = TestCudaTarget();
+  CompilationConfig config(pass_ctx, {raw_cuda_target});
+
+  Target host_target = TestDefaultCpuTarget();
+  Target cuda_target = Target::WithHost(TestCudaTarget(), host_target);
+  Target cpu_target = Target::WithHost(TestDefaultCpuTarget(), host_target);
+
+  VirtualDevice expected_default_primitive_virtual_device(kDLCUDA, 0, cuda_target);
   VirtualDevice expected_host_virtual_device(kDLCPU, 0, host_target);
 
-  ASSERT_EQ(config->legacy_target_map.size(), 2);
-  for (const auto& pair : config->legacy_target_map) {
-    if (pair.first->value == kDLCPU) {
-      EXPECT_TRUE(StructuralEqual()(pair.second, Target::WithHost(cpu_target, host_target)));
-    } else if (pair.first->value == kDLCUDA) {
-      EXPECT_TRUE(StructuralEqual()(pair.second, Target::WithHost(cuda_target, host_target)));
-    }
-  }
+  // Host is chosen as per Rule C.
   EXPECT_TRUE(config->host_target.defined());
   EXPECT_TRUE(StructuralEqual()(config->host_target, host_target));
-  ASSERT_EQ(config->primitive_targets.size(), 2);
+  EXPECT_TRUE(StructuralEqual()(config->host_virtual_device, expected_host_virtual_device));
+
+  ASSERT_EQ(config->primitive_targets.size(), 1);
+  EXPECT_TRUE(StructuralEqual()(config->primitive_targets[0], cuda_target));
+
+  // Default primitive virtual device chosen as per Rule E
   EXPECT_TRUE(StructuralEqual()(config->default_primitive_virtual_device,
                                 expected_default_primitive_virtual_device));
-  EXPECT_TRUE(StructuralEqual()(config->host_virtual_device, expected_host_virtual_device));
-  EXPECT_FALSE(config->optional_homogeneous_target.defined());
+
+  // Homogeneous case.
+  ASSERT_TRUE(config->optional_homogeneous_target.defined());
+  EXPECT_TRUE(StructuralEqual()(config->optional_homogeneous_target, cuda_target));
+}
+
+TEST(CompilationConfig, Constructor_Heterogeneous_CorrectOrdering) {
+  transform::PassContext pass_ctx = transform::PassContext::Create();
+
+  Target host_target = TestDefaultCpuTarget();
+  Target cuda_target = Target::WithHost(TestCudaTarget(), host_target);
+  Target ext_codegen1_target = Target::WithHost(TestExtCodegenTarget1(), host_target);
+  Target ext_codegen2_target = Target::WithHost(TestExtCodegenTarget2(), host_target);
+  CompilationConfig config(pass_ctx, {cuda_target, ext_codegen1_target, ext_codegen2_target});
+
+  ASSERT_EQ(config->primitive_targets.size(), 3);
+  EXPECT_TRUE(StructuralEqual()(config->primitive_targets[0], cuda_target));
+  EXPECT_TRUE(StructuralEqual()(config->primitive_targets[1], ext_codegen1_target));
+  EXPECT_TRUE(StructuralEqual()(config->primitive_targets[2], ext_codegen2_target));
+}
+
+TEST(CompilationConfig, Constructor_Heterogeneous_InvalidOrdering) {
+  transform::PassContext pass_ctx = transform::PassContext::Create();
+
+  Target host_target = TestDefaultCpuTarget();
+  Target ext_codegen1_target = Target::WithHost(TestExtCodegenTarget1(), host_target);
+  Target cuda_target = Target::WithHost(TestCudaTarget(), host_target);
+  Target ext_codegen2_target = Target::WithHost(TestExtCodegenTarget2(), host_target);
+
+  EXPECT_ANY_THROW(
+      CompilationConfig(pass_ctx, {ext_codegen1_target, cuda_target, ext_codegen2_target}));
+}
+
+TEST(CompilationConfig, Constructor_NoTargets) {
+  transform::PassContext pass_ctx = transform::PassContext::Create();
+  EXPECT_ANY_THROW(CompilationConfig(pass_ctx, {}));
 }
 
 TEST(CompilationConfig, Constructor_InvalidAttribute) {
   transform::PassContext pass_ctx = transform::PassContext::Create();
   pass_ctx->config.Set("relay.fallback_device_type", Integer(static_cast(kInvalidDeviceType)));
-  TargetMap legacy_target_map;
-  legacy_target_map.Set(Integer(static_cast(kDLCUDA)), TestCudaTarget());
-  EXPECT_ANY_THROW(
-      CompilationConfig config(pass_ctx, legacy_target_map, /*optional_host_target_arg=*/{}));
+
+  Target cuda_target = Target::WithHost(TestCudaTarget(), TestDefaultCpuTarget());
+  EXPECT_ANY_THROW(CompilationConfig(pass_ctx, {cuda_target}));
 }
 
 TEST(CompilationConfig, Constructor_NoMatchingPrimitiveTarget) {
   transform::PassContext pass_ctx = transform::PassContext::Create();
   pass_ctx->config.Set("relay.fallback_device_type", Integer(static_cast(kDLMetal)));
-  TargetMap legacy_target_map;
-  legacy_target_map.Set(Integer(static_cast(kDLCUDA)), TestCudaTarget());
-  EXPECT_ANY_THROW(
-      CompilationConfig config(pass_ctx, legacy_target_map, /*optional_host_target_arg=*/{}));
+  Target host_target = TestDefaultCpuTarget();
+  Target cuda_target = Target::WithHost(TestCudaTarget(), host_target);
+  EXPECT_ANY_THROW(CompilationConfig(pass_ctx, {cuda_target}));
 }
 
 TEST(CompilationConfig, Constructor_DefaultNoMatchingPrimitiveTarget) {
   transform::PassContext pass_ctx = transform::PassContext::Create();
-  TargetMap legacy_target_map;
-  legacy_target_map.Set(Integer(static_cast(kDLCUDA)), TestCudaTarget());
-  legacy_target_map.Set(Integer(static_cast(kDLExtDev)), TestExtDevTarget());
-  EXPECT_ANY_THROW(
-      CompilationConfig config(pass_ctx, legacy_target_map, /*optional_host_target_arg=*/{}));
+  Target host_target = TestDefaultCpuTarget();
+  Target cuda_target = Target::WithHost(TestCudaTarget(), host_target);
+  Target ext_target = Target::WithHost(TestExtDevTarget(), host_target);
+  EXPECT_ANY_THROW(CompilationConfig config(pass_ctx, {cuda_target, ext_target}));
+}
+
+TEST(CompilationConfig, Constructor_Idempotent) {
+  transform::PassContext pass_ctx = transform::PassContext::Create();
+
+  Target host_target = TestDefaultCpuTarget();
+  Target cuda_target = Target::WithHost(TestCudaTarget(), host_target);
+  Target ignored_target = TestExtDevTarget();
+  Target raw_cpu_target = Target::WithHost(TestCpuTarget(), ignored_target);
+  CompilationConfig orig_config(pass_ctx, {cuda_target, raw_cpu_target});
+
+  CompilationConfig reconstructed_config(pass_ctx, orig_config->primitive_targets);
+
+  ASSERT_EQ(orig_config->primitive_targets.size(), reconstructed_config->primitive_targets.size());
+  ASSERT_TRUE(StructuralEqual()(orig_config->primitive_targets[0],
+                                reconstructed_config->primitive_targets[0]));
+  ASSERT_TRUE(StructuralEqual()(orig_config->primitive_targets[1],
+                                reconstructed_config->primitive_targets[1]));
+}
+
+TEST(CompilationConfig, FindPrimitiveTargetOrFail_Valid) {
+  CompilationConfig config = TestCompilationConfig();
+  Target cpu_target = Target::WithHost(TestCpuTarget(), TestDefaultCpuTarget());
+  ASSERT_TRUE(StructuralEqual()(config->FindPrimitiveTargetOrFail(kDLCPU), cpu_target));
+}
+
+TEST(CompilationConfig, FindPrimitiveTargetOrFail_Invalid) {
+  CompilationConfig config = TestCompilationConfig();
+  EXPECT_ANY_THROW(config->FindPrimitiveTargetOrFail(kDLMetal));
 }
 
 TEST(CompilationConfig, CanonicalVirtualDevice) {
@@ -223,7 +289,6 @@ TEST(CompilationConfig, CanonicalVirtualDevice_NoMatchingTarget) {
   VirtualDevice no_such_target(kDLMetal);
   EXPECT_ANY_THROW(config->CanonicalVirtualDevice(no_such_target));
 }
-#endif
 
 }  // namespace
 }  // namespace tvm
diff --git a/tests/python/driver/tvmc/test_target.py b/tests/python/driver/tvmc/test_target.py
index 54913e080b76..f6431a6cb20d 100644
--- a/tests/python/driver/tvmc/test_target.py
+++ b/tests/python/driver/tvmc/test_target.py
@@ -15,8 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import pytest
-
 from tvm.driver.tvmc import TVMCException
 from tvm.driver.tvmc.target import target_from_cli, tokenize_target, parse_target
 
diff --git a/tests/python/relay/test_build_module.py b/tests/python/relay/test_build_module.py
index 747062201fea..b03e760a968a 100644
--- a/tests/python/relay/test_build_module.py
+++ b/tests/python/relay/test_build_module.py
@@ -64,12 +64,11 @@ def test_build_relay_graph_():
     """Test to build a simple relay graph by using APIs directly"""
 
     def build_graph(mod, target):
-        target = relay.build_module.build_target_by_device_type_map(target)
         target, target_host = tvm.target.Target.check_and_update_host_consist(target)
-        mod, _ = relay.optimize(mod, target, None)
+        mod, _ = relay.optimize(mod, target)
         grc = graph_executor_codegen.GraphExecutorCodegen(None, target)
         _, lowered_funcs, _ = grc.codegen(mod, mod["main"])
-        _ = relay.backend._backend.build(lowered_funcs, target, target_host)
+        _ = relay.backend._backend.build(lowered_funcs, target)
 
     def add(shape, dtype):
         lhs = relay.var("A", shape=shape, dtype=dtype)
@@ -83,4 +82,6 @@ def add(shape, dtype):
 
 
 if __name__ == "__main__":
-    pytest.main()
+    import sys
+
+    sys.exit(pytest.main([__file__] + sys.argv[1:]))
diff --git a/tests/python/unittest/test_target_target.py b/tests/python/unittest/test_target_target.py
index 99cdb86314e7..9f5f62b8b991 100644
--- a/tests/python/unittest/test_target_target.py
+++ b/tests/python/unittest/test_target_target.py
@@ -48,10 +48,11 @@ def test_all_targets_device_type_verify():
     all_targets = [tvm.target.Target(t) for t in tvm.target.Target.list_kinds()]
 
     for tgt in all_targets:
-        # skip target hook
+        # skip targets with hooks or otherwise intended to be used with external codegen
         relay_to_tir = tgt.get_kind_attr("RelayToTIR")
         tir_to_runtime = tgt.get_kind_attr("TIRToRuntime")
-        if relay_to_tir is not None or tir_to_runtime is not None:
+        is_external_codegen = tgt.get_kind_attr("is_external_codegen")
+        if relay_to_tir is not None or tir_to_runtime is not None or is_external_codegen:
             continue
 
         if tgt.kind.name not in tvm._ffi.runtime_ctypes.Device.STR2MASK:
@@ -404,6 +405,44 @@ def test_check_and_update_host_consist_4():
     assert host_2.kind.name == "llvm"
 
 
+def test_canonicalize_target_and_host_0():
+    with pytest.raises(AssertionError):
+        Target.canonicalize_target_and_host(None)
+
+
+def test_canonicalize_target_and_host_1():
+    raw_targets = Target.canonicalize_target_and_host({"kind": "llvm"})
+    assert len(raw_targets) == 1
+    assert raw_targets[0].kind.name == "llvm"
+
+
+def test_canonicalize_target_and_host_2():
+    raw_targets = Target.canonicalize_target_and_host({1: "llvm", 2: "cuda"})
+    assert len(raw_targets) == 2
+    assert raw_targets[0].kind.name == "llvm"
+    assert raw_targets[1].kind.name == "cuda"
+
+
+def test_canonicalize_target_and_host_3():
+    raw_targets = Target.canonicalize_target_and_host(["llvm", "cuda"])
+    assert len(raw_targets) == 2
+    assert raw_targets[0].kind.name == "llvm"
+    assert raw_targets[1].kind.name == "cuda"
+
+
+def test_canonicalize_target_and_host_4():
+    raw_targets = Target.canonicalize_target_and_host("llvm")
+    assert len(raw_targets) == 1
+    assert raw_targets[0].kind.name == "llvm"
+
+
+def test_canonicalize_target_and_host_5():
+    raw_targets = Target.canonicalize_target_and_host("cuda", "llvm")
+    assert len(raw_targets) == 1
+    assert raw_targets[0].kind.name == "cuda"
+    assert raw_targets[0].host.kind.name == "llvm"
+
+
 def test_target_attr_bool_value():
     target0 = Target("vulkan --supports_float16=True")
     assert target0.attrs["supports_float16"] == 1

From 9082e6d4c3900e383a49c865e05b7b62be5b8e4b Mon Sep 17 00:00:00 2001
From: mbs-octoml 
Date: Mon, 2 May 2022 16:35:27 -0700
Subject: [PATCH 05/11] - Lints

---
 python/tvm/relay/backend/vm.py |  3 ---
 python/tvm/target/target.py    | 20 +++++++++++---------
 2 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/python/tvm/relay/backend/vm.py b/python/tvm/relay/backend/vm.py
index 72613421f11f..256293f6538b 100644
--- a/python/tvm/relay/backend/vm.py
+++ b/python/tvm/relay/backend/vm.py
@@ -20,11 +20,8 @@
 
 Implements a Python interface to compiling and executing on the Relay VM.
 """
-import warnings
-
 import numpy as np
 
-import tvm
 import tvm.runtime.ndarray as _nd
 import tvm.runtime.vm as vm_rt
 from tvm import autotvm
diff --git a/python/tvm/target/target.py b/python/tvm/target/target.py
index 7b38bff94006..bd43bd3b7abe 100644
--- a/python/tvm/target/target.py
+++ b/python/tvm/target/target.py
@@ -220,7 +220,8 @@ def list_kinds():
 
     @staticmethod
     def canonicalize_target(target):
-        """Given a single target-like object, returns the TVM Target object representing it. Can convert from:
+        """Given a single target-like object, returns the TVM Target object representing it.
+        Can convert from:
         - None (to None).
         - An existing TVM Target object.
         - A string.
@@ -234,11 +235,12 @@ def canonicalize_target(target):
 
     @staticmethod
     def canonicalize_multi_targets(multi_targets):
-        """Given a single or collection of target-like objects, returns a TVM Array of Target objects representing
-        then. Can convert from:
+        """Given a single or collection of target-like objects, returns a TVM Array of Target
+        objects representing then. Can convert from:
          - None (to None).
          - A single target-like object in a form recognized by canonicalize_target.
-         - A Python list or TVM Array of target-like objects in a form recognized by canonicalize_target.
+         - A Python list or TVM Array of target-like objects in a form recognized by
+           canonicalize_target.
          - A Python dict or TVM Map from TVM IntImm objects representing device types to
            a target-like object in a form recognized by canonicalize_target."""
         if multi_targets is None:
@@ -255,11 +257,11 @@ def canonicalize_multi_targets(multi_targets):
 
     @staticmethod
     def canonicalize_target_and_host(target, target_host=None):
-        """Returns a TVM Array capturing target and target_host. The given target can be in any
-        form recognized by Target.canonicalize_target or Target.canonicalize_multi_targets. If given
-        target_host can be in any form recognized by Target.canonicalize_target. If target_host is given
-        it will be set as the 'host' in each result Target object (and a warning given).
-        """
+        """Returns a TVM Array capturing target and target_host. The given target can be in
+        any form recognized by Target.canonicalize_target or Target.canonicalize_multi_targets. If
+        given target_host can be in any form recognized by Target.canonicalize_target. If
+        target_host is given it will be set as the 'host' in each result Target object (and a
+        warning given)."""
         # Convert target to Array, but not yet accounting for any host.
         raw_targets = Target.canonicalize_multi_targets(target)
         assert raw_targets is not None

From efbb2e476a68e78edede1fe21bfacabbaff98f80 Mon Sep 17 00:00:00 2001
From: mbs-octoml 
Date: Mon, 2 May 2022 17:58:58 -0700
Subject: [PATCH 06/11] - Moar Lints

---
 .../relay/backend/graph_executor_codegen.py   |  1 -
 python/tvm/target/target.py                   | 27 +++++++++----------
 2 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/python/tvm/relay/backend/graph_executor_codegen.py b/python/tvm/relay/backend/graph_executor_codegen.py
index 9ce6a056a80d..531f9f69e0e0 100644
--- a/python/tvm/relay/backend/graph_executor_codegen.py
+++ b/python/tvm/relay/backend/graph_executor_codegen.py
@@ -36,7 +36,6 @@
 from tvm.runtime.ndarray import empty
 from tvm.relay import _build_module
 from tvm.target import Target
-from tvm.tir import expr as _expr
 from .utils import mangle_module_name
 
 
diff --git a/python/tvm/target/target.py b/python/tvm/target/target.py
index bd43bd3b7abe..6c96ef94d017 100644
--- a/python/tvm/target/target.py
+++ b/python/tvm/target/target.py
@@ -228,10 +228,9 @@ def canonicalize_target(target):
         - A Python dictionary binding the target 'kind' and other attributes."""
         if target is None:
             return None
-        elif isinstance(target, Target):
+        if isinstance(target, Target):
             return target
-        else:
-            return Target(target)
+        return Target(target)
 
     @staticmethod
     def canonicalize_multi_targets(multi_targets):
@@ -245,15 +244,14 @@ def canonicalize_multi_targets(multi_targets):
            a target-like object in a form recognized by canonicalize_target."""
         if multi_targets is None:
             return None
-        elif isinstance(multi_targets, (dict, Map)) and "kind" not in multi_targets:
+        if isinstance(multi_targets, (dict, Map)) and "kind" not in multi_targets:
             # Convert legacy heterogeneous map representation to ordinary list of targets.
             return Target.canonicalize_multi_targets([t for _, t in multi_targets.items()])
-        elif isinstance(multi_targets, (list, Array)):
+        if isinstance(multi_targets, (list, Array)):
             # Multiple Target results.
             return convert([Target.canonicalize_target(t) for t in multi_targets])
-        else:
-            # Single Target result.
-            return convert([Target.canonicalize_target(multi_targets)])
+        # Single Target result.
+        return convert([Target.canonicalize_target(multi_targets)])
 
     @staticmethod
     def canonicalize_target_and_host(target, target_host=None):
@@ -269,13 +267,12 @@ def canonicalize_target_and_host(target, target_host=None):
         target_host = Target.canonicalize_target(target_host)
         if target_host is None:
             return raw_targets
-        else:
-            warnings.warn(
-                "target_host parameter is going to be deprecated. "
-                "Please pass in tvm.target.Target(target, host=target_host) instead."
-            )
-            # Make sure the (canonical) host is captured in all the (canonical) targets.
-            return convert([Target(t, target_host) for t in raw_targets])
+        warnings.warn(
+            "target_host parameter is going to be deprecated. "
+            "Please pass in tvm.target.Target(target, host=target_host) instead."
+        )
+        # Make sure the (canonical) host is captured in all the (canonical) targets.
+        return convert([Target(t, target_host) for t in raw_targets])
 
     @staticmethod
     def check_and_update_host_consist(target, host=None, target_is_dict_key=True):

From 7db601f2f3d4b6542daf887c205b146acef70a72 Mon Sep 17 00:00:00 2001
From: mbs-octoml 
Date: Tue, 3 May 2022 14:58:07 -0700
Subject: [PATCH 07/11] - Fix some unit tests

---
 python/tvm/autotvm/task/relay_integration.py  |  2 +-
 python/tvm/relay/build_module.py              |  2 +-
 python/tvm/target/compilation_config.py       | 14 +++---
 python/tvm/target/target.py                   | 43 ++++++++++---------
 src/relay/backend/aot_executor_codegen.cc     |  3 ++
 src/tir/usmp/transform/assign_pool_info.cc    |  1 +
 tests/python/driver/tvmc/test_target.py       |  7 +++
 tests/python/relay/aot/test_crt_aot_usmp.py   | 11 +++++
 .../relay/test_pass_annotate_spans_defuse.py  |  8 +++-
 tests/python/relay/test_pass_plan_devices.py  |  5 +--
 tests/python/relay/test_vm.py                 |  2 +-
 11 files changed, 62 insertions(+), 36 deletions(-)

diff --git a/python/tvm/autotvm/task/relay_integration.py b/python/tvm/autotvm/task/relay_integration.py
index 04ca333a5ea8..2643a01439e6 100644
--- a/python/tvm/autotvm/task/relay_integration.py
+++ b/python/tvm/autotvm/task/relay_integration.py
@@ -44,7 +44,7 @@ def _lower(mod, target, params, opt_level=3):
         import vta
 
         with vta.build_config(opt_level=opt_level, disabled_pass={"AlterOpLayout"}):
-            mod, _ = relay.optimize(mod, target, params)
+            mod, _ = relay.optimize(mod, target=target, params=params)
             grc = graph_executor_codegen.GraphExecutorCodegen(None, target)
             grc.codegen(mod, mod["main"])
             return
diff --git a/python/tvm/relay/build_module.py b/python/tvm/relay/build_module.py
index 10030118f6ad..06fa212ff396 100644
--- a/python/tvm/relay/build_module.py
+++ b/python/tvm/relay/build_module.py
@@ -554,7 +554,7 @@ def optimize(mod, target=None, params=None):
 
     with tophub_context:
         bld_mod = BuildModule()
-        mod, params = bld_mod.optimize(mod, target, params)
+        mod, params = bld_mod.optimize(mod, target=target, params=params)
     return mod, params
 
 
diff --git a/python/tvm/target/compilation_config.py b/python/tvm/target/compilation_config.py
index 2796ec4b5135..28bb95dc041c 100644
--- a/python/tvm/target/compilation_config.py
+++ b/python/tvm/target/compilation_config.py
@@ -15,13 +15,13 @@
 # specific language governing permissions and limitations
 # under the License.
 """Python bindings for creating CompilationConfigs."""
+import tvm
 from . import _ffi_api
 
 
-def make_compilation_config(ctxt, targets, host_target=None):
-    """Returns a CompilationConfig appropriate for targets and an optional host_target.
-    Currently intended just for unit tests and will be replaced by a Python CompilationConfig
-    class in the future. Note that targets must be a dictionary from IntImm objects to Targets
-    and we do not support any of the lighter-weight conventions used by the various build(...)
-    APIs."""
-    return _ffi_api.MakeCompilationConfig(ctxt, targets, host_target)
+def make_compilation_config(ctxt, target, target_host=None):
+    """Returns a CompilationConfig appropriate for target and target_host, using the same
+    representation conventions as for the standard build interfaces. Intended only for unit
+    testing."""
+    raw_targets=tvm.target.Target.canonicalize_target_and_host(target, target_host)
+    return _ffi_api.MakeCompilationConfig(ctxt, raw_targets)
diff --git a/python/tvm/target/target.py b/python/tvm/target/target.py
index 6c96ef94d017..9dce20902541 100644
--- a/python/tvm/target/target.py
+++ b/python/tvm/target/target.py
@@ -225,7 +225,8 @@ def canonicalize_target(target):
         - None (to None).
         - An existing TVM Target object.
         - A string.
-        - A Python dictionary binding the target 'kind' and other attributes."""
+        - A Python dictionary binding the target 'kind' and other attributes.
+        """
         if target is None:
             return None
         if isinstance(target, Target):
@@ -236,12 +237,13 @@ def canonicalize_target(target):
     def canonicalize_multi_targets(multi_targets):
         """Given a single or collection of target-like objects, returns a TVM Array of Target
         objects representing then. Can convert from:
-         - None (to None).
-         - A single target-like object in a form recognized by canonicalize_target.
-         - A Python list or TVM Array of target-like objects in a form recognized by
-           canonicalize_target.
-         - A Python dict or TVM Map from TVM IntImm objects representing device types to
-           a target-like object in a form recognized by canonicalize_target."""
+        - None (to None).
+        - A single target-like object in a form recognized by canonicalize_target.
+        - A Python list or TVM Array of target-like objects in a form recognized by
+        canonicalize_target.
+        - A Python dict or TVM Map from TVM IntImm objects representing device types to
+        a target-like object in a form recognized by canonicalize_target.
+        """
         if multi_targets is None:
             return None
         if isinstance(multi_targets, (dict, Map)) and "kind" not in multi_targets:
@@ -259,7 +261,8 @@ def canonicalize_target_and_host(target, target_host=None):
         any form recognized by Target.canonicalize_target or Target.canonicalize_multi_targets. If
         given target_host can be in any form recognized by Target.canonicalize_target. If
         target_host is given it will be set as the 'host' in each result Target object (and a
-        warning given)."""
+        warning given).
+        """
         # Convert target to Array, but not yet accounting for any host.
         raw_targets = Target.canonicalize_multi_targets(target)
         assert raw_targets is not None
@@ -616,7 +619,7 @@ def get_arch_version(cpu_ver):
     # Check for valid codegen cpu
     valid_hex = ["v65", "v66", "v67", "v67t", "v68", "v69"]
     try:
-        cpu_ver = cpu_ver[cpu_ver.index("v") :].lower()
+        cpu_ver = cpu_ver[cpu_ver.index("v"):].lower()
         assert cpu_ver in valid_hex
     except:
         msg = "{} is not a valid Hexagon version\nvalid versions include {}"
@@ -684,7 +687,7 @@ def validate_hvx_length(codegen_hvx, sim_options):
                 # If --hvx_length was specified, check HVX length of sim
                 # vs codegen
                 i = sim_options.index("hvx_length") + len("hvx_length") + 1
-                sim_hvx = sim_options[i : i + 3]
+                sim_hvx = sim_options[i: i + 3]
                 if sim_hvx != str(codegen_hvx):
                     msg = "sim hvx {} and codegen hvx {} mismatch!".format(sim_hvx, codegen_hvx)
                     # Set the stacklevel to the tvm.target.hexagon() call.
@@ -712,9 +715,9 @@ def validate_hvx_length(codegen_hvx, sim_options):
 
             # Regex match for allowed cpus
             valid_cpu_str_regex = (
-                r"(?P
--.*\s)?(--m)?"
-                + r"(?Pv6[25678])(?P[a-z])?"
-                + r"(?P_[0-9]+)?(?P_rev[0-9])?\s?(?P--.*)?"
+                    r"(?P
--.*\s)?(--m)?"
+                    + r"(?Pv6[25678])(?P[a-z])?"
+                    + r"(?P_[0-9]+)?(?P_rev[0-9])?\s?(?P--.*)?"
             )
             m = re.match(valid_cpu_str_regex, sim_options.lower())
             if not m:
@@ -723,13 +726,13 @@ def validate_hvx_length(codegen_hvx, sim_options):
             # Parse options into correct order
             cpu_attr = {x: str(m.groupdict()[x] or "") for x in m.groupdict()}
             sim_options = (
-                cpu_attr["base_version"]
-                + cpu_attr["sub_version"]
-                + cpu_attr["l2_size"]
-                + cpu_attr["rev"]
-                + " "
-                + cpu_attr["pre"]
-                + cpu_attr["post"]
+                    cpu_attr["base_version"]
+                    + cpu_attr["sub_version"]
+                    + cpu_attr["l2_size"]
+                    + cpu_attr["rev"]
+                    + " "
+                    + cpu_attr["pre"]
+                    + cpu_attr["post"]
             )
 
         return sim_cpu + " " + validate_hvx_length(hvx, sim_options)
diff --git a/src/relay/backend/aot_executor_codegen.cc b/src/relay/backend/aot_executor_codegen.cc
index 399d84594de9..c981f9d62b19 100644
--- a/src/relay/backend/aot_executor_codegen.cc
+++ b/src/relay/backend/aot_executor_codegen.cc
@@ -855,6 +855,7 @@ class AOTExecutorCodegen : public MixedModeVisitor {
    * brief Run USMP to plan memory for lowered IRModule
    */
   IRModule PlanMemoryWithUSMP(const IRModule& mod) {
+    VLOG(1) << "Planning memory with USMP for module:" << std::endl << PrettyPrint(mod);
     Executor executor_config = mod->GetAttr(tvm::attr::kExecutor).value();
     Integer workspace_byte_alignment =
         executor_config->GetAttr("workspace-byte-alignment").value_or(16);
@@ -870,6 +871,8 @@ class AOTExecutorCodegen : public MixedModeVisitor {
       for (const tir::usmp::AllocatedPoolInfo& allocated_pool_info : allocated_pool_infos.value()) {
         for (const auto& kv : allocated_pool_info->pool_info->target_access) {
           Target tgt = kv.first;
+          VLOG(1) << "USMP requires target " << tgt->ToDebugString() << " to have pool size "
+                  << allocated_pool_info->allocated_size->value;
           if (main_func_info->workspace_sizes.find(tgt) == main_func_info->workspace_sizes.end()) {
             main_func_info->workspace_sizes.Set(tgt, allocated_pool_info->allocated_size);
           } else {
diff --git a/src/tir/usmp/transform/assign_pool_info.cc b/src/tir/usmp/transform/assign_pool_info.cc
index 930299e4f039..e291eaa0519e 100644
--- a/src/tir/usmp/transform/assign_pool_info.cc
+++ b/src/tir/usmp/transform/assign_pool_info.cc
@@ -77,6 +77,7 @@ class PoolInfoAssigner : public StmtExprMutator {
 };
 
 PoolInfo PoolInfoAssigner::CreateDefaultMemoryPool(const tvm::IRModule& module) {
+  VLOG(1) << "Creating default memory pool for:" << std::endl << PrettyPrint(module);
   Map target_access;
   tir::PrimFunc tir_main_func =
       Downcast(module->Lookup(::tvm::runtime::symbol::tvm_module_main));
diff --git a/tests/python/driver/tvmc/test_target.py b/tests/python/driver/tvmc/test_target.py
index f6431a6cb20d..b02f89d2e425 100644
--- a/tests/python/driver/tvmc/test_target.py
+++ b/tests/python/driver/tvmc/test_target.py
@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import pytest
 from tvm.driver.tvmc import TVMCException
 from tvm.driver.tvmc.target import target_from_cli, tokenize_target, parse_target
 
@@ -174,3 +175,9 @@ def test_parse_multiple_target_with_opts_ethos_n78():
     assert "myopt" in targets[0]["opts"]
     assert "value" == targets[0]["opts"]["myopt"]
     assert "llvm" == targets[1]["name"]
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__] + sys.argv[1:]))
diff --git a/tests/python/relay/aot/test_crt_aot_usmp.py b/tests/python/relay/aot/test_crt_aot_usmp.py
index 23283392ee3b..371defdf518e 100644
--- a/tests/python/relay/aot/test_crt_aot_usmp.py
+++ b/tests/python/relay/aot/test_crt_aot_usmp.py
@@ -87,6 +87,9 @@ def test_memory_planning(workspace_byte_alignment, main_workspace_size):
         },
     ):
         lib = tvm.relay.build(mod, target, executor=executor, runtime=runtime, params=params)
+    print(lib.function_metadata)
+    print(sum(lib.function_metadata["__tvm_main__"].workspace_sizes.values()))
+    print(main_workspace_size)
     assert (
         sum(lib.function_metadata["__tvm_main__"].workspace_sizes.values()) == main_workspace_size
     )
@@ -634,3 +637,11 @@ def test_u4_usecase_incompatible_interface_api_errors():
             config={"tir.usmp.enable": True, "tir.usmp.use_workspace_io": True},
         ):
             tvm.relay.build(mod, target, executor=executor, runtime=runtime, params=params)
+
+
+if __name__ == "__main__":
+    import sys
+    import pytest
+
+    #sys.exit(pytest.main([__file__] + sys.argv[1:]))
+    test_memory_planning(8, 17280)
diff --git a/tests/python/relay/test_pass_annotate_spans_defuse.py b/tests/python/relay/test_pass_annotate_spans_defuse.py
index def4a1da1b55..dc3c9f99813c 100644
--- a/tests/python/relay/test_pass_annotate_spans_defuse.py
+++ b/tests/python/relay/test_pass_annotate_spans_defuse.py
@@ -42,7 +42,7 @@ def test_annotate_spans_compatibility():
 
     # Apply some simple passes to legalize the IR.
     with tvm.transform.PassContext(opt_level=0):
-        module, params = relay.optimize(module, tvm.testing.enabled_targets()[0][0], params)
+        module, params = relay.optimize(module, target=tvm.testing.enabled_targets()[0][0], params=params)
 
     seq = tvm.transform.Sequential([relay.transform.AnnotateSpans(), relay.transform.DefuseOps()])
     with tvm.transform.PassContext(opt_level=3):
@@ -50,4 +50,8 @@ def test_annotate_spans_compatibility():
 
 
 if __name__ == "__main__":
-    test_annotate_spans_compatibility()
+    import sys
+    import pytest
+
+    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+
diff --git a/tests/python/relay/test_pass_plan_devices.py b/tests/python/relay/test_pass_plan_devices.py
index f9fe9cf3555b..e485b626b4da 100644
--- a/tests/python/relay/test_pass_plan_devices.py
+++ b/tests/python/relay/test_pass_plan_devices.py
@@ -37,10 +37,7 @@
 GPU_DEVICE = tvm.device("cuda")
 GPU_TARGET = tvm.target.Target("cuda").with_host(HOST_TARGET)
 
-TARGETS = {
-    tvm.tir.IntImm("int32", CPU_DEVICE.device_type): CPU_TARGET,
-    tvm.tir.IntImm("int32", GPU_DEVICE.device_type): GPU_TARGET,
-}
+TARGETS = [CPU_TARGET, GPU_TARGET]
 
 HOST = tvm.target.VirtualDevice(HOST_DEVICE, HOST_TARGET)  # device_type=1
 CPU = tvm.target.VirtualDevice(CPU_DEVICE, CPU_TARGET)  # device_type=1
diff --git a/tests/python/relay/test_vm.py b/tests/python/relay/test_vm.py
index cde78068a7b1..e05e84d2ec35 100644
--- a/tests/python/relay/test_vm.py
+++ b/tests/python/relay/test_vm.py
@@ -1292,7 +1292,7 @@ def test_let_bound_constants():
     mod = IRModule.from_expr(f)
 
     compiler = VMCompiler()
-    compiler.optimize(mod, "llvm")
+    compiler.optimize(mod, target="llvm")
 
 
 def test_large_constants():

From 538d0358088d9c6ccadc8f8fd4685ef1f8640c9d Mon Sep 17 00:00:00 2001
From: mbs-octoml 
Date: Tue, 3 May 2022 18:01:40 -0700
Subject: [PATCH 08/11] - Fix last unit test failures

---
 include/tvm/target/compilation_config.h     | 19 +++++++++++++++----
 include/tvm/target/target.h                 |  2 +-
 src/target/target_kind.cc                   |  5 ++++-
 tests/python/relay/aot/test_crt_aot_usmp.py | 14 ++++++--------
 4 files changed, 26 insertions(+), 14 deletions(-)

diff --git a/include/tvm/target/compilation_config.h b/include/tvm/target/compilation_config.h
index bb278c911358..87b9798b20e8 100644
--- a/include/tvm/target/compilation_config.h
+++ b/include/tvm/target/compilation_config.h
@@ -63,10 +63,10 @@ class CompilationConfigNode : public Object {
   Target host_target;
 
   /*!
-   * \brief Vector of all available \p Targets for compiling primitive tensor operators (kernels).
-   * May contain a \p Target for the same device type as for the \p host_target, however the \p
-   * host_target should be used for all host computations and data. Each \p Target will have \p
-   * host_target as its 'host'.
+   * \brief Vector of all available \p Targets for partitioning or compiling primitive tensor
+   * operators (kernels). May contain a \p Target for the same device type as for the
+   * \p host_target, however the \p host_target should be used for all host computations and data.
+   * Each \p Target will have \p host_target as its 'host'.
    *
    * It is possible to have multiple primitive targets for the same device type. However given
    * primitive targets left and right where:
@@ -76,6 +76,17 @@ class CompilationConfigNode : public Object {
    *  - right.IsExternalCodegenFor(left) must be true
    * In this way the FindPrimitiveTargetOrFail method will find the 'most general' target for
    * the requested device type.
+   *
+   * In the homogeneous case primitive_targets will have just one entry, which will be pointer equal
+   * to optional_homogeneous_target.
+   *
+   * In the homogenous case where the 'host' is the same device as used for compiling kernels it
+   * is *not* the case that optional_homogenous_target == host_target. This is because all
+   * primitive always have their host field set to the host_target. Ie, it is valid to have:
+   * \code
+   *   host_target=Target("llvm")
+   *   optional_homogenous_target=Target("llvm", host=host_target)
+   * \endcode
    */
   Array primitive_targets;
 
diff --git a/include/tvm/target/target.h b/include/tvm/target/target.h
index 29103bd20866..c61db5ab50a0 100644
--- a/include/tvm/target/target.h
+++ b/include/tvm/target/target.h
@@ -185,7 +185,7 @@ class Target : public ObjectRef {
    *  - \p this and \p that have the same kind->device_type
    *
    * After partitioning, the external codegen compilation path may use \p that to guide it's
-   * compilation to a \p runtime::Module. Given on \p this, an appropriate \p that can be
+   * compilation to a \p runtime::Module. Given \p this, an appropriate \p that can be
    * found using \p CompilationConfig::FindPrimitiveTargetOrFail(this->kind->device_type).
    *
    * The \p CollagePartition pass uses this method to guide it's search over candidate partitions
diff --git a/src/target/target_kind.cc b/src/target/target_kind.cc
index 2ad75259d69b..43bcfef105ff 100644
--- a/src/target/target_kind.cc
+++ b/src/target/target_kind.cc
@@ -267,7 +267,10 @@ TVM_REGISTER_TARGET_KIND("llvm", kDLCPU)
     .add_attr_option("fast-math-contract")
     .add_attr_option("fast-math-reassoc")
     .add_attr_option("opt-level")
-    .set_default_keys({"cpu"});
+    .set_default_keys({"cpu"})
+    // Force the external codegen kind attribute to be registered, even if no external
+    // codegen targets are enabled by the TVM build.
+    .set_attr(tvm::attr::kIsExternalCodegen, Bool(false));
 
 TVM_REGISTER_TARGET_KIND("c", kDLCPU)
     .add_attr_option("system-lib")
diff --git a/tests/python/relay/aot/test_crt_aot_usmp.py b/tests/python/relay/aot/test_crt_aot_usmp.py
index 371defdf518e..b1be3e8a8317 100644
--- a/tests/python/relay/aot/test_crt_aot_usmp.py
+++ b/tests/python/relay/aot/test_crt_aot_usmp.py
@@ -87,12 +87,10 @@ def test_memory_planning(workspace_byte_alignment, main_workspace_size):
         },
     ):
         lib = tvm.relay.build(mod, target, executor=executor, runtime=runtime, params=params)
-    print(lib.function_metadata)
-    print(sum(lib.function_metadata["__tvm_main__"].workspace_sizes.values()))
-    print(main_workspace_size)
-    assert (
-        sum(lib.function_metadata["__tvm_main__"].workspace_sizes.values()) == main_workspace_size
-    )
+    # The workspace_size dictionary will have an entry for both the 'primitive' and 'host'
+    # targets, though both are identical.
+    for size in lib.function_metadata["__tvm_main__"].workspace_sizes.values():
+        assert size == main_workspace_size
 
 
 @parametrize_aot_options
@@ -643,5 +641,5 @@ def test_u4_usecase_incompatible_interface_api_errors():
     import sys
     import pytest
 
-    #sys.exit(pytest.main([__file__] + sys.argv[1:]))
-    test_memory_planning(8, 17280)
+    sys.exit(pytest.main([__file__] + sys.argv[1:]))
+

From 320caf4bb8663f830ab5cb1fa0ebdc9ee8ca9293 Mon Sep 17 00:00:00 2001
From: mbs-octoml 
Date: Tue, 3 May 2022 18:12:42 -0700
Subject: [PATCH 09/11] - whitespace

---
 python/tvm/target/compilation_config.py       |  2 +-
 python/tvm/target/target.py                   | 24 +++++++++----------
 tests/python/relay/aot/test_crt_aot_usmp.py   |  1 -
 .../relay/test_pass_annotate_spans_defuse.py  |  5 ++--
 4 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/python/tvm/target/compilation_config.py b/python/tvm/target/compilation_config.py
index 28bb95dc041c..8a59a33c1a47 100644
--- a/python/tvm/target/compilation_config.py
+++ b/python/tvm/target/compilation_config.py
@@ -23,5 +23,5 @@ def make_compilation_config(ctxt, target, target_host=None):
     """Returns a CompilationConfig appropriate for target and target_host, using the same
     representation conventions as for the standard build interfaces. Intended only for unit
     testing."""
-    raw_targets=tvm.target.Target.canonicalize_target_and_host(target, target_host)
+    raw_targets = tvm.target.Target.canonicalize_target_and_host(target, target_host)
     return _ffi_api.MakeCompilationConfig(ctxt, raw_targets)
diff --git a/python/tvm/target/target.py b/python/tvm/target/target.py
index 9dce20902541..03115612c5ce 100644
--- a/python/tvm/target/target.py
+++ b/python/tvm/target/target.py
@@ -619,7 +619,7 @@ def get_arch_version(cpu_ver):
     # Check for valid codegen cpu
     valid_hex = ["v65", "v66", "v67", "v67t", "v68", "v69"]
     try:
-        cpu_ver = cpu_ver[cpu_ver.index("v"):].lower()
+        cpu_ver = cpu_ver[cpu_ver.index("v") :].lower()
         assert cpu_ver in valid_hex
     except:
         msg = "{} is not a valid Hexagon version\nvalid versions include {}"
@@ -687,7 +687,7 @@ def validate_hvx_length(codegen_hvx, sim_options):
                 # If --hvx_length was specified, check HVX length of sim
                 # vs codegen
                 i = sim_options.index("hvx_length") + len("hvx_length") + 1
-                sim_hvx = sim_options[i: i + 3]
+                sim_hvx = sim_options[i : i + 3]
                 if sim_hvx != str(codegen_hvx):
                     msg = "sim hvx {} and codegen hvx {} mismatch!".format(sim_hvx, codegen_hvx)
                     # Set the stacklevel to the tvm.target.hexagon() call.
@@ -715,9 +715,9 @@ def validate_hvx_length(codegen_hvx, sim_options):
 
             # Regex match for allowed cpus
             valid_cpu_str_regex = (
-                    r"(?P
--.*\s)?(--m)?"
-                    + r"(?Pv6[25678])(?P[a-z])?"
-                    + r"(?P_[0-9]+)?(?P_rev[0-9])?\s?(?P--.*)?"
+                r"(?P
--.*\s)?(--m)?"
+                + r"(?Pv6[25678])(?P[a-z])?"
+                + r"(?P_[0-9]+)?(?P_rev[0-9])?\s?(?P--.*)?"
             )
             m = re.match(valid_cpu_str_regex, sim_options.lower())
             if not m:
@@ -726,13 +726,13 @@ def validate_hvx_length(codegen_hvx, sim_options):
             # Parse options into correct order
             cpu_attr = {x: str(m.groupdict()[x] or "") for x in m.groupdict()}
             sim_options = (
-                    cpu_attr["base_version"]
-                    + cpu_attr["sub_version"]
-                    + cpu_attr["l2_size"]
-                    + cpu_attr["rev"]
-                    + " "
-                    + cpu_attr["pre"]
-                    + cpu_attr["post"]
+                cpu_attr["base_version"]
+                + cpu_attr["sub_version"]
+                + cpu_attr["l2_size"]
+                + cpu_attr["rev"]
+                + " "
+                + cpu_attr["pre"]
+                + cpu_attr["post"]
             )
 
         return sim_cpu + " " + validate_hvx_length(hvx, sim_options)
diff --git a/tests/python/relay/aot/test_crt_aot_usmp.py b/tests/python/relay/aot/test_crt_aot_usmp.py
index b1be3e8a8317..86e0d18021fd 100644
--- a/tests/python/relay/aot/test_crt_aot_usmp.py
+++ b/tests/python/relay/aot/test_crt_aot_usmp.py
@@ -642,4 +642,3 @@ def test_u4_usecase_incompatible_interface_api_errors():
     import pytest
 
     sys.exit(pytest.main([__file__] + sys.argv[1:]))
-
diff --git a/tests/python/relay/test_pass_annotate_spans_defuse.py b/tests/python/relay/test_pass_annotate_spans_defuse.py
index dc3c9f99813c..d6b16e70a50a 100644
--- a/tests/python/relay/test_pass_annotate_spans_defuse.py
+++ b/tests/python/relay/test_pass_annotate_spans_defuse.py
@@ -42,7 +42,9 @@ def test_annotate_spans_compatibility():
 
     # Apply some simple passes to legalize the IR.
     with tvm.transform.PassContext(opt_level=0):
-        module, params = relay.optimize(module, target=tvm.testing.enabled_targets()[0][0], params=params)
+        module, params = relay.optimize(
+            module, target=tvm.testing.enabled_targets()[0][0], params=params
+        )
 
     seq = tvm.transform.Sequential([relay.transform.AnnotateSpans(), relay.transform.DefuseOps()])
     with tvm.transform.PassContext(opt_level=3):
@@ -54,4 +56,3 @@ def test_annotate_spans_compatibility():
     import pytest
 
     sys.exit(pytest.main([__file__] + sys.argv[1:]))
-

From 0533b4c727644403baca59e770c4ad7c782b70b0 Mon Sep 17 00:00:00 2001
From: mbs-octoml 
Date: Wed, 4 May 2022 09:29:45 -0700
Subject: [PATCH 10/11] - Address Eric's comments.   CI likely to fail due to
 stricter FindPrimitiveTargetOrFail but let's see.

---
 include/tvm/target/target.h      |  8 ++++++++
 src/target/compilation_config.cc | 25 ++++++++++++++-----------
 src/target/target.cc             | 10 +++++++---
 3 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/include/tvm/target/target.h b/include/tvm/target/target.h
index c61db5ab50a0..a9d893ff5402 100644
--- a/include/tvm/target/target.h
+++ b/include/tvm/target/target.h
@@ -177,6 +177,14 @@ class Target : public ObjectRef {
    */
   static Target WithHost(const Target& target, const Target& host);
 
+  /*!
+   * \brief Returns true if \p this target represents an external codegen. If so,
+   * \p this->kind->name can be used as the "Compiler" attribute on partitioned functions,
+   * and can be used to retrieve a partitioning pattern table using
+   * \p get_pattern_table.
+   */
+  bool IsExternalCodegen() const;
+
   /*!
    * \brief Returns true if \p this target represents an external codegen which is compatible
    * with \p that target. In particular:
diff --git a/src/target/compilation_config.cc b/src/target/compilation_config.cc
index d9d7d4b9675c..b7a7b17a712d 100644
--- a/src/target/compilation_config.cc
+++ b/src/target/compilation_config.cc
@@ -39,11 +39,7 @@ void CompilationConfigNode::VisitAttrs(AttrVisitor* v) {
 }
 
 Target CompilationConfigNode::FindPrimitiveTargetOrFail(DLDeviceType device_type) const {
-  if (device_type < 0 && primitive_targets.size() == 1) {
-    // In the homogenous case don't be fussy with device types.
-    return primitive_targets.front();
-  }
-  ICHECK_GT(device_type, 0);
+  ICHECK_GT(device_type, 0) << "Invalid device type";
   auto itr = std::find_if(
       primitive_targets.begin(), primitive_targets.end(),
       [device_type](const Target& target) { return target->kind->device_type == device_type; });
@@ -144,6 +140,8 @@ void CompilationConfigNode::Init(const transform::PassContext& pass_ctx,
 
   //
   // Check the primitive_targets are ordered correctly re Target::IsExternalCodegenFor.
+  // Note we could just sort the list, but given all the implicit defaulting for backwards
+  // compat it seems we should avoid making this any more magical than necessarny.
   //
   std::unordered_set primitive_target_device_types;
   for (const auto& target : primitive_targets) {
@@ -157,13 +155,18 @@ void CompilationConfigNode::Init(const transform::PassContext& pass_ctx,
       }
       if (!first_primitive_target.defined()) {
         first_primitive_target = current_primitive_target;
-        continue;
+        CHECK(!first_primitive_target.IsExternalCodegen())
+            << "The first given target for device type " << device_type
+            << " must not be for an external codegen, however given "
+            << first_primitive_target->ToDebugString();
+      } else {
+        CHECK(current_primitive_target.IsExternalCodegenFor(first_primitive_target))
+            << "When given multiple targets for the device type " << device_type
+            << " the first must be for non external codegen, and all subsequent must be for "
+               "external codegen. However have been given first "
+            << first_primitive_target->ToDebugString() << " and subsequent "
+            << current_primitive_target->ToDebugString();
       }
-      CHECK(current_primitive_target.IsExternalCodegenFor(first_primitive_target))
-          << "The first given target for device type " << device_type << " is "
-          << first_primitive_target->ToDebugString() << ", however a later target "
-          << current_primitive_target->ToDebugString()
-          << " for the same device type is not an external codegen target.";
     }
   }
 
diff --git a/src/target/target.cc b/src/target/target.cc
index e26fed6c74f9..75126ed11c70 100644
--- a/src/target/target.cc
+++ b/src/target/target.cc
@@ -494,10 +494,14 @@ Target::Target(TargetKind kind, Optional host, String tag, Array attr_map = TargetKind::GetAttrMap(::tvm::attr::kIsExternalCodegen);
-  return get()->kind->device_type == that->kind->device_type &&
-         attr_map.get(get()->kind, Bool(false)) && !attr_map.get(that->kind, Bool(false));
+  return attr_map.get(get()->kind, Bool(false));
+}
+
+bool Target::IsExternalCodegenFor(const Target& that) const {
+  return get()->kind->device_type == that->kind->device_type && IsExternalCodegen() &&
+         !that.IsExternalCodegen();
 }
 
 std::vector TargetNode::GetKeys() const {

From 53a8aa715e595655dc9c1b99cfc69d27f3ad20e9 Mon Sep 17 00:00:00 2001
From: mbs-octoml 
Date: Wed, 4 May 2022 11:19:58 -0700
Subject: [PATCH 11/11] - Comment adjustments. - Unit test for new Target
 members.

---
 src/target/compilation_config.cc | 21 ++++++++++++---------
 tests/cpp/target_test.cc         | 28 ++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/src/target/compilation_config.cc b/src/target/compilation_config.cc
index b7a7b17a712d..7260427bc1a1 100644
--- a/src/target/compilation_config.cc
+++ b/src/target/compilation_config.cc
@@ -78,14 +78,14 @@ void CompilationConfigNode::Init(const transform::PassContext& pass_ctx,
   // Decide on the host target.
   //
 
-  // Any CPU-like targets?
-  auto cpu_itr = std::find_if(raw_targets.begin(), raw_targets.end(), [](const Target& target) {
-    // TODO(tvm-team): AoT only works with kDLCPU device type. We can remove kDLHexagon
+  // Any targets which could act as a host?
+  auto hosting_itr = std::find_if(raw_targets.begin(), raw_targets.end(), [](const Target& target) {
+    // TODO(tvm-team): The kDLHexagon device can act as a host. We can remove kDLHexagon
     // here once we refactored kDLHexagon to kDLCPU.
     return target->kind->device_type == kDLCPU || target->kind->device_type == kDLHexagon;
   });
 
-  // Any targets with a host?
+  // Any targets with their host field set?
   auto has_host_itr = std::find_if(raw_targets.begin(), raw_targets.end(),
                                    [](const Target& target) { return target->host.defined(); });
 
@@ -95,9 +95,10 @@ void CompilationConfigNode::Init(const transform::PassContext& pass_ctx,
     host_target = Target((*has_host_itr)->GetHost().value(), /*host=*/Target());
     VLOG(1) << "The target " << (*has_host_itr)->ToDebugString() << " supplies a host target "
             << host_target->ToDebugString() << " of device type " << host_target->kind->device_type;
-  } else if (cpu_itr != raw_targets.end()) {
-    // RULE B: If any raw target is for a CPU-like device then also use that as the host.
-    host_target = Target(*cpu_itr, /*host=*/Target());
+  } else if (hosting_itr != raw_targets.end()) {
+    // RULE B: If any raw target is for a device which could be a host then use the first such as
+    // the host.
+    host_target = Target(*hosting_itr, /*host=*/Target());
     VLOG(1) << "Using target " << host_target->ToDebugString() << " of CPU-like device type "
             << host_target->kind->device_type << " as the host target";
   } else {
@@ -140,9 +141,11 @@ void CompilationConfigNode::Init(const transform::PassContext& pass_ctx,
 
   //
   // Check the primitive_targets are ordered correctly re Target::IsExternalCodegenFor.
-  // Note we could just sort the list, but given all the implicit defaulting for backwards
-  // compat it seems we should avoid making this any more magical than necessarny.
   //
+
+  // TODO(mbs): We could just sort the list, but given all the implicit defaulting for backwards
+  // compat it seems we should avoid making this any more magical than necessary. But revisit
+  // if usability suffers.
   std::unordered_set primitive_target_device_types;
   for (const auto& target : primitive_targets) {
     primitive_target_device_types.emplace(static_cast(target->kind->device_type));
diff --git a/tests/cpp/target_test.cc b/tests/cpp/target_test.cc
index 6106eb2225e1..b657ac0c5783 100644
--- a/tests/cpp/target_test.cc
+++ b/tests/cpp/target_test.cc
@@ -135,6 +135,34 @@ TEST(TargetCreationFail, TargetKindNotFound) {
   ASSERT_EQ(failed, true);
 }
 
+TVM_REGISTER_TARGET_KIND("test_external_codegen_0", kDLCUDA)
+    .set_attr(tvm::attr::kIsExternalCodegen, Bool(true));
+
+TVM_REGISTER_TARGET_KIND("test_external_codegen_1", kDLCUDA)
+    .set_attr(tvm::attr::kIsExternalCodegen, Bool(true));
+
+TVM_REGISTER_TARGET_KIND("test_external_codegen_2", kDLMetal)
+    .set_attr(tvm::attr::kIsExternalCodegen, Bool(true));
+
+TEST(Target, ExternalCodegen) {
+  Target regular("cuda");
+  Target external0("test_external_codegen_0");
+  Target external1("test_external_codegen_1");
+  Target external2("test_external_codegen_2");
+
+  ASSERT_FALSE(regular.IsExternalCodegen());
+  ASSERT_TRUE(external0.IsExternalCodegen());
+  ASSERT_TRUE(external1.IsExternalCodegen());
+  ASSERT_TRUE(external2.IsExternalCodegen());
+
+  ASSERT_TRUE(external0.IsExternalCodegenFor(regular));
+  ASSERT_FALSE(regular.IsExternalCodegenFor(external0));
+  ASSERT_TRUE(external1.IsExternalCodegenFor(regular));
+  ASSERT_FALSE(regular.IsExternalCodegenFor(external1));
+  ASSERT_FALSE(external2.IsExternalCodegenFor(regular));
+  ASSERT_FALSE(regular.IsExternalCodegenFor(external2));
+}
+
 TEST(TargetCreation, DeduplicateKeys) {
   Map config = {
       {"kind", String("llvm")},