visibility = ["//visibility:public"],
deps = select({
"//tensorflow:android": [
+ "//tensorflow/core:lib",
"//tensorflow/core:android_tensorflow_lib_lite",
],
"//conditions:default": [
// See WARNING comment below - would be nice to rework to avoid this
// subtlety.
tensorflow::tf_shared_lock l(ctx->functions_mu);
- status->status =
- tensorflow::KernelAndDevice::Init(ndef, ctx->func_lib(device), kernel);
+ status->status = tensorflow::KernelAndDevice::Init(
+ ndef, ctx->func_lib(device), &ctx->runner, kernel);
if (!status->status.ok()) {
delete kernel;
return;
#include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/common_runtime/rendezvous_mgr.h"
#include "tensorflow/core/framework/rendezvous.h"
+#include "tensorflow/core/lib/core/threadpool.h"
#include "tensorflow/core/lib/gtl/map_util.h"
#include "tensorflow/core/lib/gtl/stl_util.h"
#include "tensorflow/core/platform/mutex.h"
struct TFE_Context {
explicit TFE_Context(const TFE_ContextOptions& opts, TF_Session* s)
- : policy(opts.policy),
+ : thread_pool(new tensorflow::thread::ThreadPool(
+ opts.session_options.options.env, "EagerCompute",
+ opts.session_options.options.config
+ .inter_op_parallelism_threads() != 0
+ ? opts.session_options.options.config
+ .inter_op_parallelism_threads()
+ : tensorflow::port::NumSchedulableCPUs())),
+ runner([this](std::function<void()> f) { thread_pool->Schedule(f); }),
+ policy(opts.policy),
session(s),
rendezvous(new tensorflow::IntraProcessRendezvous(s->device_mgr)),
pflr(new tensorflow::ProcessFunctionLibraryRuntime(
log_device_placement(
opts.session_options.options.config.log_device_placement()) {}
+ const std::unique_ptr<tensorflow::thread::ThreadPool> thread_pool;
+ std::function<void(std::function<void()>)> runner;
+
const TFE_ContextDevicePlacementPolicy policy;
// Note: we cannot use C++11 thread_local here as there is no concept of a
out->device_ = device;
out->kernel_.reset(k);
out->flib_ = nullptr;
+ out->runner_ = nullptr;
+ out->default_runner_ = [](std::function<void()> f) { f(); };
return s;
}
// static
Status KernelAndDevice::Init(const NodeDef& ndef, FunctionLibraryRuntime* flib,
+ std::function<void(std::function<void()>)>* runner,
KernelAndDevice* out) {
OpKernel* k = nullptr;
Status s = flib->CreateKernel(ndef, &k);
out->device_ = flib->device();
out->kernel_.reset(k);
out->flib_ = flib;
+ out->runner_ = runner;
+ out->default_runner_ = [](std::function<void()> f) { f(); };
return s;
}
if (stats != nullptr) {
params.track_allocations = true;
}
- // TODO(apassos): use a thread pool.
- std::function<void(std::function<void()>)> runner =
- [](std::function<void()> f) { f(); };
- params.runner = &runner;
+ if (runner_ == nullptr) {
+ params.runner = &default_runner_;
+ } else {
+ params.runner = runner_;
+ }
OpKernelContext context(¶ms);
device_->Compute(kernel_.get(), &context);
// the FunctionLibraryRuntime is pushed on to the caller (see locking in
// c_api.cc).
static Status Init(const NodeDef& ndef, FunctionLibraryRuntime* flib,
+ std::function<void(std::function<void()>)>* runner,
KernelAndDevice* out);
// TODO(ashankar): Remove this
static Status InitOp(Device* device, const NodeDef& ndef,
private:
std::unique_ptr<OpKernel> kernel_;
Device* device_;
+ std::function<void(std::function<void()>)>* runner_;
+ std::function<void(std::function<void()>)> default_runner_;
FunctionLibraryRuntime* flib_;
checkpoint::TensorSliceReaderCacheWrapper slice_reader_cache_;
Rendezvous* rendez_;
.BuildNodeDef());
TestEnv env;
KernelAndDevice kernel(nullptr);
- Status s =
- KernelAndDevice::Init(ndef, env.function_library_runtime(), &kernel);
+ Status s = KernelAndDevice::Init(ndef, env.function_library_runtime(),
+ nullptr, &kernel);
ASSERT_TRUE(s.ok()) << s;
std::vector<Tensor> outputs;
s = kernel.Run(&inputs, &outputs, nullptr);
KernelAndDevice k(nullptr);
tensorflow::testing::StartTiming();
for (int i = 0; i < iters; ++i) {
- TF_CHECK_OK(
- KernelAndDevice::Init(ndef, env.function_library_runtime(), &k));
+ TF_CHECK_OK(KernelAndDevice::Init(ndef, env.function_library_runtime(),
+ nullptr, &k));
}
}
BENCHMARK(BM_KernelAndDeviceInit);
.BuildNodeDef());
TestEnv env;
KernelAndDevice kernel(nullptr);
- TF_CHECK_OK(
- KernelAndDevice::Init(ndef, env.function_library_runtime(), &kernel));
+ TF_CHECK_OK(KernelAndDevice::Init(ndef, env.function_library_runtime(),
+ nullptr, &kernel));
tensorflow::testing::StartTiming();
for (int i = 0; i < iters; ++i) {
TF_CHECK_OK(kernel.Run(&inputs, &outputs, nullptr));