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:
parent
8842777dbb
commit
374322eb6a
@ -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_);
|
||||
|
@ -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:
|
||||
|
@ -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) {
|
||||
|
@ -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.
|
||||
|
Loading…
Reference in New Issue
Block a user