[SE] Make ExecutorCache thread-safe, change ExecutorCache::Insert to ExecutorCache::GetOrCreate. Add support for creating Executors for different device ordinals in parallel.

[XLA] Create Executors in parallel.

PiperOrigin-RevId: 163734988

2 files changed, 2 insertions, 20 deletions

diff --git a/tensorflow/compiler/plugin/executor/platform.cc b/tensorflow/compiler/plugin/executor/platform.cc
index 2f339f04a7..404e1c3da3 100644
--- a/tensorflow/compiler/plugin/executor/platform.cc
+++ b/tensorflow/compiler/plugin/executor/platform.cc
@@ -64,23 +64,8 @@ ExecutorPlatform::ExecutorForDeviceWithPluginConfig(
 
 port::StatusOr<StreamExecutor*> ExecutorPlatform::GetExecutor(
     const StreamExecutorConfig& config) {
-  mutex_lock lock(executors_mutex_);
-
-  port::StatusOr<StreamExecutor*> status = executor_cache_.Get(config);
-  if (status.ok()) {
-    return status.ValueOrDie();
-  }
-
-  port::StatusOr<std::unique_ptr<StreamExecutor>> executor =
-      GetUncachedExecutor(config);
-  if (!executor.ok()) {
-    return executor.status();
-  }
-
-  StreamExecutor* naked_executor = executor.ValueOrDie().get();
-  SE_RETURN_IF_ERROR(
-      executor_cache_.Insert(config, executor.ConsumeValueOrDie()));
-  return naked_executor;
+  return executor_cache_.GetOrCreate(
+      config, [&]() { return GetUncachedExecutor(config); });
 }
 
 port::StatusOr<std::unique_ptr<StreamExecutor>>
diff --git a/tensorflow/compiler/plugin/executor/platform.h b/tensorflow/compiler/plugin/executor/platform.h
index c252a589d4..624bcd5a4e 100644
--- a/tensorflow/compiler/plugin/executor/platform.h
+++ b/tensorflow/compiler/plugin/executor/platform.h
@@ -67,9 +67,6 @@ class ExecutorPlatform : public Platform {
   // This platform's name.
   string name_;
 
-  // mutex that guards the ordinal-to-executor map.
-  mutable mutex executors_mutex_;
-
   // Cache of created StreamExecutors.
   ExecutorCache executor_cache_;