target_kind.h

Go to the documentation of this file.

/*
 * 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.
 */
 
#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/node/attr_registry_map.h>
#include <tvm/node/node.h>
 
#include <memory>
#include <unordered_map>
#include <utility>
#include <vector>
 
namespace tvm {
 
class Target;
 
using TargetFeatures = ffi::Map<ffi::String, ffi::Any>;
 
using TargetJSON = ffi::Map<ffi::String, ffi::Any>;
using FTVMTargetParser = ffi::TypedFunction<TargetJSON(TargetJSON)>;
 
namespace detail {
template <typename, typename, typename>
struct ValueTypeInfoMaker;
}
 
class TargetInternal;
 
template <typename>
class TargetKindAttrMap;
 
class TargetKindNode : public Object {
 public:
  ffi::String name;
  int default_device_type;
  ffi::Array<ffi::String> default_keys;
  ffi::Function preprocessor;
  FTVMTargetParser target_parser;
 
  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, Object);
 
 private:
  uint32_t AttrRegistryIndex() const { return index_; }
  ffi::String AttrRegistryName() const { return name; }
  struct ValueTypeInfo {
    ffi::String type_key;
    int32_t type_index;
    std::unique_ptr<ValueTypeInfo> key;
    std::unique_ptr<ValueTypeInfo> val;
  };
  std::unordered_map<ffi::String, ValueTypeInfo> key2vtype_;
  std::unordered_map<ffi::String, ffi::Any> key2default_;
  uint32_t index_;
 
  template <typename, typename, typename>
  friend struct detail::ValueTypeInfoMaker;
  template <typename, typename>
  friend class AttrRegistry;
  template <typename>
  friend class AttrRegistryMapContainerMap;
  friend class TargetKindRegEntry;
  friend class TargetInternal;
};
 
class TargetKind : public ObjectRef {
 public:
  TargetKind() = default;
  explicit TargetKind(ObjectPtr<TargetKindNode> data) : ObjectRef(data) {
    TVM_FFI_ICHECK(data != nullptr);
  }
  template <typename ValueType>
  static inline TargetKindAttrMap<ValueType> GetAttrMap(const ffi::String& attr_name);
  TVM_DLL static ffi::Optional<TargetKind> Get(const ffi::String& target_kind_name);
  TargetKindNode* operator->() { return static_cast<TargetKindNode*>(data_.get()); }
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TargetKind, ObjectRef, TargetKindNode);
 
 private:
  TVM_DLL static const AttrRegistryMapContainerMap<TargetKind>& GetAttrMapContainer(
      const ffi::String& attr_name);
  friend class TargetKindRegEntry;
  friend class TargetInternal;
};
 
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) {}
};
 
static constexpr const char* kTvmRuntimeCpp = "c++";
 
static constexpr const char* kTvmRuntimeCrt = "c";
 
class TargetKindRegEntry {
 public:
  template <typename ValueType>
  inline TargetKindRegEntry& set_attr(const ffi::String& attr_name, const ValueType& value,
                                      int plevel = 10);
  inline TargetKindRegEntry& set_default_device_type(int device_type);
  inline TargetKindRegEntry& set_default_keys(std::vector<ffi::String> keys);
  template <typename FLambda>
  inline TargetKindRegEntry& set_attrs_preprocessor(FLambda f);
  inline TargetKindRegEntry& set_target_parser(FTVMTargetParser parser);
  template <typename ValueType>
  inline TargetKindRegEntry& add_attr_option(const ffi::String& key);
  template <typename ValueType>
  inline TargetKindRegEntry& add_attr_option(const ffi::String& key, ffi::Any default_value);
  inline TargetKindRegEntry& set_name();
  TVM_DLL static ffi::Array<ffi::String> ListTargetKinds();
  TVM_DLL static ffi::Map<ffi::String, ffi::String> ListTargetKindOptions(const TargetKind& kind);
 
  TVM_DLL static TargetKindRegEntry& RegisterOrGet(const ffi::String& target_kind_name);
 
 private:
  TargetKind kind_;
  ffi::String name;
 
  explicit TargetKindRegEntry(uint32_t reg_index) : kind_(ffi::make_object<TargetKindNode>()) {
    kind_->index_ = reg_index;
  }
  TVM_DLL void UpdateAttr(const ffi::String& key, ffi::Any value, int plevel);
  template <typename, typename>
  friend class AttrRegistry;
  friend class TargetKind;
};
 
namespace detail {
template <typename Type, template <typename...> class Container>
struct is_specialized : std::false_type {
  using type = std::false_type;
};
 
template <template <typename...> class Container, typename... Args>
struct is_specialized<Container<Args...>, Container> : std::true_type {
  using type = std::true_type;
};
 
template <typename ValueType,
          typename IsArray = typename is_specialized<ValueType, ffi::Array>::type,
          typename IsMap = typename is_specialized<ValueType, ffi::Map>::type>
struct ValueTypeInfoMaker {};
 
template <typename ValueType>
struct ValueTypeInfoMaker<ValueType, std::false_type, std::false_type> {
  using ValueTypeInfo = TargetKindNode::ValueTypeInfo;
 
