A series of changes to significantly reduce the number of allocations

done by models distributed across many devices.  A small
microbenchmark model that runs two banks (A and B) of 30 nodes with a
30x30 full shuffle between them, where each of the nodes in A and in B
run with one node on each of the 30 devices (so 30*29+30+30, or ~930
separate RPCs) was showing ~111,000 allocations per iteration of the graph.

With the changes here, this is now down to ~64,300 allocations per iteration.

Changes include:

o DeviceContext::CopyDeviceTensorToCPU and related helper routines:
use StringPiece instead of const string& for the tensor name (avoids
creating a string in some cases where the caller only has a
StringPiece available).

o Change some Rendezvous and BaseRemoteRendezvous interfaces to
take a 'const Rendezvous::ParsedKey& key', rather than 'const string& key'.
In many cases, the callers were already having to parse the key
into a ParsedKey, and so we were doing the parsing multiple times at
different levels as we processed receiving or sending of a tensor.  This
reduces the number of times that we parse a key as it flows from a Send
node through to a Recv node on another worker.

o Changed Rendezvous::ParsedKey so that it makes a copy of the underlying
full key, and then uses StringPiece objects to point into this copy for
the src_device, dst_device, and edge_name pieces.  This turns 3 string
allocations into 1 per Rendezvous::ParseKey call.

o Added new StringPiece Rendezvous::ParsedKey::FullKey() accessor to
return a StringPiece for the underlying full key, and used that in a
few places (mostly logging) where that is useful.

o In many places, used std::move(function_variable) when assigning to
  an instance variable.  This eliminates a very large number of excess
  std::function allocations/initializations (~56000 of the baseline
  allocations were related to std::function setup or cloning, and this
  is now down to ~11000 after this cl).

o In the RPC-based remote workers (StubbyRemoteWorker and
GrpcRemoteWorker), changed the code path in RecvTensorAsync to avoid
creation of a std::function with 6 arguments unless necessary.  There
are three cases now handled separately:

  (a) We're not logging, and we didn't make a copy of the request that we
      need to free: just use the passed in 'StatusCallback done' object
      directly, without creating a wrapper std::function object at all

  (b) We're not logging, but we made a copy of the request that we
      need to free: we create a simple wrapper std::function that
      invokes the passed in 'done' callback, and then frees the
      req_copy request copy object.

  (c) We're logging: we create the std::function object with all the
      necessary state to log when the recv has finished.

o Changed DeviceMgr::LookupDevice to take a StringPiece, rather than a
const string&, and changed the hash table to use StringPiece keys.
This allows clients that just have a StringPiece device name in their
hand to avoid a string creation to lookup the Device* object.

o Changed ExecutorState to use a specialized TaggedNodeReadyQueue that
internally uses a gtl::InlinedVector<TaggedNode, 16>, rather than
using a std::deque<TaggedNode> for keeping track of nodes ready to
execute.  This is faster because it avoids allocations entirely if the
ready node queue doesn't get bigger than 16, and inlined vectors are
generally faster than std::deque, at a minor risk of using more memory
if this queue grows to very large numbers of ready nodes (mostly imaginable
only in pathological graphs).

o In ExecutorState::Process, allocated a single ExecutorState::AsyncState
object to keep track of all the state we need to preserve for an asynchronously
executed node, rather than keeping this state implicitly via a very large
number of arguments to a lamda function.

o Added new atomic std::atomic<bool> status_is_ok_ in
BaseRemoteRendezvous.  This allows us to avoid acquiring the lock when
we just want to check if the status is non-OK in
BaseRemoteRendezvous::Send and BaseRemoteRendezvous::ValidateDevices.

