/*! * Copyright (c) 2016 by Contributors * \file scope.h * \brief attribute scope data structure, * defines attributes on current domain */ #ifndef TVM_PASS_SCOPE_H_ #define TVM_PASS_SCOPE_H_ #include <tvm/ir.h> #include <unordered_map> #include <vector> #include <string> namespace tvm { namespace ir { /*! * \brief Attribute scope of Nodes in the IR. * \tparam ValueType The value of of the scope. */ template<typename K, typename V> class Scope { public: /*! * \brief Push value to scope * \param key the key to be pushed. * \param v The value to be pushed. */ inline void Push(const K& key, V v) { data_[key].emplace_back(v); } /*! * \brief Pop value from scope. * \param key the key to be poped */ inline void Pop(const K& key) { auto& v = data_[key]; CHECK_NE(v.size(), 0U); v.pop_back(); } /*! * \brief Get value from the scope * \param key the key to fetch. * \return The value to be fetched. */ inline V operator[](const K& key) const { const auto it = data_.find(key); CHECK(it != data_.end() && it->second.size() != 0) << "cannot find value in scope"; return it->second.back(); } private: std::unordered_map<K, std::vector<V> > data_; }; /*! \brief Attribute key for specific attribute */ struct AttrKey { /*! \brief The node of the attribute */ NodeRef node; /*! \brief The type key of the attribute. */ std::string type_key; // overload operator == inline bool operator==(const AttrKey& other) const { return node == other.node && type_key == other.type_key; } }; } // namespace ir } // namespace tvm namespace std { template <> struct hash<::tvm::ir::AttrKey> { std::size_t operator()(const ::tvm::ir::AttrKey& k) const { size_t lhs = k.node.hash(); size_t rhs = std::hash<std::string>()(k.type_key); lhs ^= rhs + 0x9e3779b9 + (lhs << 6) + (lhs >> 2); return lhs; } }; } // namespace std #endif // TVM_PASS_SCOPE_H_