From a4110fedcf72eaede76324bb5c21a76589d75849 Mon Sep 17 00:00:00 2001
From: Lakshay Garg <lakshayg@nvidia.com>
Date: Wed, 8 Oct 2025 19:13:49 +0000
Subject: [PATCH] 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
---
 aten/src/ATen/native/cudnn/Conv_v8.cpp                      | 3 +--
 aten/src/ATen/native/cudnn/MHA.cpp                          | 3 +--
 torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.cpp    | 3 +--
 torch/csrc/profiler/standalone/execution_trace_observer.cpp | 3 +--
 4 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/aten/src/ATen/native/cudnn/Conv_v8.cpp b/aten/src/ATen/native/cudnn/Conv_v8.cpp
index f9837ccc79a2..c973d91e6674 100644
--- a/aten/src/ATen/native/cudnn/Conv_v8.cpp
+++ b/aten/src/ATen/native/cudnn/Conv_v8.cpp
@@ -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
     }
diff --git a/aten/src/ATen/native/cudnn/MHA.cpp b/aten/src/ATen/native/cudnn/MHA.cpp
index 2c5acf8fd3cb..f7b4360d88fa 100644
--- a/aten/src/ATen/native/cudnn/MHA.cpp
+++ b/aten/src/ATen/native/cudnn/MHA.cpp
@@ -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));
   }
 };
 
diff --git a/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.cpp b/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.cpp
index b0e08b49d374..0316a8d85862 100644
--- a/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.cpp
+++ b/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.cpp
@@ -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;
 }
diff --git a/torch/csrc/profiler/standalone/execution_trace_observer.cpp b/torch/csrc/profiler/standalone/execution_trace_observer.cpp
index 5afb3936fd78..918cc554c5b1 100644
--- a/torch/csrc/profiler/standalone/execution_trace_observer.cpp
+++ b/torch/csrc/profiler/standalone/execution_trace_observer.cpp
@@ -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 {