o In GraphMgr::RunAllDone, changed assignment of args.runner to avoid
one extra level of std::function indirection (binding the function directly
to the ThreadPool::Schedule routine, rather than creating an intermediate
lambda function that invokes this inside the body of the lambda.

o Added freelist of RpcRecvTensorCall objects in
third_party/tensorflow/core/distributed_runtime/rpc/rpc_rendezvous_mgr.cc

o Changed third_party/tensorflow/core/framework/rendezvous.cc to keep the
hashtable of Item* objects keyed by uint64 (hash of the tensor name), rather
than the full-string tensor name.  Collisions in the 64-bit hash space
should basically never happen.

o Sped up DeviceNameUtils::ParseFullName by optimizing for the common
ordering of parts of /job, /replica, /task, /device.  The parsing code
was general enough to handle any order, but did so by comparing the
prefixes 4, 3, 2, and 1 times, respectively, rather than 1, 1, 1, and 1 times.

o Sped up DeviceNameUtils::SplitDeviceName to avoid extra string copies.
Change: 125991891

1 files changed, 38 insertions, 13 deletions

diff --git a/tensorflow/core/util/device_name_utils.cc b/tensorflow/core/util/device_name_utils.cc
index 25edfd1405..5816dbd40c 100644
--- a/tensorflow/core/util/device_name_utils.cc
+++ b/tensorflow/core/util/device_name_utils.cc
@@ -104,22 +104,29 @@ bool DeviceNameUtils::ParseFullName(StringPiece fullname, ParsedName* p) {
   }
   StringPiece tmp;
   while (!fullname.empty()) {
+    bool progress = false;
     if (str_util::ConsumePrefix(&fullname, "/job:")) {
       p->has_job = !str_util::ConsumePrefix(&fullname, "*");
       if (p->has_job && !ConsumeJobName(&fullname, &p->job)) {
         return false;
       }
-    } else if (str_util::ConsumePrefix(&fullname, "/replica:")) {
+      progress = true;
+    }
+    if (str_util::ConsumePrefix(&fullname, "/replica:")) {
       p->has_replica = !str_util::ConsumePrefix(&fullname, "*");
       if (p->has_replica && !ConsumeNumber(&fullname, &p->replica)) {
         return false;
       }
-    } else if (str_util::ConsumePrefix(&fullname, "/task:")) {
+      progress = true;
+    }
+    if (str_util::ConsumePrefix(&fullname, "/task:")) {
       p->has_task = !str_util::ConsumePrefix(&fullname, "*");
       if (p->has_task && !ConsumeNumber(&fullname, &p->task)) {
         return false;
       }
-    } else if (str_util::ConsumePrefix(&fullname, "/device:")) {
+      progress = true;
+    }
+    if (str_util::ConsumePrefix(&fullname, "/device:")) {
       p->has_type = !str_util::ConsumePrefix(&fullname, "*");
       if (p->has_type && !ConsumeDeviceType(&fullname, &p->type)) {
         return false;
@@ -132,24 +139,31 @@ bool DeviceNameUtils::ParseFullName(StringPiece fullname, ParsedName* p) {
           return false;
         }
       }
+      progress = true;
+    }
 
-    } else if (str_util::ConsumePrefix(&fullname, "/cpu:") ||
-               str_util::ConsumePrefix(&fullname, "/CPU:")) {
+    if (str_util::ConsumePrefix(&fullname, "/cpu:") ||
+        str_util::ConsumePrefix(&fullname, "/CPU:")) {
       p->has_type = true;
       p->type = "CPU";  // Treat '/cpu:..' as uppercase '/device:CPU:...'
       p->has_id = !str_util::ConsumePrefix(&fullname, "*");
       if (p->has_id && !ConsumeNumber(&fullname, &p->id)) {
         return false;
       }
-    } else if (str_util::ConsumePrefix(&fullname, "/gpu:") ||
-               str_util::ConsumePrefix(&fullname, "/GPU:")) {
+      progress = true;
+    }
+    if (str_util::ConsumePrefix(&fullname, "/gpu:") ||
+        str_util::ConsumePrefix(&fullname, "/GPU:")) {
       p->has_type = true;
       p->type = "GPU";  // Treat '/gpu:..' as uppercase '/device:GPU:...'
       p->has_id = !str_util::ConsumePrefix(&fullname, "*");
       if (p->has_id && !ConsumeNumber(&fullname, &p->id)) {
         return false;
       }
-    } else {
+      progress = true;
+    }
+
+    if (!progress) {
       return false;
     }
   }
@@ -340,11 +354,22 @@ bool DeviceNameUtils::SplitDeviceName(StringPiece name, string* task,
                                       string* device) {
   ParsedName pn;
   if (ParseFullName(name, &pn) && pn.has_type && pn.has_id) {
-    *task = strings::StrCat(
-        (pn.has_job ? strings::StrCat("/job:", pn.job) : ""),
-        (pn.has_replica ? strings::StrCat("/replica:", pn.replica) : ""),
-        (pn.has_task ? strings::StrCat("/task:", pn.task) : ""));
-    *device = strings::StrCat(pn.type, ":", pn.id);
+    task->clear();
+    task->reserve(
+        (pn.has_job ? (5 + pn.job.size()) : 0) +
+        (pn.has_replica ? (9 + 4 /*estimated UB for # replica digits*/) : 0) +
+        (pn.has_task ? (6 + 4 /*estimated UB for # task digits*/) : 0));
+    if (pn.has_job) {
+      strings::StrAppend(task, "/job:", pn.job);
+    }
+    if (pn.has_replica) {
+      strings::StrAppend(task, "/replica:", pn.replica);
+    }
+    if (pn.has_task) {
+      strings::StrAppend(task, "/task:", pn.task);
+    }
+    device->clear();
+    strings::StrAppend(device, pn.type, ":", pn.id);
     return true;
   }
   return false;