Use insert_or_assign instead of erase+emplace (#164868)

insert_or_assign does effectively the same thing as
erase+emplace but more efficiently since the search
does not need to be repeated

Pull Request resolved: https://github.com/pytorch/pytorch/pull/164868
Approved by: https://github.com/eqy
This commit is contained in:
Lakshay Garg
2025-10-08 19:13:49 +00:00
committed by PyTorch MergeBot
parent 37c6087334
commit a4110fedcf
4 changed files with 4 additions and 8 deletions

View File

@ -337,8 +337,7 @@ struct BenchmarkCache {
engine_cache_order.begin(), engine_cache_order, it->second.second);
}
} else {
engine_cache.erase(key);
engine_cache.emplace(
engine_cache.insert_or_assign(
key,
std::make_pair(results, engine_cache_order.end())); // dummy iterator
}

View File

@ -371,8 +371,7 @@ struct MHAGraphCache {
}
void update(const KeyType& key, T& results) {
engine_cache.erase(key);
engine_cache.emplace(key, std::move(results));
engine_cache.insert_or_assign(key, std::move(results));
}
};

View File

@ -152,8 +152,7 @@ static at::Tensor empty_strided_p2p_persistent(
auto allocated = at::from_blob(dev_ptr, size, stride, options);
// Track the allocation's activeness
alloc_id_to_storage.erase(alloc_id);
alloc_id_to_storage.emplace(
alloc_id_to_storage.insert_or_assign(
alloc_id, allocated.storage().getWeakStorageImpl());
return allocated;
}

View File

@ -138,8 +138,7 @@ struct TORCH_API ExecutionTraceObserver { // NOLINT
// So we need to remove the key and insert the key with the new value.
data_ptr_to_storage_id.erase(raw_data_ptr);
data_ptr_to_storage_id[raw_data_ptr] = id;
data_ptr_to_weak_storage_ptr.erase(raw_data_ptr);
data_ptr_to_weak_storage_ptr.emplace(
data_ptr_to_weak_storage_ptr.insert_or_assign(
raw_data_ptr, t_storage.getWeakStorageImpl());
return id;
} else {