- Masa & Andrew's comments

Author: mbs-octoml

include/tvm/ir/module.h

@@ -519,12 +519,12 @@ constexpr const char* kConstantMemoryPools = "constant_memory_pools";
 
 /*
  * \brief All the runtime::NDArrays extracted from PrimFunc tir::AllocateConst nodes. The
- * node will record the index into this array. See also kConstNameToNDArray below, which is
+ * node will record the index into this array. See also kConstNameToConstant below, which is
  * the analog for Realy Functions.
  *
  * Type: Array<runtime::NDArray>
  */
-constexpr const char* kConstantsArray = "Constants";
+constexpr const char* kConstants = "constants";
 
 /*!
  * \brief All the runtime::Modules accumulated during compilation by external codegen. These
@@ -542,7 +542,7 @@ constexpr const char* kExternalMods = "external_mods";
  *
  * Type: Map<String, runtime::NDArray>
  */
-constexpr const char* kConstNameToNDArray = "const_name_to_ndarray";
+constexpr const char* kConstNameToConstant = "const_name_to_constant";
 
 }  // namespace attr
 }  // namespace tvm

python/tvm/relay/backend/interpreter.py

@@ -195,7 +195,7 @@ class Interpreter(Executor):
         The runtime device to run the code on.
 
     target : tvm.Target
-        The target option to build the function.
+        The target option to build the function. Only homogeneous execution is supported.
 
     CAUTION: Despite the API the module is prepared upon each call to evaluate
     rather than once in create_executor.

python/tvm/relay/backend/vm.py

@@ -198,8 +198,9 @@ class VMExecutor(Executor):
     device : :py:class:`~tvm.runtime.Device`
         The runtime device to run the code on.
 
-    target : :py:class:`Target`
-        The target option to build the function.
+    target : any multi-target like object, see Target.canon_multi_target
+        For homogeneous compilation, the unique build target.
+        For heterogeneous compilation, a dictionary or list of possible build targets.
     """
 
     def __init__(self, mod, device, target):

python/tvm/relay/build_module.py

@@ -570,15 +570,16 @@ class GraphExecutor(_interpreter.Executor):
     device : :py:class:`Device`
         The runtime device to run the code on.
 
-    raw_targets : Array[tvm.target.Target]
-        The available targets.
+    target : any multi-target like object, see Target.canon_multi_target
+        For homogeneous compilation, the unique build target.
+        For heterogeneous compilation, a dictionary or list of possible build targets.
     """
 
-    def __init__(self, mod, device, raw_targets):
+    def __init__(self, mod, device, target):
         assert mod is not None
         self.mod = mod
         self.device = device
-        self.raw_targets = raw_targets
+        self.target = target
 
     def _make_executor(self, expr=None):
         if expr:
@@ -589,7 +590,7 @@ def _make_executor(self, expr=None):
             raise ValueError(
                 "Graph Executor only supports static graphs, got output type", ret_type
             )
-        mod = build(self.mod, target=self.raw_targets)
+        mod = build(self.mod, target=self.target)
         gmodule = _graph_executor.GraphModule(mod["default"](self.device))
 
         def _unflatten(flat_iter, cur_type):
@@ -630,16 +631,16 @@ class AotExecutor(_interpreter.Executor):
     device : :py:class:`Device`
         The runtime device to run the code on.
 
