attr_registry_map.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_NODE_ATTR_REGISTRY_MAP_H_
#define TVM_NODE_ATTR_REGISTRY_MAP_H_
 
#include <tvm/ffi/string.h>
 
#include <utility>
#include <vector>
 
namespace tvm {
 
template <typename KeyType>
class AttrRegistryMapContainerMap {
 public:
  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;
    }
  }
  const ffi::Any& operator[](const KeyType& key) const {
    ICHECK(key.defined());
    const uint32_t idx = key->AttrRegistryIndex();
    ICHECK(idx < data_.size() && data_[idx].second != 0)
        << "Attribute " << attr_name_ << " has not been registered for " << key->name;
    return data_[idx].first;
  }
  template <typename ValueType>
  ValueType get(const KeyType& key, ValueType def_value) const {
    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:
  String attr_name_;
  std::vector<std::pair<ffi::Any, int>> data_;
  AttrRegistryMapContainerMap() = default;
  template <typename, typename>
  friend class AttrRegistry;
  friend class OpRegEntry;
};
 
template <typename KeyType, typename ValueType>
class AttrRegistryMap {
 public:
  explicit AttrRegistryMap(const AttrRegistryMapContainerMap<KeyType>& map) : map_(map) {}
  int count(const KeyType& key) const { return map_.count(key); }
  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>();
    }
  }
  ValueType get(const KeyType& key, ValueType def_value) const { return map_.get(key, def_value); }
 
 protected:
  const AttrRegistryMapContainerMap<KeyType>& map_;
};
 
}  // namespace tvm
#endif  // TVM_NODE_ATTR_REGISTRY_MAP_H_