  ValueTypeInfo operator()() const {
    ValueTypeInfo info;
    info.key = nullptr;
    info.val = nullptr;
    if constexpr (std::is_base_of_v<ObjectRef, ValueType>) {
      int32_t tindex = ffi::TypeToRuntimeTypeIndex<ValueType>::v();
      info.type_index = tindex;
      info.type_key = runtime::Object::TypeIndex2Key(tindex);
      return info;
    } else if constexpr (std::is_same_v<ValueType, ffi::String>) {
      // special handle string since it can be backed by multiple types.
      info.type_index = ffi::TypeIndex::kTVMFFIStr;
      info.type_key = ffi::TypeTraits<ValueType>::TypeStr();
      return info;
    } else {
      // TODO(tqchen) consider upgrade to leverage any system to support union type
      constexpr int32_t tindex = ffi::TypeToFieldStaticTypeIndex<ValueType>::value;
      static_assert(tindex != ffi::TypeIndex::kTVMFFIAny, "Do not support union type for now");
      info.type_index = tindex;
      info.type_key = runtime::Object::TypeIndex2Key(tindex);
      return info;
    }
  }
};
 
template <typename ValueType>
struct ValueTypeInfoMaker<ValueType, std::true_type, std::false_type> {
  using ValueTypeInfo = TargetKindNode::ValueTypeInfo;
 
  ValueTypeInfo operator()() const {
    using key_type = ValueTypeInfoMaker<typename ValueType::value_type>;
    uint32_t tindex = ValueType::ContainerType::_GetOrAllocRuntimeTypeIndex();
    ValueTypeInfo info;
    info.type_index = tindex;
    info.type_key = runtime::Object::TypeIndex2Key(tindex);
    info.key = std::make_unique<ValueTypeInfo>(key_type()());
    info.val = nullptr;
    return info;
  }
};
 
template <typename ValueType>
struct ValueTypeInfoMaker<ValueType, std::false_type, std::true_type> {
  using ValueTypeInfo = TargetKindNode::ValueTypeInfo;
  ValueTypeInfo operator()() const {
    using key_type = ValueTypeInfoMaker<typename ValueType::key_type>;
    using val_type = ValueTypeInfoMaker<typename ValueType::mapped_type>;
    uint32_t tindex = ValueType::ContainerType::_GetOrAllocRuntimeTypeIndex();
    ValueTypeInfo info;
    info.type_index = tindex;
    info.type_key = runtime::Object::TypeIndex2Key(tindex);
    info.key = std::make_unique<ValueTypeInfo>(key_type()());
    info.val = std::make_unique<ValueTypeInfo>(val_type()());
    return info;
  }
};
 
}  // namespace detail
 
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) {
  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;
}
 
template <typename FLambda>
inline TargetKindRegEntry& TargetKindRegEntry::set_attrs_preprocessor(FLambda f) {
  LOG(WARNING) << "set_attrs_preprocessor is deprecated please use set_target_parser instead";
  kind_->preprocessor = ffi::Function::FromTyped(std::move(f));
  return *this;
}
 
inline TargetKindRegEntry& TargetKindRegEntry::set_target_parser(FTVMTargetParser parser) {
  kind_->target_parser = parser;
  return *this;
}
 
template <typename ValueType>
inline TargetKindRegEntry& TargetKindRegEntry::add_attr_option(const ffi::String& key) {
  ICHECK(!kind_->key2vtype_.count(key))
      << "AttributeError: add_attr_option failed because '" << key << "' has been set once";
  kind_->key2vtype_[key] = detail::ValueTypeInfoMaker<ValueType>()();
  return *this;
}
 
template <typename ValueType>
inline TargetKindRegEntry& TargetKindRegEntry::add_attr_option(const ffi::String& key,
                                                               Any default_value) {
  add_attr_option<ValueType>(key);
  kind_->key2default_[key] = default_value;
  return *this;
}
 
inline TargetKindRegEntry& TargetKindRegEntry::set_name() {
  if (kind_->name.empty()) {
    kind_->name = name;
  }
  return *this;
}
 
#define TVM_TARGET_KIND_REGISTER_VAR_DEF \
  static DMLC_ATTRIBUTE_UNUSED ::tvm::TargetKindRegEntry& __make_##TargetKind
 
#define TVM_REGISTER_TARGET_KIND(TargetKindName, DeviceType)      \
  TVM_STR_CONCAT(TVM_TARGET_KIND_REGISTER_VAR_DEF, __COUNTER__) = \
      ::tvm::TargetKindRegEntry::RegisterOrGet(TargetKindName)    \
          .set_name()                                             \
          .set_default_device_type(DeviceType)                    \
          .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_