chore: import upstream snapshot with attribution

include/tvm/target/codegen.h

@@ -0,0 +1,94 @@
+/*
+ * 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 tvm/target/codegen.h
+ * \brief Translates IRModule to runtime::Module.
+ */
+#ifndef TVM_TARGET_CODEGEN_H_
+#define TVM_TARGET_CODEGEN_H_
+
+#include <tvm/ffi/extra/module.h>
+#include <tvm/ir/module.h>
+#include <tvm/target/target.h>
+#include <tvm/tirx/expr.h>
+
+#include <string>
+
+namespace tvm {
+/*! \brief namespace for target translation and codegen. */
+namespace codegen {
+
+/*!
+ * \brief Build a module from array of lowered function.
+ * \param mod The Module to be built
+ * \param target The target to be built.
+ * \return The result runtime::Module.
+ */
+ffi::Module Build(IRModule mod, Target target);
+
+/*!
+ * \brief Serialize runtime module including its submodules
+ * \param mod The runtime module to serialize including its import tree.
+ * \param export_dso By default, include the info of DSOExportable modules. If disabled, an error
+ * will be raised when encountering DSO modules.
+ */
+std::string SerializeModuleToBytes(const ffi::Module& mod, bool export_dso = true);
+
+/*!
+ * \brief Deserialize runtime module including its submodules
+ * \param blob byte stream, which are generated by `SerializeModuleToBytes`.
+ * \return runtime::Module runtime module constructed from the given stream
+ */
+ffi::Module DeserializeModuleFromBytes(std::string blob);
+
+/*!
+ * \brief Pack imported device library to a C file.
+ *  Compile the C file and link with the host library
+ *  will allow the DSO loader to automatically discover and import
+ *  the dependency from the shared library.
+ *
+ * \param m The host module with the imports.
+ * \param system_lib Whether expose as system library.
+ * \param c_symbol_prefix Optional symbol prefix of the blob symbol.
+ * \return cstr The C string representation of the file.
+ */
+std::string PackImportsToC(const ffi::Module& m, bool system_lib,
+                           const std::string& c_symbol_prefix = "");
+
+/*!
+ * \brief Pack imported device library to a LLVM module.
+ *  Compile the LLVM module and link with the host library
+ *  will allow the DSO loader to automatically discover and import
+ *  the dependency from the shared library.
+ *
+ * \param m The host module with the imports.
+ * \param system_lib Whether expose as system library.
+ * \param target_triple LLVM target triple
+ * \param c_symbol_prefix Optional symbol prefix of the blob symbol.
+ *
+ * \return runtime::Module The generated LLVM module.
+ */
+ffi::Module PackImportsToLLVM(const ffi::Module& m, bool system_lib,
+                              const std::string& target_triple,
+                              const std::string& c_symbol_prefix = "");
+
+}  // namespace codegen
+}  // namespace tvm
+#endif  // TVM_TARGET_CODEGEN_H_

include/tvm/target/tag.h

