diff --git a/tensorflow/compiler/tf2xla/xla_compiler.cc b/tensorflow/compiler/tf2xla/xla_compiler.cc
index 17adb9b1fdd..d56d6518b7a 100644
--- a/tensorflow/compiler/tf2xla/xla_compiler.cc
+++ b/tensorflow/compiler/tf2xla/xla_compiler.cc
@@ -256,24 +256,12 @@ Status ExecuteGraph(XlaContext* xla_context, std::unique_ptr<Graph> graph,
std::unique_ptr<Executor> exec(exec_ptr);
// At this point ownership of the graph has been transferred to exec.
- auto runner = [](Executor::Args::Closure c) {
- // TODO(misard) Temporarily just schedule c eagerly while we
- // decide what to do about the fact that the ComputationBuilder is
- // thread-compatible, but we don't really want Op writers to have
- // to remember to acquire a lock around every call to
- // ComputationBuilder. One possibility is to add the (generally
- // useful) ability to run a single-threaded Executor based on an
- // option in LocalExecutorParams. Another is to automagically
- // acquire a lock around ComputationBuilder calls using some
- // wrapper or RAII funny business.
- c();
- };
-
// Run the graph symbolically, turning the graph into an XLA computation.
Executor::Args exec_args;
exec_args.step_id = step_id;
exec_args.step_container = step_container.get();
- exec_args.runner = runner;
+ // Run all compilation kernels on the main thread.
+ exec_args.runner = [](Executor::Args::Closure c) { c(); };
TF_RETURN_WITH_CONTEXT_IF_ERROR(
exec->Run(exec_args),
"Conversion from TensorFlow graph to XLA computation failed.");
diff --git a/tensorflow/compiler/tf2xla/xla_context.cc b/tensorflow/compiler/tf2xla/xla_context.cc
index a770271628c..05876f86dca 100644
--- a/tensorflow/compiler/tf2xla/xla_context.cc
+++ b/tensorflow/compiler/tf2xla/xla_context.cc
@@ -131,8 +131,6 @@ Status XlaContext::CollectResults(
xla::Computation* computation, bool* requires_runtime_context,
std::vector<ConstRetVal>* compile_time_constants,
int* num_nonconst_outputs) {
- mutex_lock l(mu_);
-
xla::ComputationDataHandle handle;
if (retval_.empty() && has_side_effects_) {
// Build a empty tuple return value for computations that have side effects
@@ -200,7 +198,6 @@ XlaContext::XlaContext(XlaCompiler* compiler, xla::Client* client,
const xla::ComputationDataHandle&
XlaContext::GetOrCreateRuntimeContextParameter() {
- mutex_lock lock(mu_);
CHECK(allow_cpu_custom_calls_);
CHECK(!use_tuple_arg_);
if (has_context_parameter_) return context_parameter_;
@@ -220,7 +217,6 @@ void XlaContext::AddRetval(int retval_index,
// Add the return value to the list being built up. The executor
// is multi-threaded so this has to happen under the
// lock.
- mutex_lock l(mu_);
retval_.emplace_back(retval_index, handle);
}
@@ -232,17 +228,14 @@ Status XlaContext::AddConstRetval(int retval_index, DataType dtype,
ConstRetVal value;
value.index = retval_index;
TF_RETURN_IF_ERROR(LiteralToHostTensor(literal, dtype, &value.value));
- mutex_lock l(mu_);
compile_time_constant_.push_back(std::move(value));
} else {
- mutex_lock l(mu_);
retval_.emplace_back(retval_index, xla_builder_.ConstantLiteral(literal));
}
return Status::OK();
}
void XlaContext::AddSideEffects() {
- mutex_lock lock(mu_);
has_side_effects_ = true;
}
@@ -323,7 +316,6 @@ const xla::Computation* XlaContext::LookupOrCreate(
DataType type, ComputationMap* out,
const std::function<xla::Computation()>& create) {
{
- mutex_lock l(mu_);
const auto& entry = (*out)[type];
if (!entry.IsNull()) {
return &entry;
@@ -331,7 +323,6 @@ const xla::Computation* XlaContext::LookupOrCreate(
}
auto new_entry = create();
{
- mutex_lock l(mu_);
// Somebody else might have made one concurrently.
auto& entry = (*out)[type];
if (entry.IsNull()) {
diff --git a/tensorflow/compiler/tf2xla/xla_context.h b/tensorflow/compiler/tf2xla/xla_context.h
index 8ece3d37984..53009bf8809 100644
--- a/tensorflow/compiler/tf2xla/xla_context.h
+++ b/tensorflow/compiler/tf2xla/xla_context.h
@@ -29,8 +29,6 @@ limitations under the License.
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/graph/graph.h"
#include "tensorflow/core/platform/macros.h"
-#include "tensorflow/core/platform/mutex.h"
-#include "tensorflow/core/platform/thread_annotations.h"
namespace tensorflow {
@@ -223,11 +221,9 @@ class XlaContext : public ResourceBase {
XlaCompiler* const compiler_;
- mutable mutex mu_;
-
// The ComputationBuilder used to construct the subgraph's compiled
// representation.
- xla::ComputationBuilder xla_builder_ GUARDED_BY(mu_);
+ xla::ComputationBuilder xla_builder_;
// Number of XLA Parameters, not counting the context parameter, if any.
int num_parameters_;
@@ -252,18 +248,17 @@ class XlaContext : public ResourceBase {
// for an additional final parameter to the computation, through which will be
// passed a XlaLocalRuntimeContext* at runtime. Created on demand by
// GetOrCreateRuntimeContextParameter().
- bool has_context_parameter_ GUARDED_BY(mu_) = false;
- xla::ComputationDataHandle context_parameter_ GUARDED_BY(mu_);
+ bool has_context_parameter_ = false;
+ xla::ComputationDataHandle context_parameter_;
// The data-dependent return values of the computation.
- std::vector<std::pair<int, xla::ComputationDataHandle>> retval_
- GUARDED_BY(mu_);
+ std::vector<std::pair<int, xla::ComputationDataHandle>> retval_;
// The non-data-dependent return values of the computation.
- std::vector<ConstRetVal> compile_time_constant_ GUARDED_BY(mu_);
+ std::vector<ConstRetVal> compile_time_constant_;
// Does the computation have side effects, i.e., Send() calls?
- bool has_side_effects_ GUARDED_BY(mu_) = false;
+ bool has_side_effects_ = false;
// Cache of prebuilt computations indexed by their type.
using ComputationMap = std::map<DataType, xla::Computation>;
@@ -273,16 +268,16 @@ class XlaContext : public ResourceBase {
// map. The returned value != nullptr and is owned by the map.
const xla::Computation* LookupOrCreate(
DataType type, ComputationMap* out,
- const std::function<xla::Computation()>& create) LOCKS_EXCLUDED(mu_);
+ const std::function<xla::Computation()>& create);
// Cached computation to compute Max of two elements, specialized by type.
- ComputationMap max_func_ GUARDED_BY(mu_);
+ ComputationMap max_func_;
// Cached computation to compute Sum of two elements, specialized by type.
- ComputationMap add_func_ GUARDED_BY(mu_);
+ ComputationMap add_func_;
// Cached computation to compute Sigmoid of an element, specialized by type.
- ComputationMap sigmoid_func_ GUARDED_BY(mu_);
+ ComputationMap sigmoid_func_;
TF_DISALLOW_COPY_AND_ASSIGN(XlaContext);
};