-    raw_targets : Array[tvm.target.Target]
-        The available targets.
+    target : any multi-target like object, see Target.canon_multi_target
+        For homogeneous compilation, the unique build target.
+        For heterogeneous compilation, a dictionary or list of possible build targets.
     """
 
-    def __init__(self, mod, device, raw_targets):
+    def __init__(self, mod, device, target):
         assert mod is not None
         self.mod = mod
         self.device = device
-        self.raw_targets = raw_targets
-        assert raw_targets[0].attrs.get("executor", "graph") == "aot"
+        self.target = target
 
     def _make_executor(self, expr=None):
         if expr:
@@ -648,7 +649,7 @@ def _make_executor(self, expr=None):
         ret_type = self.mod["main"].checked_type.ret_type
         if _ty.is_dynamic(ret_type):
             raise ValueError("AOT Executor only supports static graphs, got output type", ret_type)
-        mod = build(self.mod, target=self.raw_targets)
+        mod = build(self.mod, target=self.target)
 
         # NOTE: Given AOT requires use of the "c" backend, must export/import to compile the
         # generated code.
@@ -722,6 +723,8 @@ def create_executor(kind="debug", mod=None, device=None, target="llvm", params=N
     target : any multi-target like object, see Target.canon_multi_target
         For homogeneous compilation, the unique build target.
         For heterogeneous compilation, a dictionary or list of possible build targets.
+        CAUTION: Though this API allows multiple targets, it does not allow multiple devices, so
+        heterogenous compilation is not yet supported.
 
     params : dict of str to NDArray
          Input parameters to the graph that do not change
@@ -737,11 +740,14 @@ def create_executor(kind="debug", mod=None, device=None, target="llvm", params=N
     if device is not None:
         assert device.device_type == raw_targets[0].kind.device_type
     else:
+        # Use the first target as the device.
         device = _nd.device(raw_targets[0].kind.device_type, 0)
 
     if params is not None:
         mod = IRModule.from_expr(bind_params_by_name(mod["main"], params))
 
+    assert raw_targets[0].attrs.get("executor") == kind
+
     if kind == "debug":
         assert len(raw_targets) == 1, "The interpreter currently only supports a single target"
         return _interpreter.Interpreter(mod, device, raw_targets[0])

src/relay/backend/aot_executor_codegen.cc

@@ -1170,12 +1170,10 @@ class AOTExecutorCodegen : public MixedModeVisitor {
 
     // Collect any constants extracted by external codegen.
     ret.params = std::unordered_map<std::string, tvm::runtime::NDArray>();
-    Map<String, runtime::NDArray> const_name_to_ndarray =
-        lowered_mod
-            ->GetAttr<Map<String, runtime::NDArray>>(tvm::attr::kConstNameToNDArray,
-                                                     Map<String, runtime::NDArray>())
-            .value();
-    for (const auto& kv : const_name_to_ndarray) {
+    Map<String, runtime::NDArray> const_name_to_constant =
+        lowered_mod->GetAttr<Map<String, runtime::NDArray>>(tvm::attr::kConstNameToConstant)
+            .value_or({});
+    for (const auto& kv : const_name_to_constant) {
       ICHECK(ret.params.emplace(kv.first, kv.second).second);
     }
 
@@ -1223,10 +1221,8 @@ class AOTExecutorCodegen : public MixedModeVisitor {
     }
 
     // Collect any runtime modules generated by external codegen.
-    ret.external_mods = lowered_mod
-                            ->GetAttr<Array<tvm::runtime::Module>>(tvm::attr::kExternalMods,
-                                                                   Array<tvm::runtime::Module>())
-                            .value();
+    ret.external_mods =
+        lowered_mod->GetAttr<Array<tvm::runtime::Module>>(tvm::attr::kExternalMods).value_or({});
 
     // This is the point where we separate the functions in the module by target
     VLOG(1) << "lowered module:" << std::endl << PrettyPrint(lowered_mod);

src/relay/backend/contrib/arm_compute_lib/codegen.cc

@@ -393,7 +393,7 @@ runtime::Module ACLCompiler(const ObjectRef& ref) {
   serializer.serialize();
   std::string graph_json = serializer.GetJSON();
 
-  // Note that serializer.const_name_to_ndarray() is ignored. Instead the TECompiler invokes
+  // Note that serializer.const_name_to_constant() is ignored. Instead the TECompiler invokes
   // a callback which calls backend::UpdateConstants to capture the map before the function
   // 'disappears' into lowered form, on the assumption the visit order and thus constant
   // names match those generated by the JSONSerializer.

src/relay/backend/contrib/bnns/codegen.cc

@@ -137,7 +137,7 @@ runtime::Module BNNSCompiler(const ObjectRef& ref) {
   serializer.serialize();
   std::string graph_json = serializer.GetJSON();
 
-  // Note that serializer.const_name_to_ndarray() is ignored. Instead the TECompiler invokes
+  // Note that serializer.const_name_to_constant() is ignored. Instead the TECompiler invokes
   // a callback which calls backend::UpdateConstants to capture the map before the function
   // 'disappears' into lowered form, on the assumption the visit order and thus constant
   // names match those generated by the JSONSerializer.

src/relay/backend/contrib/codegen_json/codegen_json.h

@@ -169,8 +169,8 @@ class JSONSerializer : public MemoizedExprTranslator<std::vector<JSONGraphNodeEn
    * \brief Returns the accumulated map from constant names to the NDArray they must be bound to
    * at runtime. Also referred to a 'params' elsewhere in the code.
    */
-  const std::unordered_map<std::string, runtime::NDArray>& const_name_to_ndarray() const {
-    return const_name_to_ndarray_;
+  const std::unordered_map<std::string, runtime::NDArray>& const_name_to_constant() const {
+    return const_name_to_constant_;
   }
 
   /*!
@@ -260,9 +260,10 @@ class JSONSerializer : public MemoizedExprTranslator<std::vector<JSONGraphNodeEn
 
   std::vector<JSONGraphNodeEntry> VisitExpr_(const ConstantNode* constant_node) {
     std::string name = symbol_ + "_const_" + std::to_string(const_names_.size());
-    VLOG(1) << "Will require parameter '" << name << "' to be available at runtime";
-    ICHECK_EQ(const_name_to_ndarray_.count(name), 0);
-    const_name_to_ndarray_.emplace(name, constant_node->data);
+    VLOG(1) << "Will require parameter '" << name
+            << "' to be supplied by the ConstLoaderModule at runtime";
+    ICHECK_EQ(const_name_to_constant_.count(name), 0);
+    const_name_to_constant_.emplace(name, constant_node->data);
     const_names_.push_back(name);
     auto node = std::make_shared<JSONGraphNode>(name, /*op_type=*/"const");
     return AddNode(node, GetRef<Expr>(constant_node));
@@ -361,7 +362,7 @@ class JSONSerializer : public MemoizedExprTranslator<std::vector<JSONGraphNodeEn
    * translation to JSON. The JSON will record only the constant name. The actual NDArray must
    * be made available at runtime from a ConstLoaderModule.
    */
-  std::unordered_map<std::string, runtime::NDArray> const_name_to_ndarray_;
+  std::unordered_map<std::string, runtime::NDArray> const_name_to_constant_;
   /*!
    * \brief The domain of the above map, but in order the constants were encountered during
    * translation.

src/relay/backend/contrib/cutlass/codegen.cc

@@ -47,7 +47,9 @@ namespace {
 Target GetCutlassTarget() {
   Target target = Target::Current(/*allow_not_defined=*/true);
   if (!target.defined() || target->kind->name != "cutlass") {
-    // Use the default CUTLASS compilation options.
+    // Use the default CUTLASS compilation options if no specific "cutlass" target was given
+    // in the overall targets list. In that case target_hooks.cc will invoke the custom pass
+    // without pushing any target instance onto the implicit target stack.
     target = Target("cutlass");
   }
   return target;
@@ -912,8 +914,7 @@ transform::Pass CompileForCutlassImpl() {
     Target target = GetCutlassTarget();
     runtime::Module runtime_mod = (*pf)(mod, target);
     Array<runtime::Module> external_mods =
-        mod->GetAttr<Array<runtime::Module>>(tvm::attr::kExternalMods, Array<runtime::Module>())
-            .value();
+        mod->GetAttr<Array<runtime::Module>>(tvm::attr::kExternalMods).value_or({});
     external_mods.push_back(runtime_mod);
     return WithAttr(mod, tvm::attr::kExternalMods, external_mods);
   };

src/relay/backend/contrib/dnnl/codegen.cc

@@ -586,7 +586,7 @@ runtime::Module DNNLCompiler(const ObjectRef& ref) {
   serializer.serialize();
   std::string graph_json = serializer.GetJSON();
 
-  // Note that serializer.const_name_to_ndarray() is ignored. Instead the TECompiler invokes
+  // Note that serializer.const_name_to_constant() is ignored. Instead the TECompiler invokes
   // a callback which calls backend::UpdateConstants to capture the map before the function
   // 'disappears' into lowered form, on the assumption the visit order and thus constant
   // names match those generated by the JSONSerializer.