[Relay] Support 'external codegen targets'. (apache#11173)

* [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<Target> of
   'raw targets' instead of a Map<Integer, Target>. 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<Target> 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.

* - Working on unit tests

* - Fix two Debug-only failures

* - Use Array<Target> in GraphExecutorCodegen/AOTExecutorCodegen ifaces instead
  of CompilationConfig (don't want to bake it into any official APIs).
- Started unit tests.

* - Lints

* - Moar Lints

* - Fix some unit tests

* - Fix last unit test failures

* - whitespace

* - Address Eric's comments.
  CI likely to fail due to stricter FindPrimitiveTargetOrFail but let's see.

* - Comment adjustments.
- Unit test for new Target members.

Author: mbs-octoml

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 $<TARGET_PROPERTY:tvm,INTERFACE_COMPILE_DEFINITIONS>)
   gtest_discover_tests(cpptest)

cmake/modules/CUDA.cmake

@@ -41,13 +41,17 @@ 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})
   endif(USE_CUDNN)
 
   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})

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,30 @@ 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.
+ * The construction of \p CompilationConfig is idempotent, in that given the same \p PassContext
+ * \p ctx and an arbitrary \p Array<Target> \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.
  */
 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 +63,37 @@ 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 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:
+   *  - 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.
+   *
+   * 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<Target> 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 +104,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 +139,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<Target>& 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 +170,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<Target>& raw_targets);
 
   TVM_DEFINE_OBJECT_REF_METHODS(CompilationConfig, ObjectRef, CompilationConfigNode);
 };

include/tvm/target/target.h

@@ -177,7 +177,34 @@ 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:
+   *  - \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 \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<ObjectRef> host, String tag, Array<String> keys,
+         Map<String, ObjectRef> attrs);
+
   // enable with syntax.
   friend class TargetInternal;
   friend class With<Target>;
@@ -194,8 +221,6 @@ class Target : public ObjectRef {
   TVM_DLL void ExitWithScope();
 };
 
-using TargetMap = Map<Integer, Target>;
-
 /*!
  * \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 +230,6 @@ using TargetMap = Map<Integer, Target>;
  */
 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,

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.

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

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):

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
 
 
@@ -54,15 +53,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.