@@ -0,0 +1,173 @@
+/*
+ * 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 tvm/target/tag.h
+ * \brief Target tag registry
+ */
+#ifndef TVM_TARGET_TAG_H_
+#define TVM_TARGET_TAG_H_
+
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/attr_registry_map.h>
+#include <tvm/target/target.h>
+
+#include <utility>
+
+namespace tvm {
+
+/*! \brief A target tag */
+class TargetTagNode : public ffi::Object {
+ public:
+  /*! \brief Name of the target */
+  ffi::String name;
+  /*! \brief Config map to generate the target */
+  ffi::Map<ffi::String, Any> config;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<TargetTagNode>()
+        .def_ro("name", &TargetTagNode::name)
+        .def_ro("config", &TargetTagNode::config);
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("target.TargetTag", TargetTagNode, ffi::Object);
+
+ private:
+  /*! \brief Return the index stored in attr registry */
+  uint32_t AttrRegistryIndex() const { return index_; }
+  /*! \brief Return the name stored in attr registry */
+  ffi::String AttrRegistryName() const { return name; }
+  /*! \brief Index used for internal lookup of attribute registry */
+  uint32_t index_;
+
+  template <typename, typename>
+  friend class AttrRegistry;
+  template <typename>
+  friend class AttrRegistryMapContainerMap;
+  friend class TargetTagRegEntry;
+};
+
+/*!
+ * \brief Managed reference class to TargetTagNode
+ * \sa TargetTagNode
+ */
+class TargetTag : public ffi::ObjectRef {
+ public:
+  /*!
+   * \brief Retrieve the Target given it the name of target tag
+   * \param target_tag_name Name of the target tag
+   * \return The Target requested
+   */
+  TVM_DLL static ffi::Optional<Target> Get(const ffi::String& target_tag_name);
+  /*!
+   * \brief List all names of the existing target tags
+   * \return A dictionary that maps tag name to the concrete target it corresponds to
+   */
+  TVM_DLL static ffi::Map<ffi::String, Target> ListTags();
+  /*!
+   * \brief Retrieve the raw config dict for a target tag
+   * \param target_tag_name Name of the target tag
+   * \return The config dict if the tag exists, nullopt otherwise
+   */
+  TVM_DLL static ffi::Optional<ffi::Map<ffi::String, Any>> GetConfig(
+      const ffi::String& target_tag_name);
+  /*!
+   * \brief Add a tag into the registry
+   * \param name Name of the tag
+   * \param config The target config corresponding to the tag
+   * \param override Allow overriding existing tags
+   * \return Target created with the tag
+   */
+  TVM_DLL static Target AddTag(ffi::String name, ffi::Map<ffi::String, Any> config, bool override);
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TargetTag, ffi::ObjectRef, TargetTagNode);
+
+ private:
+  /*! \brief Mutable access to the container class  */
+  TargetTagNode* operator->() { return static_cast<TargetTagNode*>(data_.get()); }
+  friend class TargetTagRegEntry;
+};
+
+class TargetTagRegEntry {
+ public:
+  /*!
+   * \brief Set the config dict corresponding to the target tag
+   * \param config The config dict for target creation
+   */
+  inline TargetTagRegEntry& set_config(ffi::Map<ffi::String, Any> config);
+  /*!
+   * \brief Add a key-value pair to the config dict
+   * \param key The attribute name
+   * \param value The attribute value
+   */
+  inline TargetTagRegEntry& with_config(ffi::String key, Any value);
+  /*! \brief Set name of the TargetTag to be the same as registry if it is empty */
+  inline TargetTagRegEntry& set_name();
+  /*!
+   * \brief Register or get a new entry.
+   * \param target_tag_name The name of the TargetTag.
+   * \return the corresponding entry.
+   */
+  TVM_DLL static TargetTagRegEntry& RegisterOrGet(const ffi::String& target_tag_name);
+
+ private:
+  TargetTag tag_;
+  ffi::String name;
+
+  /*! \brief private constructor */
+  explicit TargetTagRegEntry(uint32_t reg_index) : tag_(ffi::make_object<TargetTagNode>()) {
+    tag_->index_ = reg_index;
+  }
+  template <typename, typename>
+  friend class AttrRegistry;
+  friend class TargetTag;
+};
+
+inline TargetTagRegEntry& TargetTagRegEntry::set_config(ffi::Map<ffi::String, Any> config) {
+  tag_->config = std::move(config);
+  return *this;
+}
+
+inline TargetTagRegEntry& TargetTagRegEntry::with_config(ffi::String key, ffi::Any value) {
+  tag_->config.Set(key, value);
+  return *this;
+}
+
+inline TargetTagRegEntry& TargetTagRegEntry::set_name() {
+  if (tag_->name.empty()) {
+    tag_->name = name;
+  }
+  return *this;
+}
+
+#define TVM_TARGET_TAG_REGISTER_VAR_DEF \
+  [[maybe_unused]] static ::tvm::TargetTagRegEntry& __make_##TargetTag
+
+/*!
+ * \def TVM_REGISTER_TARGET_TAG
+ * \brief Register a new target tag, or set attribute of the corresponding target tag.
+ * \param TargetTagName The name of target tag
+ */
+#define TVM_REGISTER_TARGET_TAG(TargetTagName)                       \
+  TVM_FFI_STR_CONCAT(TVM_TARGET_TAG_REGISTER_VAR_DEF, __COUNTER__) = \
+      ::tvm::TargetTagRegEntry::RegisterOrGet(TargetTagName).set_name()
+
+}  // namespace tvm
+
+#endif  // TVM_TARGET_TAG_H_

include/tvm/target/target.h

