Use insert_or_assign instead of erase+emplace (pytorch#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: pytorch#164868
Approved by: https://github.com/eqy

Author: lakshayg

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
     }

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));
   }
 };
 

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;
 }

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 {