[PyTorch] Reduce heap allocations in OperatorName::setNamespaceIfNotSet (#64673)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/64673

We are now guaranteed to allocate at most one time in this function.
ghstack-source-id: 137786392

Test Plan: Previous diff adds test coverage for this function.

Reviewed By: dhruvbird

Differential Revision: D30813014

fbshipit-source-id: 17d844a1cc8c30574afcc6b0b41b219e62c0b723

diff --git a/aten/src/ATen/core/operator_name.h b/aten/src/ATen/core/operator_name.h
index 2a92697..6440a69 100644 (file)
--- a/aten/src/ATen/core/operator_name.h
+++ b/aten/src/ATen/core/operator_name.h
@@ -34,10 +34,15 @@ struct OperatorName final {
 
   // Returns true if we successfully set the namespace
   bool setNamespaceIfNotSet(const char* ns) {
-    std::ostringstream oss;
     if (!getNamespace().has_value()) {
-      oss << ns << "::" << name;
-      name = oss.str();
+      const auto ns_len = strlen(ns);
+      const auto old_name_size = name.size();
+      name.resize(ns_len + 2 + old_name_size);
+      // Shift current value of name to the end of the new space.
+      name.replace(name.size() - old_name_size, old_name_size, name, 0, old_name_size);
+      name.replace(0, ns_len, ns, ns_len);
+      name[ns_len] = ':';
+      name[ns_len + 1] = ':';
       return true;
     } else {
       return false;