@@ -0,0 +1,200 @@
+/*
+ * 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 tvm/target/target.h
+ * \brief Compilation target object.
+ */
+#ifndef TVM_TARGET_TARGET_H_
+#define TVM_TARGET_TARGET_H_
+
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/expr.h>
+#include <tvm/ir/function.h>
+#include <tvm/ir/with_context.h>
+#include <tvm/runtime/device_api.h>
+#include <tvm/target/target_kind.h>
+
+#include <string>
+
+namespace tvm {
+
+class TargetInternal;
+class Target;
+
+/*!
+ * \brief Compilation target.
+ * \sa Target
+ */
+class TargetNode : public ffi::Object {
+ public:
+  /*! \brief The kind of the target device */
+  TargetKind kind;
+  /*! \brief Target host information, must be Target type */
+  ffi::Optional<ffi::ObjectRef> host;
+  /*! \brief Tag of the target, can be empty */
+  ffi::String tag;
+  /*! \brief Keys for this target */
+  ffi::Array<ffi::String> keys;
+  /*! \brief Collection of attributes (includes feature.* keys set by canonicalizer) */
+  ffi::Map<ffi::String, Any> attrs;
+
+  /*!
+   * \brief The JSON string representation of the target
+   * \return JSON string of the target configuration (e.g. {"kind": "llvm", "mcpu": "cortex-a53"})
+   */
+  TVM_DLL const std::string& str() const;
+  /*! \return Export target to JSON-like configuration */
+  TVM_DLL ffi::Map<ffi::String, ffi::Any> ToConfig() const;
+  /*! \return The ffi::Optional<Target> typed target host of the TargetNode */
+  TVM_DLL ffi::Optional<Target> GetHost() const;
+  /*! \return The device type for this target */
+  TVM_DLL int GetTargetDeviceType() const;
+
+  /*!
+   * \brief Check if the target contains a key
+   *
+   * \param query_key The string name of the key to be checked
+   *
+   * \return True if the target's `TargetNode::keys` contains the
+   * specified key, False otherwise.
+   */
+  TVM_DLL bool HasKey(const std::string& query_key) const;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<TargetNode>()
+        .def_ro("kind", &TargetNode::kind)
+        .def_ro("tag", &TargetNode::tag)
+        .def_ro("keys", &TargetNode::keys)
+        .def_ro("attrs", &TargetNode::attrs)
+        .def_ro("host", &TargetNode::host);
+  }
+
+  /*!
+   * \brief Get an entry from attrs of the target
+   * \tparam TObjectRef Type of the attribute
+   * \param attr_key The name of the attribute key
+   * \param default_value The value returned if the key is not present
+   * \return An optional, std::nullopt if not found, otherwise the value found
+   */
+  template <typename TObjectRef>
+  ffi::Optional<TObjectRef> GetAttr(
+      const std::string& attr_key,
+      ffi::Optional<TObjectRef> default_value = ffi::Optional<TObjectRef>(std::nullopt)) const {
+    auto it = attrs.find(attr_key);
+    if (it != attrs.end()) {
+      return (*it).second.as_or_throw<ffi::Optional<TObjectRef>>();
+    } else {
+      return default_value;
+    }
+  }
+  /*!
+   * \brief Get an entry from attrs of the target
+   * \tparam TObjectRef Type of the attribute
+   * \param attr_key The name of the attribute key
+   * \param default_value The value returned if the key is not present
+   * \return An optional, std::nullopt if not found, otherwise the value found
+   */
+  template <typename TObjectRef>
+  ffi::Optional<TObjectRef> GetAttr(const std::string& attr_key, TObjectRef default_value) const {
+    return GetAttr<TObjectRef>(attr_key, ffi::Optional<TObjectRef>(default_value));
+  }
+
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("target.Target", TargetNode, ffi::Object);
+
+ private:
+  /*! \brief Internal string repr. */
+  mutable std::string str_repr_;
+
+  friend class TargetInternal;
+};
+
+/*!
+ * \brief Managed reference class to TargetNode.
+ * \sa TargetNode
+ */
+class Target : public ffi::ObjectRef {
+ public:
+  /*! \brief Construct a null Target */
+  TVM_DLL explicit Target(std::nullptr_t) { data_ = nullptr; }
+  /*!
+   * \brief Construct a Target given a string
+   * \param tag_or_config_or_target_str the string to parse for target
+   */
+  TVM_DLL explicit Target(const ffi::String& tag_or_config_or_target_str);
+  /*!
+   * \brief Construct a Target using a JSON-like configuration
+   * \param config The JSON-like configuration for target
+   */
+  TVM_DLL explicit Target(const ffi::Map<ffi::String, ffi::Any>& config);
+  /*!
+   * \brief Get the current target context from thread local storage.
+   * \param allow_not_defined If the context stack is empty and this is set to true, an
+   *   undefined Target will be returned. Otherwise, an empty context stack will cause a
+   *   runtime error.
+   * \return The target that is the current context. The target may not be defined if
+   * allow_not_defined is true.
+   */
+  TVM_DLL static tvm::Target Current(bool allow_not_defined = true);
+  /*!
+   * \brief Construct a Target given target and host
+   * \param target The Target typed object with host field undefined for target
+   * \param host The Target typed object for target host
+   */
+  TVM_DLL explicit Target(Target target, Target host);
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Target, ffi::ObjectRef, TargetNode);
+
+  static Target WithHost(const Target& target, const Target& host);
+
+  /*! \return The target with the host stripped out */
+  Target WithoutHost() const;
+
+ private:
+  Target(TargetKind kind, ffi::Optional<ffi::ObjectRef> host, ffi::String tag,
+         ffi::Array<ffi::String> keys, ffi::Map<ffi::String, ffi::Any> attrs);
+
+  // enable with syntax.
+  friend class TargetInternal;
+  friend class With<Target>;
+  /*!
+   * \brief Push a new target context onto the thread local stack.
+   *  The Target on top of the stack is used to determine which
+   *  specialization to use when invoking a GenericFunc.
+   */
+  TVM_DLL void EnterWithScope();
+  /*!
+   * \brief Pop a target off the thread local context stack,
+   *  restoring the previous target as the current context.
+   */
+  TVM_DLL void ExitWithScope();
+};
+
+/*!
+ * \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
+ *  recommended for other use.
+ * \param target The pointer to a Target typed object with host field to be updated
+ * \param host The pointer to a Target typed object for target host to be updated
+ */
+void CheckAndUpdateHostConsistency(Target* target, Target* host);
+
+}  // namespace tvm
+#endif  // TVM_TARGET_TARGET_H_

include/tvm/target/target_kind.h

