wehub-resource-sync
145 lines
4.5 KiB
C++
145 lines
4.5 KiB
C++
/*
|
|
* 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/ir/attr_registry_map.h
|
|
* \brief Attribute map used in registry.
|
|
*/
|
|
#ifndef TVM_IR_ATTR_REGISTRY_MAP_H_
|
|
#define TVM_IR_ATTR_REGISTRY_MAP_H_
|
|
|
|
#include <tvm/ffi/string.h>
|
|
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
namespace tvm {
|
|
|
|
/*!
|
|
* \brief Generic attribute map.
|
|
* \tparam KeyType the type of the key.
|
|
*/
|
|
template <typename KeyType>
|
|
class AttrRegistryMapContainerMap {
|
|
public:
|
|
/*!
|
|
* \brief Check if the map has key.
|
|
* \param key The key to the map
|
|
* \return 1 if key is contained in map, 0 otherwise.
|
|
*/
|
|
int count(const KeyType& key) const {
|
|
if (key.defined()) {
|
|
const uint32_t idx = key->AttrRegistryIndex();
|
|
return idx < data_.size() ? (data_[idx].second != 0) : 0;
|
|
} else {
|
|
return 0;
|
|
}
|
|
}
|
|
/*!
|
|
* \brief get the corresponding value element at key.
|
|
* \param key The key to the map
|
|
* \return the const reference to the content value.
|
|
*/
|
|
const ffi::Any& operator[](const KeyType& key) const {
|
|
TVM_FFI_ICHECK(key.defined());
|
|
const uint32_t idx = key->AttrRegistryIndex();
|
|
TVM_FFI_ICHECK(idx < data_.size() && data_[idx].second != 0)
|
|
<< "Attribute " << attr_name_ << " has not been registered for " << key->name;
|
|
return data_[idx].first;
|
|
}
|
|
/*!
|
|
* \brief get the corresponding value element at key with default value.
|
|
* \param key The key to the map
|
|
* \param def_value The default value when the key does not exist.
|
|
* \return the const reference to the content value.
|
|
* \tparam ValueType The content value type.
|
|
*/
|
|
template <typename ValueType>
|
|
ValueType get(const KeyType& key, ValueType def_value) const {
|
|
TVM_FFI_ICHECK(key.defined());
|
|
const uint32_t idx = key->AttrRegistryIndex();
|
|
if (idx < data_.size() && data_[idx].second != 0) {
|
|
if constexpr (std::is_same_v<ValueType, ffi::Any>) {
|
|
return data_[idx].first;
|
|
} else {
|
|
return data_[idx].first.template cast<ValueType>();
|
|
}
|
|
} else {
|
|
return def_value;
|
|
}
|
|
}
|
|
|
|
private:
|
|
/*! \brief The name of the attr field */
|
|
ffi::String attr_name_;
|
|
/*! \brief The internal data. */
|
|
std::vector<std::pair<ffi::Any, int>> data_;
|
|
/*! \brief The constructor */
|
|
AttrRegistryMapContainerMap() = default;
|
|
template <typename, typename>
|
|
friend class AttrRegistry;
|
|
friend class OpRegEntry;
|
|
};
|
|
|
|
/*!
|
|
* \brief ffi::Map<Key, ValueType> used to store meta-data.
|
|
* \tparam KeyType The type of the key
|
|
* \tparam ValueType The type of the value stored in map.
|
|
*/
|
|
template <typename KeyType, typename ValueType>
|
|
class AttrRegistryMap {
|
|
public:
|
|
/*!
|
|
* \brief constructor
|
|
* \param map The internal map.
|
|
*/
|
|
explicit AttrRegistryMap(const AttrRegistryMapContainerMap<KeyType>& map) : map_(map) {}
|
|
/*!
|
|
* \brief Check if the map has op as key.
|
|
* \param key The key to the map
|
|
* \return 1 if op is contained in map, 0 otherwise.
|
|
*/
|
|
int count(const KeyType& key) const { return map_.count(key); }
|
|
/*!
|
|
* \brief get the corresponding value element at key.
|
|
* \param key The key to the map
|
|
* \return the const reference to the content value.
|
|
*/
|
|
ValueType operator[](const KeyType& key) const {
|
|
if constexpr (std::is_same_v<ValueType, ffi::Any>) {
|
|
return map_[key];
|
|
} else {
|
|
return map_[key].template cast<ValueType>();
|
|
}
|
|
}
|
|
/*!
|
|
* \brief get the corresponding value element at key with default value.
|
|
* \param key The key to the map
|
|
* \param def_value The default value when the key does not exist.
|
|
* \return the const reference to the content value.
|
|
*/
|
|
ValueType get(const KeyType& key, ValueType def_value) const { return map_.get(key, def_value); }
|
|
|
|
protected:
|
|
/*! \brief The internal map field */
|
|
const AttrRegistryMapContainerMap<KeyType>& map_;
|
|
};
|
|
|
|
} // namespace tvm
|
|
#endif // TVM_IR_ATTR_REGISTRY_MAP_H_
|