Fix a bug that in cost model manager, graph ptr is used as key of map but memory allocator returns same addresses for graph objects when the graph mgr is being used repeatedly.

Change: 138251588
This commit is contained in:
Yuefeng Zhou 2016-11-04 16:04:16 -08:00 committed by TensorFlower Gardener
parent 8842777dbb
commit 374322eb6a
4 changed files with 21 additions and 0 deletions

View File

@ -42,6 +42,17 @@ CostModel* CostModelManager::FindOrCreateCostModel(const Graph* graph) {
return cost_model;
}
bool CostModelManager::RemoveCostModelForGraph(const Graph* graph) {
mutex_lock l(mu_);
auto itr = cost_models_.find(graph);
if (itr == cost_models_.end()) {
return false;
}
delete itr->second;
cost_models_.erase(graph);
return true;
}
Status CostModelManager::AddToCostGraphDef(const Graph* graph,
CostGraphDef* cost_graph) {
mutex_lock l(mu_);

View File

@ -41,6 +41,8 @@ class CostModelManager {
CostModel* FindOrCreateCostModel(const Graph* graph);
bool RemoveCostModelForGraph(const Graph* graph);
Status AddToCostGraphDef(const Graph* graph, CostGraphDef* cost_graph);
private:

View File

@ -53,6 +53,9 @@ GraphMgr::~GraphMgr() {
GraphMgr::Item::~Item() {
for (const auto& unit : this->units) {
CHECK_NOTNULL(unit.device);
if (!graph_mgr->skip_cost_models_) {
graph_mgr->cost_model_manager_.RemoveCostModelForGraph(unit.graph);
}
delete unit.root;
delete unit.lib;
unit.device->op_segment()->RemoveHold(this->session);
@ -139,6 +142,7 @@ Status GraphMgr::InitItem(const string& session, const GraphDef& gdef,
Status s;
item->units.reserve(partitions.size());
item->graph_mgr = this;
const auto& optimizer_opts = graph_options.optimizer_options();
GraphOptimizer optimizer(optimizer_opts);
for (auto&& p : partitions) {

View File

@ -118,6 +118,10 @@ class GraphMgr {
// A graph is partitioned over multiple devices. Each partition
// has a root executor which may call into the runtime library.
std::vector<ExecutionUnit> units;
// Used to deresgister a cost model when cost model is requried in graph
// manager.
GraphMgr* graph_mgr;
};
// Not owned.