@@ -0,0 +1,319 @@
+/*
+ * 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 tvm/target/target_kind.h
+ * \brief Target kind registry
+ */
+#ifndef TVM_TARGET_TARGET_KIND_H_
+#define TVM_TARGET_TARGET_KIND_H_
+
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/attr_registry_map.h>
+#include <tvm/ir/config_schema.h>
+#include <tvm/runtime/base.h>
+
+#include <memory>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+
+class Target;
+
+/*!
+ * \brief Target canonicalizer applied on instantiation of a given TargetKind.
+ *
+ * \param target_json Target in JSON format to be transformed during canonicalization.
+ * \return The transformed Target JSON object.
+ */
+using FTargetCanonicalizer =
+    ffi::TypedFunction<ffi::Map<ffi::String, ffi::Any>(ffi::Map<ffi::String, ffi::Any>)>;
+
+class TargetInternal;
+
+template <typename>
+class TargetKindAttrMap;
+
+/*! \brief Target kind, specifies the kind of the target */
+class TargetKindNode : public ffi::Object {
+ public:
+  /*! \brief Name of the target kind */
+  ffi::String name;
+  /*! \brief Device type of target kind */
+  int default_device_type;
+  /*! \brief Default keys of the target */
+  ffi::Array<ffi::String> default_keys;
+  /*! \brief Function used to canonicalize a JSON target during creation */
+  FTargetCanonicalizer target_canonicalizer;
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<TargetKindNode>()
+        .def_ro("name", &TargetKindNode::name)
+        .def_ro("default_device_type", &TargetKindNode::default_device_type,
+                refl::AttachFieldFlag::SEqHashIgnore())
+        .def_ro("default_keys", &TargetKindNode::default_keys,
+                refl::AttachFieldFlag::SEqHashIgnore());
+  }
+
+  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindUniqueInstance;
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("target.TargetKind", TargetKindNode, ffi::Object);
+
+ private:
+  /*! \brief Return the index stored in attr registry */
+  uint32_t AttrRegistryIndex() const { return index_; }
+  /*! \brief Return the name stored in attr registry */
+  ffi::String AttrRegistryName() const { return name; }
+  /*! \brief ConfigSchema for validating and resolving target attributes */
+  ir::ConfigSchema schema_;
+  /*! \brief Index used for internal lookup of attribute registry */
+  uint32_t index_;
+
+  template <typename, typename>
+  friend class AttrRegistry;
+  template <typename>
+  friend class AttrRegistryMapContainerMap;
+  friend class TargetKindRegEntry;
+  friend class TargetInternal;
+};
+
+/*!
+ * \brief Managed reference class to TargetKindNode
+ * \sa TargetKindNode
+ */
+class TargetKind : public ffi::ObjectRef {
+ public:
+  TargetKind() = default;
+  explicit TargetKind(ffi::ObjectPtr<TargetKindNode> data) : ffi::ObjectRef(data) {
+    TVM_FFI_ICHECK(data != nullptr);
+  }
+  /*! \brief Get the attribute map given the attribute name */
+  template <typename ValueType>
+  static inline TargetKindAttrMap<ValueType> GetAttrMap(const ffi::String& attr_name);
+  /*!
+   * \brief Retrieve the TargetKind given its name
+   * \param target_kind_name Name of the target kind
+   * \return The TargetKind requested
+   */
+  TVM_DLL static ffi::Optional<TargetKind> Get(const ffi::String& target_kind_name);
+  /*! \brief Mutable access to the container class  */
+  TargetKindNode* operator->() { return static_cast<TargetKindNode*>(data_.get()); }
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TargetKind, ffi::ObjectRef, TargetKindNode);
+
+ private:
+  TVM_DLL static const AttrRegistryMapContainerMap<TargetKind>& GetAttrMapContainer(
+      const ffi::String& attr_name);
+  friend class TargetKindRegEntry;
+  friend class TargetInternal;
+};
+
+/*!
+ * \brief ffi::Map<TargetKind, ValueType> used to store meta-information about TargetKind
+ * \tparam ValueType The type of the value stored in map
+ */
+template <typename ValueType>
+class TargetKindAttrMap : public AttrRegistryMap<TargetKind, ValueType> {
+ public:
+  using TParent = AttrRegistryMap<TargetKind, ValueType>;
+  using TParent::count;
+  using TParent::get;
+  using TParent::operator[];
+  explicit TargetKindAttrMap(const AttrRegistryMapContainerMap<TargetKind>& map) : TParent(map) {}
+};
+
+/*! \brief Value used with --runtime in target specs to indicate the C++ runtime. */
+static constexpr const char* kTvmRuntimeCpp = "c++";
+
+/*! \brief Value used with --runtime in target specs to indicate the C runtime. */
+static constexpr const char* kTvmRuntimeCrt = "c";
+
+/*!
+ * \brief Helper structure to register TargetKind
+ * \sa TVM_REGISTER_TARGET_KIND
+ */
+class TargetKindRegEntry {
+ public:
+  /*!
+   * \brief Register additional attributes to target_kind.
+   * \param attr_name The name of the attribute.
+   * \param value The value to be set.
+   * \param plevel The priority level of this attribute,
+   *  an higher priority level attribute
+   *  will replace lower priority level attribute.
+   *  Must be bigger than 0.
+   *
+   *  Cannot set with same plevel twice in the code.
+   *
+   * \tparam ValueType The type of the value to be set.
+   */
+  template <typename ValueType>
+  inline TargetKindRegEntry& set_attr(const ffi::String& attr_name, const ValueType& value,
+                                      int plevel = 10);
+  /*!
+   * \brief Set DLPack's device_type the target
+   * \param device_type Device type
+   */
+  inline TargetKindRegEntry& set_default_device_type(int device_type);
+  /*!
+   * \brief Set DLPack's device_type the target
+   * \param keys The default keys
+   */
+  inline TargetKindRegEntry& set_default_keys(std::vector<ffi::String> keys);
+  /*!
+   * \brief Set the canonicalizer function applied upon target creation.
+   * \param canonicalizer The target canonicalizer function.
+   */
+  inline TargetKindRegEntry& set_target_canonicalizer(FTargetCanonicalizer canonicalizer);
+  /*!
+   * \brief Register a valid configuration option and its ValueType for validation
+   * \param key The configuration key
+   * \param traits Optional traits (e.g. refl::DefaultValue, doc string, or raw default value)
+   * \tparam ValueType The value type to be registered
+   * \tparam Traits Optional trait types
+   */
+  template <typename ValueType, typename... Traits>
+  inline TargetKindRegEntry& add_attr_option(const ffi::String& key, Traits&&... traits);
+  /*! \brief Set name of the TargetKind to be the same as registry if it is empty */
+  inline TargetKindRegEntry& set_name();
+  /*!
+   * \brief List all the entry names in the registry.
+   * \return The entry names.
+   */
+  TVM_DLL static ffi::Array<ffi::String> ListTargetKinds();
+  /*!
+   * \brief Get all supported option names and types for a given Target kind.
+   * \return Map of option name to type
+   */
+  TVM_DLL static ffi::Map<ffi::String, ffi::String> ListTargetKindOptions(const TargetKind& kind);
+
+  /*!
+   * \brief Register or get a new entry.
+   * \param target_kind_name The name of the TargetKind.
+   * \return the corresponding entry.
+   */
+  TVM_DLL static TargetKindRegEntry& RegisterOrGet(const ffi::String& target_kind_name);
+
+ private:
+  TargetKind kind_;
+  ffi::String name;
+
+  /*! \brief private constructor */
+  explicit TargetKindRegEntry(uint32_t reg_index) : kind_(ffi::make_object<TargetKindNode>()) {
+    kind_->index_ = reg_index;
+  }
+  /*!
+   * \brief update the attribute TargetKindAttrMap
+   * \param key The name of the attribute
+   * \param value The value to be set
+   * \param plevel The priority level
+   */
+  TVM_DLL void UpdateAttr(const ffi::String& key, ffi::Any value, int plevel);
+  template <typename, typename>
+  friend class AttrRegistry;
+  friend class TargetKind;
+};
+
+template <typename ValueType>
+inline TargetKindAttrMap<ValueType> TargetKind::GetAttrMap(const ffi::String& attr_name) {
+  return TargetKindAttrMap<ValueType>(GetAttrMapContainer(attr_name));
+}
+
+template <typename ValueType>
+inline TargetKindRegEntry& TargetKindRegEntry::set_attr(const ffi::String& attr_name,
+                                                        const ValueType& value, int plevel) {
+  TVM_FFI_ICHECK_GT(plevel, 0) << "plevel in set_attr must be greater than 0";
+  ffi::Any rv;
+  rv = value;
+  UpdateAttr(attr_name, rv, plevel);
+  return *this;
+}
+
+inline TargetKindRegEntry& TargetKindRegEntry::set_default_device_type(int device_type) {
+  kind_->default_device_type = device_type;
+  return *this;
+}
+
+inline TargetKindRegEntry& TargetKindRegEntry::set_default_keys(std::vector<ffi::String> keys) {
+  kind_->default_keys = keys;
+  return *this;
+}
+
+inline TargetKindRegEntry& TargetKindRegEntry::set_target_canonicalizer(
+    FTargetCanonicalizer canonicalizer) {
+  kind_->target_canonicalizer = canonicalizer;
+  kind_->schema_.set_canonicalizer(canonicalizer);
+  return *this;
+}
+
+template <typename ValueType, typename... Traits>
+inline TargetKindRegEntry& TargetKindRegEntry::add_attr_option(const ffi::String& key,
+                                                               Traits&&... traits) {
+  kind_->schema_.def_option<ValueType>(key, std::forward<Traits>(traits)...);
+  return *this;
+}
+
+inline TargetKindRegEntry& TargetKindRegEntry::set_name() {
+  if (kind_->name.empty()) {
+    kind_->name = name;
+  }
+  return *this;
+}
+
+#define TVM_TARGET_KIND_REGISTER_VAR_DEF \
+  [[maybe_unused]] static ::tvm::TargetKindRegEntry& __make_##TargetKind
+
+/*!
+ * \def TVM_REGISTER_TARGET_KIND
+ * \brief Register a new target kind, or set attribute of the corresponding target kind.
+ *
+ * \param TargetKindName The name of target kind
+ * \param DeviceType The DLDeviceType of the target kind
+ *
+ * \code
+ *
+ *  TVM_REGISTER_TARGET_KIND("llvm")
+ *  .set_attr<TPreCodegenPass>("TPreCodegenPass", a-pre-codegen-pass)
+ *  .add_attr_option<Bool>("system_lib")
+ *  .add_attr_option<ffi::String>("mtriple")
+ *  .add_attr_option<ffi::String>("mattr");
+ *
+ * \endcode
+ */
+#define TVM_REGISTER_TARGET_KIND(TargetKindName, DeviceType)          \
+  TVM_FFI_STR_CONCAT(TVM_TARGET_KIND_REGISTER_VAR_DEF, __COUNTER__) = \
+      ::tvm::TargetKindRegEntry::RegisterOrGet(TargetKindName)        \
+          .set_name()                                                 \
+          .set_default_device_type(DeviceType)                        \
+          .add_attr_option<ffi::String>("kind")                       \
+          .add_attr_option<ffi::Array<ffi::String>>("keys")           \
+          .add_attr_option<ffi::String>("tag")                        \
+          .add_attr_option<ffi::String>("device")                     \
+          .add_attr_option<ffi::String>("model")                      \
+          .add_attr_option<ffi::Array<ffi::String>>("libs")           \
+          .add_attr_option<Target>("host")                            \
+          .add_attr_option<int64_t>("from_device")                    \
+          .add_attr_option<int64_t>("target_device_type")
+
+}  // namespace tvm
+
+#endif  // TVM_TARGET_TARGET_KIND_H_

