Reserve input/output edgeset sizes when copying graphs. PiperOrigin-RevId: 355055758 Change-Id: Id78260cda6f8bf9ed30663ecc819b5936fff26a8
246 lines
6.4 KiB
C++
246 lines
6.4 KiB
C++
/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
|
|
|
|
Licensed 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 TENSORFLOW_GRAPH_EDGESET_H_
|
|
#define TENSORFLOW_GRAPH_EDGESET_H_
|
|
|
|
#include <stddef.h>
|
|
|
|
#include "tensorflow/core/lib/gtl/flatset.h"
|
|
#include "tensorflow/core/platform/logging.h"
|
|
#include "tensorflow/core/platform/macros.h"
|
|
#include "tensorflow/core/platform/types.h"
|
|
namespace tensorflow {
|
|
|
|
class Edge;
|
|
|
|
// An unordered set of edges. Uses very little memory for small sets.
|
|
// Unlike gtl::FlatSet, EdgeSet does NOT allow mutations during
|
|
// iteration.
|
|
class EdgeSet {
|
|
public:
|
|
EdgeSet();
|
|
~EdgeSet();
|
|
|
|
typedef const Edge* key_type;
|
|
typedef const Edge* value_type;
|
|
typedef size_t size_type;
|
|
typedef ptrdiff_t difference_type;
|
|
|
|
class const_iterator;
|
|
typedef const_iterator iterator;
|
|
|
|
bool empty() const;
|
|
size_type size() const;
|
|
void clear();
|
|
std::pair<iterator, bool> insert(value_type value);
|
|
size_type erase(key_type key);
|
|
void reserve(size_type new_size) {
|
|
if (new_size > kInline) {
|
|
auto s = new gtl::FlatSet<const Edge*>(new_size);
|
|
s->insert(reinterpret_cast<const Edge**>(std::begin(ptrs_)),
|
|
reinterpret_cast<const Edge**>(&ptrs_[0] + size()));
|
|
ptrs_[0] = this;
|
|
ptrs_[1] = s;
|
|
}
|
|
}
|
|
|
|
// Caller is not allowed to mutate the EdgeSet while iterating.
|
|
const_iterator begin() const;
|
|
const_iterator end() const;
|
|
|
|
private:
|
|
// Up to kInline elements are stored directly in ptrs_ (nullptr means none).
|
|
// If ptrs_[0] == this then ptrs_[1] points to a set<const Edge*>.
|
|
// kInline must be >= 2, and is chosen such that ptrs_ fills a 64 byte
|
|
// cacheline.
|
|
static constexpr int kInline = 64 / sizeof(const void*);
|
|
const void* ptrs_[kInline];
|
|
|
|
gtl::FlatSet<const Edge*>* get_set() const {
|
|
if (ptrs_[0] == this) {
|
|
return static_cast<gtl::FlatSet<const Edge*>*>(
|
|
const_cast<void*>(ptrs_[1]));
|
|
} else {
|
|
return nullptr;
|
|
}
|
|
}
|
|
|
|
// To detect mutations while iterating.
|
|
#ifdef NDEBUG
|
|
void RegisterMutation() {}
|
|
#else
|
|
uint32 mutations_ = 0;
|
|
void RegisterMutation() { mutations_++; }
|
|
#endif
|
|
|
|
TF_DISALLOW_COPY_AND_ASSIGN(EdgeSet);
|
|
};
|
|
|
|
class EdgeSet::const_iterator {
|
|
public:
|
|
typedef typename EdgeSet::value_type value_type;
|
|
typedef const typename EdgeSet::value_type& reference;
|
|
typedef const typename EdgeSet::value_type* pointer;
|
|
typedef typename EdgeSet::difference_type difference_type;
|
|
typedef std::forward_iterator_tag iterator_category;
|
|
|
|
const_iterator() {}
|
|
|
|
const_iterator& operator++();
|
|
const_iterator operator++(int /*unused*/);
|
|
const value_type* operator->() const;
|
|
value_type operator*() const;
|
|
bool operator==(const const_iterator& other) const;
|
|
bool operator!=(const const_iterator& other) const {
|
|
return !(*this == other);
|
|
}
|
|
|
|
private:
|
|
friend class EdgeSet;
|
|
|
|
void const* const* array_iter_ = nullptr;
|
|
typename gtl::FlatSet<const Edge*>::const_iterator tree_iter_;
|
|
|
|
#ifdef NDEBUG
|
|
inline void Init(const EdgeSet* e) {}
|
|
inline void CheckNoMutations() const {}
|
|
#else
|
|
inline void Init(const EdgeSet* e) {
|
|
owner_ = e;
|
|
init_mutations_ = e->mutations_;
|
|
}
|
|
inline void CheckNoMutations() const {
|
|
CHECK_EQ(init_mutations_, owner_->mutations_);
|
|
}
|
|
const EdgeSet* owner_ = nullptr;
|
|
uint32 init_mutations_ = 0;
|
|
#endif
|
|
};
|
|
|
|
inline EdgeSet::EdgeSet() {
|
|
for (int i = 0; i < kInline; i++) {
|
|
ptrs_[i] = nullptr;
|
|
}
|
|
}
|
|
|
|
inline EdgeSet::~EdgeSet() { delete get_set(); }
|
|
|
|
inline bool EdgeSet::empty() const { return size() == 0; }
|
|
|
|
inline EdgeSet::size_type EdgeSet::size() const {
|
|
auto s = get_set();
|
|
if (s) {
|
|
return s->size();
|
|
} else {
|
|
size_t result = 0;
|
|
for (int i = 0; i < kInline; i++) {
|
|
if (ptrs_[i]) result++;
|
|
}
|
|
return result;
|
|
}
|
|
}
|
|
|
|
inline void EdgeSet::clear() {
|
|
RegisterMutation();
|
|
delete get_set();
|
|
for (int i = 0; i < kInline; i++) {
|
|
ptrs_[i] = nullptr;
|
|
}
|
|
}
|
|
|
|
inline EdgeSet::const_iterator EdgeSet::begin() const {
|
|
const_iterator ci;
|
|
ci.Init(this);
|
|
auto s = get_set();
|
|
if (s) {
|
|
ci.tree_iter_ = s->begin();
|
|
} else {
|
|
ci.array_iter_ = &ptrs_[0];
|
|
}
|
|
return ci;
|
|
}
|
|
|
|
inline EdgeSet::const_iterator EdgeSet::end() const {
|
|
const_iterator ci;
|
|
ci.Init(this);
|
|
auto s = get_set();
|
|
if (s) {
|
|
ci.tree_iter_ = s->end();
|
|
} else {
|
|
ci.array_iter_ = &ptrs_[size()];
|
|
}
|
|
return ci;
|
|
}
|
|
|
|
inline EdgeSet::const_iterator& EdgeSet::const_iterator::operator++() {
|
|
CheckNoMutations();
|
|
if (array_iter_ != nullptr) {
|
|
++array_iter_;
|
|
} else {
|
|
++tree_iter_;
|
|
}
|
|
return *this;
|
|
}
|
|
|
|
inline EdgeSet::const_iterator EdgeSet::const_iterator::operator++(
|
|
int /*unused*/) {
|
|
CheckNoMutations();
|
|
const_iterator tmp = *this;
|
|
operator++();
|
|
return tmp;
|
|
}
|
|
|
|
// gcc's set and multiset always use const_iterator since it will otherwise
|
|
// allow modification of keys.
|
|
inline const EdgeSet::const_iterator::value_type* EdgeSet::const_iterator::
|
|
operator->() const {
|
|
CheckNoMutations();
|
|
if (array_iter_ != nullptr) {
|
|
return reinterpret_cast<const value_type*>(array_iter_);
|
|
} else {
|
|
return tree_iter_.operator->();
|
|
}
|
|
}
|
|
|
|
// gcc's set and multiset always use const_iterator since it will otherwise
|
|
// allow modification of keys.
|
|
inline EdgeSet::const_iterator::value_type EdgeSet::const_iterator::operator*()
|
|
const {
|
|
CheckNoMutations();
|
|
if (array_iter_ != nullptr) {
|
|
return static_cast<value_type>(*array_iter_);
|
|
} else {
|
|
return *tree_iter_;
|
|
}
|
|
}
|
|
|
|
inline bool EdgeSet::const_iterator::operator==(
|
|
const const_iterator& other) const {
|
|
DCHECK((array_iter_ == nullptr) == (other.array_iter_ == nullptr))
|
|
<< "Iterators being compared must be from same set that has not "
|
|
<< "been modified since the iterator was constructed";
|
|
CheckNoMutations();
|
|
if (array_iter_ != nullptr) {
|
|
return array_iter_ == other.array_iter_;
|
|
} else {
|
|
return other.array_iter_ == nullptr && tree_iter_ == other.tree_iter_;
|
|
}
|
|
}
|
|
|
|
} // namespace tensorflow
|
|
|
|
#endif // TENSORFLOW_GRAPH_EDGESET_H_
|