diff --git a/tensorflow/core/framework/op_kernel.cc b/tensorflow/core/framework/op_kernel.cc index 3e34bf04188..c26214c8579 100644 --- a/tensorflow/core/framework/op_kernel.cc +++ b/tensorflow/core/framework/op_kernel.cc @@ -15,6 +15,7 @@ limitations under the License. #include "tensorflow/core/framework/op_kernel.h" +#include #include #include #include @@ -38,6 +39,7 @@ limitations under the License. #include "tensorflow/core/lib/io/path.h" #include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/lib/strings/strcat.h" +#include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/logging.h" #include "tensorflow/core/platform/mutex.h" #include "tensorflow/core/platform/types.h" @@ -924,12 +926,52 @@ struct KernelRegistration { // KernelDef. typedef std::unordered_multimap KernelRegistry; +#if defined(_WIN32) +static const char kKernelLibPattern[] = "libtfkernel*.dll"; +#elif defined(__APPLE__) +static const char kKernelLibPattern[] = "libtfkernel*.dylib"; +#else +static const char kKernelLibPattern[] = "libtfkernel*.so"; +#endif + +void LoadDynamicKernelsInternal() { + Env* env = Env::Default(); + string bazel_kernel_dir = io::JoinPath(env->GetRunfilesDir(), + "tensorflow", + "core", + "kernels"); + std::vector files; + Status s_kernel_dir = env->GetChildren(bazel_kernel_dir, &files); + if (s_kernel_dir.ok()) { + string dll_spec = io::JoinPath(bazel_kernel_dir, kKernelLibPattern); + for (const auto& file : files) { + string fullpath = io::JoinPath(bazel_kernel_dir, file); + if (env->MatchPath(fullpath, dll_spec)) { + // TODO(gunan): Store the handles to the opened files. + void* unused_filehandle; + TF_CHECK_OK(env->LoadLibrary(fullpath.c_str(), &unused_filehandle)); + } + } + } +} + +// Mechanism for loading existing kernel libraries. +void LoadDynamicKernels() { + // TODO(gunan): As more features are available, add intelligent kernel + // selection, and dropping unsuitable kernel logic here. + static std::once_flag dll_loader_flag; + std::call_once(dll_loader_flag, LoadDynamicKernelsInternal); +} + void* GlobalKernelRegistry() { static KernelRegistry* global_kernel_registry = new KernelRegistry; return global_kernel_registry; } static KernelRegistry* GlobalKernelRegistryTyped() { +#ifdef AUTOLOAD_DYNAMIC_KERNELS + LoadDynamicKernels(); +#endif // AUTOLOAD_DYNAMIC_KERNELS return reinterpret_cast(GlobalKernelRegistry()); } @@ -949,8 +991,17 @@ void OpKernelRegistrar::InitInternal(const KernelDef* kernel_def, const string key = Key(kernel_def->op(), DeviceType(kernel_def->device_type()), kernel_def->label()); - GlobalKernelRegistryTyped()->insert(std::make_pair( - key, KernelRegistration(*kernel_def, kernel_class_name, factory))); + + // To avoid calling LoadDynamicKernels DO NOT CALL GlobalKernelRegistryTyped + // here. + // InitInternal gets called by static initializers, so it ends up executing + // before main. This causes LoadKernelLibraries function to get called + // before some file libraries can initialize, which in turn crashes the + // program flakily. Until we get rid of static initializers in kernel + // registration mechanism, we have this workaround here. + reinterpret_cast(GlobalKernelRegistry()) + ->insert(std::make_pair( + key, KernelRegistration(*kernel_def, kernel_class_name, factory))); } delete kernel_def; } diff --git a/tools/bazel.rc b/tools/bazel.rc index cee0b0b064d..a9c86fd4854 100644 --- a/tools/bazel.rc +++ b/tools/bazel.rc @@ -29,7 +29,7 @@ build:mkl -c opt # This config option is used to enable MKL-DNN open source library only, # without depending on MKL binary version. -build:mkl_open_source_only --define=build_with_mkl_dnn_only=true +build:mkl_open_source_only --define=build_with_mkl_dnn_only=true build:mkl_open_source_only --define=build_with_mkl=true --define=enable_mkl=true build:download_clang --crosstool_top=@local_config_download_clang//:toolchain @@ -86,6 +86,7 @@ build --define=grpc_no_ares=true # Modular TF build options build:dynamic_kernels --define=dynamic_loaded_kernels=true +build:dynamic_kernels -DAUTOLOAD_DYNAMIC_KERNELS # Default paths for TF_SYSTEM_LIBS build --define=PREFIX=/usr