include/tvm/target/virtual_device.h

@@ -0,0 +1,382 @@
+/*
+ * 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 tvm/target/virtual_device.h
+ * \brief A compile time representation for where data is to be stored at runtime, and how to
+ * compile code to compute it.
+ */
+
+#ifndef TVM_TARGET_VIRTUAL_DEVICE_H_
+#define TVM_TARGET_VIRTUAL_DEVICE_H_
+
+#include <tvm/ir/transform.h>
+#include <tvm/target/target.h>
+
+#include <string>
+#include <unordered_set>
+#include <utility>
+
+namespace tvm {
+
+/*!
+ * \brief Abstract label for an area of memory.
+ *
+ * Currently uninterpreted and arbitrary. Likely to be replaced by a structured representation
+ * of a memory pool in the future. Please try to use this alias instead of ffi::String to aid future
+ * code migration.
+ */
+using MemoryScope = ffi::String;
+
+// NOTE: cannot use enum as they are out of bound of the original enum
+// and results in an undefined behavior
+// A 'null' device type, does not correspond to any DLDeviceType enum.
+// TODO(mbs): This is to help us as we transition away from representing the 'homogenous' case
+// as a singleton target map indexed by the invalid DLDeviceType '0'.
+constexpr int kNullDeviceType = 0;
+
+// An 'invalid' device type, does not correspond to any DLDeviceType enum.
+constexpr int kInvalidDeviceType = -1;
+
+/*!
+ * \brief Describes at compile time the constraints on where data is to be stored at runtime
+ * down to the (virtual) device and memory scope level, and how to compile code to compute that
+ * data. Used by the \p PlanDevices pass to collect and solve (virtual) device constraints for
+ * the whole Relax program.
+ *
+ * Is a quadruple of:
+ * - A \p device_type (\p DLDeviceType). May be \p kInvalidDeviceType if unconstrained.
+ * - A \p virtual_device_id (\p int). This allows us to distinguish distinct devices
+ *   with the same \p Target, for example in a multi-GPU system. May be -1 if unconstrained.
+ *   See "Virtual Devices" below.
+ * - A \p target (\p Target) describing how to compile code for the intended device. May be null
+ *   if unconstrained.
+ * - A \p memory_scope (\p MemoryScope, which is currently just \p String) describing which memory
+ *   area is to be used to hold data. May be "" if unconstrained. See "Memory Scopes and Devices"
+ *   below.
+ *
+ * Some or all of these fields may be unconstrained, signaling that device planning is free to
+ * choose a value consistent with the whole program. However if a \p target is given then the \p
+ * device_type must equal \p target->GetTargetDeviceType().
+ *
+ * Note that currently we assume if a function returns its result on a particular (virtual) device
+ * then the function body is also executed on that device.
+ *
+ *
+ * By 'execution' we include both (fused) primitive operators, and all the Relax expressions
+ * surrounding them which coordinates data and control flow. Again, typically non-primitive
+ * operators must be executed on a 'CPU'-like device with good support for control flow.
+ *
+ * Since TVM targets such a wide range of systems it is not possible for \p VirtualDevice to impose
+ * much semantics on these fields, particularly for \p virtual_device_id and \p memory_scope.
+ * Instead we assume downstream passes and codegen will interpret an validate these fields
+ * appropriately.
+ *
+ * Targets vs Devices
+ * ------------------
+ * Generally \p Targets (a compile-time only datastructue) describe compiler options for a specific
+ * microarchitecture and toolchain, while \p Devices (a runtime datastructure also available at
+ * compile time) describe a physical device on the target system. Obviously the target must agree
+ * with the device's microarchitecture, but we otherwise don't impose any constraints between them:
+ *  - It's ok to use different \p Targets for the same \p Device, eg to squeeze some extra perf
+ *    out of a particular primitive using particular compiler flags.
+ *  - It's ok to use the same \p Target for multiple \p Devices, eg if we have multiple CPUs.
+ *
+ * Traditionally TVM assumes at most one \p Target per \p DLDeviceType. We are moving away from that
+ * assumption.
+ *
+ * Virtual vs Physical Devices
+ * ---------------------------
+ * The \p virtual_device_id may be used by downstream passes or the runtime to help decide which
+ * \p device_id to use for a particular physical runtime \p Device. For example:
+ *  - Some runtimes may support passing in an array of actual `device` specifications, and the
+ *    \p virtual_device_id can be used at runtime as an index into that array.
+ *  - Some runtimes may support dynamically allocating computations to physical devices. On these
+ *    systems a large space of \p virtual_device_ids could be used at compile time, even though
+ *    at runtime only a few physical devices will be present.
+ *
+ * The \p virtual_device_id may also be left unconstrained if not needed.
+ *
+ * Memory Scopes and Devices
+ * -------------------------
+ * Multi-device systems can have complex memory hierarchies. For example
+ * \code
+ * (kDLCPU, 0, "llvm", "global")
+ * \endcode
+ * and
+ * \code
+ * (kDLCPU, 1, "llvm", "global")
+ * \endcode
+ * could denote:
+ * - The same memory area accessible from two separate CPUs without any CPU affinity;
+ * - Distinct memory areas in a NUMA architecture for which cross-device access is handled
+ *   by the memory system;
+ * - Outright distinct memory areas, where one device cannot directly address the memory of
+ *   another.
+ *
+ * Similarly:
+ * \code
+ * (kDLCPU, 0, "llvm", "global")
+ * \endcode
+ * and
+ * \code
+ * (kDLCUDA, 0, "cuda", "host")
+ * \endcode
+ * could denote the same memory area, but with very different access costs.
+ *
+ * Furthermore, not all memory scopes are accessible to all devices, and it is possible for
+ * a memory scope to only be accessible to a device when code is compiled with particular
+ * \p Target options.
+ *
+ * \p VirtualDevices themselves have no system-level understanding. Currently the \p PlanDevices
+ * pass will simply insert "device_copy" operators wherever \p VirtualDevices are not exactly
+ * pointwise equal. We may revisit this in the future as the work on memory pools matures.
+ *
+ * Joining and Defaulting
+ * ----------------------
+ * It is possible to 'join' two \p VirtualDevices to yield the most constrained \p VirtualDevice
+ * which agrees with both join arguments. Eg:
+ * \code
+ * Join((kDLCPU, -1, "llvm", ""), (kInvalidDeviceType, 3, null, "global))
+ *   => (kDLCPU, 3, "llvm", "global")
+ * Join((kDLCPU, -1, "llvm", ""), (kInvalidDeviceType, 3, null, "local))
+ *   => null (no join possible)
+ * \endcode
+ *
+ * Related to 'join' is 'default', which only takes constrained fields from the rhs when the
+ * lhs is unconstrained:
+ * \code
+ * Default(kDLCPU, -1, "llvm", "local"), (kDLCPU, 3, null, "global"))
+ *   => (kDLCPU, 3, "llvm", "local")
+ * \endcode
+ *
+ * These operations are needed during device planning.
+ */
+
+class VirtualDeviceNode : public AttrsNode {
+ private:
+  /*!
+   * \brief The \p DLDeviceType (represented as an int) of the virtual device. If \p target is
+   * known then this will be equal to \p target->GetTargetDeviceType(). If \p target is null then
+   * the target is to be determined later.
+   *
+   * This is needed to support the legacy "on_device" and "device_copy" calls which only allow
+   * a \p DLDeviceTypes (as an integer) to be given.
+   *
+   * kInvalidDeviceType denotes unconstrained. An int since the DLDeviceType enum representation
+   * is not fixed. Private to discourage further int vs DLDeviceType confusion.
+   */
+  int /* actually DLDeviceType */ device_type_int;
+
+ public:
+  DLDeviceType device_type() const { return static_cast<DLDeviceType>(device_type_int); }
+
+  /*!
+   * \brief The device identifier for the virtual device. This must be resolved to a physical
+   * device identifier either during compilation or at runtime.
+   *
+   * -1 denotes unconstrained.
+   */
+  int virtual_device_id;
+
+  /*!
+   * \brief The \p Target describing how to compile for the virtual device.
+   *
+   * Null denotes unconstrained. Note that if a target later becomes known for this \p VirtualDevice
+   * then it must be consistent with the \p device_type if already known. This is enforced by the
+   * Join and Default methods.
+   */
+  Target target;
+
+  /*!
+   * \brief The scope of memory w.r.t. the virtual device which holds data.
+   *
+   * Empty denotes unconstrained.
+   */
+  MemoryScope memory_scope;
+
+  /*!
+   * \brief Returns true if virtual device is 'fully unconstrained', ie no target/device type,
+   * device id or memory scope is specified.
+   */
+  bool IsFullyUnconstrained() const {
+    return !target.defined() && device_type() == kInvalidDeviceType && virtual_device_id == -1 &&
+           memory_scope.empty();
+  }
+
+  /*!
+   * \brief Returns true if virtual device is 'fully constrained', ie target, device id and memory
+   * scope are all specified.
+   */
+  bool IsFullyConstrained() const {
+    return target.defined() && virtual_device_id != -1 && !memory_scope.empty();
+  }
+
+  /*!
+   * \brief Returns the (virtual) \p Device implied by this \p VirtualDevice. Both the \p
+   * device_type and \p virtual_device_must be constrained. The returned \p Device may not
+   * correspond to any physical device available at compile time or even runtime: see "Virtual vs
+   * Physical Devices" above.
+   */
+  Device ToDevice() const {
+    TVM_FFI_ICHECK(device_type_int != kInvalidDeviceType);
+    TVM_FFI_ICHECK(virtual_device_id != -1);
+    Device device;
+    device.device_type = device_type();
+    device.device_id = virtual_device_id;
+    return device;
+  }
+
+  static void RegisterReflection() {
+    namespace refl = tvm::ffi::reflection;
+    refl::ObjectDef<VirtualDeviceNode>()
+        .def_ro("device_type_int", &VirtualDeviceNode::device_type_int,
+                "The type of the virtual device.", refl::DefaultValue(kInvalidDeviceType))
+        .def_ro("virtual_device_id", &VirtualDeviceNode::virtual_device_id,
+                "The device id of the virtual device.", refl::DefaultValue(-1))
+        .def_ro("target", &VirtualDeviceNode::target,
+                "The target describing how to compile for the virtual device.",
+                refl::DefaultValue(Target()))
+        .def_ro("memory_scope", &VirtualDeviceNode::memory_scope,
+                "The area of memory w.r.t. the virtual device where data is stored.",
+                refl::DefaultValue(""));
+  }
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("target.VirtualDevice", VirtualDeviceNode, AttrsNode);
+
+  friend class VirtualDevice;
+};
+
+/*!
+ * \brief Managed reference class to \p VirtualDeviceNode.
+ */
+class VirtualDevice : public ffi::ObjectRef {
+ public:
+  /*!
+   * \brief Construct a virtual device.
+   * \param device_type_int The device type for the virtual device, or \p kInvalidDeviceType if
+   * unconstrained.  If \p target is defined then must match its \p target->GetTargetDeviceType().
+   * \param virtual_device_id The device id for the virtual device, or -1 if unconstrained.
+   * \param target The target describing how to compile for the virtual device, or null if
+   * unconstrained.
+   * \param memory_scope The memory scope w.r.t. the virtual device which holds data, or "" if
+   * unconstrained.
+   */
+  TVM_DLL explicit VirtualDevice(int device_type_int = kInvalidDeviceType,
+                                 int virtual_device_id = -1, Target target = {},
+                                 MemoryScope memory_scope = {});
+
+  /*! \brief Returns the unique fully unconstrained \p VirtualDevice. */
+  static VirtualDevice FullyUnconstrained();
+
+  /*!
+   * \brief Returns the \p VirtualDevice for \p device_type and (if not -1) \p virtual_device_id.
+   * The target and memory scope will be unconstrained.
+   */
+  static VirtualDevice ForDeviceType(DLDeviceType device_type, int virtual_device_id = -1) {
+    TVM_FFI_ICHECK_GT(device_type, 0);
+    return VirtualDevice(device_type, virtual_device_id);
+  }
+  static VirtualDevice ForDeviceType(int device_type, int virtual_device_id = -1) {
+    return ForDeviceType(static_cast<DLDeviceType>(device_type), virtual_device_id);
+  }
+  static VirtualDevice ForDeviceType(int64_t device_type, int virtual_device_id = -1) {
+    return ForDeviceType(static_cast<int>(device_type), virtual_device_id);
+  }
+
+  /*! \brief Returns the \p VirtualDevice for \p device. */
+  static VirtualDevice ForDevice(const Device& device) {
+    return ForDeviceType(device.device_type, device.device_id);
+  }
+
+  /*! \brief Returns the \p VirtualDevice for \p device and \p target. */
+  static VirtualDevice ForDeviceAndTarget(const Device& device, Target target) {
+    return VirtualDevice(device.device_type, device.device_id, std::move(target));
+  }
+
+  /*! \brief Returns the \p VirtualDevice for \p target. */
+  static VirtualDevice ForTarget(Target target) {
+    DLDeviceType device_type = static_cast<DLDeviceType>(target->GetTargetDeviceType());
+    return VirtualDevice(device_type, /*virtual_device_id=*/0, std::move(target));
+  }
+
+  /*! \brief Returns the \p VirtualDevice for \p memory_scope alone. */
+  static VirtualDevice ForMemoryScope(MemoryScope memory_scope) {
+    return VirtualDevice(kInvalidDeviceType, -1, {}, std::move(memory_scope));
+  }
+
+  /*! \brief Returns the \p VirtualDevice for \p device, \p target and \p memory_scope. */
+  TVM_DLL static VirtualDevice ForDeviceTargetAndMemoryScope(const Device& device, Target target,
+                                                             MemoryScope memory_scope) {
+    return VirtualDevice(device.device_type, device.device_id, std::move(target),
+                         std::move(memory_scope));
+  }
+
+  /*!
+   * \brief Returns the 'join' of \p lhs and \p rhs. The result will agree pointwise with
+   * \p lhs and \p rhs on all their constrained fields. Returns the null optional if no such
+   * join exists, ie there's disagreement on at least one constrained field.
+   */
+  TVM_DLL static ffi::Optional<VirtualDevice> Join(const VirtualDevice& lhs,
+                                                   const VirtualDevice& rhs);
+
+  /*!
+   * \brief Returns the 'default' of \p lhs and \p rhs. The result will be \p lhs, except any
+   * unconstrained fields in \p lhs will take their value from \p rhs. Always well-defined.
+   */
+  TVM_DLL static VirtualDevice Default(const VirtualDevice& lhs, const VirtualDevice& rhs);
+
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VirtualDevice, ffi::ObjectRef, VirtualDeviceNode);
+
+  friend class VirtualDeviceCache;  // Private implementation helper.
+};
+
+/*!
+ * \brief A cache of \p VirtualDevices. This can be used:
+ *  - To avoid ending up with lots of identical instances, since the space of VirtualDevices for any
+ *    one compilation is very small but the number of points they need to be constructed can
+ *    be very large (eg during device planning).
+ *  - So we can assume \p VirtualDevices are pointer equal if and only if they are structurally
+ * equal. This simplifies the unification of 'device domains' which are built on \p VirtualDevices.
+ */
+class TVM_DLL VirtualDeviceCache {
+ public:
+  /*! \brief Returns the unique \p VirtualDevice representing given fields. */
+  VirtualDevice Make(int device_type = kInvalidDeviceType, int virtual_device_id = -1,
+                     Target target = {}, MemoryScope memory_scope = {});
+
+  /*!
+   * \brief Returns the unique \p VirtualDevice structurally equal to the given \p virtual_device.
+   */
+  VirtualDevice Unique(const VirtualDevice& virtual_device);
+
+ private:
+  /*! \brief Already constructed VirtualDevices. */
+  std::unordered_set<VirtualDevice, ffi::StructuralHash, ffi::StructuralEqual> cache_;
+};
+
+/*! brief The attribute key for the virtual device. This key will be promoted to first class on
+ * functions. For use in the parser and printer only.
+ *
+ * Type: VirtualDevice
+ */
+constexpr const char* kVirtualDevice = "virtual_device";
+
+}  // namespace tvm
+
+#endif  //  TVM_TARGET_VIRTUAL_DEVICE_H_